diff --git a/third_party/transformers/src/transformers/models/apertus/__init__.py b/third_party/transformers/src/transformers/models/apertus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..726e3d61c3bdcd7bb1ee85b0dfbaebe8e4815ecf --- /dev/null +++ b/third_party/transformers/src/transformers/models/apertus/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2025 The HuggingFace Inc. team and the Swiss AI Initiative. All rights reserved. +# +# This code is based on HuggingFace's LLaMA implementation in this library. +# It has been modified from its original forms to accommodate the architectural +# differences made by the Swiss AI Initiative 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_apertus import * + from .modeling_apertus 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/apertus/configuration_apertus.py b/third_party/transformers/src/transformers/models/apertus/configuration_apertus.py new file mode 100644 index 0000000000000000000000000000000000000000..1e0122160b307a63a33c7079c6f321cc4ac96a33 --- /dev/null +++ b/third_party/transformers/src/transformers/models/apertus/configuration_apertus.py @@ -0,0 +1,99 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/apertus/modular_apertus.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_apertus.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team and the Swiss AI Initiative. 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="swiss-ai/Apertus-8B-Instruct-2509") +@strict +class ApertusConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import ApertusModel, ApertusConfig + + >>> # Initializing a Apertus-8B style configuration + >>> configuration = ApertusConfig() + + >>> # Initializing a model from the Apertus-8B style configuration + >>> model = ApertusModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "apertus" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 12000000.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.q_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise", + "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 = 131072 + hidden_size: int = 4096 + intermediate_size: int = 14336 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "xielu" + max_position_embeddings: int = 65536 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = 3 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.rope_parameters is None: + self.rope_parameters = { + "rope_type": "llama3", + "rope_theta": 12000000.0, + "factor": 8.0, + "original_max_position_embeddings": 8192, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + } + super().__post_init__(**kwargs) + + +__all__ = ["ApertusConfig"] diff --git a/third_party/transformers/src/transformers/models/apertus/modeling_apertus.py b/third_party/transformers/src/transformers/models/apertus/modeling_apertus.py new file mode 100644 index 0000000000000000000000000000000000000000..7d14dd3d14c8c3909df49f1e0dd7057a4e3a4194 --- /dev/null +++ b/third_party/transformers/src/transformers/models/apertus/modeling_apertus.py @@ -0,0 +1,505 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/apertus/modular_apertus.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_apertus.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team and the Swiss AI Initiative. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from ...activations import ACT2CLS, ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_apertus import ApertusConfig + + +class ApertusMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + 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] + if config.hidden_act == "xielu": + self.act_fn = ACT2CLS["xielu"](dtype=config.dtype) + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +@use_kernel_forward_from_hub("RMSNorm") +class ApertusRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + ApertusRMSNorm 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 ApertusRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: ApertusConfig, 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: ApertusConfig | 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 ApertusAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: ApertusConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = ApertusRMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = ApertusRMSNorm(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]: + 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) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(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 ApertusDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: ApertusConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = ApertusAttention(config=config, layer_idx=layer_idx) + + self.mlp = ApertusMLP(config) + self.attention_layernorm = ApertusRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feedforward_layernorm = ApertusRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.attention_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class ApertusPreTrainedModel(PreTrainedModel): + config: ApertusConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["ApertusDecoderLayer"] + _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": ApertusDecoderLayer, + "attentions": ApertusAttention, + } + + +@auto_docstring +class ApertusModel(ApertusPreTrainedModel): + def __init__(self, config: ApertusConfig): + 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( + [ApertusDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = ApertusRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = ApertusRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class ApertusForCausalLM(ApertusPreTrainedModel, 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 = ApertusModel(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""" + 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, ApertusForCausalLM + + >>> model = ApertusForCausalLM.from_pretrained("swiss-ai/Apertus-8B-Instruct-2509") + >>> tokenizer = AutoTokenizer.from_pretrained("swiss-ai/Apertus-8B-Instruct-2509") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class ApertusForTokenClassification(GenericForTokenClassification, ApertusPreTrainedModel): + pass + + +__all__ = ["ApertusModel", "ApertusForCausalLM", "ApertusForTokenClassification", "ApertusPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/apertus/modular_apertus.py b/third_party/transformers/src/transformers/models/apertus/modular_apertus.py new file mode 100644 index 0000000000000000000000000000000000000000..3c9eb6d8b6eaab3e3f31fb22cc369bad2a41d2bc --- /dev/null +++ b/third_party/transformers/src/transformers/models/apertus/modular_apertus.py @@ -0,0 +1,268 @@ +# Copyright 2025 the HuggingFace Inc. team and the Swiss AI Initiative. 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 +from huggingface_hub.dataclasses import strict +from torch import nn + +from ...activations import ACT2CLS +from ...cache_utils import Cache +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaForTokenClassification, + LlamaModel, + LlamaPreTrainedModel, + LlamaRMSNorm, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from ..nemotron.modeling_nemotron import NemotronMLP + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="swiss-ai/Apertus-8B-Instruct-2509") +@strict +class ApertusConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import ApertusModel, ApertusConfig + + >>> # Initializing a Apertus-8B style configuration + >>> configuration = ApertusConfig() + + >>> # Initializing a model from the Apertus-8B style configuration + >>> model = ApertusModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "apertus" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 12000000.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.q_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise", + "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 = 131072 + hidden_size: int = 4096 + intermediate_size: int = 14336 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "xielu" + max_position_embeddings: int = 65536 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = 3 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.rope_parameters is None: + self.rope_parameters = { + "rope_type": "llama3", + "rope_theta": 12000000.0, + "factor": 8.0, + "original_max_position_embeddings": 8192, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + } + super().__post_init__(**kwargs) + + +class ApertusMLP(NemotronMLP): + def __init__(self, config): + super().__init__(config) + 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) + if config.hidden_act == "xielu": + self.act_fn = ACT2CLS["xielu"](dtype=config.dtype) + + +class ApertusRMSNorm(LlamaRMSNorm): + pass + + +class ApertusRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class ApertusAttention(LlamaAttention): + def __init__(self, config: ApertusConfig, layer_idx: int | None = None): + super().__init__(config, layer_idx) + self.q_norm = ApertusRMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = ApertusRMSNorm(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]: + 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) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(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 ApertusDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: ApertusConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.attention_layernorm = ApertusRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feedforward_layernorm = ApertusRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + del self.input_layernorm + del self.post_attention_layernorm + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.attention_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class ApertusPreTrainedModel(LlamaPreTrainedModel): + pass + + +class ApertusModel(LlamaModel): + pass + + +class ApertusForCausalLM(LlamaForCausalLM): + def forward(self, **super_kwargs): + 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, ApertusForCausalLM + + >>> model = ApertusForCausalLM.from_pretrained("swiss-ai/Apertus-8B-Instruct-2509") + >>> tokenizer = AutoTokenizer.from_pretrained("swiss-ai/Apertus-8B-Instruct-2509") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + return super().forward(**super_kwargs) + + +class ApertusForTokenClassification(LlamaForTokenClassification): + pass + + +__all__ = [ + "ApertusConfig", + "ApertusModel", + "ApertusForCausalLM", + "ApertusForTokenClassification", + "ApertusPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/arcee/__init__.py b/third_party/transformers/src/transformers/models/arcee/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c3df45b2a3b14a72b36362c87833298be8fce48 --- /dev/null +++ b/third_party/transformers/src/transformers/models/arcee/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 Arcee 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_arcee import * + from .modeling_arcee 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/arcee/configuration_arcee.py b/third_party/transformers/src/transformers/models/arcee/configuration_arcee.py new file mode 100644 index 0000000000000000000000000000000000000000..ee711d608204e048d280922de813e493bccb1bfb --- /dev/null +++ b/third_party/transformers/src/transformers/models/arcee/configuration_arcee.py @@ -0,0 +1,100 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/arcee/modular_arcee.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_arcee.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Arcee 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 huggingface_hub.dataclasses import strict + +from transformers.utils import auto_docstring + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +@strict +class ArceeConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import ArceeModel, ArceeConfig + + >>> # Initializing an Arcee AFM-4.5B-Base style configuration + >>> configuration = ArceeConfig() + + >>> # Initializing a model from the AFM-4.5B-Base style configuration + >>> model = ArceeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "arcee" + 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.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 = 32000 + hidden_size: int = 2560 + intermediate_size: int = 18432 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "relu2" + max_position_embeddings: int = 4096 + 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 = 128000 + eos_token_id: int | list[int] | None = 128001 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + + def __post_init__(self, **kwargs): + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +__all__ = ["ArceeConfig"] diff --git a/third_party/transformers/src/transformers/models/arcee/modeling_arcee.py b/third_party/transformers/src/transformers/models/arcee/modeling_arcee.py new file mode 100644 index 0000000000000000000000000000000000000000..8d2d05bf2952ba39cc04b856a901c19eaa810d1b --- /dev/null +++ b/third_party/transformers/src/transformers/models/arcee/modeling_arcee.py @@ -0,0 +1,520 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/arcee/modular_arcee.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_arcee.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Arcee 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 collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from transformers.utils import auto_docstring + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import ( + GenericForQuestionAnswering, + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_arcee import ArceeConfig + + +class ArceeMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = 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.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +@use_kernel_forward_from_hub("RMSNorm") +class ArceeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + ArceeRMSNorm 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 ArceeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: ArceeConfig, 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: ArceeConfig | 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 ArceeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: ArceeConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ArceeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: ArceeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = ArceeAttention(config=config, layer_idx=layer_idx) + + self.mlp = ArceeMLP(config) + self.input_layernorm = ArceeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = ArceeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class ArceePreTrainedModel(PreTrainedModel): + config: ArceeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["ArceeDecoderLayer"] + _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": ArceeDecoderLayer, + "attentions": ArceeAttention, + } + + +@auto_docstring +class ArceeModel(ArceePreTrainedModel): + def __init__(self, config: ArceeConfig): + 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( + [ArceeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = ArceeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = ArceeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForCausalLM(ArceePreTrainedModel, 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 = ArceeModel(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, ArceeForCausalLM + + >>> model = ArceeForCausalLM.from_pretrained("meta-arcee/Arcee-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-arcee/Arcee-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, + ) + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForSequenceClassification(GenericForSequenceClassification, ArceePreTrainedModel): + pass + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForQuestionAnswering(GenericForQuestionAnswering, ArceePreTrainedModel): + base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model` + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForTokenClassification(GenericForTokenClassification, ArceePreTrainedModel): + pass + + +__all__ = [ + "ArceeForCausalLM", + "ArceeForQuestionAnswering", + "ArceeForSequenceClassification", + "ArceeForTokenClassification", + "ArceeModel", + "ArceePreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/arcee/modular_arcee.py b/third_party/transformers/src/transformers/models/arcee/modular_arcee.py new file mode 100644 index 0000000000000000000000000000000000000000..91cd8e13f1eddc424ed156835909c7d76e062412 --- /dev/null +++ b/third_party/transformers/src/transformers/models/arcee/modular_arcee.py @@ -0,0 +1,117 @@ +# Copyright 2025 Arcee AI and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Arcee model.""" + +from huggingface_hub.dataclasses import strict + +from transformers.utils import auto_docstring, logging + +from ...modeling_rope_utils import RopeParameters +from ..llama.configuration_llama import LlamaConfig +from ..llama.modeling_llama import ( + LlamaForCausalLM, + LlamaForQuestionAnswering, + LlamaForSequenceClassification, + LlamaForTokenClassification, +) +from ..nemotron.modeling_nemotron import NemotronMLP + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +@strict +class ArceeConfig(LlamaConfig): + r""" + ```python + >>> from transformers import ArceeModel, ArceeConfig + + >>> # Initializing an Arcee AFM-4.5B-Base style configuration + >>> configuration = ArceeConfig() + + >>> # Initializing a model from the AFM-4.5B-Base style configuration + >>> model = ArceeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "arcee" + 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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + + vocab_size: int = 32000 + hidden_size: int = 2560 + intermediate_size: int = 18432 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "relu2" + max_position_embeddings: int = 4096 + 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 = 128000 + eos_token_id: int | list[int] | None = 128001 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + + pretraining_tp = AttributeError() + + +class ArceeMLP(NemotronMLP): + pass + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForCausalLM(LlamaForCausalLM): + pass + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForSequenceClassification(LlamaForSequenceClassification): + pass + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForQuestionAnswering(LlamaForQuestionAnswering): + pass + + +@auto_docstring(checkpoint="arcee-ai/AFM-4.5B") +class ArceeForTokenClassification(LlamaForTokenClassification): + pass + + +__all__ = [ + "ArceeConfig", + "ArceeForCausalLM", + "ArceeForQuestionAnswering", + "ArceeForSequenceClassification", + "ArceeForTokenClassification", + "ArceeModel", # noqa: F822 + "ArceePreTrainedModel", # noqa: F822 +] diff --git a/third_party/transformers/src/transformers/models/aya_vision/__init__.py b/third_party/transformers/src/transformers/models/aya_vision/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f8be47cb228b19f02b87d195747e78c2a87de752 --- /dev/null +++ b/third_party/transformers/src/transformers/models/aya_vision/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_aya_vision import * + from .modeling_aya_vision import * + from .processing_aya_vision 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/aya_vision/configuration_aya_vision.py b/third_party/transformers/src/transformers/models/aya_vision/configuration_aya_vision.py new file mode 100644 index 0000000000000000000000000000000000000000..442ee717e772c7d415adb6a836760dfc017c64a5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/aya_vision/configuration_aya_vision.py @@ -0,0 +1,80 @@ +# Copyright 2025 Cohere 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. +"""AyaVision model configuration""" + +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="CohereForAI/aya-vision-8b") +@strict +class AyaVisionConfig(PreTrainedConfig): + r""" + downsample_factor (`int`, *optional*, defaults to 2): + The downsample factor to apply to the vision features. + adapter_layer_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon value used for layer normalization in the adapter. + """ + + model_type = "aya_vision" + attribute_map = { + "image_token_id": "image_token_index", + } + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + vision_feature_select_strategy: str = "full" + vision_feature_layer: int | list[int] = -1 + downsample_factor: int = 2 + adapter_layer_norm_eps: float = 1e-6 + image_token_index: int = 255036 + 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", "siglip_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["siglip_vision_model"]( + hidden_size=1152, + intermediate_size=4304, + patch_size=14, + image_size=384, + num_hidden_layers=26, + num_attention_heads=14, + vision_use_head=False, + ) + + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "cohere2") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["cohere2"]() + + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.vision_feature_select_strategy not in ["default", "full"]: + raise ValueError( + "vision_feature_select_strategy should be one of 'default', 'full'." + f"Got: {self.vision_feature_select_strategy}" + ) + + +__all__ = ["AyaVisionConfig"] diff --git a/third_party/transformers/src/transformers/models/aya_vision/modeling_aya_vision.py b/third_party/transformers/src/transformers/models/aya_vision/modeling_aya_vision.py new file mode 100644 index 0000000000000000000000000000000000000000..38f434d405a56284277df84ed11cfbcd4400ecc1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/aya_vision/modeling_aya_vision.py @@ -0,0 +1,446 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/aya_vision/modular_aya_vision.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_aya_vision.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the Cohere 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, torch_compilable_check +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ..auto import AutoModel +from .configuration_aya_vision import AyaVisionConfig + + +class AyaVisionMultiModalProjector(nn.Module): + def __init__(self, config: AyaVisionConfig): + super().__init__() + self.config = config + self.downsample_factor = config.downsample_factor + self.alignment_intermediate_size = getattr( + config, "alignment_intermediate_size", config.text_config.hidden_size + ) + self.layernorm = nn.LayerNorm( + config.vision_config.hidden_size * (config.downsample_factor**2), eps=config.adapter_layer_norm_eps + ) + + self.linear_1 = nn.Linear( + config.vision_config.hidden_size * (config.downsample_factor**2), + self.alignment_intermediate_size, + bias=True, + ) + + self.act = ACT2FN["silu"] # SwiGLU uses SiLU activation + # For SwiGLU, project down to half size since we split intermediate dim + self.linear_2 = nn.Linear(self.alignment_intermediate_size // 2, config.text_config.hidden_size, bias=True) + + def forward(self, image_features): + image_features = self.pixel_shuffle(image_features) + image_features = self.layernorm(image_features) + hidden_states = self.linear_1(image_features) + + # Split along last dimension and apply SwiGLU + x, gate = hidden_states.chunk(2, dim=-1) + hidden_states = self.act(gate) * x + + hidden_states = self.linear_2(hidden_states) + return hidden_states + + def pixel_shuffle(self, image_features): # B, S, D + batch_size, seq_length, feature_dim = image_features.shape + height = width = int(seq_length**0.5) + image_features = image_features.reshape(image_features.shape[0], width, height, -1) + channels = image_features.shape[-1] + image_features = image_features.reshape( + batch_size, width, int(height / self.downsample_factor), int(channels * self.downsample_factor) + ) + image_features = image_features.permute(0, 2, 1, 3) + image_features = image_features.reshape( + batch_size, int(height / self.downsample_factor), int(width / self.downsample_factor), -1 + ) + image_features = image_features.permute(0, 2, 1, 3) + return image_features + + +@auto_docstring +class AyaVisionPreTrainedModel(PreTrainedModel): + config: AyaVisionConfig + 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 AyaVision causal language model (or autoregressive) outputs. + """ +) +class AyaVisionCausalLMOutputWithPast(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 AyaVision outputs, with hidden states and attentions. + """ +) +class AyaVisionModelOutputWithPast(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 AyaVision model which consists of a vision backbone and a language model, without a language modeling head. + """ +) +class AyaVisionModel(AyaVisionPreTrainedModel): + def __init__(self, config: AyaVisionConfig): + super().__init__(config) + self.vision_tower = AutoModel.from_config(config.vision_config) + + self.multi_modal_projector = AyaVisionMultiModalProjector(config) + self.language_model = AutoModel.from_config(config.text_config) + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + @merge_with_config_defaults + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + kwargs = {k: v for k, v in kwargs.items() if v is not None} + # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) + + # If we have one vision feature layer, return the corresponding hidden states, + # otherwise, select the hidden states of each feature layer and concatenate them + if isinstance(vision_feature_layer, int): + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] + if vision_feature_select_strategy == "default": + selected_image_feature = selected_image_feature[:, 1:] + else: + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + # For default; crop CLS from each hidden state in the hidden state pool + if vision_feature_select_strategy == "default": + hs_pool = [hs[:, 1:] for hs in hs_pool] + selected_image_feature = torch.cat(hs_pool, dim=-1) + + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) + + return image_outputs + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + + n_image_tokens = special_image_mask.sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", + ) + return special_image_mask + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | AyaVisionModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + return_dict=True, + ).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return AyaVisionModelOutputWithPast( + 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 AYA_VISION model which consists of a vision backbone and a language model. + """ +) +class AyaVisionForConditionalGeneration(AyaVisionPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + + def __init__(self, config: AyaVisionConfig): + super().__init__(config) + self.model = AyaVisionModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + def get_output_embeddings(self) -> nn.Module: + return self.lm_head + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, + ) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + image_sizes: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | AyaVisionCausalLMOutputWithPast: + 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 AutoProcessor, AyaVisionForConditionalGeneration + >>> import torch + + >>> torch_device = "cuda:0" + >>> processor = AutoProcessor.from_pretrained("CohereForAI/aya-vision-8b", use_fast=True) + >>> model = AyaVisionForConditionalGeneration.from_pretrained("CohereForAI/aya-vision-8b", device_map=torch_device) + + >>> messages = [ + ... { + ... "role": "user", + ... "content": [ + ... { + ... "type": "image", + ... "url": "https://pbs.twimg.com/media/Fx7YvfQWYAIp6rZ?format=jpg&name=medium", + ... }, + ... {"type": "text", "text": "चित्र में लिखा पाठ क्या कहता है?"}, + ... ], + ... } + ... ] + + >>> inputs = processor.apply_chat_template( + ... messages, padding=True, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", device=torch_device + ... ).to(model.device) + + >>> gen_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.3) + >>> processor.tokenizer.decode(gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + image_sizes=image_sizes, + **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 AyaVisionCausalLMOutputWithPast( + 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__ = ["AyaVisionForConditionalGeneration", "AyaVisionPreTrainedModel", "AyaVisionModel"] diff --git a/third_party/transformers/src/transformers/models/aya_vision/modular_aya_vision.py b/third_party/transformers/src/transformers/models/aya_vision/modular_aya_vision.py new file mode 100644 index 0000000000000000000000000000000000000000..d1ab6548b4fa3f195af1e5678b3529582a89e930 --- /dev/null +++ b/third_party/transformers/src/transformers/models/aya_vision/modular_aya_vision.py @@ -0,0 +1,262 @@ +# Copyright 2025 the Cohere 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 AyaVision model.""" + +import torch +from torch import nn + +from transformers.models.llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, + LlavaForConditionalGeneration, + LlavaModel, + LlavaModelOutputWithPast, + LlavaPreTrainedModel, +) + +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, logging +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from .configuration_aya_vision import AyaVisionConfig + + +logger = logging.get_logger(__name__) + + +class AyaVisionMultiModalProjector(nn.Module): + def __init__(self, config: AyaVisionConfig): + super().__init__() + self.config = config + self.downsample_factor = config.downsample_factor + self.alignment_intermediate_size = getattr( + config, "alignment_intermediate_size", config.text_config.hidden_size + ) + self.layernorm = nn.LayerNorm( + config.vision_config.hidden_size * (config.downsample_factor**2), eps=config.adapter_layer_norm_eps + ) + + self.linear_1 = nn.Linear( + config.vision_config.hidden_size * (config.downsample_factor**2), + self.alignment_intermediate_size, + bias=True, + ) + + self.act = ACT2FN["silu"] # SwiGLU uses SiLU activation + # For SwiGLU, project down to half size since we split intermediate dim + self.linear_2 = nn.Linear(self.alignment_intermediate_size // 2, config.text_config.hidden_size, bias=True) + + def forward(self, image_features): + image_features = self.pixel_shuffle(image_features) + image_features = self.layernorm(image_features) + hidden_states = self.linear_1(image_features) + + # Split along last dimension and apply SwiGLU + x, gate = hidden_states.chunk(2, dim=-1) + hidden_states = self.act(gate) * x + + hidden_states = self.linear_2(hidden_states) + return hidden_states + + def pixel_shuffle(self, image_features): # B, S, D + batch_size, seq_length, feature_dim = image_features.shape + height = width = int(seq_length**0.5) + image_features = image_features.reshape(image_features.shape[0], width, height, -1) + channels = image_features.shape[-1] + image_features = image_features.reshape( + batch_size, width, int(height / self.downsample_factor), int(channels * self.downsample_factor) + ) + image_features = image_features.permute(0, 2, 1, 3) + image_features = image_features.reshape( + batch_size, int(height / self.downsample_factor), int(width / self.downsample_factor), -1 + ) + image_features = image_features.permute(0, 2, 1, 3) + return image_features + + +class AyaVisionPreTrainedModel(LlavaPreTrainedModel): + _can_compile_fullgraph = False + + +class AyaVisionCausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass + + +class AyaVisionModelOutputWithPast(LlavaModelOutputWithPast): + pass + + +class AyaVisionModel(LlavaModel): + # Unlike LLaVA, the model doesn't have to deal with Pixtral-style image states + @merge_with_config_defaults + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + kwargs = {k: v for k, v in kwargs.items() if v is not None} + # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) + + # If we have one vision feature layer, return the corresponding hidden states, + # otherwise, select the hidden states of each feature layer and concatenate them + if isinstance(vision_feature_layer, int): + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] + if vision_feature_select_strategy == "default": + selected_image_feature = selected_image_feature[:, 1:] + else: + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + # For default; crop CLS from each hidden state in the hidden state pool + if vision_feature_select_strategy == "default": + hs_pool = [hs[:, 1:] for hs in hs_pool] + selected_image_feature = torch.cat(hs_pool, dim=-1) + + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) + + return image_outputs + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | AyaVisionModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + return_dict=True, + ).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return AyaVisionModelOutputWithPast( + 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 AyaVisionForConditionalGeneration(LlavaForConditionalGeneration): + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + image_sizes: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | AyaVisionCausalLMOutputWithPast: + 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 AutoProcessor, AyaVisionForConditionalGeneration + >>> import torch + + >>> torch_device = "cuda:0" + >>> processor = AutoProcessor.from_pretrained("CohereForAI/aya-vision-8b", use_fast=True) + >>> model = AyaVisionForConditionalGeneration.from_pretrained("CohereForAI/aya-vision-8b", device_map=torch_device) + + >>> messages = [ + ... { + ... "role": "user", + ... "content": [ + ... { + ... "type": "image", + ... "url": "https://pbs.twimg.com/media/Fx7YvfQWYAIp6rZ?format=jpg&name=medium", + ... }, + ... {"type": "text", "text": "चित्र में लिखा पाठ क्या कहता है?"}, + ... ], + ... } + ... ] + + >>> inputs = processor.apply_chat_template( + ... messages, padding=True, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", device=torch_device + ... ).to(model.device) + + >>> gen_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.3) + >>> processor.tokenizer.decode(gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) + ```""" + super().forward( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + labels=labels, + logits_to_keep=logits_to_keep, + image_sizes=image_sizes, + **kwargs, + ) + + +__all__ = ["AyaVisionForConditionalGeneration", "AyaVisionPreTrainedModel", "AyaVisionModel"] diff --git a/third_party/transformers/src/transformers/models/aya_vision/processing_aya_vision.py b/third_party/transformers/src/transformers/models/aya_vision/processing_aya_vision.py new file mode 100644 index 0000000000000000000000000000000000000000..90188519aba7a4b0d4907ff2830bc9c6885dc618 --- /dev/null +++ b/third_party/transformers/src/transformers/models/aya_vision/processing_aya_vision.py @@ -0,0 +1,208 @@ +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...image_processing_utils import BatchFeature +from ...image_utils import ImageInput, make_flat_list_of_images +from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +class AyaVisionProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding_side": "left", + "padding": True, + "return_mm_token_type_ids": False, + }, + "images_kwargs": { + "crop_to_patches": True, + }, + } + + +@auto_docstring +class AyaVisionProcessor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + patch_size: int = 28, + img_size: int = 364, + image_token="", # set the default and let users change if they have peculiar special tokens in rare cases + downsample_factor: int = 1, + start_of_img_token="<|START_OF_IMG|>", + end_of_img_token="<|END_OF_IMG|>", + img_patch_token="<|IMG_PATCH|>", + img_line_break_token="<|IMG_LINE_BREAK|>", + tile_token="TILE", + tile_global_token="TILE_GLOBAL", + chat_template=None, + **kwargs, + ): + r""" + patch_size (`int`, *optional*, defaults to 28): + The size of image patches for tokenization. + img_size (`int`, *optional*, defaults to 364): + The size of the image to be tokenized. This should correspond to the size given to the image processor. + image_token (`str`, *optional*, defaults to `""`): + The token to be used to represent an image in the text. + downsample_factor (`int`, *optional*, defaults to 1): + The factor by which to scale the patch size. + start_of_img_token (`str`, *optional*, defaults to `"<|START_OF_IMG|>"`): + The token to be used to represent the start of an image in the text. + end_of_img_token (`str`, *optional*, defaults to `"<|END_OF_IMG|>"`): + The token to be used to represent the end of an image in the text. + img_patch_token (`str`, *optional*, defaults to `"<|IMG_PATCH|>"`): + The token to be used to represent an image patch in the text. + img_line_break_token (`str`, *optional*, defaults to `"<|IMG_LINE_BREAK|>"`): + The token to be used to represent a line break in the text. + tile_token (`str`, *optional*, defaults to `"TILE"`): + The token to be used to represent an image patch in the text. + tile_global_token (`str`, *optional*, defaults to `"TILE_GLOBAL"`): + The token to be used to represent the cover image in the text. + """ + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + self.image_token = image_token + self.patch_size = patch_size * downsample_factor + self.img_size = img_size + + self.start_of_img_token = start_of_img_token + self.end_of_img_token = end_of_img_token + self.img_patch_token = img_patch_token + self.img_line_break_token = img_line_break_token + self.tile_token = tile_token + self.tile_global_token = tile_global_token + self.image_token_id = tokenizer.convert_tokens_to_ids(self.img_patch_token) + self.image_ids = tokenizer.convert_tokens_to_ids( + [img_patch_token, tile_token, tile_global_token, start_of_img_token, end_of_img_token] + ) + + def _prompt_split_image(self, num_patches): + """ + Create a structured string representation of image tokens + + Args: + num_patches: Number of patches in the image + + Returns: + String with appropriate image tokens + """ + + img_patches_per_tile = (self.img_size // self.patch_size) ** 2 + img_string = f"{self.start_of_img_token}" + if num_patches > 1: + for idx in range(1, num_patches): + img_string += f"{self.tile_token}_{idx}" + f"{self.img_patch_token}" * img_patches_per_tile + + img_string += f"{self.tile_global_token}" + f"{self.img_patch_token}" * img_patches_per_tile + img_string += f"{self.end_of_img_token}" + return img_string + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + **kwargs: Unpack[AyaVisionProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + if text is None: + raise ValueError("You have to specify text.") + + output_kwargs = self._merge_kwargs( + AyaVisionProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if not isinstance(text, (list, tuple)): + text = [text] + + # Process images + image_inputs = {} + if images is not None: + images = self.image_processor.fetch_images(images) + images = make_flat_list_of_images(images) + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + num_patches = image_inputs.pop("num_patches") + image_index = 0 + processed_text = [] + for prompt in text: + new_prompt = prompt + while "" in new_prompt: + # Replace the image placeholder with structured image tokens + image_tokens = self._prompt_split_image(num_patches[image_index]) + new_prompt = new_prompt.replace("", image_tokens, 1) + image_index += 1 + processed_text.append(new_prompt) + + if image_index != len(images): + raise ValueError("Number of image placeholders in the prompt does not match the number of images.") + + text = processed_text + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + + vision_data = {} + if image_sizes is not None: + images_kwargs = AyaVisionProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + + num_image_patches = [ + self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) + for image_size in image_sizes + ] + + token_per_patch = (self.img_size // self.patch_size) ** 2 + num_image_tokens = [ + token_per_patch + 3 + sum(token_per_patch + 1 for _ in range(1, num_patches)) + for num_patches in num_image_patches + ] # Add +3 and +1 for BOI/EOI and image tile tokens + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + +__all__ = ["AyaVisionProcessor"] diff --git a/third_party/transformers/src/transformers/models/blt/__init__.py b/third_party/transformers/src/transformers/models/blt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..703b81ecdd09dda47a97c641f7e440bcb5e81119 --- /dev/null +++ b/third_party/transformers/src/transformers/models/blt/__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_blt import * + from .modeling_blt import * + from .tokenization_blt 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/blt/configuration_blt.py b/third_party/transformers/src/transformers/models/blt/configuration_blt.py new file mode 100644 index 0000000000000000000000000000000000000000..d685bff3ae2da0c0907641083975e02d073a1971 --- /dev/null +++ b/third_party/transformers/src/transformers/models/blt/configuration_blt.py @@ -0,0 +1,286 @@ +# 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. +"""Blt model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="itazap/blt-1b-hf") +@strict +class BltLocalEncoderConfig(PreTrainedConfig): + r""" + cross_attn_all_layers (`bool`, *optional*, defaults to `True`): + Whether all attention layers have cross attention. + cross_attn_k (`int`, *optional*, defaults to 2): + Number of cross-attention heads used in the model. + hidden_size_global (`int`, *int*, defaults to 2048): + Hidden size of the global transformer layer. + """ + + model_type = "blt_local_encoder" + default_theta = 500000.0 + + vocab_size: int = 260 + cross_attn_all_layers: bool | None = False + cross_attn_k: int | None = 2 + hidden_size_global: int | None = 2048 + hidden_size: int = 1024 + num_attention_heads: int = 16 + num_key_value_heads: int | None = None + num_hidden_layers: int = 1 + rms_norm_eps: float = 1e-5 + dropout: float | int | None = 0.0 + max_position_embeddings: int = 24576 + rope_parameters: RopeParameters | dict | None = None + hidden_act: str = "silu" + intermediate_size: int | None = None + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + self.num_key_value_heads = self.num_key_value_heads or self.num_attention_heads + self.intermediate_size = self.intermediate_size or int(8 * self.hidden_size / 3) + self.tie_word_embeddings = False + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="itazap/blt-1b-hf") +@strict +class BltLocalDecoderConfig(PreTrainedConfig): + r""" + cross_attn_all_layers (`bool`, *optional*, defaults to `True`): + Whether all attention layers have cross attention. + cross_attn_k (`int`, *optional*, defaults to 2): + Number of cross-attention heads used in the model. + hidden_size_global (`int`, *int*, defaults to 2048): + Hidden size of the global transformer layer. + """ + + model_type = "blt_local_decoder" + default_theta = 500000.0 + + vocab_size: int = 260 + cross_attn_all_layers: bool | None = True + cross_attn_k: int | None = 2 + hidden_size_global: int | None = 2048 + hidden_size: int = 1024 + num_attention_heads: int = 16 + num_key_value_heads: int | None = None + num_hidden_layers: int = 9 + rms_norm_eps: float = 1e-5 + dropout: float | int | None = 0.0 + max_position_embeddings: int = 24576 + rope_parameters: RopeParameters | dict | None = None + hidden_act: str = "silu" + intermediate_size: int = 2816 + initializer_range: float = 0.02 + pad_token_id: int | None = None + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + self.num_key_value_heads = self.num_key_value_heads or self.num_attention_heads + self.head_dim = self.hidden_size // self.num_attention_heads + self.intermediate_size = self.intermediate_size or int(8 * self.hidden_size / 3) + self.tie_word_embeddings = False # Force-set to False for BC + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="itazap/blt-1b-hf") +@strict +class BltGlobalTransformerConfig(PreTrainedConfig): + model_type = "blt_global_transformer" + default_theta = 500000.0 + + hidden_size: int = 2048 + num_attention_heads: int = 16 + num_key_value_heads: int | None = None + num_hidden_layers: int = 25 + rms_norm_eps: float = 1e-5 + dropout: float | int | None = 0.0 + max_position_embeddings: int = 4096 + rope_parameters: RopeParameters | dict | None = None + hidden_act: str = "silu" + intermediate_size: int = 5632 + initializer_range: float = 0.02 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + self.num_key_value_heads = self.num_key_value_heads or self.num_attention_heads + self.head_dim = self.hidden_size // self.num_attention_heads + self.intermediate_size = self.intermediate_size or int(8 * self.hidden_size / 3) + self.tie_word_embeddings = False + + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="itazap/blt-1b-hf") +@strict +class BltPatcherConfig(PreTrainedConfig): + model_type = "blt_patcher" + + vocab_size: int = 260 + hidden_size: int = 768 + num_hidden_layers: int = 14 + num_attention_heads: int = 12 + num_key_value_heads: int | None = None + max_position_embeddings: int = 8192 + rms_norm_eps: float = 1e-5 + dropout: float | int | None = 0.0 + intermediate_size: int = 2048 + rope_parameters: RopeParameters | dict | None = None + initializer_range: float = 0.02 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + self.num_key_value_heads = self.num_key_value_heads or self.num_attention_heads + self.head_dim = self.hidden_size // self.num_attention_heads + self.intermediate_size = self.intermediate_size or int(8 * self.hidden_size / 3) + self.tie_word_embeddings = False + self.hidden_act = "silu" # Blt uses silu activation + + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="itazap/blt-1b-hf") +@strict +class BltConfig(PreTrainedConfig): + r""" + patch_in_forward (`bool`, *optional*, defaults to `True`): + Whether to perform patching during the forward pass. + patch_size (`int`, *optional*, defaults to 4): + Size of the patches used in the patching mechanism. + patching_mode (`str`, *optional*, defaults to `"entropy"`): + The mode used for patching, such as entropy-based patching. + patching_threshold (`float`, *optional*, defaults to 1.34): + Threshold value used for determining when to apply patches. + patching_batch_size (`int`, *optional*, defaults to 1): + Batch size used during the patching process. + max_patch_length (`int`, *optional*): + Maximum length of patches that can be generated. + cross_attn_k (`int`, *optional*, defaults to 2): + Number of cross-attention heads used in the model. + encoder_hash_byte_group_size (`list`, *optional*): + List of byte group sizes used in the encoder hash function. + encoder_hash_byte_group_vocab (`int`, *optional*, defaults to 500002): + Vocabulary size for the encoder hash byte groups. + encoder_hash_byte_group_nb_functions (`int`, *optional*, defaults to 1): + Number of hash functions used in the encoder byte grouping. + patcher_config (`BltPatcherConfig`, *optional*): + Configuration for the patcher component of the model. + global_config (`BltGlobalTransformerConfig`, *optional*): + Configuration for the global transformer component of the model. + + Example: + ```python + >>> from transformers import BltModel, BltConfig + + >>> # Initializing a Blt configuration + >>> configuration = BltConfig() + + >>> # Initializing a model from the configuration + >>> model = BltModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "blt" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 500000.0 + sub_configs = { + "patcher_config": BltPatcherConfig, + "encoder_config": BltLocalEncoderConfig, + "decoder_config": BltLocalDecoderConfig, + "global_config": BltGlobalTransformerConfig, + } + + vocab_size: int = 260 + max_position_embeddings: int = 4096 + patch_in_forward: bool | None = True + patch_size: int | None = 4 + patching_mode: str | None = "entropy" + patching_threshold: float | None = 1.335442066192627 + patching_batch_size: int | None = 1 + max_patch_length: int | None = None + cross_attn_k: int | None = 2 + encoder_hash_byte_group_size: list[int] | None = None + encoder_hash_byte_group_vocab: int | None = 500002 + encoder_hash_byte_group_nb_functions: int | None = 1 + patcher_config: dict | PreTrainedConfig | None = None + encoder_config: dict | PreTrainedConfig | None = None + decoder_config: dict | PreTrainedConfig | None = None + global_config: dict | PreTrainedConfig | None = None + tie_word_embeddings: bool = False + pad_token_id: int | None = None + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + initializer_range: float = 0.02 + rope_parameters: RopeParameters | dict | None = None + + def __post_init__(self, **kwargs): + self.encoder_hash_byte_group_size = self.encoder_hash_byte_group_size or [3, 4, 5, 6, 7, 8] + + # Initialize component configurations + if self.patcher_config is None: + self.patcher_config = BltPatcherConfig(initializer_range=self.initializer_range) + logger.info("patcher_config is None, using default Blt patcher config") + elif isinstance(self.patcher_config, dict): + self.patcher_config.setdefault("initializer_range", self.initializer_range) + self.patcher_config = BltPatcherConfig(**self.patcher_config) + + if self.encoder_config is None: + self.encoder_config = BltLocalEncoderConfig(initializer_range=self.initializer_range) + logger.info("encoder_config is None, using default Blt encoder config") + elif isinstance(self.encoder_config, dict): + self.encoder_config.setdefault("initializer_range", self.initializer_range) + self.encoder_config = BltLocalEncoderConfig(**self.encoder_config) + + if self.decoder_config is None: + self.decoder_config = BltLocalDecoderConfig(initializer_range=self.initializer_range) + logger.info("decoder_config is None, using default Blt decoder config") + elif isinstance(self.decoder_config, dict): + self.decoder_config.setdefault("initializer_range", self.initializer_range) + self.decoder_config = BltLocalDecoderConfig(**self.decoder_config) + + if self.global_config is None: + self.global_config = BltGlobalTransformerConfig(initializer_range=self.initializer_range) + logger.info("global_config is None, using default Blt global config") + elif isinstance(self.global_config, dict): + self.global_config.setdefault("initializer_range", self.initializer_range) + self.global_config = BltGlobalTransformerConfig(**self.global_config) + + # Determine if token embedding projection is needed based on dimension mismatch (7b) + encoder_cross_output_size = self.encoder_config.hidden_size * self.cross_attn_k + self.global_config.encoder_cross_output_size = ( + encoder_cross_output_size if encoder_cross_output_size != self.global_config.hidden_size else None + ) + + super().__post_init__(**kwargs) + + +__all__ = [ + "BltConfig", + "BltPatcherConfig", + "BltLocalEncoderConfig", + "BltLocalDecoderConfig", + "BltGlobalTransformerConfig", +] diff --git a/third_party/transformers/src/transformers/models/blt/convert_blt_weights_to_hf.py b/third_party/transformers/src/transformers/models/blt/convert_blt_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..054ca9f84d164d0f2cbe620c5287e484a53de6ed --- /dev/null +++ b/third_party/transformers/src/transformers/models/blt/convert_blt_weights_to_hf.py @@ -0,0 +1,487 @@ +import argparse +import json +import logging +import os +from typing import Any + +import torch +from huggingface_hub import hf_hub_download, upload_folder +from safetensors.torch import load_file, save_file +from tokenizers import Tokenizer, decoders, pre_tokenizers, processors +from tokenizers.models import BPE + +from transformers import PreTrainedTokenizerFast +from transformers.convert_slow_tokenizer import bytes_to_unicode +from transformers.utils import logging as transformers_logging + + +logger = transformers_logging.get_logger(__name__) +transformers_logging.set_verbosity_info() + + +def merge_configurations(config_path: str, entropy_params_path: str) -> dict[str, Any]: + logger.info("Merging configurations") + + with open(config_path, "r") as f: + main_config = json.load(f) + + with open(entropy_params_path, "r") as f: + entropy_data = json.load(f) + + entropy_model_params = entropy_data.get("entropy_model", {}) + patcher_args = entropy_data.get("data", {}).get("patcher_args", {}) + + unified_config = main_config.copy()["args"] + + for key in ["vocab_size", "dim", "n_layers", "n_heads", "max_seqlen"]: + if key in unified_config and not isinstance(unified_config[key], int): + unified_config[key] = int(unified_config[key]) + + patch_size = patcher_args.get("patch_size", 8) + if isinstance(patch_size, float): + patch_size = int(patch_size) + + # Create patcher config + patcher_hidden_size = int(entropy_model_params.get("dim", 512)) + patcher_multiple_of = int(entropy_model_params.get("multiple_of", 256)) + patcher_intermediate_size = patcher_multiple_of * ( + (int(8 * patcher_hidden_size / 3) + patcher_multiple_of - 1) // patcher_multiple_of + ) + + patcher_config = { + "vocab_size": int(entropy_model_params.get("vocab_size", 256)), + "hidden_size": patcher_hidden_size, + "num_hidden_layers": int(entropy_model_params.get("n_layers", 8)), + "num_attention_heads": int(entropy_model_params.get("n_heads", 8)), + "num_key_value_heads": int(entropy_model_params.get("n_kv_heads")) + if entropy_model_params.get("n_kv_heads") is not None + else None, + "max_position_embeddings": int(entropy_model_params.get("max_seqlen", 1024)), + "norm_eps": entropy_model_params.get("norm_eps", 1e-5), + "dropout": entropy_model_params.get("dropout", 0.0), + "rope_theta": entropy_model_params.get("rope_theta", 10000.0), + "attn_impl": entropy_model_params.get("attn_impl", "sdpa"), + "attn_bias_type": entropy_model_params.get("attn_bias_type", "causal"), + "intermediate_size": patcher_intermediate_size, + } + + # Create encoder config + encoder_hidden_size = unified_config.get("dim_local_encoder", 1024) + encoder_multiple_of = unified_config.get("multiple_of", 256) + encoder_intermediate_size = encoder_multiple_of * ( + (int(8 * encoder_hidden_size / 3) + encoder_multiple_of - 1) // encoder_multiple_of + ) + + encoder_config = { + "vocab_size": unified_config.get("vocab_size", 256), + "cross_attn_all_layers": unified_config.get("cross_attn_all_layers_encoder", False), + "cross_attn_k": unified_config.get("cross_attn_k", 2), + "hidden_size_global": unified_config.get("dim_global", 2048), + "pm_size": unified_config.get("pm_size", 0), + "hidden_size": encoder_hidden_size, + "num_attention_heads": unified_config.get("n_heads_local_encoder", 16), + "num_key_value_heads": unified_config.get("n_kv_heads"), + "num_hidden_layers": unified_config.get("n_layers_local_encoder", 1), + "norm_eps": unified_config.get("norm_eps", 1e-5), + "dropout": unified_config.get("dropout", 0.0), + "max_position_embeddings": unified_config.get("max_encoder_seq_length") + or unified_config.get("max_seqlen", 1024), + "rope_theta": unified_config.get("rope_theta", 10000.0), + "rope_parameters": {"rope_type": "default"}, + "hidden_act": unified_config.get("hidden_act", "silu"), + "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"), + "intermediate_size": encoder_intermediate_size, + } + + # Create decoder config + decoder_hidden_size = unified_config.get("dim_local_decoder", 1024) + decoder_multiple_of = unified_config.get("multiple_of", 256) + decoder_intermediate_size = decoder_multiple_of * ( + (int(8 * decoder_hidden_size / 3) + decoder_multiple_of - 1) // decoder_multiple_of + ) + + decoder_config = { + "vocab_size": unified_config.get("vocab_size", 256), + "cross_attn_all_layers": unified_config.get("cross_attn_all_layers_decoder", False), + "cross_attn_k": unified_config.get("cross_attn_k", 2), + "hidden_size_global": unified_config.get("dim_global", 2048), + "hidden_size": decoder_hidden_size, + "num_attention_heads": unified_config.get("n_heads_local_decoder", 16), + "num_key_value_heads": unified_config.get("n_kv_heads"), + "num_hidden_layers": unified_config.get("n_layers_local_decoder", 9), + "norm_eps": unified_config.get("norm_eps", 1e-5), + "dropout": unified_config.get("dropout", 0.0), + "max_position_embeddings": unified_config.get("max_encoder_seq_length") + or unified_config.get("max_seqlen", 1024), + "rope_theta": unified_config.get("rope_theta", 10000.0), + "rope_parameters": {"rope_type": "default"}, + "hidden_act": unified_config.get("hidden_act", "silu"), + "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"), + "intermediate_size": decoder_intermediate_size, + } + + # Create global transformer config + global_hidden_size = unified_config.get("dim_global", 2048) + global_multiple_of = unified_config.get("multiple_of", 256) + global_intermediate_size = global_multiple_of * ( + (int(8 * global_hidden_size / 3) + global_multiple_of - 1) // global_multiple_of + ) + + global_config = { + "hidden_size": global_hidden_size, + "num_attention_heads": unified_config.get("n_heads_global", 16), + "num_key_value_heads": unified_config.get("n_kv_heads_global"), + "num_hidden_layers": unified_config.get("n_layers_global", 25), + "norm_eps": unified_config.get("norm_eps", 1e-5), + "dropout": unified_config.get("dropout", 0.0), + "max_position_embeddings": unified_config.get("max_seqlen", 1024), + "rope_theta": unified_config.get("rope_theta", 10000.0), + "rope_parameters": {"rope_type": "default"}, + "hidden_act": unified_config.get("hidden_act", "silu"), + "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"), + "intermediate_size": global_intermediate_size, + } + + # Create main config with sub-configs + main_config_dict = { + "model_type": "blt", + "vocab_size": unified_config.get("vocab_size", 256), + "max_position_embeddings": unified_config.get("max_seqlen", 1024), + "patch_in_forward": True, + "realtime_patching": True, + "patching_mode": "entropy", + "patch_size": patch_size, + "patching_threshold": patcher_args.get("threshold", 0.5), + "patching_threshold_add": patcher_args.get("threshold_add", 0.0), + "max_patch_length": patcher_args.get("max_patch_length"), + "patching_batch_size": patcher_args.get("patching_batch_size", 1), + "patching_device": patcher_args.get("patching_device", "cuda"), + "monotonicity": patcher_args.get("monotonicity", False), + "cross_attn_k": unified_config.get("cross_attn_k", 2), + "encoder_hash_byte_group_size": unified_config.get("encoder_hash_byte_group_size"), + "encoder_hash_byte_group_vocab": unified_config.get("encoder_hash_byte_group_vocab", 30000), + "encoder_hash_byte_group_nb_functions": unified_config.get("encoder_hash_byte_group_nb_functions", 3), + "pm_size": unified_config.get("pm_size", 0), + "patcher_config": patcher_config, + "encoder_config": encoder_config, + "decoder_config": decoder_config, + "global_config": global_config, + } + + main_config_dict["tie_word_embeddings"] = False + + logger.info(f"Merged configuration with {len(main_config_dict)} parameters") + return main_config_dict + + +def apply_weight_mapping(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + component_mappings = { + ".attention.": ".self_attn.", + ".feed_forward.": ".mlp.", + ".attention_norm.": ".input_layernorm.", + ".ffn_norm.": ".post_attention_layernorm.", + ".tok_embeddings.": ".embed_tokens.", + ".cross_attn_norm_q.": ".q_norm.", + ".cross_attn_norm_kv.": ".k_norm.", + ".w1.": ".gate_proj.", + ".w2.": ".down_proj.", + ".w3.": ".up_proj.", + ".wq.": ".q_proj.", + ".wk.": ".k_proj.", + ".wv.": ".v_proj.", + ".wo.": ".o_proj.", + ".output.": ".lm_head.", + } + + new_state_dict = {} + + for old_key, tensor in state_dict.items(): + new_key = old_key + + for old_pattern, new_pattern in component_mappings.items(): + if old_pattern in new_key: + new_key = new_key.replace(old_pattern, new_pattern) + + new_state_dict[new_key] = tensor + + return new_state_dict + + +def convert_hash_embeddings_to_fused( + unified_weights: dict[str, torch.Tensor], config: dict[str, Any] +) -> dict[str, torch.Tensor]: + """Convert ModuleList hash embeddings to nn.embedding format""" + original_keys_format = [ + key + for key in unified_weights.keys() + if "encoder_hash_tok_embedding." in key and ".weight" in key and key.split(".")[-2].isdigit() + ] + + num_embeddings = config.get("encoder_hash_byte_group_nb_functions", 1) * len( + config.get("encoder_hash_byte_group_size", [3, 4, 5, 6, 7, 8]) + ) + vocab_size = config.get("encoder_hash_byte_group_vocab", 500002) + hidden_size = config.get("encoder_config", {}).get("hidden_size", 1024) + + fused_weight = torch.zeros(vocab_size * num_embeddings, hidden_size) + + sorted_keys = sorted(original_keys_format, key=lambda k: int(k.split(".")[-2])) + + for i, old_key in enumerate(sorted_keys): + start_idx = i * vocab_size + end_idx = (i + 1) * vocab_size + fused_weight[start_idx:end_idx] = unified_weights[old_key] + logger.info(f"Copied {old_key} to indices {start_idx}:{end_idx}") + del unified_weights[old_key] + + fused_key = "model.encoder_hash_tok_embedding.weight" + unified_weights[fused_key] = fused_weight + + return unified_weights + + +def merge_weights(weights_path: str, entropy_weights_path: str) -> dict[str, torch.Tensor]: + main_weights = load_file(weights_path) + + entropy_weights = torch.load(entropy_weights_path, map_location="cpu", weights_only=True) + + if "model" in entropy_weights: + entropy_weights = entropy_weights["model"] + elif "state_dict" in entropy_weights: + entropy_weights = entropy_weights["state_dict"] + + unified_weights = main_weights.copy() + + for key, tensor in entropy_weights.items(): + patcher_key = f"patcher.{key}" + unified_weights[patcher_key] = tensor + + unified_weights = apply_weight_mapping(unified_weights) + + decoder_lm_head_key = "local_decoder.lm_head.weight" + top_lm_head_key = "lm_head.weight" + unified_weights[top_lm_head_key] = unified_weights[decoder_lm_head_key] + del unified_weights[decoder_lm_head_key] + + prefixed_weights = {} + for key, tensor in unified_weights.items(): + if key == top_lm_head_key: + prefixed_weights[key] = tensor + elif not key.startswith("model."): + prefixed_weights[f"model.{key}"] = tensor + else: + prefixed_weights[key] = tensor + + unified_weights = prefixed_weights + + return unified_weights + + +def create_tokenizer_config(output_dir: str, config: dict[str, Any]): + tokenizer_config = { + "tokenizer_class": "PreTrainedTokenizerFast", + "vocab_size": config.get("vocab_size", 256), + "model_max_length": config.get("max_seqlen", 1024), + "model_input_names": ["input_ids", "attention_mask"], + "add_bos_token": True, + "add_eos_token": True, + "bos_token": "", + "eos_token": "", + "pad_token": "", + "unk_token": "", + } + + tokenizer_path = os.path.join(output_dir, "tokenizer_config.json") + with open(tokenizer_path, "w") as f: + json.dump(tokenizer_config, f, indent=2) + + +def create_tokenizer_json(output_dir: str, config: dict[str, Any]): + byte_encoder = bytes_to_unicode() + + vocab: dict[str, int] = {} + vocab[""] = 0 + vocab[""] = 1 + vocab[""] = 2 + vocab[""] = 3 + + offset = 4 + for byte_val, unicode_char in byte_encoder.items(): + vocab[unicode_char] = byte_val + offset + + backend = Tokenizer( + BPE(vocab=vocab, merges=[], continuing_subword_prefix="", end_of_word_suffix="", fuse_unk=False) + ) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + backend.decoder = decoders.ByteLevel() + + bos = config.get("bos_token", "") + backend.post_processor = processors.TemplateProcessing( + single=f"{bos}:0 $A:0", + pair=f"{bos}:0 $A:0 $B:1", + special_tokens=[(bos, 1)], + ) + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + bos_token=config.get("bos_token", ""), + eos_token=config.get("eos_token", ""), + pad_token=config.get("pad_token", ""), + unk_token=config.get("unk_token", ""), + ) + + tokenizer.add_bos_token = bool(config.get("add_bos_token", True)) + tokenizer.add_eos_token = bool(config.get("add_eos_token", True)) + + tokenizer.save_pretrained(output_dir) + logger.info(f"Saved tokenizer.json to {os.path.join(output_dir, 'tokenizer.json')}") + + +def push_to_hub( + local_dir: str, + repo_id: str, + commit_message: str = "Upload converted Blt model", + private: bool = False, + token: str | None = None, +) -> None: + try: + upload_folder( + folder_path=local_dir, + repo_id=repo_id, + commit_message=commit_message, + repo_type="model", + token=token, + ) + logger.info(f"Successfully pushed model to {repo_id}") + + except Exception as e: + logger.error(f"Failed to push model to Hub: {e}") + raise + + +def convert_hf_blt_to_unified( + model_id: str, + output_dir: str, + config_name: str = "config.json", + weights_name: str = "model.bin", + cache_dir: str | None = None, + push_to_hub_repo: str | None = None, + hub_private: bool = False, + hub_token: str | None = None, +) -> None: + # Download model files + config_path = hf_hub_download(repo_id=model_id, filename="config.json", cache_dir=cache_dir) + weights_path = hf_hub_download(repo_id=model_id, filename="model.safetensors", cache_dir=cache_dir) + entropy_params_path = hf_hub_download(repo_id=model_id, filename="entropy_model/params.json", cache_dir=cache_dir) + entropy_weights_path = hf_hub_download( + repo_id=model_id, filename="entropy_model/consolidated.pth", cache_dir=cache_dir + ) + + unified_config = merge_configurations(config_path, entropy_params_path) + unified_weights = merge_weights(weights_path, entropy_weights_path) + + unified_weights = convert_hash_embeddings_to_fused(unified_weights, unified_config) + + os.makedirs(output_dir, exist_ok=True) + + config_path = os.path.join(output_dir, config_name) + with open(config_path, "w") as f: + json.dump(unified_config, f, indent=2) + + if weights_name.endswith(".bin"): + weights_name = weights_name.replace(".bin", ".safetensors") + + weights_path = os.path.join(output_dir, weights_name) + save_file(unified_weights, weights_path) + + create_tokenizer_json(output_dir=output_dir, config=unified_config) + + create_tokenizer_config(output_dir, unified_config) + + logger.info(f"Conversion completed, model saved to: {output_dir}") + + if push_to_hub_repo: + push_to_hub( + local_dir=output_dir, + repo_id=push_to_hub_repo, + commit_message="Upload Blt model converted", + private=hub_private, + token=hub_token, + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Convert Blt models from HuggingFace Hub format to unified format", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + parser.add_argument( + "--model_id", + type=str, + default="facebook/blt-7b", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./blt_converted", + ) + parser.add_argument( + "--config_name", + type=str, + default="config.json", + ) + parser.add_argument( + "--weights_name", + type=str, + default="model.bin", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + ) + parser.add_argument( + "--debug", + action="store_true", + default=True, + ) + parser.add_argument( + "--push_to_hub", + type=str, + default=None, + ) + parser.add_argument( + "--hub_private", + action="store_true", + default=False, + ) + parser.add_argument( + "--hub_token", + type=str, + default="hf_token", + ) + + args = parser.parse_args() + + transformers_logging.set_verbosity_debug() + logging.basicConfig(level=logging.DEBUG) + + try: + convert_hf_blt_to_unified( + model_id=args.model_id, + output_dir=args.output_dir, + config_name=args.config_name, + weights_name=args.weights_name, + cache_dir=args.cache_dir, + push_to_hub_repo=False, # args.push_to_hub, + hub_private=args.hub_private, + hub_token=args.hub_token, + ) + except Exception as e: + logger.error(f"Conversion failed: {e}") + raise + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/blt/modeling_blt.py b/third_party/transformers/src/transformers/models/blt/modeling_blt.py new file mode 100644 index 0000000000000000000000000000000000000000..778f7ba80cf62618100cd0980e6634c310856f6f --- /dev/null +++ b/third_party/transformers/src/transformers/models/blt/modeling_blt.py @@ -0,0 +1,1463 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/blt/modular_blt.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_blt.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.distributions +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import 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.deprecation import deprecate_kwarg +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_blt import ( + BltConfig, + BltGlobalTransformerConfig, + BltLocalDecoderConfig, + BltLocalEncoderConfig, + BltPatcherConfig, +) + + +class BltMLP(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) + # Ignore copy + 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 BltRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + BltRMSNorm 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 BltRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: BltConfig, 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: BltConfig | 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) + + +# Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer +class BltTransformerLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = BltSelfAttention(config=config, layer_idx=layer_idx) + self.mlp = BltMLP(config) + self.input_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.layer_idx = layer_idx + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + """ + 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. + + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +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. + """ + 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 + + +class BltSelfAttention(nn.Module): + def __init__(self, config: BltConfig, layer_idx: int): + super().__init__() + self.config = config + self.num_heads = config.num_attention_heads + self.dropout = config.dropout + self.hidden_size = config.hidden_size + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.hidden_size // self.num_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + + self.layer_idx = layer_idx + self.is_causal = True + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_embeddings: torch.Tensor, + past_key_values=None, + **kwargs, + ): + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + 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.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class BltCrossAttention(nn.Module): + """Cross-attention module for Blt, following transformers style""" + + def __init__(self, config: BltConfig, layer_idx: int, hidden_size: int | None = None): + super().__init__() + self.config = config + self.num_heads = self.config.num_attention_heads + self.num_key_value_heads = self.config.num_key_value_heads + self.dropout = config.dropout + self.hidden_size = config.hidden_size + self.head_dim = config.hidden_size // self.num_heads + self.layer_idx = layer_idx + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.q_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps) + self.k_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps) + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + bsz, q_len, _ = hidden_states.size() + query_states = self.q_norm(hidden_states) + query_states = self.q_proj(query_states) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + cross_attention_states = self.k_norm(cross_attention_states) + key_states = self.k_proj(cross_attention_states) + value_states = self.v_proj(cross_attention_states) + key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, -1, self.num_key_value_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, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + attn_output = attn_output + hidden_states + return attn_output, attn_weights + + +@auto_docstring +class BltPreTrainedModel(PreTrainedModel): + config: BltConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["BltTransformerLayer"] + _can_compile_fullgraph = False # static cache cannot have different shapes for each layer + _supports_sdpa = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_attention_backend = False + _can_record_outputs = { + "hidden_states": OutputRecorder(BltTransformerLayer, index=0), + "attentions": OutputRecorder(BltSelfAttention, index=1), + } + + @torch.no_grad() + def _init_weights(self, module): + """ + Initialize BLT weights following the original ByteLatentTransformer: + + - Most weights are drawn from a truncated normal. + - Scale is ~ 1 / sqrt(model_dim) (or 1 / sqrt(hidden_dim) for FFN outputs). + - Norm layers are set to weight = 1, bias = 0. + """ + class_name = module.__class__.__name__ + + # Norms: RMSNorm / LayerNorm + if isinstance(module, (BltRMSNorm, nn.LayerNorm)) or "RMSNorm" in class_name or "LayerNorm" in class_name: + if getattr(module, "weight", None) is not None: + init.ones_(module.weight) + if getattr(module, "bias", None) is not None: + init.zeros_(module.bias) + return + + # Embeddings (encoder / patcher / hash embeddings) + if isinstance(module, nn.Embedding): + hidden_size = getattr(self.config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "encoder_config"): + hidden_size = getattr(self.config.encoder_config, "hidden_size", None) + if hidden_size is None: + hidden_size = module.embedding_dim + + std = hidden_size**-0.5 + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if module.padding_idx is not None: + init.zeros_(module.weight[module.padding_idx]) + return + + # Self-attention / cross-attention projections + if isinstance(module, (BltSelfAttention, BltCrossAttention)) or class_name in ( + "MllamaTextSelfAttention", + "MllamaTextCrossAttention", + ): + dim = getattr(self.config, "hidden_size", None) + if dim is None and hasattr(module, "hidden_size"): + dim = module.hidden_size + if dim is None: + for name in ("q_proj", "k_proj", "v_proj", "o_proj", "dense"): + proj = getattr(module, name, None) + if proj is not None and hasattr(proj, "weight"): + dim = proj.weight.shape[-1] + break + if dim is None: + return + + std = dim**-0.5 + + # Input projections (q, k, v) + for proj_name in ("q_proj", "k_proj", "v_proj"): + proj = getattr(module, proj_name, None) + if proj is not None and hasattr(proj, "weight"): + init.trunc_normal_( + proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(proj, "bias", None) is not None: + init.zeros_(proj.bias) + + # Output projection: o_proj or dense + o_proj = getattr(module, "o_proj", getattr(module, "dense", None)) + if o_proj is not None and hasattr(o_proj, "weight"): + init.trunc_normal_( + o_proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(o_proj, "bias", None) is not None: + init.zeros_(o_proj.bias) + return + + # MLP / FFN blocks + if isinstance(module, BltMLP) or class_name == "MllamaTextMLP": + hidden_size = getattr(self.config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "decoder_config"): + hidden_size = getattr(self.config.decoder_config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "encoder_config"): + hidden_size = getattr(self.config.encoder_config, "hidden_size", None) + + # Input-side std + in_std = None + if hidden_size is not None: + in_std = hidden_size**-0.5 + + gate_proj = getattr(module, "gate_proj", getattr(module, "fc1", None)) + up_proj = getattr(module, "up_proj", None) + down_proj = getattr(module, "down_proj", getattr(module, "fc2", None)) + + # gate / input projections + for proj in (gate_proj, up_proj): + if proj is not None and hasattr(proj, "weight"): + std = in_std or (proj.weight.shape[1] ** -0.5) + init.trunc_normal_( + proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(proj, "bias", None) is not None: + init.zeros_(proj.bias) + + # output/ down projections + if down_proj is not None and hasattr(down_proj, "weight"): + hidden_dim = down_proj.weight.shape[1] + out_std = hidden_dim**-0.5 + init.trunc_normal_( + down_proj.weight, + mean=0.0, + std=out_std, + a=-3 * out_std, + b=3 * out_std, + ) + if getattr(down_proj, "bias", None) is not None: + init.zeros_(down_proj.bias) + return + + # Generic Linear layers (projections, lm_head, etc.) + if isinstance(module, nn.Linear): + fan_in = module.in_features + std = fan_in**-0.5 + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if module.bias is not None: + init.zeros_(module.bias) + return + + if isinstance(module, BltRotaryEmbedding): + rope_fn = ( + ROPE_INIT_FUNCTIONS[module.rope_type] + if module.rope_type != "default" + else module.compute_default_rope_parameters + ) + buffer_value, _ = rope_fn(module.config) + init.copy_(module.inv_freq, buffer_value) + init.copy_(module.original_inv_freq, buffer_value) + + +class BltLocalEncoder(BltPreTrainedModel): + config: BltLocalEncoderConfig + _can_record_outputs = { + "encoder_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_encoder"), + } + + def __init__(self, config: BltLocalEncoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + self.config = config + self.layers = nn.ModuleList( + [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = BltRotaryEmbedding(config=config) + self.patch_embedding_projection = nn.Linear( + in_features=config.hidden_size, + out_features=config.hidden_size * config.cross_attn_k, + bias=False, + ) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.cross_attn_layers = nn.ModuleList() + layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1 + for layer_idx in range(layers_to_add): + self.cross_attn_layers.append( + BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size) + ) + + self.post_init() + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + patch_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + encoder_attention_mask: torch.Tensor | None = None, + num_patches: int | None = None, + patch_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + batch_size = inputs_embeds.shape[0] + hidden_states = F.dropout(inputs_embeds, p=self.config.dropout, training=self.training) + + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + + for idx, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + if idx == len(self.layers) - 1 or self.config.cross_attn_all_layers: + patch_embeds = self.patch_reduce(hidden_states, num_patches, patch_ids) + patch_embeds = self.patch_embedding_projection(patch_embeds) + patch_embeds = patch_embeds.reshape( + batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size + ) + layer_idx = idx if self.config.cross_attn_all_layers else 0 + cross_attention_output, _ = self.cross_attn_layers[layer_idx]( + hidden_states=patch_embeds, + cross_attention_states=hidden_states, + attention_mask=encoder_attention_mask, + **kwargs, + ) + patch_embeds = patch_embeds + cross_attention_output + encoder_cross_states = patch_embeds + return hidden_states, encoder_cross_states + + def patch_reduce(self, hidden_states, max_num_patches, patch_ids): + """ + Reduce variable length patches to single embedding per patch + Note: this works with variable number of patches for different sequences in the batch + It handles variable length patches by assuming that patch_lengths will be 0 for any + extra patches on the *right*. Since there can be a variable number of patches + this function also return the number of patches for each sequence in the batch. + Any embeddings on the right that are not allocated to a patch + (i.e. if the sum(patch_lengths[i]) < seq_len for any i) + will be sent to a dummy patch, which is trimmed before returning. + """ + batch_size = hidden_states.shape[0] + embedding_dim = hidden_states.shape[-1] + + patch_ids = patch_ids.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]) + + reduced_embeddings = torch.zeros( + (batch_size, max_num_patches, embedding_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + reduced_embeddings = reduced_embeddings.scatter_reduce( + src=hidden_states, + dim=1, + index=patch_ids, + reduce="amax", + include_self=False, + ) + reduced_embeddings = reduced_embeddings[:, :max_num_patches, :] + + return reduced_embeddings + + +class BltLocalDecoder(BltPreTrainedModel): + config: BltLocalDecoderConfig + + def __init__(self, config: BltLocalDecoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + self.config = config + self.cross_attn_decoder = True + self.layers = nn.ModuleList( + [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = BltRotaryEmbedding(config=config) + self.patch_embedding_projection = nn.Linear( + in_features=config.hidden_size_global, + out_features=config.hidden_size * config.cross_attn_k, + bias=False, + ) + self.norm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn_layers = nn.ModuleList() + layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1 + for layer_idx in range(layers_to_add): + self.cross_attn_layers.append( + BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size) + ) + + self.post_init() + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + patch_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size = inputs_embeds.shape[0] + hidden_states = inputs_embeds + patch_embeds = self.patch_embedding_projection(patch_embeds) + patch_embeds = patch_embeds.reshape( + batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size + ) + + if patch_embeds is not None and not self.cross_attn_decoder: + hidden_states = hidden_states + patch_embeds + + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + + for i, layer in enumerate(self.layers): + if i == 0 or self.config.cross_attn_all_layers: + cross_attention_output, _ = self.cross_attn_layers[i]( + hidden_states=hidden_states, + cross_attention_states=patch_embeds, + attention_mask=encoder_attention_mask, + **kwargs, + ) + hidden_states = hidden_states + cross_attention_output + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + logits = self.norm(hidden_states) + return logits + + +class BltGlobalTransformer(BltPreTrainedModel): + config: BltGlobalTransformerConfig + _can_record_outputs = { + "global_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="global_transformer"), + } + + def __init__(self, config: BltGlobalTransformerConfig): + super().__init__(config) + self.config = config + self.layers = nn.ModuleList() + for layer_idx in range(config.num_hidden_layers): + self.layers.append(BltTransformerLayer(config, layer_idx)) + self.rotary_emb = BltRotaryEmbedding(config=config) + + # Create token embedding projection (use nn.Identity() when no projection needed) + if getattr(config, "encoder_cross_output_size", None) is not None: + self.token_embedding_projection = nn.Linear( + config.encoder_cross_output_size, config.hidden_size, bias=False + ) + else: + self.token_embedding_projection = nn.Identity() + + self.post_init() + + @deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size, seq_len, _ = inputs_embeds.shape + hidden_states = self.token_embedding_projection(inputs_embeds) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for i, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + return hidden_states + + +def process_patch_lengths(patch_lengths: torch.Tensor, max_patch_length: int | None) -> torch.Tensor: + """ + Splits patch lengths into smaller segments if they exceed `max_patch_length`. + Pads the result to uniform length across the batch. + + Args: + patch_lengths (torch.Tensor): [batch_size, num_patches] tensor of patch lengths. + max_patch_length (int, optional): Maximum allowed length per patch. + + Returns: + torch.Tensor: [batch_size, max_len] tensor of split and padded patch lengths. + """ + if max_patch_length is None: + return patch_lengths + + batch_size = patch_lengths.size(0) + processed = [] + + for seq in patch_lengths: + splits = [] + for length in seq[seq > 0]: + length = length.item() + full_chunks, remainder = divmod(length, max_patch_length) + splits.extend([max_patch_length] * full_chunks) + if remainder: + splits.append(remainder) + processed.append(splits) + + # Find max length to pad to + max_len = max(len(splits) for splits in processed) + padded = torch.zeros((batch_size, max_len), dtype=patch_lengths.dtype, device=patch_lengths.device) + + for i, splits in enumerate(processed): + if splits: + padded[i, : len(splits)] = torch.tensor(splits, dtype=patch_lengths.dtype, device=patch_lengths.device) + + # Trim zero columns + if (padded != 0).any(dim=0).sum() < padded.shape[1]: + last_nonzero = (padded != 0).any(dim=0).nonzero().max().item() + 1 + padded = padded[:, :last_nonzero] + + return padded + + +class BltPatcher(BltPreTrainedModel): + config: BltPatcherConfig + + def __init__(self, config: BltPatcherConfig): + super().__init__(config) + self.rotary_emb = BltRotaryEmbedding(config=self.config) + self.layers = nn.ModuleList() + for layer_idx in range(self.config.num_hidden_layers): + self.layers.append(BltTransformerLayer(self.config, layer_idx)) + self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size) + self.norm = BltRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.lm_head = nn.Linear( + self.config.hidden_size, + self.config.vocab_size, + bias=False, + ) + + self.post_init() + + 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, + patch_size: int | None = None, + threshold: float | None = None, + max_patch_length: int | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + 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, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings=position_embeddings, attention_mask=causal_mask) + + logits = self.lm_head(self.norm(hidden_states)) + prediction_entropies = torch.distributions.Categorical(logits=logits).entropy() + + batch_size, sequence_length = inputs_embeds.shape[:2] + if patch_size is not None: + patch_lengths = self.patch_lengths_from_entropies( + entropies=prediction_entropies, + sequence_length=sequence_length, + patch_size=patch_size, + threshold=threshold, + ) + else: + patch_lengths = torch.ones( + (batch_size, sequence_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + patch_lengths = process_patch_lengths(patch_lengths, max_patch_length) + return prediction_entropies, patch_lengths, logits + + @staticmethod + def patch_lengths_from_entropies( + entropies, + sequence_length, + patch_size=None, + threshold=None, + ): + """ + Computes patch lengths from token entropies. + + Depending on whether a threshold is provided, the function uses either: + - Thresholding the entropy values (when `threshold` is set). + """ + + batch_size = entropies.shape[0] + + # Always include token 0 and 1 as starting tokens + init_tokens = ( + torch.tensor([0, 1], dtype=torch.long, device=entropies.device).unsqueeze(0).repeat(batch_size, 1) + ) + offset = init_tokens.shape[1] + + # Ignore first token entropy (BOS) + entropies = entropies[:, 1:] + + # Threshold the entropy values to define patch start points + patch_mask = entropies > threshold + + seq_len = patch_mask.shape[1] + + # Create patch IDs (token indices), and add a sentinel to ensure alignment + token_indices = torch.arange(seq_len, device=entropies.device).unsqueeze(0).expand(batch_size, -1) + sentinel = torch.full_like(token_indices, seq_len) + padded_indices = torch.cat([token_indices, sentinel], dim=1) + + # Pad mask with inverse to align sentinel correctly + padded_mask = torch.cat([patch_mask, ~patch_mask], dim=1) + + # Select indices where mask is True + patch_starts = padded_indices[padded_mask].reshape(batch_size, seq_len) + max_valid_patches = patch_mask.sum(dim=1).max() + patch_starts = patch_starts[:, :max_valid_patches] + + # Offset patch starts to account for the two initial tokens + patch_start_ids = torch.cat((init_tokens, patch_starts + offset), dim=1) + + # Compute patch end positions by shifting start positions + last_token = torch.full_like(patch_start_ids[:, :1], sequence_length - 1) + patch_ends = torch.cat((patch_start_ids[:, 1:] - 1, last_token), dim=1) + + patch_lengths = patch_ends - patch_start_ids + 1 + + return patch_lengths + + +def rolling_polynomial_hash(token_tensor, prime: int = 1000000007): + """ + A polynomial rolling hash algorithm that converts sequences + of tokens into hash values. The hash is computed as: + hash = (token_0 * prime^0 + token_1 * prime^1 + ... + token_n * prime^n) + + The rolling hash allows the model to efficiently + identify and encode recurring byte-level patterns in the input text. + + Args: + token_tensor (torch.Tensor): [batch_size, seq_len, group_size] containing token IDs to hash + prime (int): Prime number used as the base for the polynomial hash. + + Returns: + torch.Tensor: Hash values of shape [batch_size, seq_len] where each value + represents the hash of the corresponding token group + + Example: + >>> tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) + >>> hashes = rolling_polynomial_hash(tokens, prime=31) + >>> # hash[0] = 1*31^0 + 2*31^1 + 3*31^2 + >>> # hash[1] = 4*31^0 + 5*31^1 + 6*31^2 + """ + prime_tensor = torch.tensor(prime, dtype=torch.int64, device=token_tensor.device) + powers = torch.arange(token_tensor.shape[-1], device=token_tensor.device) + prime_powers = prime_tensor**powers + return torch.sum(token_tensor * prime_powers, dim=-1) + + +def byte_group_hash_function( + token_ids: torch.Tensor, group_size: int = 2, prime: int = 1000000007, max_hash: int = 30000 +): + """Hash token groups and map to range [0, max_hash].""" + with torch.no_grad(): + batch_size, seq_len = token_ids.shape + # Add padding for sliding window + padding = torch.zeros(batch_size, group_size - 1, dtype=torch.int64, device=token_ids.device) + padded_tokens = torch.cat([padding, token_ids], dim=1) + + # Create sliding windows and compute hashes + windows = padded_tokens.unfold(1, group_size, 1) + hashes = rolling_polynomial_hash(windows, prime) + hash_values = hashes % max_hash + + return hash_values + + +def compute_hash_embeddings( + local_encoder_tokens: torch.Tensor, + local_encoder, + encoder_hash_tok_embedding: nn.Embedding, + encoder_hash_byte_group_nb_functions: int, + encoder_hash_byte_group_size: list, + encoder_hash_byte_group_vocab: int, +) -> torch.Tensor: + """Compute token embeddings enhanced with hash-based embeddings.""" + # Available primes for hash functions + primes = [ + 1000000007, + 5915587277, + 1500450271, + 3267000013, + 5754853343, + 4093082899, + 9576890767, + 3628273133, + 2860486313, + 5463458053, + 3367900313, + ] + + embeddings = local_encoder.embed_tokens(local_encoder_tokens) + embedding_idx = 0 + for func_nb in range(encoder_hash_byte_group_nb_functions): + prime = primes[func_nb % len(primes)] # Cycle through primes if more functions than primes + for group_size in encoder_hash_byte_group_size: + hash_ids = byte_group_hash_function(local_encoder_tokens, group_size, prime, encoder_hash_byte_group_vocab) + # Apply offset to get the correct slice of the fused embedding + offset_hash_ids = hash_ids + embedding_idx * encoder_hash_byte_group_vocab + embeddings += encoder_hash_tok_embedding(offset_hash_ids).to(embeddings.device) + embedding_idx += 1 + + return embeddings + + +def _prepare_patch_cross_attention_mask( + patch_ids: torch.Tensor, + num_patches: int, + sequence_length: int, + patches_as_queries: bool = False, + cross_attn_k: int = 1, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Prepare cross-attention mask for patch-based attention, following mllama's robust approach. + + This function creates masks that control which patches can attend to which other patches, + with support for query/key role swapping and cross-attention multipliers. + + Args: + patch_ids (torch.Tensor): Tensor of shape [batch_size, seq_len] containing patch ids. + num_patches (int): Total number of patches. + sequence_length (int): Length of the sequence. + patches_as_queries (bool): If True, patches are used as queries, otherwise as keys. + cross_attn_k (int): Cross-attention multiplier for repeating patches. + dtype (torch.dtype): Data type for the output mask. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + - cross_attention_mask: 4D tensor [batch_size, 1, q_len, kv_len] + """ + batch_size, seq_len = patch_ids.shape + device = patch_ids.device + + # Determine query and key lengths based on configuration + if patches_as_queries: + q_len = num_patches * cross_attn_k + kv_len = sequence_length + # Create patch-to-sequence mapping + q_patch_ids = ( + torch.arange(num_patches, device=device) + .unsqueeze(0) + .unsqueeze(-1) + .expand(batch_size, num_patches, seq_len) + ) + kv_patch_ids = patch_ids.unsqueeze(1).expand(batch_size, num_patches, seq_len) + else: + q_len = sequence_length + kv_len = num_patches * cross_attn_k + # Create sequence-to-patch mapping + q_patch_ids = patch_ids.unsqueeze(-1).expand(batch_size, seq_len, num_patches) + kv_patch_ids = ( + torch.arange(num_patches, device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, seq_len, num_patches) + ) + + # Create base attention mask - boolean mask where True means "should attend" + # Exact patch matching + cross_attention_mask = q_patch_ids == kv_patch_ids + + # Handle cross_attn_k multiplier by repeating along appropriate dimension + repeat_dim = 1 if patches_as_queries else -1 + cross_attention_mask = cross_attention_mask.repeat_interleave(cross_attn_k, dim=repeat_dim) + + # Validate dimensions + expected_shape = (batch_size, q_len, kv_len) + if cross_attention_mask.shape != expected_shape: + raise ValueError( + f"Cross attention mask shape {cross_attention_mask.shape} doesn't match expected {expected_shape}" + ) + + # Reshape so it can be used by attn module - add head dimension + cross_attention_mask = cross_attention_mask.unsqueeze(1) # [batch_size, 1, q_len, kv_len] + + # Invert the mask (following mllama pattern exactly) + # True -> 0.0 (attend), False -> 1.0 (will become -inf) + inverted_cross_attn_mask = 1.0 - cross_attention_mask.to(dtype) + cross_attention_mask = inverted_cross_attn_mask.masked_fill( + inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min + ) + + return cross_attention_mask + + +class BltModel(BltPreTrainedModel): + def __init__(self, config: BltConfig): + super().__init__(config) + self.gradient_checkpointing = False + + self.config = config + self.local_encoder = BltLocalEncoder(config.encoder_config) + self.global_transformer = BltGlobalTransformer(config.global_config) + self.local_decoder = BltLocalDecoder(config.decoder_config) + num_embeddings = config.encoder_hash_byte_group_nb_functions * len(config.encoder_hash_byte_group_size) + total_vocab_size = config.encoder_hash_byte_group_vocab * num_embeddings + self.encoder_hash_tok_embedding = nn.Embedding(total_vocab_size, config.encoder_config.hidden_size) + if self.config.patch_in_forward: + self.patcher = BltPatcher(config.patcher_config) + self.patcher.eval() + for param in self.patcher.parameters(): + param.requires_grad = False + else: + self.patcher = None + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor | None = None, + patch_lengths: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | 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 use_cache: + if past_key_values is None: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + elif not isinstance(past_key_values, EncoderDecoderCache): + # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache + # if not yet created by the user + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config)) + + # Extract input embeddings as early as possible + if inputs_embeds is not None: + encoder_embeds = inputs_embeds + batch_size, sequence_length, _ = inputs_embeds.shape + else: + batch_size, sequence_length = input_ids.shape + encoder_embeds = compute_hash_embeddings( + input_ids, + self.local_encoder, + self.encoder_hash_tok_embedding, + self.config.encoder_hash_byte_group_nb_functions, + self.config.encoder_hash_byte_group_size, + self.config.encoder_hash_byte_group_vocab, + ) + + if patch_lengths is None: + if self.config.patching_mode == "entropy" and self.patcher is not None: + if input_ids is None: + raise ValueError("input_ids is required for entropy-based patching") + _, patch_lengths, _ = self.patcher( + input_ids, + patch_size=self.config.patch_size, + threshold=self.config.patching_threshold, + max_patch_length=self.config.max_patch_length, + patching_batch_size=self.config.patching_batch_size, + device=input_ids.device, + ) + else: + device = input_ids.device if input_ids is not None else inputs_embeds.device + dtype = input_ids.dtype if input_ids is not None else inputs_embeds.dtype + patch_lengths = process_patch_lengths( + torch.ones((batch_size, sequence_length + 1), dtype=dtype, device=device), + self.config.max_patch_length, + ) + patch_ids = self._patch_ids_from_lengths(patch_lengths, sequence_length) + + 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(encoder_embeds.shape[1], device=encoder_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=encoder_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values.self_attention_cache if past_key_values is not None else None, + position_ids=position_ids, + ) + + cross_attn_mask_enc = _prepare_patch_cross_attention_mask( + patch_ids=patch_ids, + num_patches=patch_lengths.shape[1], + sequence_length=sequence_length, + patches_as_queries=True, + cross_attn_k=self.config.cross_attn_k, + dtype=encoder_embeds.dtype, + ) + encoder_hidden_states, encoder_cross_states = self.local_encoder( + input_ids=input_ids, + inputs_embeds=encoder_embeds, + attention_mask=causal_mask, + position_ids=position_ids, + encoder_attention_mask=cross_attn_mask_enc, + num_patches=patch_lengths.shape[1], + patch_ids=patch_ids, + past_key_values=past_key_values.self_attention_cache if past_key_values is not None else None, + **kwargs, + ) + encoder_cross_states = encoder_cross_states.view(batch_size, patch_lengths.shape[1], -1) + global_position_ids = torch.arange(0, encoder_cross_states.shape[1], device=encoder_cross_states.device) + global_position_ids = global_position_ids.unsqueeze(0) + global_causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=encoder_cross_states, + attention_mask=None, + past_key_values=None, + position_ids=None, + ) + + global_hidden_states = self.global_transformer( + inputs_embeds=encoder_cross_states, + attention_mask=global_causal_mask, + position_ids=global_position_ids, + **kwargs, + ) + decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length) + cross_attn_mask_dec = _prepare_patch_cross_attention_mask( + patch_ids=decoder_patch_ids, + num_patches=patch_lengths.shape[1], + sequence_length=sequence_length, + patches_as_queries=False, + cross_attn_k=self.config.cross_attn_k, + dtype=encoder_embeds.dtype, + ) + output = self.local_decoder( + input_ids=input_ids, + inputs_embeds=encoder_hidden_states, + patch_embeds=global_hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values.cross_attention_cache if past_key_values is not None else None, + encoder_attention_mask=cross_attn_mask_dec, + **kwargs, + ) + return BaseModelOutputWithPast( + last_hidden_state=output, + past_key_values=past_key_values, + ) + + def get_input_embeddings(self): + return self.local_encoder.embed_tokens + + def set_input_embeddings(self, value): + self.local_encoder.embed_tokens = value + + def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) -> torch.Tensor: + batch_size = patch_lengths.shape[0] + patch_starts = torch.cat( + [ + torch.zeros(batch_size, 1, dtype=patch_lengths.dtype, device=patch_lengths.device), + patch_lengths.cumsum(dim=-1)[:, :-1], + ], + dim=-1, + ) + token_positions = torch.arange(seq_len, device=patch_lengths.device) + return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1 + + +@auto_docstring( + custom_intro=""" + The Blt Text Model with a language modeling head on top. + """ +) +class BltForCausalLM(BltPreTrainedModel, GenerationMixin): + config: BltConfig + _can_compile_fullgraph = False + base_model_prefix = "model" + _tied_weights_keys = {"model.local_encoder.embed_tokens.weight": "lm_head.weight"} + + def __init__(self, config: BltConfig): + super().__init__(config) + self.text_config = config.get_text_config() + self.vocab_size = config.vocab_size + self.model = BltModel(config) + self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False) + + 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, + cross_attention_states: torch.LongTensor | None = None, # Keep for compatibility + cross_attention_mask: torch.LongTensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*): + A tuple containing two tensors that mask out rows in the cross-attention mechanism: + - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1. + A value of 0 indicates that the corresponding text token's entire row in the cross-attention + matrix should be masked out (all image tokens ignored). + - The second tensor has the same shape and is used internally to apply the masking during + the forward pass of cross-attention layers. + This mask is derived from the cross_attention_mask and is used to handle cases where a text token + should not attend to any image token. + 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, BltForCausalLM + + >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf") + + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ + # Call parent forward but exclude cross_attention_states from model call + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + 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, :]).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__ = ["BltPreTrainedModel", "BltModel", "BltPatcher", "BltForCausalLM"] diff --git a/third_party/transformers/src/transformers/models/blt/modular_blt.py b/third_party/transformers/src/transformers/models/blt/modular_blt.py new file mode 100644 index 0000000000000000000000000000000000000000..4dfff77263db97cd451918f7c825abc446a6e799 --- /dev/null +++ b/third_party/transformers/src/transformers/models/blt/modular_blt.py @@ -0,0 +1,1195 @@ +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Blt modular model, inheriting from Mllama where appropriate.""" + +from collections.abc import Callable + +import torch +import torch.distributions +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.deprecation import deprecate_kwarg +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..cohere2.modeling_cohere2 import rotate_half # noqa: F401 +from ..llama.modeling_llama import LlamaRotaryEmbedding +from ..mllama.modeling_mllama import ( + MllamaPreTrainedModel, + MllamaSelfAttentionDecoderLayer, + MllamaTextCrossAttention, + MllamaTextMLP, + MllamaTextRMSNorm, + MllamaTextSelfAttention, + eager_attention_forward, +) +from .configuration_blt import ( + BltConfig, + BltGlobalTransformerConfig, + BltLocalDecoderConfig, + BltLocalEncoderConfig, + BltPatcherConfig, +) + + +logger = logging.get_logger(__name__) + + +def rolling_polynomial_hash(token_tensor, prime: int = 1000000007): + """ + A polynomial rolling hash algorithm that converts sequences + of tokens into hash values. The hash is computed as: + hash = (token_0 * prime^0 + token_1 * prime^1 + ... + token_n * prime^n) + + The rolling hash allows the model to efficiently + identify and encode recurring byte-level patterns in the input text. + + Args: + token_tensor (torch.Tensor): [batch_size, seq_len, group_size] containing token IDs to hash + prime (int): Prime number used as the base for the polynomial hash. + + Returns: + torch.Tensor: Hash values of shape [batch_size, seq_len] where each value + represents the hash of the corresponding token group + + Example: + >>> tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) + >>> hashes = rolling_polynomial_hash(tokens, prime=31) + >>> # hash[0] = 1*31^0 + 2*31^1 + 3*31^2 + >>> # hash[1] = 4*31^0 + 5*31^1 + 6*31^2 + """ + prime_tensor = torch.tensor(prime, dtype=torch.int64, device=token_tensor.device) + powers = torch.arange(token_tensor.shape[-1], device=token_tensor.device) + prime_powers = prime_tensor**powers + return torch.sum(token_tensor * prime_powers, dim=-1) + + +def byte_group_hash_function( + token_ids: torch.Tensor, group_size: int = 2, prime: int = 1000000007, max_hash: int = 30000 +): + """Hash token groups and map to range [0, max_hash].""" + with torch.no_grad(): + batch_size, seq_len = token_ids.shape + # Add padding for sliding window + padding = torch.zeros(batch_size, group_size - 1, dtype=torch.int64, device=token_ids.device) + padded_tokens = torch.cat([padding, token_ids], dim=1) + + # Create sliding windows and compute hashes + windows = padded_tokens.unfold(1, group_size, 1) + hashes = rolling_polynomial_hash(windows, prime) + hash_values = hashes % max_hash + + return hash_values + + +def compute_hash_embeddings( + local_encoder_tokens: torch.Tensor, + local_encoder, + encoder_hash_tok_embedding: nn.Embedding, + encoder_hash_byte_group_nb_functions: int, + encoder_hash_byte_group_size: list, + encoder_hash_byte_group_vocab: int, +) -> torch.Tensor: + """Compute token embeddings enhanced with hash-based embeddings.""" + # Available primes for hash functions + primes = [ + 1000000007, + 5915587277, + 1500450271, + 3267000013, + 5754853343, + 4093082899, + 9576890767, + 3628273133, + 2860486313, + 5463458053, + 3367900313, + ] + + embeddings = local_encoder.embed_tokens(local_encoder_tokens) + embedding_idx = 0 + for func_nb in range(encoder_hash_byte_group_nb_functions): + prime = primes[func_nb % len(primes)] # Cycle through primes if more functions than primes + for group_size in encoder_hash_byte_group_size: + hash_ids = byte_group_hash_function(local_encoder_tokens, group_size, prime, encoder_hash_byte_group_vocab) + # Apply offset to get the correct slice of the fused embedding + offset_hash_ids = hash_ids + embedding_idx * encoder_hash_byte_group_vocab + embeddings += encoder_hash_tok_embedding(offset_hash_ids).to(embeddings.device) + embedding_idx += 1 + + return embeddings + + +def _prepare_patch_cross_attention_mask( + patch_ids: torch.Tensor, + num_patches: int, + sequence_length: int, + patches_as_queries: bool = False, + cross_attn_k: int = 1, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Prepare cross-attention mask for patch-based attention, following mllama's robust approach. + + This function creates masks that control which patches can attend to which other patches, + with support for query/key role swapping and cross-attention multipliers. + + Args: + patch_ids (torch.Tensor): Tensor of shape [batch_size, seq_len] containing patch ids. + num_patches (int): Total number of patches. + sequence_length (int): Length of the sequence. + patches_as_queries (bool): If True, patches are used as queries, otherwise as keys. + cross_attn_k (int): Cross-attention multiplier for repeating patches. + dtype (torch.dtype): Data type for the output mask. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + - cross_attention_mask: 4D tensor [batch_size, 1, q_len, kv_len] + """ + batch_size, seq_len = patch_ids.shape + device = patch_ids.device + + # Determine query and key lengths based on configuration + if patches_as_queries: + q_len = num_patches * cross_attn_k + kv_len = sequence_length + # Create patch-to-sequence mapping + q_patch_ids = ( + torch.arange(num_patches, device=device) + .unsqueeze(0) + .unsqueeze(-1) + .expand(batch_size, num_patches, seq_len) + ) + kv_patch_ids = patch_ids.unsqueeze(1).expand(batch_size, num_patches, seq_len) + else: + q_len = sequence_length + kv_len = num_patches * cross_attn_k + # Create sequence-to-patch mapping + q_patch_ids = patch_ids.unsqueeze(-1).expand(batch_size, seq_len, num_patches) + kv_patch_ids = ( + torch.arange(num_patches, device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, seq_len, num_patches) + ) + + # Create base attention mask - boolean mask where True means "should attend" + # Exact patch matching + cross_attention_mask = q_patch_ids == kv_patch_ids + + # Handle cross_attn_k multiplier by repeating along appropriate dimension + repeat_dim = 1 if patches_as_queries else -1 + cross_attention_mask = cross_attention_mask.repeat_interleave(cross_attn_k, dim=repeat_dim) + + # Validate dimensions + expected_shape = (batch_size, q_len, kv_len) + if cross_attention_mask.shape != expected_shape: + raise ValueError( + f"Cross attention mask shape {cross_attention_mask.shape} doesn't match expected {expected_shape}" + ) + + # Reshape so it can be used by attn module - add head dimension + cross_attention_mask = cross_attention_mask.unsqueeze(1) # [batch_size, 1, q_len, kv_len] + + # Invert the mask (following mllama pattern exactly) + # True -> 0.0 (attend), False -> 1.0 (will become -inf) + inverted_cross_attn_mask = 1.0 - cross_attention_mask.to(dtype) + cross_attention_mask = inverted_cross_attn_mask.masked_fill( + inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min + ) + + return cross_attention_mask + + +def process_patch_lengths(patch_lengths: torch.Tensor, max_patch_length: int | None) -> torch.Tensor: + """ + Splits patch lengths into smaller segments if they exceed `max_patch_length`. + Pads the result to uniform length across the batch. + + Args: + patch_lengths (torch.Tensor): [batch_size, num_patches] tensor of patch lengths. + max_patch_length (int, optional): Maximum allowed length per patch. + + Returns: + torch.Tensor: [batch_size, max_len] tensor of split and padded patch lengths. + """ + if max_patch_length is None: + return patch_lengths + + batch_size = patch_lengths.size(0) + processed = [] + + for seq in patch_lengths: + splits = [] + for length in seq[seq > 0]: + length = length.item() + full_chunks, remainder = divmod(length, max_patch_length) + splits.extend([max_patch_length] * full_chunks) + if remainder: + splits.append(remainder) + processed.append(splits) + + # Find max length to pad to + max_len = max(len(splits) for splits in processed) + padded = torch.zeros((batch_size, max_len), dtype=patch_lengths.dtype, device=patch_lengths.device) + + for i, splits in enumerate(processed): + if splits: + padded[i, : len(splits)] = torch.tensor(splits, dtype=patch_lengths.dtype, device=patch_lengths.device) + + # Trim zero columns + if (padded != 0).any(dim=0).sum() < padded.shape[1]: + last_nonzero = (padded != 0).any(dim=0).nonzero().max().item() + 1 + padded = padded[:, :last_nonzero] + + return padded + + +class BltMLP(MllamaTextMLP): + pass + + +class BltRMSNorm(MllamaTextRMSNorm): + pass + + +class BltRotaryEmbedding(LlamaRotaryEmbedding): + @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 BltTransformerLayer(MllamaSelfAttentionDecoderLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + + self.self_attn = BltSelfAttention(config=config, layer_idx=layer_idx) + self.mlp = BltMLP(config) + self.input_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class BltSelfAttention(MllamaTextSelfAttention): + def __init__(self, config: BltConfig, layer_idx: int): + super().__init__(config, layer_idx) + + +class BltCrossAttention(MllamaTextCrossAttention): + """Cross-attention module for Blt, following transformers style""" + + def __init__(self, config: BltConfig, layer_idx: int, hidden_size: int | None = None): + super().__init__() + self.is_causal = False + self.q_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps) + self.k_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + bsz, q_len, _ = hidden_states.size() + query_states = self.q_norm(hidden_states) + query_states = self.q_proj(query_states) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + cross_attention_states = self.k_norm(cross_attention_states) + key_states = self.k_proj(cross_attention_states) + value_states = self.v_proj(cross_attention_states) + key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, -1, self.num_key_value_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, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + attn_output = attn_output + hidden_states + return attn_output, attn_weights + + +@auto_docstring +class BltPreTrainedModel(MllamaPreTrainedModel): + config: BltConfig + _supports_attention_backend = False + _supports_flash_attn = False + _supports_flex_attn = False + _no_split_modules = ["BltTransformerLayer"] + _can_record_outputs = { + "hidden_states": OutputRecorder(BltTransformerLayer, index=0), + "attentions": OutputRecorder(BltSelfAttention, index=1), + } + + # Weight initialization is adapted from: + # - https://github.com/facebookresearch/blt/blob/main/bytelatent/model/blt.py + # - https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/model/model.py + # + # Both implementations use truncated normal initialization with std ~ 1 / sqrt(d_model) + # (or 1 / sqrt(hidden_dim) for FFN outputs), and unit initialization for normalization layers. + # We follow the same scheme here, but expressed in the Transformers APIs. + + @torch.no_grad() + def _init_weights(self, module): + """ + Initialize BLT weights following the original ByteLatentTransformer: + + - Most weights are drawn from a truncated normal. + - Scale is ~ 1 / sqrt(model_dim) (or 1 / sqrt(hidden_dim) for FFN outputs). + - Norm layers are set to weight = 1, bias = 0. + """ + class_name = module.__class__.__name__ + + # Norms: RMSNorm / LayerNorm + if isinstance(module, (BltRMSNorm, nn.LayerNorm)) or "RMSNorm" in class_name or "LayerNorm" in class_name: + if getattr(module, "weight", None) is not None: + init.ones_(module.weight) + if getattr(module, "bias", None) is not None: + init.zeros_(module.bias) + return + + # Embeddings (encoder / patcher / hash embeddings) + if isinstance(module, nn.Embedding): + hidden_size = getattr(self.config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "encoder_config"): + hidden_size = getattr(self.config.encoder_config, "hidden_size", None) + if hidden_size is None: + hidden_size = module.embedding_dim + + std = hidden_size**-0.5 + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if module.padding_idx is not None: + init.zeros_(module.weight[module.padding_idx]) + return + + # Self-attention / cross-attention projections + if isinstance(module, (BltSelfAttention, BltCrossAttention)) or class_name in ( + "MllamaTextSelfAttention", + "MllamaTextCrossAttention", + ): + dim = getattr(self.config, "hidden_size", None) + if dim is None and hasattr(module, "hidden_size"): + dim = module.hidden_size + if dim is None: + for name in ("q_proj", "k_proj", "v_proj", "o_proj", "dense"): + proj = getattr(module, name, None) + if proj is not None and hasattr(proj, "weight"): + dim = proj.weight.shape[-1] + break + if dim is None: + return + + std = dim**-0.5 + + # Input projections (q, k, v) + for proj_name in ("q_proj", "k_proj", "v_proj"): + proj = getattr(module, proj_name, None) + if proj is not None and hasattr(proj, "weight"): + init.trunc_normal_( + proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(proj, "bias", None) is not None: + init.zeros_(proj.bias) + + # Output projection: o_proj or dense + o_proj = getattr(module, "o_proj", getattr(module, "dense", None)) + if o_proj is not None and hasattr(o_proj, "weight"): + init.trunc_normal_( + o_proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(o_proj, "bias", None) is not None: + init.zeros_(o_proj.bias) + return + + # MLP / FFN blocks + if isinstance(module, BltMLP) or class_name == "MllamaTextMLP": + hidden_size = getattr(self.config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "decoder_config"): + hidden_size = getattr(self.config.decoder_config, "hidden_size", None) + if hidden_size is None and hasattr(self.config, "encoder_config"): + hidden_size = getattr(self.config.encoder_config, "hidden_size", None) + + # Input-side std + in_std = None + if hidden_size is not None: + in_std = hidden_size**-0.5 + + gate_proj = getattr(module, "gate_proj", getattr(module, "fc1", None)) + up_proj = getattr(module, "up_proj", None) + down_proj = getattr(module, "down_proj", getattr(module, "fc2", None)) + + # gate / input projections + for proj in (gate_proj, up_proj): + if proj is not None and hasattr(proj, "weight"): + std = in_std or (proj.weight.shape[1] ** -0.5) + init.trunc_normal_( + proj.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if getattr(proj, "bias", None) is not None: + init.zeros_(proj.bias) + + # output/ down projections + if down_proj is not None and hasattr(down_proj, "weight"): + hidden_dim = down_proj.weight.shape[1] + out_std = hidden_dim**-0.5 + init.trunc_normal_( + down_proj.weight, + mean=0.0, + std=out_std, + a=-3 * out_std, + b=3 * out_std, + ) + if getattr(down_proj, "bias", None) is not None: + init.zeros_(down_proj.bias) + return + + # Generic Linear layers (projections, lm_head, etc.) + if isinstance(module, nn.Linear): + fan_in = module.in_features + std = fan_in**-0.5 + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-3 * std, + b=3 * std, + ) + if module.bias is not None: + init.zeros_(module.bias) + return + + if isinstance(module, BltRotaryEmbedding): + rope_fn = ( + ROPE_INIT_FUNCTIONS[module.rope_type] + if module.rope_type != "default" + else module.compute_default_rope_parameters + ) + buffer_value, _ = rope_fn(module.config) + init.copy_(module.inv_freq, buffer_value) + init.copy_(module.original_inv_freq, buffer_value) + + +class BltLocalEncoder(BltPreTrainedModel): + config: BltLocalEncoderConfig + _can_record_outputs = { + "encoder_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_encoder"), + } + + def __init__(self, config: BltLocalEncoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + self.config = config + self.layers = nn.ModuleList( + [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = BltRotaryEmbedding(config=config) + self.patch_embedding_projection = nn.Linear( + in_features=config.hidden_size, + out_features=config.hidden_size * config.cross_attn_k, + bias=False, + ) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.cross_attn_layers = nn.ModuleList() + layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1 + for layer_idx in range(layers_to_add): + self.cross_attn_layers.append( + BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size) + ) + + self.post_init() + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + patch_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + encoder_attention_mask: torch.Tensor | None = None, + num_patches: int | None = None, + patch_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + batch_size = inputs_embeds.shape[0] + hidden_states = F.dropout(inputs_embeds, p=self.config.dropout, training=self.training) + + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + + for idx, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + if idx == len(self.layers) - 1 or self.config.cross_attn_all_layers: + patch_embeds = self.patch_reduce(hidden_states, num_patches, patch_ids) + patch_embeds = self.patch_embedding_projection(patch_embeds) + patch_embeds = patch_embeds.reshape( + batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size + ) + layer_idx = idx if self.config.cross_attn_all_layers else 0 + cross_attention_output, _ = self.cross_attn_layers[layer_idx]( + hidden_states=patch_embeds, + cross_attention_states=hidden_states, + attention_mask=encoder_attention_mask, + **kwargs, + ) + patch_embeds = patch_embeds + cross_attention_output + encoder_cross_states = patch_embeds + return hidden_states, encoder_cross_states + + def patch_reduce(self, hidden_states, max_num_patches, patch_ids): + """ + Reduce variable length patches to single embedding per patch + Note: this works with variable number of patches for different sequences in the batch + It handles variable length patches by assuming that patch_lengths will be 0 for any + extra patches on the *right*. Since there can be a variable number of patches + this function also return the number of patches for each sequence in the batch. + Any embeddings on the right that are not allocated to a patch + (i.e. if the sum(patch_lengths[i]) < seq_len for any i) + will be sent to a dummy patch, which is trimmed before returning. + """ + batch_size = hidden_states.shape[0] + embedding_dim = hidden_states.shape[-1] + + patch_ids = patch_ids.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]) + + reduced_embeddings = torch.zeros( + (batch_size, max_num_patches, embedding_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + reduced_embeddings = reduced_embeddings.scatter_reduce( + src=hidden_states, + dim=1, + index=patch_ids, + reduce="amax", + include_self=False, + ) + reduced_embeddings = reduced_embeddings[:, :max_num_patches, :] + + return reduced_embeddings + + +class BltLocalDecoder(BltPreTrainedModel): + config: BltLocalDecoderConfig + + def __init__(self, config: BltLocalDecoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + self.config = config + self.cross_attn_decoder = True + self.layers = nn.ModuleList( + [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = BltRotaryEmbedding(config=config) + self.patch_embedding_projection = nn.Linear( + in_features=config.hidden_size_global, + out_features=config.hidden_size * config.cross_attn_k, + bias=False, + ) + self.norm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn_layers = nn.ModuleList() + layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1 + for layer_idx in range(layers_to_add): + self.cross_attn_layers.append( + BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size) + ) + + self.post_init() + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + patch_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size = inputs_embeds.shape[0] + hidden_states = inputs_embeds + patch_embeds = self.patch_embedding_projection(patch_embeds) + patch_embeds = patch_embeds.reshape( + batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size + ) + + if patch_embeds is not None and not self.cross_attn_decoder: + hidden_states = hidden_states + patch_embeds + + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + + for i, layer in enumerate(self.layers): + if i == 0 or self.config.cross_attn_all_layers: + cross_attention_output, _ = self.cross_attn_layers[i]( + hidden_states=hidden_states, + cross_attention_states=patch_embeds, + attention_mask=encoder_attention_mask, + **kwargs, + ) + hidden_states = hidden_states + cross_attention_output + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + logits = self.norm(hidden_states) + return logits + + +class BltGlobalTransformer(BltPreTrainedModel): + config: BltGlobalTransformerConfig + _can_record_outputs = { + "global_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="global_transformer"), + } + + def __init__(self, config: BltGlobalTransformerConfig): + super().__init__(config) + self.config = config + self.layers = nn.ModuleList() + for layer_idx in range(config.num_hidden_layers): + self.layers.append(BltTransformerLayer(config, layer_idx)) + self.rotary_emb = BltRotaryEmbedding(config=config) + + # Create token embedding projection (use nn.Identity() when no projection needed) + if getattr(config, "encoder_cross_output_size", None) is not None: + self.token_embedding_projection = nn.Linear( + config.encoder_cross_output_size, config.hidden_size, bias=False + ) + else: + self.token_embedding_projection = nn.Identity() + + self.post_init() + + @deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size, seq_len, _ = inputs_embeds.shape + hidden_states = self.token_embedding_projection(inputs_embeds) + hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training) + if position_ids is None: + position_ids = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1) + ) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for i, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + return hidden_states + + +class BltPatcher(BltPreTrainedModel): + config: BltPatcherConfig + + def __init__(self, config: BltPatcherConfig): + super().__init__(config) + self.rotary_emb = BltRotaryEmbedding(config=self.config) + self.layers = nn.ModuleList() + for layer_idx in range(self.config.num_hidden_layers): + self.layers.append(BltTransformerLayer(self.config, layer_idx)) + self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size) + self.norm = BltRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.lm_head = nn.Linear( + self.config.hidden_size, + self.config.vocab_size, + bias=False, + ) + + self.post_init() + + 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, + patch_size: int | None = None, + threshold: float | None = None, + max_patch_length: int | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + 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, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings=position_embeddings, attention_mask=causal_mask) + + logits = self.lm_head(self.norm(hidden_states)) + prediction_entropies = torch.distributions.Categorical(logits=logits).entropy() + + batch_size, sequence_length = inputs_embeds.shape[:2] + if patch_size is not None: + patch_lengths = self.patch_lengths_from_entropies( + entropies=prediction_entropies, + sequence_length=sequence_length, + patch_size=patch_size, + threshold=threshold, + ) + else: + patch_lengths = torch.ones( + (batch_size, sequence_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + patch_lengths = process_patch_lengths(patch_lengths, max_patch_length) + return prediction_entropies, patch_lengths, logits + + @staticmethod + def patch_lengths_from_entropies( + entropies, + sequence_length, + patch_size=None, + threshold=None, + ): + """ + Computes patch lengths from token entropies. + + Depending on whether a threshold is provided, the function uses either: + - Thresholding the entropy values (when `threshold` is set). + """ + + batch_size = entropies.shape[0] + + # Always include token 0 and 1 as starting tokens + init_tokens = ( + torch.tensor([0, 1], dtype=torch.long, device=entropies.device).unsqueeze(0).repeat(batch_size, 1) + ) + offset = init_tokens.shape[1] + + # Ignore first token entropy (BOS) + entropies = entropies[:, 1:] + + # Threshold the entropy values to define patch start points + patch_mask = entropies > threshold + + seq_len = patch_mask.shape[1] + + # Create patch IDs (token indices), and add a sentinel to ensure alignment + token_indices = torch.arange(seq_len, device=entropies.device).unsqueeze(0).expand(batch_size, -1) + sentinel = torch.full_like(token_indices, seq_len) + padded_indices = torch.cat([token_indices, sentinel], dim=1) + + # Pad mask with inverse to align sentinel correctly + padded_mask = torch.cat([patch_mask, ~patch_mask], dim=1) + + # Select indices where mask is True + patch_starts = padded_indices[padded_mask].reshape(batch_size, seq_len) + max_valid_patches = patch_mask.sum(dim=1).max() + patch_starts = patch_starts[:, :max_valid_patches] + + # Offset patch starts to account for the two initial tokens + patch_start_ids = torch.cat((init_tokens, patch_starts + offset), dim=1) + + # Compute patch end positions by shifting start positions + last_token = torch.full_like(patch_start_ids[:, :1], sequence_length - 1) + patch_ends = torch.cat((patch_start_ids[:, 1:] - 1, last_token), dim=1) + + patch_lengths = patch_ends - patch_start_ids + 1 + + return patch_lengths + + +class BltModel(BltPreTrainedModel): + def __init__(self, config: BltConfig): + super().__init__(config) + self.gradient_checkpointing = False + + self.config = config + self.local_encoder = BltLocalEncoder(config.encoder_config) + self.global_transformer = BltGlobalTransformer(config.global_config) + self.local_decoder = BltLocalDecoder(config.decoder_config) + num_embeddings = config.encoder_hash_byte_group_nb_functions * len(config.encoder_hash_byte_group_size) + total_vocab_size = config.encoder_hash_byte_group_vocab * num_embeddings + self.encoder_hash_tok_embedding = nn.Embedding(total_vocab_size, config.encoder_config.hidden_size) + if self.config.patch_in_forward: + self.patcher = BltPatcher(config.patcher_config) + self.patcher.eval() + for param in self.patcher.parameters(): + param.requires_grad = False + else: + self.patcher = None + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor | None = None, + patch_lengths: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | 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 use_cache: + if past_key_values is None: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + elif not isinstance(past_key_values, EncoderDecoderCache): + # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache + # if not yet created by the user + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config)) + + # Extract input embeddings as early as possible + if inputs_embeds is not None: + encoder_embeds = inputs_embeds + batch_size, sequence_length, _ = inputs_embeds.shape + else: + batch_size, sequence_length = input_ids.shape + encoder_embeds = compute_hash_embeddings( + input_ids, + self.local_encoder, + self.encoder_hash_tok_embedding, + self.config.encoder_hash_byte_group_nb_functions, + self.config.encoder_hash_byte_group_size, + self.config.encoder_hash_byte_group_vocab, + ) + + if patch_lengths is None: + if self.config.patching_mode == "entropy" and self.patcher is not None: + if input_ids is None: + raise ValueError("input_ids is required for entropy-based patching") + _, patch_lengths, _ = self.patcher( + input_ids, + patch_size=self.config.patch_size, + threshold=self.config.patching_threshold, + max_patch_length=self.config.max_patch_length, + patching_batch_size=self.config.patching_batch_size, + device=input_ids.device, + ) + else: + device = input_ids.device if input_ids is not None else inputs_embeds.device + dtype = input_ids.dtype if input_ids is not None else inputs_embeds.dtype + patch_lengths = process_patch_lengths( + torch.ones((batch_size, sequence_length + 1), dtype=dtype, device=device), + self.config.max_patch_length, + ) + patch_ids = self._patch_ids_from_lengths(patch_lengths, sequence_length) + + 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(encoder_embeds.shape[1], device=encoder_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=encoder_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values.self_attention_cache if past_key_values is not None else None, + position_ids=position_ids, + ) + + cross_attn_mask_enc = _prepare_patch_cross_attention_mask( + patch_ids=patch_ids, + num_patches=patch_lengths.shape[1], + sequence_length=sequence_length, + patches_as_queries=True, + cross_attn_k=self.config.cross_attn_k, + dtype=encoder_embeds.dtype, + ) + encoder_hidden_states, encoder_cross_states = self.local_encoder( + input_ids=input_ids, + inputs_embeds=encoder_embeds, + attention_mask=causal_mask, + position_ids=position_ids, + encoder_attention_mask=cross_attn_mask_enc, + num_patches=patch_lengths.shape[1], + patch_ids=patch_ids, + past_key_values=past_key_values.self_attention_cache if past_key_values is not None else None, + **kwargs, + ) + encoder_cross_states = encoder_cross_states.view(batch_size, patch_lengths.shape[1], -1) + global_position_ids = torch.arange(0, encoder_cross_states.shape[1], device=encoder_cross_states.device) + global_position_ids = global_position_ids.unsqueeze(0) + global_causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=encoder_cross_states, + attention_mask=None, + past_key_values=None, + position_ids=None, + ) + + global_hidden_states = self.global_transformer( + inputs_embeds=encoder_cross_states, + attention_mask=global_causal_mask, + position_ids=global_position_ids, + **kwargs, + ) + decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length) + cross_attn_mask_dec = _prepare_patch_cross_attention_mask( + patch_ids=decoder_patch_ids, + num_patches=patch_lengths.shape[1], + sequence_length=sequence_length, + patches_as_queries=False, + cross_attn_k=self.config.cross_attn_k, + dtype=encoder_embeds.dtype, + ) + output = self.local_decoder( + input_ids=input_ids, + inputs_embeds=encoder_hidden_states, + patch_embeds=global_hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values.cross_attention_cache if past_key_values is not None else None, + encoder_attention_mask=cross_attn_mask_dec, + **kwargs, + ) + return BaseModelOutputWithPast( + last_hidden_state=output, + past_key_values=past_key_values, + ) + + def get_input_embeddings(self): + return self.local_encoder.embed_tokens + + def set_input_embeddings(self, value): + self.local_encoder.embed_tokens = value + + def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) -> torch.Tensor: + batch_size = patch_lengths.shape[0] + patch_starts = torch.cat( + [ + torch.zeros(batch_size, 1, dtype=patch_lengths.dtype, device=patch_lengths.device), + patch_lengths.cumsum(dim=-1)[:, :-1], + ], + dim=-1, + ) + token_positions = torch.arange(seq_len, device=patch_lengths.device) + return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1 + + +@auto_docstring( + custom_intro=""" + The Blt Text Model with a language modeling head on top. + """ +) +class BltForCausalLM(BltPreTrainedModel, GenerationMixin): + config: BltConfig + _can_compile_fullgraph = False + base_model_prefix = "model" + _tied_weights_keys = {"model.local_encoder.embed_tokens.weight": "lm_head.weight"} + + def __init__(self, config: BltConfig): + super().__init__(config) + self.text_config = config.get_text_config() + self.vocab_size = config.vocab_size + self.model = BltModel(config) + self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False) + + 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, + cross_attention_states: torch.LongTensor | None = None, # Keep for compatibility + cross_attention_mask: torch.LongTensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*): + A tuple containing two tensors that mask out rows in the cross-attention mechanism: + - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1. + A value of 0 indicates that the corresponding text token's entire row in the cross-attention + matrix should be masked out (all image tokens ignored). + - The second tensor has the same shape and is used internally to apply the masking during + the forward pass of cross-attention layers. + This mask is derived from the cross_attention_mask and is used to handle cases where a text token + should not attend to any image token. + 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, BltForCausalLM + + >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf") + + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ + # Call parent forward but exclude cross_attention_states from model call + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + 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, :]).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__ = [ + "BltPreTrainedModel", + "BltModel", + "BltPatcher", + "BltForCausalLM", +] diff --git a/third_party/transformers/src/transformers/models/byt5/__init__.py b/third_party/transformers/src/transformers/models/byt5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cb726942b0f16105f8a5a5f7c661485951d7ccc7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/byt5/__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_byt5 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/byt5/convert_byt5_original_tf_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/byt5/convert_byt5_original_tf_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..b73bc5aaf75e5ed89904c314b834125923562c77 --- /dev/null +++ b/third_party/transformers/src/transformers/models/byt5/convert_byt5_original_tf_checkpoint_to_pytorch.py @@ -0,0 +1,167 @@ +# Copyright 2018 The T5 authors and 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 T5 checkpoint.""" + +import argparse +import os + +import torch + +from transformers import T5Config, T5ForConditionalGeneration +from transformers.utils import logging + + +logger = logging.get_logger(__name__) +logging.set_verbosity_info() + + +def load_tf_weights_in_t5(model, config, tf_checkpoint_path): + """Load tf checkpoints in a pytorch model.""" + try: + import re + + import numpy as np + import tensorflow as tf + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions." + ) + raise + tf_path = os.path.abspath(tf_checkpoint_path) + logger.info(f"Converting TensorFlow checkpoint from {tf_path}") + # Load weights from TF model + init_vars = tf.train.list_variables(tf_path) + names = [] + tf_weights = {} + for name, shape in init_vars: + logger.info(f"Loading TF weight {name} with shape {shape}") + array = tf.train.load_variable(tf_path, name) + names.append(name) + tf_weights[name] = array + + for txt_name in names: + name = txt_name.split("/") + # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v + # which are not required for using pretrained model + if any( + n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] + for n in name + ): + logger.info(f"Skipping {'/'.join(name)}") + tf_weights.pop(txt_name, None) + continue + if "_slot_" in name[-1]: + logger.info(f"Skipping {'/'.join(name)}") + tf_weights.pop(txt_name, None) + continue + pointer = model + array = tf_weights[txt_name] + + for m_name in name: + if re.fullmatch(r"[A-Za-z]+_\d+", m_name): + scope_names = re.split(r"_(\d+)", m_name) + else: + scope_names = [m_name] + if scope_names[0] in ["kernel", "scale", "embedding"]: + pointer = getattr(pointer, "weight") + elif scope_names[0] == "self_attention": + pointer = getattr(pointer, "layer") + pointer = pointer[0] + elif scope_names[0] == "enc_dec_attention": + pointer = getattr(pointer, "layer") + pointer = pointer[1] + elif scope_names[0] == "dense_relu_dense": + pointer = getattr(pointer, "layer") + pointer = pointer[2] + elif scope_names[0] == "rms_norm": + if hasattr(pointer, "layer_norm"): + pointer = getattr(pointer, "layer_norm") + elif hasattr(pointer, "final_layer_norm"): + pointer = getattr(pointer, "final_layer_norm") + elif scope_names[0] == "scale": + pointer = getattr(pointer, "weight") + elif scope_names[0] == "output_bias" or scope_names[0] == "beta": + pointer = getattr(pointer, "bias") + elif scope_names[0] == "squad": + pointer = getattr(pointer, "classifier") + elif scope_names[0] == "decoder" and name[1] == "logits": + continue + elif scope_names[0] == "logits": + pointer = getattr(pointer, "lm_head") + elif scope_names[0] == "wi" and len(scope_names) > 1 and scope_names[1].isdigit(): + pointer = getattr(pointer, f"wi_{scope_names[1]}") + continue + else: + try: + pointer = getattr(pointer, scope_names[0]) + except AttributeError: + logger.info(f"Skipping {'/'.join(name)}") + continue + if len(scope_names) >= 2: + num = int(scope_names[1]) + pointer = pointer[num] + if scope_names[0] not in ["kernel", "scale", "embedding"]: + pointer = getattr(pointer, "weight") + if scope_names[0] != "embedding": + logger.info(f"Transposing numpy weight of shape {array.shape} for {name}") + array = np.transpose(array) + try: + if pointer.shape != array.shape: + raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched") + except AssertionError as e: + e.args += (pointer.shape, array.shape) + raise + logger.info(f"Initialize PyTorch weight {name}") + pointer.data = torch.from_numpy(array.astype(np.float32)) + tf_weights.pop(txt_name, None) + + logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.") + return model + + +def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path): + # Initialise PyTorch model + config = T5Config.from_json_file(config_file) + print(f"Building PyTorch model from configuration: {config}") + model = T5ForConditionalGeneration(config) + + # Load weights from tf checkpoint + load_tf_weights_in_t5(model, config, tf_checkpoint_path) + + # Save pytorch-model + print(f"Save PyTorch model to {pytorch_dump_path}") + model.save_pretrained(pytorch_dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path." + ) + parser.add_argument( + "--config_file", + default=None, + type=str, + required=True, + help=( + "The config json file corresponding to the pre-trained T5 model. \nThis specifies the model architecture." + ), + ) + parser.add_argument( + "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model." + ) + args = parser.parse_args() + convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path) diff --git a/third_party/transformers/src/transformers/models/byt5/tokenization_byt5.py b/third_party/transformers/src/transformers/models/byt5/tokenization_byt5.py new file mode 100644 index 0000000000000000000000000000000000000000..ce4ed37a4ec387b157dd0e862602b7c40c1704b2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/byt5/tokenization_byt5.py @@ -0,0 +1,234 @@ +# Copyright 2021 T5 Authors and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization class for model ByT5.""" + +import warnings + +from ...tokenization_python import AddedToken, PreTrainedTokenizer +from ...utils import logging + + +logger = logging.get_logger(__name__) + + +class ByT5Tokenizer(PreTrainedTokenizer): + """ + Construct a ByT5 tokenizer. ByT5 simply uses raw bytes utf-8 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: + 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`. + + + + 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. + extra_ids (`int`, *optional*, defaults to 125): + Add a number of extra ids added to the end of the vocabulary for use as sentinels. These tokens are + accessible as "" where "{%d}" is a number between 0 and extra_ids-1. Extra tokens are + indexed from the end of the vocabulary up to beginning ("" is the last token in the vocabulary + like in ByT5 preprocessing see + [here](https://github.com/google-research/text-to-text-transfer-transformer/blob/9fd7b14a769417be33bc6c850f9598764913c833/t5/data/preprocessors.py#L2117)). + additional_special_tokens (`list[str]`, *optional*): + Additional special tokens used by the tokenizer. + """ + + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + eos_token="", + unk_token="", + pad_token="", + extra_ids=125, + additional_special_tokens=None, + **kwargs, + ) -> None: + # Add extra_ids to the special token list + if extra_ids > 0 and additional_special_tokens is None: + additional_special_tokens = [f"" for i in range(extra_ids)] + elif extra_ids > 0 and additional_special_tokens is not None and len(additional_special_tokens) > 0: + # Check that we have the right number of extra_id special tokens + extra_tokens = len(set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens))) + if extra_tokens != extra_ids: + raise ValueError( + f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are" + " provided to ByT5Tokenizer. In this case the additional_special_tokens must include the" + " extra_ids tokens" + ) + + pad_token = AddedToken(pad_token, lstrip=True, rstrip=True) if isinstance(pad_token, str) else pad_token + # we force left and right stripping for backward compatibility. The byt5tests depend on this. + eos_token = AddedToken(eos_token, lstrip=True, rstrip=True) if isinstance(eos_token, str) else eos_token + unk_token = AddedToken(unk_token, lstrip=True, rstrip=True) if isinstance(unk_token, str) else unk_token + # unk token needs to be in the vocab with correct index + self._added_tokens_decoder = {0: pad_token, 1: eos_token, 2: unk_token} + self.offset = len(self._added_tokens_decoder) + self._utf_vocab_size = 2**8 # utf is 8 bits + super().__init__( + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + extra_ids=0, + additional_special_tokens=additional_special_tokens, # TODO extra ids are not used :sweatywmile: + **kwargs, + ) + + @property + def vocab_size(self): + return self._utf_vocab_size + + def get_vocab(self): + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size + self.offset)} + vocab.update(self.added_tokens_encoder) + return vocab + + 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 + ) + + # normal case: some special tokens + if token_ids_1 is None: + return ([0] * len(token_ids_0)) + [1] + return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + + def _add_eos_if_not_present(self, token_ids: list[int]) -> list[int]: + """Do not add eos again if user already added it.""" + if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id: + warnings.warn( + f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated" + " eos tokens being added." + ) + return token_ids + else: + return token_ids + [self.eos_token_id] + + 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. ByT5 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. + """ + eos = [self.eos_token_id] + + if token_ids_1 is None: + return len(token_ids_0 + eos) * [0] + return len(token_ids_0 + eos + token_ids_1 + eos) * [0] + + 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 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. + """ + token_ids_0 = self._add_eos_if_not_present(token_ids_0) + if token_ids_1 is None: + return token_ids_0 + else: + token_ids_1 = self._add_eos_if_not_present(token_ids_1) + return token_ids_0 + token_ids_1 + + def _tokenize(self, text: str) -> list[str]: + """Take as input a string and return a list of strings (tokens) for words/sub-words""" + tokens = [chr(i) for i in text.encode("utf-8")] + return tokens + + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + + if len(token) != 1: + token_id = None + else: + token_id = ord(token) + self.offset + + return token_id + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + token = chr(index - self.offset) + return token + + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + bstring = b"" + for token in tokens: + if token in self.added_tokens_decoder: + tok_string = self.added_tokens_decoder[token].encode("utf-8") + elif token in self.added_tokens_encoder: + tok_string = token.encode("utf-8") + else: + tok_string = bytes([ord(token)]) + bstring += tok_string + string = bstring.decode("utf-8", errors="ignore") + return string + + # ByT5Tokenizer has no vocab file + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + return () + + +__all__ = ["ByT5Tokenizer"] diff --git a/third_party/transformers/src/transformers/models/clip/__init__.py b/third_party/transformers/src/transformers/models/clip/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b3554d4ac39ae7b8316c59457058222f1b38dd42 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_clip import * + from .image_processing_clip import * + from .image_processing_pil_clip import * + from .modeling_clip import * + from .processing_clip import * + from .tokenization_clip 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/clip/configuration_clip.py b/third_party/transformers/src/transformers/models/clip/configuration_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..2e6392df37cb24fba1bc85f876f39095e225366a --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/configuration_clip.py @@ -0,0 +1,248 @@ +# Copyright 2021 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. +"""CLIP 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="openai/clip-vit-base-patch32") +@strict +class CLIPTextConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import CLIPTextConfig, CLIPTextModel + + >>> # Initializing a CLIPTextConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPTextConfig() + + >>> # Initializing a CLIPTextModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPTextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "clip_text_model" + base_config_key = "text_config" + + vocab_size: int = 49408 + hidden_size: int = 512 + intermediate_size: int = 2048 + projection_dim: int | None = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 8 + max_position_embeddings: int = 77 + hidden_act: str = "quick_gelu" + layer_norm_eps: float | None = 1e-5 + attention_dropout: int | float | None = 0.0 + initializer_range: float = 0.02 + initializer_factor: float | None = 1.0 + + # This differs from `CLIPTokenizer`'s default and from openai/clip + # See https://github.com/huggingface/transformers/pull/24773#issuecomment-1632287538 + pad_token_id: int | None = 1 + bos_token_id: int | None = 49406 + eos_token_id: int | list[int] | None = 49407 + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +@auto_docstring(checkpoint="openai/clip-vit-base-patch32") +@strict +class CLIPVisionConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import CLIPVisionConfig, CLIPVisionModel + + >>> # Initializing a CLIPVisionConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPVisionConfig() + + >>> # Initializing a CLIPVisionModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPVisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "clip_vision_model" + base_config_key = "vision_config" + + hidden_size: int = 768 + intermediate_size: int = 3072 + projection_dim: int | None = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int | None = 3 + image_size: int | None = 224 + patch_size: int | None = 32 + hidden_act: str = "quick_gelu" + layer_norm_eps: float | None = 1e-5 + attention_dropout: int | float | None = 0.0 + initializer_range: float = 0.02 + initializer_factor: float | None = 1.0 + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +@auto_docstring(checkpoint="openai/clip-vit-base-patch32") +@strict +class CLIPConfig(PreTrainedConfig): + r""" + text_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`CLIPTextConfig`]. + vision_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`CLIPVisionConfig`]. + logit_scale_init_value (`float | int`, *optional*, defaults to 2.6592): + The initial value of the *logit_scale* parameter. Default is used as per the original CLIP implementation. + + Example: + + ```python + >>> from transformers import CLIPConfig, CLIPModel + + >>> # Initializing a CLIPConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPConfig() + + >>> # Initializing a CLIPModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a CLIPConfig from a CLIPTextConfig and a CLIPVisionConfig + >>> from transformers import CLIPTextConfig, CLIPVisionConfig + + >>> # Initializing a CLIPText and CLIPVision configuration + >>> config_text = CLIPTextConfig() + >>> config_vision = CLIPVisionConfig() + + >>> config = CLIPConfig(text_config=config_text, vision_config=config_vision) + ```""" + + model_type = "clip" + sub_configs = {"text_config": CLIPTextConfig, "vision_config": CLIPVisionConfig} + + text_config: dict | CLIPTextConfig | None = None + vision_config: dict | CLIPVisionConfig | None = None + projection_dim: int | None = 512 + logit_scale_init_value: float | int | None = 2.6592 + initializer_factor: float | None = 1.0 + + def __post_init__(self, **kwargs): + if self.text_config is None: + text_config = {} + logger.info("`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.") + elif isinstance(self.text_config, CLIPTextConfig): + text_config = self.text_config.to_dict() + else: + text_config = self.text_config + + if self.vision_config is None: + vision_config = {} + logger.info("`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.") + elif isinstance(self.vision_config, CLIPVisionConfig): + vision_config = self.vision_config.to_dict() + else: + vision_config = self.vision_config + + # For backward compatibility check keyword args + # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in + # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most + # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`. + text_config_dict = kwargs.pop("text_config_dict", None) + vision_config_dict = kwargs.pop("vision_config_dict", None) + + if text_config_dict is not None: + # This is the complete result when using `text_config_dict`. + _text_config_dict = CLIPTextConfig(**text_config_dict).to_dict() + + # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different. + for key, value in _text_config_dict.items(): + if key in text_config and value != text_config[key] and key != "transformers_version": + # If specified in `text_config_dict` + if key in text_config_dict: + message = ( + f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. " + f'The value `text_config_dict["{key}"]` will be used instead.' + ) + # If inferred from default argument values (just to be super careful) + else: + message = ( + f"`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The " + f'value `text_config["{key}"]` will be overridden.' + ) + logger.info(message) + + # Update all values in `text_config` with the ones in `_text_config_dict`. + text_config.update(_text_config_dict) + + if vision_config_dict is not None: + # This is the complete result when using `vision_config_dict`. + _vision_config_dict = CLIPVisionConfig(**vision_config_dict).to_dict() + # convert keys to string instead of integer + if "id2label" in _vision_config_dict: + _vision_config_dict["id2label"] = { + str(key): value for key, value in _vision_config_dict["id2label"].items() + } + + # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different. + for key, value in _vision_config_dict.items(): + if key in vision_config and value != vision_config[key] and key != "transformers_version": + # If specified in `vision_config_dict` + if key in vision_config_dict: + message = ( + f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different " + f'values. The value `vision_config_dict["{key}"]` will be used instead.' + ) + # If inferred from default argument values (just to be super careful) + else: + message = ( + f"`vision_config_dict` is provided which will be used to initialize `CLIPVisionConfig`. " + f'The value `vision_config["{key}"]` will be overridden.' + ) + logger.info(message) + + # Update all values in `vision_config` with the ones in `_vision_config_dict`. + vision_config.update(_vision_config_dict) + + # Finally we can convert back our unified text/vision configs to `PretrainedConfig` + self.text_config = CLIPTextConfig(**text_config) + self.vision_config = CLIPVisionConfig(**vision_config) + + super().__post_init__(**kwargs) + + +__all__ = ["CLIPConfig", "CLIPTextConfig", "CLIPVisionConfig"] diff --git a/third_party/transformers/src/transformers/models/clip/convert_clip_original_pytorch_to_hf.py b/third_party/transformers/src/transformers/models/clip/convert_clip_original_pytorch_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0987e7785988e186b11b72a42e89bb5920d580b5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/convert_clip_original_pytorch_to_hf.py @@ -0,0 +1,155 @@ +# Copyright 2021 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 torch +from clip import load + +from transformers import CLIPConfig, CLIPModel + + +def copy_attn_layer(hf_attn_layer, pt_attn_layer): + q_proj, k_proj, v_proj = pt_attn_layer.in_proj_weight.chunk(3, dim=0) + q_proj_bias, k_proj_bias, v_proj_bias = pt_attn_layer.in_proj_bias.chunk(3, dim=0) + + out_proj_weights = pt_attn_layer.out_proj.weight + out_proj_bias = pt_attn_layer.out_proj.bias + + hf_attn_layer.q_proj.weight.data = q_proj + hf_attn_layer.q_proj.bias.data = q_proj_bias + + hf_attn_layer.k_proj.weight.data = k_proj + hf_attn_layer.k_proj.bias.data = k_proj_bias + + hf_attn_layer.v_proj.weight.data = v_proj + hf_attn_layer.v_proj.bias.data = v_proj_bias + + hf_attn_layer.out_proj.weight = out_proj_weights + hf_attn_layer.out_proj.bias = out_proj_bias + + +def copy_mlp(hf_mlp, pt_mlp): + copy_linear(hf_mlp.fc1, pt_mlp.c_fc) + copy_linear(hf_mlp.fc2, pt_mlp.c_proj) + + +def copy_linear(hf_linear, pt_linear): + hf_linear.weight = pt_linear.weight + hf_linear.bias = pt_linear.bias + + +def copy_layer(hf_layer, pt_layer): + # copy layer norms + copy_linear(hf_layer.layer_norm1, pt_layer.ln_1) + copy_linear(hf_layer.layer_norm2, pt_layer.ln_2) + + # copy MLP + copy_mlp(hf_layer.mlp, pt_layer.mlp) + + # copy attn + copy_attn_layer(hf_layer.self_attn, pt_layer.attn) + + +def copy_layers(hf_layers, pt_layers): + for hf_layer, pt_layer in zip(hf_layers, pt_layers): + copy_layer(hf_layer, pt_layer) + + +def copy_encoder(hf_encoder, pt_model): + # copy embeds + hf_encoder.embeddings.token_embedding.weight = pt_model.token_embedding.weight + hf_encoder.embeddings.position_embedding.weight.data = pt_model.positional_embedding + + # copy layer norm + copy_linear(hf_encoder.final_layer_norm, pt_model.ln_final) + + # copy hidden layers + copy_layers(hf_encoder.encoder.layers, pt_model.transformer.resblocks) + + +def copy_text_model_and_projection(hf_model, pt_model): + # copy projection + hf_model.text_projection.weight.data = pt_model.text_projection.data.T.contiguous() + + # copy text encoder + copy_encoder(hf_model.text_model, pt_model) + + +def copy_vison_model_and_projection(hf_model, pt_model): + # copy projection + hf_model.visual_projection.weight.data = pt_model.visual.proj.data.T.contiguous() + + # copy layer norms + copy_linear(hf_model.vision_model.pre_layrnorm, pt_model.visual.ln_pre) + copy_linear(hf_model.vision_model.post_layernorm, pt_model.visual.ln_post) + + # copy embeds + hf_model.vision_model.embeddings.patch_embedding.weight.data = pt_model.visual.conv1.weight.data + hf_model.vision_model.embeddings.class_embedding = pt_model.visual.class_embedding + hf_model.vision_model.embeddings.position_embedding.weight.data = pt_model.visual.positional_embedding.data + + # copy encoder + copy_layers(hf_model.vision_model.encoder.layers, pt_model.visual.transformer.resblocks) + + +@torch.no_grad() +def convert_clip_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_path=None): + """ + Copy/paste/tweak model's weights to transformers design. + """ + if config_path is not None: + config = CLIPConfig.from_pretrained(config_path) + else: + config = CLIPConfig(projection_dim=512, text_config={}, vision_config={}) + + hf_model = CLIPModel(config).eval() + + pt_model, _ = load(checkpoint_path, device="cpu", jit=False) + pt_model = pt_model.eval() + + copy_text_model_and_projection(hf_model, pt_model) + copy_vison_model_and_projection(hf_model, pt_model) + hf_model.logit_scale = pt_model.logit_scale + + # Use `eos_token` so the example is more meaningful + input_ids = torch.tensor( + [ + [config.text_config.bos_token_id] + + list(range(3, 77)) + + [config.text_config.eos_token_id] + + [config.text_config.pad_token_id] + ] + ) + pixel_values = torch.randn(1, 3, 224, 224) + + hf_outputs = hf_model(input_ids=input_ids, pixel_values=pixel_values, return_dict=True) + hf_logits_per_image = hf_outputs.logits_per_image + hf_logits_per_text = hf_outputs.logits_per_text + pt_logits_per_image, pt_logits_per_text = pt_model(pixel_values, input_ids) + + assert torch.allclose(hf_logits_per_image, pt_logits_per_image, atol=1e-3) + assert torch.allclose(hf_logits_per_text, pt_logits_per_text, atol=1e-3) + + hf_model.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") + parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to OpenAI checkpoint") + parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") + args = parser.parse_args() + + convert_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path) diff --git a/third_party/transformers/src/transformers/models/clip/image_processing_clip.py b/third_party/transformers/src/transformers/models/clip/image_processing_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..d79b1a716b02e38e94f84e3ae7f89a47062e858d --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/image_processing_clip.py @@ -0,0 +1,45 @@ +# Copyright 2021 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 CLIP.""" + +from ...image_processing_backends import TorchvisionBackend +from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import auto_docstring + + +@auto_docstring +class CLIPImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"shortest_edge": 224} + default_to_square = False + crop_size = {"height": 224, "width": 224} + do_resize = True + do_center_crop = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + + def __init__(self, **kwargs: Unpack[ImagesKwargs]): + # for backwards compatibility of KOSMOS-2 + if "use_square_size" in kwargs and kwargs["use_square_size"]: + kwargs["size"] = {"height": self.size["shortest_edge"], "width": self.size["shortest_edge"]} + kwargs.pop("use_square_size") + + super().__init__(**kwargs) + + +__all__ = ["CLIPImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/clip/image_processing_pil_clip.py b/third_party/transformers/src/transformers/models/clip/image_processing_pil_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..228309ec0557196877d9585ed872bf036dbcc1d1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/image_processing_pil_clip.py @@ -0,0 +1,45 @@ +# Copyright 2021 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 CLIP.""" + +from ...image_processing_backends import PilBackend +from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import auto_docstring + + +@auto_docstring +class CLIPImageProcessorPil(PilBackend): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"shortest_edge": 224} + default_to_square = False + crop_size = {"height": 224, "width": 224} + do_resize = True + do_center_crop = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + + def __init__(self, **kwargs: Unpack[ImagesKwargs]): + # for backwards compatibility of KOSMOS-2 + if "use_square_size" in kwargs and kwargs["use_square_size"]: + kwargs["size"] = {"height": self.size["shortest_edge"], "width": self.size["shortest_edge"]} + kwargs.pop("use_square_size") + + super().__init__(**kwargs) + + +__all__ = ["CLIPImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/clip/modeling_clip.py b/third_party/transformers/src/transformers/models/clip/modeling_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..923c650d21583cc307d54020ece225bd6761ac13 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/modeling_clip.py @@ -0,0 +1,1148 @@ +# Copyright 2021 The OpenAI Team Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch CLIP model.""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + logging, + torch_int, +) +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig + + +logger = logging.get_logger(__name__) + + +# contrastive loss function, adapted from +# https://sachinruk.github.io/blog/2021-03-07-clip.html +def contrastive_loss(logits: torch.Tensor) -> torch.Tensor: + return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device)) + + +def clip_loss(similarity: torch.Tensor) -> torch.Tensor: + caption_loss = contrastive_loss(similarity) + image_loss = contrastive_loss(similarity.t()) + return (caption_loss + image_loss) / 2.0 + + +def _get_vector_norm(tensor: torch.Tensor) -> torch.Tensor: + """ + This method is equivalent to tensor.norm(p=2, dim=-1, keepdim=True) and used to make + model `executorch` exportable. See issue https://github.com/pytorch/executorch/issues/3566 + """ + square_tensor = torch.pow(tensor, 2) + sum_tensor = torch.sum(square_tensor, dim=-1, keepdim=True) + normed_tensor = torch.pow(sum_tensor, 0.5) + return normed_tensor + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + """ +) +class CLIPVisionModelOutput(ModelOutput): + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + """ + + image_embeds: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for text model's outputs that also contains a pooling of the last hidden states. + """ +) +class CLIPTextModelOutput(ModelOutput): + r""" + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + """ + + text_embeds: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring +class CLIPOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`CLIPTextModel`]. + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The image embeddings obtained by applying the projection layer to the pooled output of [`CLIPVisionModel`]. + text_model_output (`BaseModelOutputWithPooling`): + The output of the [`CLIPTextModel`]. + vision_model_output (`BaseModelOutputWithPooling`): + The output of the [`CLIPVisionModel`]. + """ + + loss: torch.FloatTensor | None = None + logits_per_image: torch.FloatTensor | None = None + logits_per_text: torch.FloatTensor | None = None + text_embeds: torch.FloatTensor | None = None + image_embeds: torch.FloatTensor | None = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPooling = None + + def to_tuple(self) -> tuple[Any]: + return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values()) + + +class CLIPVisionEmbeddings(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] - 1 + position_embedding = self.position_embedding.weight.unsqueeze(0) + num_positions = position_embedding.shape[1] - 1 + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + class_pos_embed = position_embedding[:, :1] + patch_pos_embed = position_embedding[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size): + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})." + ) + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +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], +): + 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 CLIPAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: CLIPVisionConfig | CLIPTextConfig): + 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 + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class CLIPMLP(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 CLIPEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: CLIPVisionConfig | CLIPTextConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = CLIPAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> 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 + + +@auto_docstring +class CLIPPreTrainedModel(PreTrainedModel): + config: CLIPConfig + base_model_prefix = "clip" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": CLIPEncoderLayer, + "attentions": CLIPAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor + if isinstance(module, CLIPTextEmbeddings): + init.normal_(module.token_embedding.weight, mean=0.0, std=factor * 0.02) + init.normal_(module.position_embedding.weight, mean=0.0, std=factor * 0.02) + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + elif isinstance(module, CLIPVisionEmbeddings): + factor = self.config.initializer_factor + init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) + elif isinstance(module, CLIPAttention): + factor = self.config.initializer_factor + in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor + out_proj_std = (module.embed_dim**-0.5) * factor + init.normal_(module.q_proj.weight, std=in_proj_std) + init.normal_(module.k_proj.weight, std=in_proj_std) + init.normal_(module.v_proj.weight, std=in_proj_std) + init.normal_(module.out_proj.weight, std=out_proj_std) + elif isinstance(module, CLIPMLP): + factor = self.config.initializer_factor + in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor + fc_std = (2 * module.config.hidden_size) ** -0.5 * factor + init.normal_(module.fc1.weight, std=fc_std) + init.normal_(module.fc2.weight, std=in_proj_std) + elif isinstance(module, CLIPModel): + init.normal_( + module.text_projection.weight, + std=module.text_embed_dim**-0.5 * self.config.initializer_factor, + ) + init.normal_( + module.visual_projection.weight, + std=module.vision_embed_dim**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPVisionModelWithProjection): + init.normal_( + module.visual_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPTextModelWithProjection): + init.normal_( + module.text_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPForImageClassification): + init.normal_( + module.classifier.weight, + std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor, + ) + + if isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + if isinstance(module, nn.Linear) and module.bias is not None: + init.zeros_(module.bias) + + +class CLIPEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__(self, config: CLIPConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + """ + hidden_states = inputs_embeds + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_states, + attention_mask, + **kwargs, + ) + + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +class CLIPTextTransformer(CLIPPreTrainedModel): + config: CLIPTextConfig + input_modalities = ("text",) + + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + self.config = config + embed_dim = config.hidden_size + self.embeddings = CLIPTextEmbeddings(config) + self.encoder = CLIPEncoder(config) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + self.post_init() + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPooling: + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + past_key_values=None, + ) + + kwargs.pop("is_causal", None) + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + is_causal=True, + **kwargs, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of extra new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id) + .int() + .argmax(dim=-1), + ] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + ) + + +@auto_docstring( + custom_intro=""" + The text model from CLIP without any head or projection on top. + """ +) +class CLIPTextModel(CLIPPreTrainedModel): + config: CLIPTextConfig + input_modalities = ("text",) + + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + self.text_model = CLIPTextTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPTextModel + + >>> model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + + +class CLIPVisionTransformer(CLIPPreTrainedModel): + config: CLIPVisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = CLIPEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.post_init() + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + interpolate_pos_encoding: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPooling: + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) + hidden_states = self.pre_layrnorm(hidden_states) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + **kwargs, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + pooled_output = last_hidden_state[:, 0, :] + pooled_output = self.post_layernorm(pooled_output) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + ) + + +@auto_docstring( + custom_intro=""" + The vision model from CLIP without any head or projection on top. + """ +) +class CLIPVisionModel(CLIPPreTrainedModel): + config: CLIPVisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + self.vision_model = CLIPVisionTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + interpolate_pos_encoding: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPooling: + r""" + Example: + + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, CLIPVisionModel + + >>> model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled CLS states + ```""" + + return self.vision_model( + pixel_values=pixel_values, + interpolate_pos_encoding=interpolate_pos_encoding, + **kwargs, + ) + + +@auto_docstring +class CLIPModel(CLIPPreTrainedModel): + config: CLIPConfig + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer", "CLIPVisionEmbeddings"] + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + if not isinstance(config.text_config, CLIPTextConfig): + raise TypeError( + "config.text_config is expected to be of type CLIPTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, CLIPVisionConfig): + raise TypeError( + "config.vision_config is expected to be of type CLIPVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + self.vision_embed_dim = vision_config.hidden_size + + text_model = CLIPTextModel._from_config(text_config) + self.text_model = text_model.text_model + + vision_model = CLIPVisionModel._from_config(vision_config) + self.vision_model = vision_model.vision_model + + self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False) + self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False) + self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value)) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> with torch.inference_mode(): + ... text_features = model.get_text_features(**inputs) + ```""" + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + **kwargs, + ) + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) + + return text_outputs + + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + interpolate_pos_encoding: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, CLIPModel + >>> from transformers.image_utils import load_image + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = load_image(url) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> with torch.inference_mode(): + ... image_features = model.get_image_features(**inputs) + ```""" + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, + ) + pooled_output = vision_outputs.pooler_output + vision_outputs.pooler_output = self.visual_projection(pooled_output) + + return vision_outputs + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + return_loss: bool | None = None, + interpolate_pos_encoding: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> CLIPOutput: + r""" + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + + Examples: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, CLIPModel + >>> from transformers.image_utils import load_image + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = load_image(url) + + >>> inputs = processor( + ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True + ... ) + + >>> with torch.inference_mode(): + ... outputs = model(**inputs) + >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score + >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities + ```""" + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + interpolate_pos_encoding=interpolate_pos_encoding, + **kwargs, + ) + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + + image_embeds = vision_outputs.pooler_output + image_embeds = self.visual_projection(image_embeds) + + text_embeds = text_outputs.pooler_output + text_embeds = self.text_projection(text_embeds) + + # normalized features + image_embeds = image_embeds / _get_vector_norm(image_embeds) + text_embeds = text_embeds / _get_vector_norm(text_embeds) + + # cosine similarity as logits + logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) + logits_per_text = logits_per_text * self.logit_scale.exp().to(text_embeds.device) + + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = clip_loss(logits_per_text) + + return CLIPOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +@auto_docstring +class CLIPTextModelWithProjection(CLIPPreTrainedModel): + config: CLIPTextConfig + input_modalities = ("text",) + + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + + text_model = CLIPTextModel._from_config(config) + self.text_model = text_model.text_model + + self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> CLIPTextModelOutput: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, CLIPTextModelWithProjection + + >>> model = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> with torch.inference_mode(): + ... outputs = model(**inputs) + >>> text_embeds = outputs.text_embeds + ```""" + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + pooled_output = text_outputs.pooler_output + text_embeds = self.text_projection(pooled_output) + + return CLIPTextModelOutput( + text_embeds=text_embeds, + last_hidden_state=text_outputs.last_hidden_state, + hidden_states=text_outputs.hidden_states, + attentions=text_outputs.attentions, + ) + + +@auto_docstring +class CLIPVisionModelWithProjection(CLIPPreTrainedModel): + config: CLIPVisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + + vision_model = CLIPVisionModel._from_config(config) + self.vision_model = vision_model.vision_model + + self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + interpolate_pos_encoding: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> CLIPVisionModelOutput: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, CLIPVisionModelWithProjection + >>> from transformers.image_utils import load_image + + >>> model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = load_image(url) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> with torch.inference_mode(): + ... outputs = model(**inputs) + >>> image_embeds = outputs.image_embeds + ```""" + + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + interpolate_pos_encoding=interpolate_pos_encoding, + **kwargs, + ) + pooled_output = vision_outputs.pooler_output + image_embeds = self.visual_projection(pooled_output) + + return CLIPVisionModelOutput( + image_embeds=image_embeds, + last_hidden_state=vision_outputs.last_hidden_state, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + CLIP vision encoder with an image classification head on top (a linear layer on top of the pooled final hidden states of + the patch tokens) e.g. for ImageNet. + """ +) +class CLIPForImageClassification(CLIPPreTrainedModel): + main_input_name = "pixel_values" + input_modalities = ("image",) + + def __init__(self, config: CLIPConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + vision_model = CLIPVisionModel._from_config(config.vision_config) + self.vision_model = vision_model.vision_model + + # Classifier head + self.classifier = ( + nn.Linear(config.vision_config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ImageClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values, + **kwargs, + ) + + sequence_output = outputs.last_hidden_state + + sequence_output = torch.mean(sequence_output[:, 1:, :], dim=1) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(labels, logits, self.config) + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "CLIPModel", + "CLIPPreTrainedModel", + "CLIPTextModel", + "CLIPTextModelWithProjection", + "CLIPVisionModel", + "CLIPVisionModelWithProjection", + "CLIPForImageClassification", +] diff --git a/third_party/transformers/src/transformers/models/clip/processing_clip.py b/third_party/transformers/src/transformers/models/clip/processing_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..f69b275c48b248f1e35419ba4f174f5a75d04f21 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/processing_clip.py @@ -0,0 +1,28 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Image/Text processor class for CLIP +""" + +from ...processing_utils import ProcessorMixin +from ...utils import auto_docstring + + +@auto_docstring +class CLIPProcessor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, **kwargs): + super().__init__(image_processor, tokenizer) + + +__all__ = ["CLIPProcessor"] diff --git a/third_party/transformers/src/transformers/models/clip/tokenization_clip.py b/third_party/transformers/src/transformers/models/clip/tokenization_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..018c630afbce1cc439477a94779044a2e724a7c5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/clip/tokenization_clip.py @@ -0,0 +1,142 @@ +# Copyright 2021 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 CLIP.""" + +from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import BPE + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"} + + +class CLIPTokenizer(TokenizersBackend): + """ + Construct a CLIP tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level + Byte-Pair-Encoding. + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab (`str`, `dict` or `list`, *optional*): + Vocabulary dict to use for the tokenizer. + merges (`str` or `list`, *optional*): + Merges list to use for the BPE tokenizer. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str`, *optional*, defaults to `"<|startoftext|>"`): + The beginning of sequence token. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The token used for padding, for example when batching sequences of different lengths. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + merges: str | list[str] | None = None, + unk_token: str = "<|endoftext|>", + bos_token: str = "<|startoftext|>", + eos_token: str = "<|endoftext|>", + pad_token: str = "<|endoftext|>", + **kwargs, + ): + _vocab = ( + vocab + if vocab is not None + else { + str(bos_token): 0, + str(eos_token): 1, + str(pad_token): 2, + } + ) + + self._merges = merges or [] + + self._tokenizer = Tokenizer( + BPE( + vocab=_vocab, + merges=self._merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + unk_token=str(unk_token), + ) + ) + + self._tokenizer.normalizer = normalizers.Sequence( + [normalizers.NFC(), normalizers.Replace(Regex(r"\s+"), " "), normalizers.Lowercase()] + ) + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split( + Regex( + r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""" + ), + behavior="removed", + invert=True, + ), + pre_tokenizers.ByteLevel(add_prefix_space=False), + ] + ) + + self._tokenizer.decoder = decoders.ByteLevel() + + super().__init__( + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + **kwargs, + ) + + self._tokenizer.post_processor = processors.RobertaProcessing( + sep=(str(eos_token), self.eos_token_id), + cls=(str(bos_token), self.bos_token_id), + add_prefix_space=False, + trim_offsets=False, + ) + + # Very ugly hack to enable padding to have a correct decoding see https://github.com/huggingface/tokenizers/issues/872 + self._wrap_decode_method_backend_tokenizer() + + def _wrap_decode_method_backend_tokenizer(self): + orig_decode_method = self.backend_tokenizer.decode + + ## define this as a local variable to avoid circular reference + ## See: https://github.com/huggingface/transformers/issues/30930 + end_of_word_suffix = self.backend_tokenizer.model.end_of_word_suffix + + def new_decode_method(*args, **kwargs): + text = orig_decode_method(*args, **kwargs) + text = text.replace(end_of_word_suffix, " ").strip() + return text + + self.backend_tokenizer.decode = new_decode_method + + +__all__ = ["CLIPTokenizer"] diff --git a/third_party/transformers/src/transformers/models/deberta_v2/__init__.py b/third_party/transformers/src/transformers/models/deberta_v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7c8e16f4a9ffc0478ab02d8d24fe76c6014f8141 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deberta_v2/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_deberta_v2 import * + from .modeling_deberta_v2 import * + from .tokenization_deberta_v2 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/deberta_v2/configuration_deberta_v2.py b/third_party/transformers/src/transformers/models/deberta_v2/configuration_deberta_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..0c5938d9f58b1d1d37486a7a0455b36ab4a0d6ec --- /dev/null +++ b/third_party/transformers/src/transformers/models/deberta_v2/configuration_deberta_v2.py @@ -0,0 +1,94 @@ +# Copyright 2020, Microsoft 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. +"""DeBERTa-v2 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/deberta-v2-xlarge") +@strict +class DebertaV2Config(PreTrainedConfig): + r""" + relative_attention (`bool`, *optional*, defaults to `True`): + Whether use relative position encoding. + max_relative_positions (`int`, *optional*, defaults to -1): + The range of relative positions `[-max_position_embeddings, max_position_embeddings]`. Use the same value + as `max_position_embeddings`. + position_biased_input (`bool`, *optional*, defaults to `True`): + Whether add absolute position embedding to content embedding. + pos_att_type (`list[str]`, *optional*): + The type of relative position attention, it can be a combination of `["p2c", "c2p"]`, e.g. `["p2c"]`, + `["p2c", "c2p"]`, `["p2c", "c2p"]`. + pooler_dropout (`float`, *optional*, defaults to `0`): + Dropout rate in the pooler module. + pooler_hidden_act (`str`, *optional*, defaults to `"gelu"`): + Activation function used in the dropout module. + legacy (`bool`, *optional*, defaults to `True`): + Whether or not the model should use the legacy `LegacyDebertaOnlyMLMHead`, which does not work properly + for mask infilling tasks. + + Example: + + ```python + >>> from transformers import DebertaV2Config, DebertaV2Model + + >>> # Initializing a DeBERTa-v2 microsoft/deberta-v2-xlarge style configuration + >>> configuration = DebertaV2Config() + + >>> # Initializing a model (with random weights) from the microsoft/deberta-v2-xlarge style configuration + >>> model = DebertaV2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deberta-v2" + + vocab_size: int = 128100 + hidden_size: int = 1536 + num_hidden_layers: int = 24 + num_attention_heads: int = 24 + intermediate_size: int = 6144 + 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 = 0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-7 + relative_attention: bool = False + max_relative_positions: int = -1 + pad_token_id: int | None = 0 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + position_biased_input: bool = True + pos_att_type: str | list[str] | None = None + pooler_dropout: float | int = 0.0 + pooler_hidden_act: str = "gelu" + legacy: bool = True + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + # Backwards compatibility + if isinstance(self.pos_att_type, str): + self.pos_att_type = [x.strip() for x in self.pos_att_type.lower().split("|")] + + self.pooler_hidden_size = kwargs.get("pooler_hidden_size", self.hidden_size) + super().__post_init__(**kwargs) + + +__all__ = ["DebertaV2Config"] diff --git a/third_party/transformers/src/transformers/models/deberta_v2/modeling_deberta_v2.py b/third_party/transformers/src/transformers/models/deberta_v2/modeling_deberta_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..b30cbd6342dc692f1733a86b59ed734acf2af568 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deberta_v2/modeling_deberta_v2.py @@ -0,0 +1,1361 @@ +# Copyright 2020 Microsoft and the Hugging Face 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 DeBERTa-v2 model.""" + +from collections.abc import Sequence + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + MaskedLMOutput, + MultipleChoiceModelOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_deberta_v2 import DebertaV2Config + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm +class DebertaV2SelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +@torch.jit.script +def make_log_bucket_position(relative_pos, bucket_size: int, max_position: int): + sign = torch.sign(relative_pos) + mid = bucket_size // 2 + abs_pos = torch.where( + (relative_pos < mid) & (relative_pos > -mid), + torch.tensor(mid - 1).type_as(relative_pos), + torch.abs(relative_pos), + ) + log_pos = ( + torch.ceil(torch.log(abs_pos / mid) / torch.log(torch.tensor((max_position - 1) / mid)) * (mid - 1)) + mid + ) + bucket_pos = torch.where(abs_pos <= mid, relative_pos.type_as(log_pos), log_pos * sign) + return bucket_pos + + +def build_relative_position(query_layer, key_layer, bucket_size: int = -1, max_position: int = -1): + """ + Build relative position according to the query and key + + We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key + \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q - + P_k\\) + + Args: + query_size (int): the length of query + key_size (int): the length of key + bucket_size (int): the size of position bucket + max_position (int): the maximum allowed absolute position + device (`torch.device`): the device on which tensors will be created. + + Return: + `torch.LongTensor`: A tensor with shape [1, query_size, key_size] + """ + query_size = query_layer.size(-2) + key_size = key_layer.size(-2) + + q_ids = torch.arange(query_size, dtype=torch.long, device=query_layer.device) + k_ids = torch.arange(key_size, dtype=torch.long, device=key_layer.device) + rel_pos_ids = q_ids[:, None] - k_ids[None, :] + if bucket_size > 0 and max_position > 0: + rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position) + rel_pos_ids = rel_pos_ids.to(torch.long) + rel_pos_ids = rel_pos_ids[:query_size, :] + rel_pos_ids = rel_pos_ids.unsqueeze(0) + return rel_pos_ids + + +@torch.jit.script +def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos): + return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)]) + + +@torch.jit.script +def p2c_dynamic_expand(c2p_pos, query_layer, key_layer): + return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)]) + + +@torch.jit.script +def pos_dynamic_expand(pos_index, p2c_att, key_layer): + return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2))) + + +@torch.jit.script +def scaled_size_sqrt(query_layer: torch.Tensor, scale_factor: int): + return torch.sqrt(torch.tensor(query_layer.size(-1), dtype=torch.float) * scale_factor) + + +@torch.jit.script +def build_rpos(query_layer, key_layer, relative_pos, position_buckets: int, max_relative_positions: int): + if key_layer.size(-2) != query_layer.size(-2): + return build_relative_position( + key_layer, + key_layer, + bucket_size=position_buckets, + max_position=max_relative_positions, + ) + else: + return relative_pos + + +class DisentangledSelfAttention(nn.Module): + """ + Disentangled self-attention module + + Parameters: + config (`DebertaV2Config`): + A model config class instance with the configuration to build a new model. The schema is similar to + *BertConfig*, for more details, please refer [`DebertaV2Config`] + + """ + + def __init__(self, config): + 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_attention_heads = config.num_attention_heads + _attention_head_size = config.hidden_size // config.num_attention_heads + self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True) + self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True) + self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True) + + self.share_att_key = getattr(config, "share_att_key", False) + self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else [] + self.relative_attention = getattr(config, "relative_attention", False) + + if self.relative_attention: + self.position_buckets = getattr(config, "position_buckets", -1) + self.max_relative_positions = getattr(config, "max_relative_positions", -1) + if self.max_relative_positions < 1: + self.max_relative_positions = config.max_position_embeddings + self.pos_ebd_size = self.max_relative_positions + if self.position_buckets > 0: + self.pos_ebd_size = self.position_buckets + + self.pos_dropout = nn.Dropout(config.hidden_dropout_prob) + + if not self.share_att_key: + if "c2p" in self.pos_att_type: + self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True) + if "p2c" in self.pos_att_type: + self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def transpose_for_scores(self, x, attention_heads) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (attention_heads, -1) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1)) + + def forward( + self, + hidden_states, + attention_mask, + output_attentions=False, + query_states=None, + relative_pos=None, + rel_embeddings=None, + ): + """ + Call the module + + Args: + hidden_states (`torch.FloatTensor`): + Input states to the module usually the output from previous layer, it will be the Q,K and V in + *Attention(Q,K,V)* + + attention_mask (`torch.BoolTensor`): + An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum + sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j* + th token. + + output_attentions (`bool`, *optional*): + Whether return the attention matrix. + + query_states (`torch.FloatTensor`, *optional*): + The *Q* state in *Attention(Q,K,V)*. + + relative_pos (`torch.LongTensor`): + The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with + values ranging in [*-max_relative_positions*, *max_relative_positions*]. + + rel_embeddings (`torch.FloatTensor`): + The embedding of relative distances. It's a tensor of shape [\\(2 \\times + \\text{max_relative_positions}\\), *hidden_size*]. + + + """ + if query_states is None: + query_states = hidden_states + query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads) + key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads) + value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads) + + rel_att = None + # Take the dot product between "query" and "key" to get the raw attention scores. + scale_factor = 1 + if "c2p" in self.pos_att_type: + scale_factor += 1 + if "p2c" in self.pos_att_type: + scale_factor += 1 + scale = scaled_size_sqrt(query_layer, scale_factor) + attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2) / scale.to(dtype=query_layer.dtype)) + if self.relative_attention: + rel_embeddings = self.pos_dropout(rel_embeddings) + rel_att = self.disentangled_attention_bias( + query_layer, key_layer, relative_pos, rel_embeddings, scale_factor + ) + + if rel_att is not None: + attention_scores = attention_scores + rel_att + attention_scores = attention_scores.view( + -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1) + ) + + attention_mask = attention_mask.bool() + attention_scores = attention_scores.masked_fill(~(attention_mask), torch.finfo(query_layer.dtype).min) + # bsz x height x length x dimension + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + attention_probs = self.dropout(attention_probs) + context_layer = torch.bmm( + attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer + ) + context_layer = ( + context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1)) + .permute(0, 2, 1, 3) + .contiguous() + ) + new_context_layer_shape = context_layer.size()[:-2] + (-1,) + context_layer = context_layer.view(new_context_layer_shape) + if not output_attentions: + return (context_layer, None) + return (context_layer, attention_probs) + + def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor): + if relative_pos is None: + relative_pos = build_relative_position( + query_layer, + key_layer, + bucket_size=self.position_buckets, + max_position=self.max_relative_positions, + ) + if relative_pos.dim() == 2: + relative_pos = relative_pos.unsqueeze(0).unsqueeze(0) + elif relative_pos.dim() == 3: + relative_pos = relative_pos.unsqueeze(1) + # bsz x height x query x key + elif relative_pos.dim() != 4: + raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}") + + att_span = self.pos_ebd_size + relative_pos = relative_pos.to(device=query_layer.device, dtype=torch.long) + + rel_embeddings = rel_embeddings[0 : att_span * 2, :].unsqueeze(0) + if self.share_att_key: + pos_query_layer = self.transpose_for_scores( + self.query_proj(rel_embeddings), self.num_attention_heads + ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) + pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat( + query_layer.size(0) // self.num_attention_heads, 1, 1 + ) + else: + if "c2p" in self.pos_att_type: + pos_key_layer = self.transpose_for_scores( + self.pos_key_proj(rel_embeddings), self.num_attention_heads + ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) # .split(self.all_head_size, dim=-1) + if "p2c" in self.pos_att_type: + pos_query_layer = self.transpose_for_scores( + self.pos_query_proj(rel_embeddings), self.num_attention_heads + ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) # .split(self.all_head_size, dim=-1) + + score = 0 + # content->position + if "c2p" in self.pos_att_type: + scale = scaled_size_sqrt(pos_key_layer, scale_factor) + c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2)) + c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1) + c2p_att = torch.gather( + c2p_att, + dim=-1, + index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]), + ) + score += c2p_att / scale.to(dtype=c2p_att.dtype) + + # position->content + if "p2c" in self.pos_att_type: + scale = scaled_size_sqrt(pos_query_layer, scale_factor) + r_pos = build_rpos( + query_layer, + key_layer, + relative_pos, + self.max_relative_positions, + self.position_buckets, + ) + p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1) + p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2)) + p2c_att = torch.gather( + p2c_att, + dim=-1, + index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]), + ).transpose(-1, -2) + score += p2c_att / scale.to(dtype=p2c_att.dtype) + + return score + + +# Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2 +class DebertaV2Attention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = DisentangledSelfAttention(config) + self.output = DebertaV2SelfOutput(config) + self.config = config + + def forward( + self, + hidden_states, + attention_mask, + output_attentions: bool = False, + query_states=None, + relative_pos=None, + rel_embeddings=None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self_output, att_matrix = self.self( + hidden_states, + attention_mask, + output_attentions, + query_states=query_states, + relative_pos=relative_pos, + rel_embeddings=rel_embeddings, + ) + if query_states is None: + query_states = hidden_states + attention_output = self.output(self_output, query_states) + + if output_attentions: + return (attention_output, att_matrix) + else: + return (attention_output, None) + + +# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2 +class DebertaV2Intermediate(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.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm +class DebertaV2Output(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def forward(self, hidden_states, input_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.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2 +class DebertaV2Layer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.attention = DebertaV2Attention(config) + self.intermediate = DebertaV2Intermediate(config) + self.output = DebertaV2Output(config) + + def forward( + self, + hidden_states, + attention_mask, + query_states=None, + relative_pos=None, + rel_embeddings=None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + attention_output, att_matrix = self.attention( + hidden_states, + attention_mask, + output_attentions=output_attentions, + query_states=query_states, + relative_pos=relative_pos, + rel_embeddings=rel_embeddings, + ) + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + + if output_attentions: + return (layer_output, att_matrix) + else: + return (layer_output, None) + + +class ConvLayer(nn.Module): + def __init__(self, config): + super().__init__() + kernel_size = getattr(config, "conv_kernel_size", 3) + groups = getattr(config, "conv_groups", 1) + self.conv_act = getattr(config, "conv_act", "tanh") + self.conv = nn.Conv1d( + config.hidden_size, config.hidden_size, kernel_size, padding=(kernel_size - 1) // 2, groups=groups + ) + self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def forward(self, hidden_states, residual_states, input_mask): + out = self.conv(hidden_states.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous() + rmask = (1 - input_mask).bool() + out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0) + out = ACT2FN[self.conv_act](self.dropout(out)) + + layer_norm_input = residual_states + out + output = self.LayerNorm(layer_norm_input).to(layer_norm_input) + + if input_mask is None: + output_states = output + else: + if input_mask.dim() != layer_norm_input.dim(): + if input_mask.dim() == 4: + input_mask = input_mask.squeeze(1).squeeze(1) + input_mask = input_mask.unsqueeze(2) + + input_mask = input_mask.to(output.dtype) + output_states = output * input_mask + + return output_states + + +# Copied from transformers.models.deberta.modeling_deberta.DebertaEmbeddings with DebertaLayerNorm->LayerNorm,Deberta->DebertaV2 +class DebertaV2Embeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + pad_token_id = getattr(config, "pad_token_id", 0) + self.embedding_size = getattr(config, "embedding_size", config.hidden_size) + self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id) + + self.position_biased_input = getattr(config, "position_biased_input", True) + if not self.position_biased_input: + self.position_embeddings = None + else: + self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size) + + if config.type_vocab_size > 0: + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size) + else: + self.token_type_embeddings = None + + if self.embedding_size != config.hidden_size: + self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False) + else: + self.embed_proj = None + + self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=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] + + 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 = self.word_embeddings(input_ids) + + if self.position_embeddings is not None: + position_embeddings = self.position_embeddings(position_ids.long()) + else: + position_embeddings = torch.zeros_like(inputs_embeds) + + embeddings = inputs_embeds + if self.position_biased_input: + embeddings = embeddings + position_embeddings + if self.token_type_embeddings is not None: + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = embeddings + token_type_embeddings + + if self.embed_proj is not None: + embeddings = self.embed_proj(embeddings) + + embeddings = self.LayerNorm(embeddings) + + if mask is not None: + if mask.dim() != embeddings.dim(): + if mask.dim() == 4: + mask = mask.squeeze(1).squeeze(1) + mask = mask.unsqueeze(2) + mask = mask.to(embeddings.dtype) + + embeddings = embeddings * mask + + embeddings = self.dropout(embeddings) + return embeddings + + +class DebertaV2Encoder(nn.Module): + """Modified BertEncoder with relative position bias support""" + + def __init__(self, config): + super().__init__() + + self.layer = nn.ModuleList([DebertaV2Layer(config) for _ in range(config.num_hidden_layers)]) + self.relative_attention = getattr(config, "relative_attention", False) + + if self.relative_attention: + self.max_relative_positions = getattr(config, "max_relative_positions", -1) + if self.max_relative_positions < 1: + self.max_relative_positions = config.max_position_embeddings + + self.position_buckets = getattr(config, "position_buckets", -1) + pos_ebd_size = self.max_relative_positions * 2 + + if self.position_buckets > 0: + pos_ebd_size = self.position_buckets * 2 + + self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size) + + self.norm_rel_ebd = [x.strip() for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")] + + if "layer_norm" in self.norm_rel_ebd: + self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True) + + self.conv = ConvLayer(config) if getattr(config, "conv_kernel_size", 0) > 0 else None + self.gradient_checkpointing = False + + def get_rel_embedding(self): + rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None + if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd): + rel_embeddings = self.LayerNorm(rel_embeddings) + return rel_embeddings + + def get_attention_mask(self, attention_mask): + if attention_mask.dim() <= 2: + extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) + attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1) + elif attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + + return attention_mask + + def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None): + if self.relative_attention and relative_pos is None: + if query_states is not None: + relative_pos = build_relative_position( + query_states, + hidden_states, + bucket_size=self.position_buckets, + max_position=self.max_relative_positions, + ) + else: + relative_pos = build_relative_position( + hidden_states, + hidden_states, + bucket_size=self.position_buckets, + max_position=self.max_relative_positions, + ) + return relative_pos + + def forward( + self, + hidden_states, + attention_mask, + output_hidden_states=True, + output_attentions=False, + query_states=None, + relative_pos=None, + return_dict=True, + ): + if attention_mask.dim() <= 2: + input_mask = attention_mask + else: + input_mask = attention_mask.sum(-2) > 0 + attention_mask = self.get_attention_mask(attention_mask) + relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos) + + all_hidden_states: tuple[torch.Tensor] | None = (hidden_states,) if output_hidden_states else None + all_attentions = () if output_attentions else None + + next_kv = hidden_states + rel_embeddings = self.get_rel_embedding() + for i, layer_module in enumerate(self.layer): + output_states, attn_weights = layer_module( + next_kv, + attention_mask, + query_states=query_states, + relative_pos=relative_pos, + rel_embeddings=rel_embeddings, + output_attentions=output_attentions, + ) + + if output_attentions: + all_attentions = all_attentions + (attn_weights,) + + if i == 0 and self.conv is not None: + output_states = self.conv(hidden_states, output_states, input_mask) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (output_states,) + + if query_states is not None: + query_states = output_states + if isinstance(hidden_states, Sequence): + next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None + else: + next_kv = output_states + + if not return_dict: + return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions + ) + + +@auto_docstring +class DebertaV2PreTrainedModel(PreTrainedModel): + config: DebertaV2Config + base_model_prefix = "deberta" + _keys_to_ignore_on_load_unexpected = ["position_embeddings"] + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, (LegacyDebertaV2LMPredictionHead, DebertaV2LMPredictionHead)): + init.zeros_(module.bias) + elif isinstance(module, DebertaV2Embeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@auto_docstring +# Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2 +class DebertaV2Model(DebertaV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = DebertaV2Embeddings(config) + self.encoder = DebertaV2Encoder(config) + self.z_steps = 0 + self.config = 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, new_embeddings): + self.embeddings.word_embeddings = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = torch.ones(input_shape, device=device) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + embedding_output = self.embeddings( + input_ids=input_ids, + token_type_ids=token_type_ids, + position_ids=position_ids, + mask=attention_mask, + inputs_embeds=inputs_embeds, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask, + output_hidden_states=True, + output_attentions=output_attentions, + return_dict=return_dict, + ) + encoded_layers = encoder_outputs[1] + + if self.z_steps > 1: + hidden_states = encoded_layers[-2] + layers = [self.encoder.layer[-1] for _ in range(self.z_steps)] + query_states = encoded_layers[-1] + rel_embeddings = self.encoder.get_rel_embedding() + attention_mask = self.encoder.get_attention_mask(attention_mask) + rel_pos = self.encoder.get_rel_pos(embedding_output) + for layer in layers[1:]: + query_states = layer( + hidden_states, + attention_mask, + output_attentions=False, + query_states=query_states, + relative_pos=rel_pos, + rel_embeddings=rel_embeddings, + ) + encoded_layers.append(query_states) + + sequence_output = encoded_layers[-1] + + if not return_dict: + return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :] + + return BaseModelOutput( + last_hidden_state=sequence_output, + hidden_states=encoder_outputs.hidden_states if output_hidden_states else None, + attentions=encoder_outputs.attentions, + ) + + +# Copied from transformers.models.deberta.modeling_deberta.LegacyDebertaPredictionHeadTransform with Deberta->DebertaV2 +class LegacyDebertaV2PredictionHeadTransform(nn.Module): + def __init__(self, config): + super().__init__() + self.embedding_size = getattr(config, "embedding_size", config.hidden_size) + + self.dense = nn.Linear(config.hidden_size, self.embedding_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(self.embedding_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states): + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +class LegacyDebertaV2LMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = LegacyDebertaV2PredictionHeadTransform(config) + + self.embedding_size = getattr(config, "embedding_size", config.hidden_size) + # The output weights are the same as the input embeddings, but there is + # an output-only bias for each token. + self.decoder = nn.Linear(self.embedding_size, config.vocab_size) + + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + def forward(self, hidden_states): + hidden_states = self.transform(hidden_states) + hidden_states = self.decoder(hidden_states) + return hidden_states + + +class LegacyDebertaV2OnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = LegacyDebertaV2LMPredictionHead(config) + + def forward(self, sequence_output): + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +class DebertaV2LMPredictionHead(nn.Module): + """https://github.com/microsoft/DeBERTa/blob/master/DeBERTa/deberta/bert.py#L270""" + + 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, elementwise_affine=True) + + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + # note that the input embeddings must be passed as an argument + def forward(self, hidden_states, word_embeddings): + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + hidden_states = torch.matmul(hidden_states, word_embeddings.weight.t()) + self.bias + return hidden_states + + +class DebertaV2OnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.lm_head = DebertaV2LMPredictionHead(config) + + # note that the input embeddings must be passed as an argument + def forward(self, sequence_output, word_embeddings): + prediction_scores = self.lm_head(sequence_output, word_embeddings) + return prediction_scores + + +@auto_docstring +class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.bias": "cls.predictions.bias", + "cls.predictions.decoder.weight": "deberta.embeddings.word_embeddings.weight", + } + _keys_to_ignore_on_load_unexpected = [r"mask_predictions.*"] + + def __init__(self, config): + super().__init__(config) + self.legacy = config.legacy + self.deberta = DebertaV2Model(config) + if self.legacy: + self.cls = LegacyDebertaV2OnlyMLMHead(config) + else: + self._tied_weights_keys = { + "lm_predictions.lm_head.weight": "deberta.embeddings.word_embeddings.weight", + } + self.lm_predictions = DebertaV2OnlyMLMHead(config) + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + if self.legacy: + return self.cls.predictions.decoder + else: + return self.lm_predictions.lm_head.dense + + def set_output_embeddings(self, new_embeddings): + if self.legacy: + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + else: + self.lm_predictions.lm_head.dense = new_embeddings + self.lm_predictions.lm_head.bias = new_embeddings.bias + + @auto_docstring + # Copied from transformers.models.deberta.modeling_deberta.DebertaForMaskedLM.forward with Deberta->DebertaV2 + 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_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MaskedLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.deberta( + 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] + if self.legacy: + prediction_scores = self.cls(sequence_output) + else: + prediction_scores = self.lm_predictions(sequence_output, self.deberta.embeddings.word_embeddings) + + 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.deberta.modeling_deberta.ContextPooler +class ContextPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size) + self.dropout = nn.Dropout(config.pooler_dropout) + self.config = config + + def forward(self, hidden_states): + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + + context_token = hidden_states[:, 0] + context_token = self.dropout(context_token) + pooled_output = self.dense(context_token) + pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output) + return pooled_output + + @property + def output_dim(self): + return self.config.hidden_size + + +@auto_docstring( + custom_intro=""" + DeBERTa 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 DebertaV2ForSequenceClassification(DebertaV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + num_labels = getattr(config, "num_labels", 2) + self.num_labels = num_labels + + self.deberta = DebertaV2Model(config) + self.pooler = ContextPooler(config) + output_dim = self.pooler.output_dim + + self.classifier = nn.Linear(output_dim, num_labels) + drop_out = getattr(config, "cls_dropout", None) + drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out + self.dropout = nn.Dropout(drop_out) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.deberta.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + self.deberta.set_input_embeddings(new_embeddings) + + @auto_docstring + # Copied from transformers.models.deberta.modeling_deberta.DebertaForSequenceClassification.forward with Deberta->DebertaV2 + 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_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.deberta( + input_ids, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + encoder_layer = outputs[0] + pooled_output = self.pooler(encoder_layer) + 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: + # regression task + loss_fn = nn.MSELoss() + logits = logits.view(-1).to(labels.dtype) + loss = loss_fn(logits, labels.view(-1)) + elif labels.dim() == 1 or labels.size(-1) == 1: + label_index = (labels >= 0).nonzero() + labels = labels.long() + if label_index.size(0) > 0: + labeled_logits = torch.gather( + logits, 0, label_index.expand(label_index.size(0), logits.size(1)) + ) + labels = torch.gather(labels, 0, label_index.view(-1)) + loss_fct = CrossEntropyLoss() + loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1)) + else: + loss = torch.tensor(0).to(logits) + else: + log_softmax = nn.LogSoftmax(-1) + loss = -((log_softmax(logits) * labels).sum(-1)).mean() + elif 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 +# Copied from transformers.models.deberta.modeling_deberta.DebertaForTokenClassification with Deberta->DebertaV2 +class DebertaV2ForTokenClassification(DebertaV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.deberta = DebertaV2Model(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_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.deberta( + 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[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 DebertaV2ForQuestionAnswering(DebertaV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.deberta = DebertaV2Model(config) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + # Copied from transformers.models.deberta.modeling_deberta.DebertaForQuestionAnswering.forward with Deberta->DebertaV2 + 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_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | QuestionAnsweringModelOutput: + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.deberta( + 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[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, + ) + + +@auto_docstring +class DebertaV2ForMultipleChoice(DebertaV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + num_labels = getattr(config, "num_labels", 2) + self.num_labels = num_labels + + self.deberta = DebertaV2Model(config) + self.pooler = ContextPooler(config) + output_dim = self.pooler.output_dim + + self.classifier = nn.Linear(output_dim, 1) + drop_out = getattr(config, "cls_dropout", None) + drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out + self.dropout = nn.Dropout(drop_out) + + self.post_init() + + def get_input_embeddings(self): + return self.deberta.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + self.deberta.set_input_embeddings(new_embeddings) + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MultipleChoiceModelOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + 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.deberta( + 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, + ) + + encoder_layer = outputs[0] + pooled_output = self.pooler(encoder_layer) + 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[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, + ) + + +__all__ = [ + "DebertaV2ForMaskedLM", + "DebertaV2ForMultipleChoice", + "DebertaV2ForQuestionAnswering", + "DebertaV2ForSequenceClassification", + "DebertaV2ForTokenClassification", + "DebertaV2Model", + "DebertaV2PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/deberta_v2/tokenization_deberta_v2.py b/third_party/transformers/src/transformers/models/deberta_v2/tokenization_deberta_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..9452fa2e921dbfff914a3f80b300392f62173bba --- /dev/null +++ b/third_party/transformers/src/transformers/models/deberta_v2/tokenization_deberta_v2.py @@ -0,0 +1,169 @@ +# Copyright 2020 Microsoft 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 class for model DeBERTa-v2.""" + +from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import Unigram + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "spm.model", "tokenizer_file": "tokenizer.json"} + + +class DebertaV2Tokenizer(TokenizersBackend): + """ + Construct a DeBERTa-v2 tokenizer (backed by HuggingFace's *tokenizers* library). Based on Unigram tokenization. + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + Path to the vocabulary file (SentencePiece model file). Not used directly but kept for compatibility. + vocab (`str`, `dict` or `list`, *optional*): + List of tuples (piece, score) for the vocabulary. + precompiled_charsmap (`bytes`, *optional*): + Precompiled character map for normalization. + do_lower_case (`bool`, *optional*, defaults to `False`): + Whether or not to lowercase the input when tokenizing. + split_by_punct (`bool`, *optional*, defaults to `False`): + Whether to split by punctuation. + bos_token (`str`, *optional*, defaults to `"[CLS]"`): + The beginning of sequence token. + eos_token (`str`, *optional*, defaults to `"[SEP]"`): + The end of sequence token. + unk_token (`str`, *optional*, defaults to `"[UNK]"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `"[SEP]"`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + pad_token (`str`, *optional*, defaults to `"[PAD]"`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `"[CLS]"`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + mask_token (`str`, *optional*, defaults to `"[MASK]"`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + add_prefix_space (`bool`, *optional*, defaults to `True`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. + unk_id (`int`, *optional*, defaults to index of `unk_token` in vocab): + The ID of the unknown token in the vocabulary. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask", "token_type_ids"] + model = Unigram + + def __init__( + self, + vocab: str | dict | list | None = None, + do_lower_case=False, + split_by_punct=False, + bos_token="[CLS]", + eos_token="[SEP]", + unk_token="[UNK]", + sep_token="[SEP]", + pad_token="[PAD]", + cls_token="[CLS]", + mask_token="[MASK]", + add_prefix_space=True, + unk_id=1, + **kwargs, + ): + self.do_lower_case = do_lower_case + self.split_by_punct = split_by_punct + self.add_prefix_space = add_prefix_space + + if vocab is None: + vocab = [ + (str(pad_token), 0.0), + (str(unk_token), 0.0), + (str(bos_token), 0.0), + (str(eos_token), 0.0), + (str(sep_token), 0.0), + (str(cls_token), 0.0), + (str(mask_token), 0.0), + ] + unk_id = 1 + elif isinstance(vocab, list): + unk_id = vocab.index((str(unk_token), 0.0)) if (str(unk_token), 0.0) in vocab else unk_id + + self._vocab = vocab + self._tokenizer = Tokenizer( + Unigram( + self._vocab, + unk_id=unk_id, + byte_fallback=False, + ) + ) + + list_normalizers = [] + if do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + + list_normalizers.extend( + [ + normalizers.Replace(Regex(r"\s{2,}|[\n\r\t]"), " "), + normalizers.NFC(), + normalizers.Strip(left=False, right=True), + ] + ) + self._tokenizer.normalizer = normalizers.Sequence(list_normalizers) + + list_pretokenizers = [] + if split_by_punct: + list_pretokenizers.append(pre_tokenizers.Punctuation(behavior="isolated")) + + prepend_scheme = "always" if add_prefix_space else "first" + list_pretokenizers.append(pre_tokenizers.Metaspace(replacement="▁", prepend_scheme=prepend_scheme)) + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(list_pretokenizers) + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme=prepend_scheme) + 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, + unk_id=unk_id, + do_lower_case=do_lower_case, + split_by_punct=split_by_punct, + add_prefix_space=add_prefix_space, + **kwargs, + ) + + cls_token_id = self.cls_token_id if self.cls_token_id is not None else 0 + sep_token_id = self.sep_token_id if self.sep_token_id is not None else 0 + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=f"{str(self.cls_token)}:0 $A:0 {str(self.sep_token)}:0", + pair=f"{str(self.cls_token)}:0 $A:0 {str(self.sep_token)}:0 $B:1 {str(self.sep_token)}:1", + special_tokens=[ + (str(self.cls_token), cls_token_id), + (str(self.sep_token), sep_token_id), + ], + ) + + +__all__ = ["DebertaV2Tokenizer"] diff --git a/third_party/transformers/src/transformers/models/depth_anything/__init__.py b/third_party/transformers/src/transformers/models/depth_anything/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7425e37e0399c792155a88488045176fb3b5e7a5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/depth_anything/__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_depth_anything import * + from .modeling_depth_anything 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/depth_anything/configuration_depth_anything.py b/third_party/transformers/src/transformers/models/depth_anything/configuration_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..bdc68ec4fe02ea3d7b8ab68f7cdac9e6dab4d0a8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/depth_anything/configuration_depth_anything.py @@ -0,0 +1,99 @@ +# 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. +"""DepthAnything 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.configuration_auto import AutoConfig + + +@auto_docstring(checkpoint="LiheYoung/depth-anything-small-hf") +@strict +class DepthAnythingConfig(PreTrainedConfig): + r""" + reassemble_hidden_size (`int`, *optional*, defaults to 384): + The number of input channels of the reassemble layers. + reassemble_factors (`list[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`): + The up/downsampling factors of the reassemble layers. + neck_hidden_sizes (`list[str]`, *optional*, defaults to `[48, 96, 192, 384]`): + The hidden sizes to project to for the feature maps of the backbone. + fusion_hidden_size (`int`, *optional*, defaults to 64): + The number of channels before fusion. + head_in_index (`int`, *optional*, defaults to -1): + The index of the features to use in the depth estimation head. + head_hidden_size (`int`, *optional*, defaults to 32): + The number of output channels in the second convolution of the depth estimation head. + depth_estimation_type (`str`, *optional*, defaults to `"relative"`): + The type of depth estimation to use. Can be one of `["relative", "metric"]`. + max_depth (`float`, *optional*): + The maximum depth to use for the "metric" depth estimation head. 20 should be used for indoor models + and 80 for outdoor models. For "relative" depth estimation, this value is ignored. + + Example: + + ```python + >>> from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation + + >>> # Initializing a DepthAnything small style configuration + >>> configuration = DepthAnythingConfig() + + >>> # Initializing a model from the DepthAnything small style configuration + >>> model = DepthAnythingForDepthEstimation(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "depth_anything" + sub_configs = {"backbone_config": AutoConfig} + + backbone_config: dict | PreTrainedConfig | None = None + patch_size: int | list[int] | tuple[int, int] = 14 + initializer_range: float = 0.02 + reassemble_hidden_size: int = 384 + reassemble_factors: list[int | float] | tuple[int | float, ...] = (4, 2, 1, 0.5) + neck_hidden_sizes: list[int] | tuple[int, ...] = (48, 96, 192, 384) + fusion_hidden_size: int = 64 + head_in_index: int = -1 + head_hidden_size: int = 32 + depth_estimation_type: str = "relative" + max_depth: int | None = None + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="dinov2", + default_config_kwargs={ + "image_size": 518, + "hidden_size": 384, + "num_attention_heads": 6, + "out_indices": [9, 10, 11, 12], + "reshape_hidden_states": False, + }, + **kwargs, + ) + + self.max_depth = self.max_depth if self.max_depth else 1 + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.depth_estimation_type not in ["relative", "metric"]: + raise ValueError("depth_estimation_type must be one of ['relative', 'metric']") + + +__all__ = ["DepthAnythingConfig"] diff --git a/third_party/transformers/src/transformers/models/depth_anything/convert_depth_anything_to_hf.py b/third_party/transformers/src/transformers/models/depth_anything/convert_depth_anything_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..805a4976dce51290b5f45457d074fda9f78f205b --- /dev/null +++ b/third_party/transformers/src/transformers/models/depth_anything/convert_depth_anything_to_hf.py @@ -0,0 +1,370 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert Depth Anything checkpoints from the original repository. URL: +https://github.com/LiheYoung/Depth-Anything""" + +import argparse +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 DepthAnythingConfig, DepthAnythingForDepthEstimation, Dinov2Config, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(model_name): + if "small" in model_name: + out_indices = [3, 6, 9, 12] if "v2" in model_name else [9, 10, 11, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-small", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 64 + neck_hidden_sizes = [48, 96, 192, 384] + elif "base" in model_name: + out_indices = [3, 6, 9, 12] if "v2" in model_name else [9, 10, 11, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-base", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 128 + neck_hidden_sizes = [96, 192, 384, 768] + elif "large" in model_name: + out_indices = [5, 12, 18, 24] if "v2" in model_name else [21, 22, 23, 24] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-large", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 256 + neck_hidden_sizes = [256, 512, 1024, 1024] + else: + raise NotImplementedError(f"Model not supported: {model_name}") + + if "metric" in model_name: + depth_estimation_type = "metric" + max_depth = 20 if "indoor" in model_name else 80 + else: + depth_estimation_type = "relative" + max_depth = None + + config = DepthAnythingConfig( + reassemble_hidden_size=backbone_config.hidden_size, + patch_size=backbone_config.patch_size, + backbone_config=backbone_config, + fusion_hidden_size=fusion_hidden_size, + neck_hidden_sizes=neck_hidden_sizes, + depth_estimation_type=depth_estimation_type, + max_depth=max_depth, + ) + + return config + + +def create_rename_keys(config): + rename_keys = [] + + # fmt: off + # stem + rename_keys.append(("pretrained.cls_token", "backbone.embeddings.cls_token")) + rename_keys.append(("pretrained.mask_token", "backbone.embeddings.mask_token")) + rename_keys.append(("pretrained.pos_embed", "backbone.embeddings.position_embeddings")) + rename_keys.append(("pretrained.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight")) + rename_keys.append(("pretrained.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias")) + + # Transformer encoder + for i in range(config.backbone_config.num_hidden_layers): + rename_keys.append((f"pretrained.blocks.{i}.ls1.gamma", f"backbone.encoder.layer.{i}.layer_scale1.lambda1")) + rename_keys.append((f"pretrained.blocks.{i}.ls2.gamma", f"backbone.encoder.layer.{i}.layer_scale2.lambda1")) + rename_keys.append((f"pretrained.blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.norm1.weight")) + rename_keys.append((f"pretrained.blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.norm1.bias")) + rename_keys.append((f"pretrained.blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.norm2.weight")) + rename_keys.append((f"pretrained.blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.norm2.bias")) + rename_keys.append((f"pretrained.blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.mlp.fc1.weight")) + rename_keys.append((f"pretrained.blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.mlp.fc1.bias")) + rename_keys.append((f"pretrained.blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.mlp.fc2.weight")) + rename_keys.append((f"pretrained.blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.mlp.fc2.bias")) + rename_keys.append((f"pretrained.blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight")) + rename_keys.append((f"pretrained.blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias")) + + # Head + rename_keys.append(("pretrained.norm.weight", "backbone.layernorm.weight")) + rename_keys.append(("pretrained.norm.bias", "backbone.layernorm.bias")) + + # activation postprocessing (readout projections + resize blocks) + # Depth Anything does not use CLS token => readout_projects not required + + for i in range(4): + rename_keys.append((f"depth_head.projects.{i}.weight", f"neck.reassemble_stage.layers.{i}.projection.weight")) + rename_keys.append((f"depth_head.projects.{i}.bias", f"neck.reassemble_stage.layers.{i}.projection.bias")) + + if i != 2: + rename_keys.append((f"depth_head.resize_layers.{i}.weight", f"neck.reassemble_stage.layers.{i}.resize.weight")) + rename_keys.append((f"depth_head.resize_layers.{i}.bias", f"neck.reassemble_stage.layers.{i}.resize.bias")) + + # refinenet (tricky here) + mapping = {1:3, 2:2, 3:1, 4:0} + + for i in range(1, 5): + j = mapping[i] + rename_keys.append((f"depth_head.scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight")) + rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias")) + + # scratch convolutions + for i in range(4): + rename_keys.append((f"depth_head.scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight")) + + # head + rename_keys.append(("depth_head.scratch.output_conv1.weight", "head.conv1.weight")) + rename_keys.append(("depth_head.scratch.output_conv1.bias", "head.conv1.bias")) + rename_keys.append(("depth_head.scratch.output_conv2.0.weight", "head.conv2.weight")) + rename_keys.append(("depth_head.scratch.output_conv2.0.bias", "head.conv2.bias")) + rename_keys.append(("depth_head.scratch.output_conv2.2.weight", "head.conv3.weight")) + rename_keys.append(("depth_head.scratch.output_conv2.2.bias", "head.conv3.bias")) + + return rename_keys + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + hidden_size = config.backbone_config.hidden_size + for i in range(config.backbone_config.num_hidden_layers): + # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"pretrained.blocks.{i}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"pretrained.blocks.{i}.attn.qkv.bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[:hidden_size] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + hidden_size : hidden_size * 2, : + ] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[ + hidden_size : hidden_size * 2 + ] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-hidden_size:] + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +name_to_checkpoint = { + "depth-anything-small": "pytorch_model.bin", + "depth-anything-base": "pytorch_model.bin", + "depth-anything-large": "pytorch_model.bin", + "depth-anything-v2-small": "depth_anything_v2_vits.pth", + "depth-anything-v2-base": "depth_anything_v2_vitb.pth", + "depth-anything-v2-large": "depth_anything_v2_vitl.pth", + "depth-anything-v2-metric-indoor-small": "depth_anything_v2_metric_hypersim_vits.pth", + "depth-anything-v2-metric-indoor-base": "depth_anything_v2_metric_hypersim_vitb.pth", + "depth-anything-v2-metric-indoor-large": "depth_anything_v2_metric_hypersim_vitl.pth", + "depth-anything-v2-metric-outdoor-small": "depth_anything_v2_metric_vkitti_vits.pth", + "depth-anything-v2-metric-outdoor-base": "depth_anything_v2_metric_vkitti_vitb.pth", + "depth-anything-v2-metric-outdoor-large": "depth_anything_v2_metric_vkitti_vitl.pth", + # v2-giant pending +} + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + # define DPT configuration + config = get_dpt_config(model_name) + + model_name_to_repo = { + "depth-anything-small": "LiheYoung/depth_anything_vits14", + "depth-anything-base": "LiheYoung/depth_anything_vitb14", + "depth-anything-large": "LiheYoung/depth_anything_vitl14", + "depth-anything-v2-small": "depth-anything/Depth-Anything-V2-Small", + "depth-anything-v2-base": "depth-anything/Depth-Anything-V2-Base", + "depth-anything-v2-large": "depth-anything/Depth-Anything-V2-Large", + "depth-anything-v2-metric-indoor-small": "depth-anything/Depth-Anything-V2-Metric-Hypersim-Small", + "depth-anything-v2-metric-indoor-base": "depth-anything/Depth-Anything-V2-Metric-Hypersim-Base", + "depth-anything-v2-metric-indoor-large": "depth-anything/Depth-Anything-V2-Metric-Hypersim-Large", + "depth-anything-v2-metric-outdoor-small": "depth-anything/Depth-Anything-V2-Metric-VKITTI-Small", + "depth-anything-v2-metric-outdoor-base": "depth-anything/Depth-Anything-V2-Metric-VKITTI-Base", + "depth-anything-v2-metric-outdoor-large": "depth-anything/Depth-Anything-V2-Metric-VKITTI-Large", + } + + # load original state_dict + repo_id = model_name_to_repo[model_name] + filename = name_to_checkpoint[model_name] + filepath = hf_hub_download( + repo_id=repo_id, + filename=f"{filename}", + ) + + state_dict = torch.load(filepath, map_location="cpu", weights_only=True) + # rename keys + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + # read in qkv matrices + read_in_q_k_v(state_dict, config) + + # load HuggingFace model + model = DepthAnythingForDepthEstimation(config) + model.load_state_dict(state_dict) + model.eval() + + processor = DPTImageProcessor( + do_resize=True, + size={"height": 518, "width": 518}, + ensure_multiple_of=14, + keep_aspect_ratio=True, + do_rescale=True, + do_normalize=True, + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + + pixel_values = processor(image, return_tensors="pt").pixel_values + + # Verify forward pass + with torch.no_grad(): + outputs = model(pixel_values) + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values:", predicted_depth[0, :3, :3]) + + # assert logits + if verify_logits: + expected_shape = torch.Size([1, 518, 686]) + if model_name == "depth-anything-small": + expected_slice = torch.tensor( + [[8.8204, 8.6468, 8.6195], [8.3313, 8.6027, 8.7526], [8.6526, 8.6866, 8.7453]], + ) + elif model_name == "depth-anything-base": + expected_slice = torch.tensor( + [[26.3997, 26.3004, 26.3928], [26.2260, 26.2092, 26.3427], [26.0719, 26.0483, 26.1254]], + ) + elif model_name == "depth-anything-large": + expected_slice = torch.tensor( + [[87.9968, 87.7493, 88.2704], [87.1927, 87.6611, 87.3640], [86.7789, 86.9469, 86.7991]] + ) + elif model_name == "depth-anything-v2-small": + expected_slice = torch.tensor( + [[2.6751, 2.6211, 2.6571], [2.5820, 2.6138, 2.6271], [2.6160, 2.6141, 2.6306]] + ) + elif model_name == "depth-anything-v2-base": + expected_slice = torch.tensor( + [[4.3576, 4.3723, 4.3908], [4.3231, 4.3146, 4.3611], [4.3016, 4.3170, 4.3121]] + ) + elif model_name == "depth-anything-v2-large": + expected_slice = torch.tensor( + [[162.2751, 161.8504, 162.8788], [160.3138, 160.8050, 161.9835], [159.3812, 159.9884, 160.0768]] + ) + elif model_name == "depth-anything-v2-metric-indoor-small": + expected_slice = torch.tensor( + [[1.3349, 1.2946, 1.2801], [1.2793, 1.2337, 1.2899], [1.2629, 1.2218, 1.2476]] + ) + elif model_name == "depth-anything-v2-metric-indoor-base": + expected_slice = torch.tensor( + [[1.4601, 1.3824, 1.4904], [1.5031, 1.4349, 1.4274], [1.4570, 1.4578, 1.4200]] + ) + elif model_name == "depth-anything-v2-metric-indoor-large": + expected_slice = torch.tensor( + [[1.5040, 1.5019, 1.5218], [1.5087, 1.5195, 1.5149], [1.5437, 1.5128, 1.5252]] + ) + elif model_name == "depth-anything-v2-metric-outdoor-small": + expected_slice = torch.tensor( + [[9.5804, 8.0339, 7.7386], [7.9890, 7.2464, 7.7149], [7.7021, 7.2330, 7.3304]] + ) + elif model_name == "depth-anything-v2-metric-outdoor-base": + expected_slice = torch.tensor( + [[10.2916, 9.0933, 8.8622], [9.1964, 9.3393, 9.0644], [8.9618, 9.4201, 9.2262]] + ) + elif model_name == "depth-anything-v2-metric-outdoor-large": + expected_slice = torch.tensor( + [[14.0137, 13.3627, 13.1080], [13.2522, 13.3943, 13.3705], [13.0581, 13.4505, 13.3925]] + ) + else: + raise ValueError("Not supported") + + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"{model_name.title()}-hf") + processor.push_to_hub(repo_id=f"{model_name.title()}-hf") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="depth-anything-small", + type=str, + choices=name_to_checkpoint.keys(), + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + parser.add_argument( + "--verify_logits", + action="store_false", + required=False, + help="Whether to verify the logits after conversion.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits) diff --git a/third_party/transformers/src/transformers/models/depth_anything/convert_distill_any_depth_to_hf.py b/third_party/transformers/src/transformers/models/depth_anything/convert_distill_any_depth_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..205b6a4ffc96b31c0d46c582ee2f845a55148225 --- /dev/null +++ b/third_party/transformers/src/transformers/models/depth_anything/convert_distill_any_depth_to_hf.py @@ -0,0 +1,249 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert Distill Any Depth checkpoints from the original repository. URL: +https://github.com/Westlake-AGI-Lab/Distill-Any-Depth""" + +import argparse +import re +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from huggingface_hub import hf_hub_download +from PIL import Image +from safetensors.torch import load_file + +from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation, Dinov2Config, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + r"(backbone|pretrained)\.cls_token": r"backbone.embeddings.cls_token", + r"(backbone|pretrained)\.mask_token": r"backbone.embeddings.mask_token", + r"(backbone|pretrained)\.pos_embed": r"backbone.embeddings.position_embeddings", + r"(backbone|pretrained)\.patch_embed\.proj\.(weight|bias)": r"backbone.embeddings.patch_embeddings.projection.\2", + r"(backbone|pretrained)\.norm\.(weight|bias)": r"backbone.layernorm.\2", + r"(backbone|pretrained)(\.blocks(\.\d+)?)?\.(\d+)\.attn\.proj\.(weight|bias)": r"backbone.encoder.layer.\4.attention.output.dense.\5", + r"(backbone|pretrained)(\.blocks(\.\d+)?)?\.(\d+)\.ls(1|2)\.gamma": r"backbone.encoder.layer.\4.layer_scale\5.lambda1", + r"(backbone|pretrained)(\.blocks(\.\d+)?)?\.(\d+)\.mlp\.fc(1|2)\.(weight|bias)": r"backbone.encoder.layer.\4.mlp.fc\5.\6", + r"(backbone|pretrained)(\.blocks(\.\d+)?)?\.(\d+)\.norm(1|2)\.(weight|bias)": r"backbone.encoder.layer.\4.norm\5.\6", + r"depth_head\.projects\.(\d+)\.(weight|bias)": r"neck.reassemble_stage.layers.\1.projection.\2", + r"depth_head\.resize_layers\.(?!2)(\d+)\.(weight|bias)": r"neck.reassemble_stage.layers.\1.resize.\2", + r"depth_head\.scratch\.layer(\d+)_rn\.weight": lambda m: f"neck.convs.{int(m[1]) - 1}.weight", + r"depth_head\.scratch\.output_conv(\d+)(?:\.(\d+))?\.(weight|bias)": lambda m: ( + f"head.conv{int(m[1]) + (int(m[2]) // 2 if m[2] else 0)}.{m[3]}" if m[1] == "2" else f"head.conv{m[1]}.{m[3]}" + ), + r"depth_head\.scratch\.refinenet(\d+)\.out_conv\.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{3 - (int(m[1]) - 1)}.projection.{m[2]}", + r"depth_head\.scratch\.refinenet(\d+)\.resConfUnit(\d+)\.conv(\d+)\.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{3 - (int(m[1]) - 1)}.residual_layer{m[2]}.convolution{m[3]}.{m[4]}", +} + + +def get_dpt_config(model_name): + if "small" in model_name: + out_indices = [3, 6, 9, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-small", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 64 + neck_hidden_sizes = [48, 96, 192, 384] + elif "base" in model_name: + out_indices = [3, 6, 9, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-base", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 128 + neck_hidden_sizes = [96, 192, 384, 768] + elif "large" in model_name: + out_indices = [5, 12, 18, 24] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-large", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 256 + neck_hidden_sizes = [256, 512, 1024, 1024] + else: + raise NotImplementedError(f"Model not supported: {model_name}") + + depth_estimation_type = "relative" + max_depth = None + + config = DepthAnythingConfig( + reassemble_hidden_size=backbone_config.hidden_size, + patch_size=backbone_config.patch_size, + backbone_config=backbone_config, + fusion_hidden_size=fusion_hidden_size, + neck_hidden_sizes=neck_hidden_sizes, + depth_estimation_type=depth_estimation_type, + max_depth=max_depth, + ) + + return config + + +def convert_key_pattern(key, mapping): + for pattern, replacement in mapping.items(): + match = re.fullmatch(pattern, key) + if match: + if callable(replacement): + return replacement(match) + return re.sub(pattern, replacement, key) + return None + + +def convert_keys(state_dict, config): + new_state_dict = {} + qkv_pattern = r"(backbone|pretrained)(\.blocks(\.\d+)?)?\.(\d+)\.attn\.qkv\.(weight|bias)" + qkv_keys = [k for k in list(state_dict.keys()) if re.match(qkv_pattern, k)] + for old_key in qkv_keys: + value = state_dict.pop(old_key) + match = re.match(qkv_pattern, old_key) + _, _, _, layer, attr = match.groups() + hidden_size = config.backbone_config.hidden_size + q = value[:hidden_size] + k = value[hidden_size : hidden_size * 2] + v = value[-hidden_size:] + + for proj, tensor in zip(["query", "key", "value"], [q, k, v]): + new_key = f"backbone.encoder.layer.{layer}.attention.attention.{proj}.{attr}" + new_state_dict[new_key] = tensor + + for old_key in list(state_dict.keys()): + value = state_dict.pop(old_key) + new_key = convert_key_pattern(old_key, ORIGINAL_TO_CONVERTED_KEY_MAPPING) + + new_state_dict[new_key] = value + + return new_state_dict + + +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 + + +name_to_checkpoint = { + "distill-any-depth-small": "small/model.safetensors", + "distill-any-depth-base": "base/model.safetensors", + "distill-any-depth-large": "large/model.safetensors", +} + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits): + config = get_dpt_config(model_name) + + repo_id = "xingyang1/Distill-Any-Depth" + filepath = hf_hub_download(repo_id=repo_id, filename=name_to_checkpoint[model_name]) + state_dict = load_file(filepath) + + converted_state_dict = convert_keys(state_dict, config) + + model = DepthAnythingForDepthEstimation(config) + model.load_state_dict(converted_state_dict) + model.eval() + + processor = DPTImageProcessor( + do_resize=True, + size={"height": 518, "width": 518}, + ensure_multiple_of=14, + keep_aspect_ratio=True, + do_rescale=True, + do_normalize=True, + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + + pixel_values = processor(image, return_tensors="pt").pixel_values + + with torch.no_grad(): + outputs = model(pixel_values) + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values:", predicted_depth[0, :3, :3]) + + if verify_logits: + print("Verifying logits...") + expected_shape = torch.Size([1, 518, 686]) + + if model_name == "distill-any-depth-small": + expected_slice = torch.tensor( + [[2.5653, 2.5249, 2.5570], [2.4897, 2.5235, 2.5355], [2.5255, 2.5261, 2.5422]] + ) + elif model_name == "distill-any-depth-base": + expected_slice = torch.tensor( + [[4.8976, 4.9075, 4.9403], [4.8872, 4.8906, 4.9448], [4.8712, 4.8898, 4.8838]] + ) + elif model_name == "distill-any-depth-large": + expected_slice = torch.tensor( + [[55.1067, 51.1828, 51.6803], [51.9098, 50.7529, 51.4494], [50.1745, 50.5491, 50.8818]] + ) + else: + raise ValueError("Not supported") + + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"{model_name.title()}-hf") + processor.push_to_hub(repo_id=f"{model_name.title()}-hf") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_name", + default="distill-any-depth-small", + type=str, + choices=name_to_checkpoint.keys(), + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + parser.add_argument( + "--verify_logits", + action="store_true", + required=False, + help="Whether to verify the logits after conversion.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits) diff --git a/third_party/transformers/src/transformers/models/depth_anything/modeling_depth_anything.py b/third_party/transformers/src/transformers/models/depth_anything/modeling_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..3da9cef51e4a61f93cbecf66908616c127877534 --- /dev/null +++ b/third_party/transformers/src/transformers/models/depth_anything/modeling_depth_anything.py @@ -0,0 +1,416 @@ +# Copyright 2024 TikTok 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 Depth Anything model.""" + +import torch +from torch import nn + +from ...backbone_utils import load_backbone +from ...modeling_outputs import DepthEstimatorOutput +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_depth_anything import DepthAnythingConfig + + +logger = logging.get_logger(__name__) + +# General docstring + + +class DepthAnythingReassembleLayer(nn.Module): + def __init__(self, config, channels, factor): + super().__init__() + self.projection = nn.Conv2d(in_channels=config.reassemble_hidden_size, out_channels=channels, kernel_size=1) + + # up/down sampling depending on factor + if factor > 1: + self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0) + elif factor == 1: + self.resize = nn.Identity() + elif factor < 1: + # so should downsample + self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=int(1 / factor), padding=1) + + # Copied from transformers.models.dpt.modeling_dpt.DPTReassembleLayer.forward + def forward(self, hidden_state): + hidden_state = self.projection(hidden_state) + hidden_state = self.resize(hidden_state) + + return hidden_state + + +class DepthAnythingReassembleStage(nn.Module): + """ + This class reassembles the hidden states of the backbone into image-like feature representations at various + resolutions. + + This happens in 3 stages: + 1. Take the patch embeddings and reshape them to image-like feature representations. + 2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`. + 3. Resizing the spatial dimensions (height, width). + + Args: + config (`[DepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.config = config + self.layers = nn.ModuleList() + for channels, factor in zip(config.neck_hidden_sizes, config.reassemble_factors): + self.layers.append(DepthAnythingReassembleLayer(config, channels=channels, factor=factor)) + + def forward(self, hidden_states: list[torch.Tensor], patch_height=None, patch_width=None) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length + 1, hidden_size)`): + List of hidden states from the backbone. + """ + out = [] + + for i, hidden_state in enumerate(hidden_states): + # reshape to (batch_size, num_channels, height, width) + hidden_state = hidden_state[:, 1:] + batch_size, _, num_channels = hidden_state.shape + hidden_state = hidden_state.reshape(batch_size, patch_height, patch_width, num_channels) + hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() + hidden_state = self.layers[i](hidden_state) + out.append(hidden_state) + + return out + + +class DepthAnythingPreActResidualLayer(nn.Module): + """ + ResidualConvUnit, pre-activate residual unit. + + Args: + config (`[DepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.activation1 = nn.ReLU() + self.convolution1 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + self.activation2 = nn.ReLU() + self.convolution2 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + residual = hidden_state + hidden_state = self.activation1(hidden_state) + hidden_state = self.convolution1(hidden_state) + hidden_state = self.activation2(hidden_state) + hidden_state = self.convolution2(hidden_state) + + return hidden_state + residual + + +class DepthAnythingFeatureFusionLayer(nn.Module): + """Feature fusion layer, merges feature maps from different stages. + + Args: + config (`[DepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True) + + self.residual_layer1 = DepthAnythingPreActResidualLayer(config) + self.residual_layer2 = DepthAnythingPreActResidualLayer(config) + + def forward(self, hidden_state, residual=None, size=None): + if residual is not None: + if hidden_state.shape != residual.shape: + residual = nn.functional.interpolate( + residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False + ) + hidden_state = hidden_state + self.residual_layer1(residual) + + hidden_state = self.residual_layer2(hidden_state) + + modifier = {"scale_factor": 2} if size is None else {"size": size} + + hidden_state = nn.functional.interpolate( + hidden_state, + **modifier, + mode="bilinear", + align_corners=True, + ) + hidden_state = self.projection(hidden_state) + + return hidden_state + + +class DepthAnythingFeatureFusionStage(nn.Module): + # Copied from transformers.models.dpt.modeling_dpt.DPTFeatureFusionStage.__init__ with DPT->DepthAnything + def __init__(self, config: DepthAnythingConfig): + super().__init__() + self.layers = nn.ModuleList() + for _ in range(len(config.neck_hidden_sizes)): + self.layers.append(DepthAnythingFeatureFusionLayer(config)) + + def forward(self, hidden_states, size=None): + # reversing the hidden_states, we start from the last + hidden_states = hidden_states[::-1] + + fused_hidden_states = [] + fused_hidden_state = None + + for idx, (hidden_state, layer) in enumerate(zip(hidden_states, self.layers)): + size = hidden_states[idx + 1].shape[2:] if idx != (len(hidden_states) - 1) else None + + if fused_hidden_state is None: + # first layer only uses the last hidden_state + fused_hidden_state = layer(hidden_state, size=size) + else: + fused_hidden_state = layer(fused_hidden_state, hidden_state, size=size) + + fused_hidden_states.append(fused_hidden_state) + + return fused_hidden_states + + +# Modified from transformers.models.dpt.modeling_dpt.DPTPreTrainedModel with DPT->DepthAnything,dpt->depth_anything +# avoiding sdpa and flash_attn_2 support, it's done in the backend +@auto_docstring +class DepthAnythingPreTrainedModel(PreTrainedModel): + config: DepthAnythingConfig + base_model_prefix = "depth_anything" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + + +class DepthAnythingNeck(nn.Module): + """ + DepthAnythingNeck. A neck is a module that is normally used between the backbone and the head. It takes a list of tensors as + input and produces another list of tensors as output. For DepthAnything, it includes 2 stages: + + * DepthAnythingReassembleStage + * DepthAnythingFeatureFusionStage. + + Args: + config (dict): config dict. + """ + + def __init__(self, config): + super().__init__() + self.config = config + + self.reassemble_stage = DepthAnythingReassembleStage(config) + + self.convs = nn.ModuleList() + for channel in config.neck_hidden_sizes: + self.convs.append(nn.Conv2d(channel, config.fusion_hidden_size, kernel_size=3, padding=1, bias=False)) + + # fusion + self.fusion_stage = DepthAnythingFeatureFusionStage(config) + + def forward(self, hidden_states: list[torch.Tensor], patch_height=None, patch_width=None) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`): + List of hidden states from the backbone. + """ + if not isinstance(hidden_states, (tuple, list)): + raise TypeError("hidden_states should be a tuple or list of tensors") + + if len(hidden_states) != len(self.config.neck_hidden_sizes): + raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.") + + # postprocess hidden states + hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width) + + features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)] + + # fusion blocks + output = self.fusion_stage(features) + + return output + + +class DepthAnythingDepthEstimationHead(nn.Module): + """ + Output head consisting of 3 convolutional layers. It progressively halves the feature dimension and upsamples + the predictions to the input resolution after the first convolutional layer (details can be found in the DPT paper's + supplementary material). The final activation function is either ReLU or Sigmoid, depending on the depth estimation + type (relative or metric). For metric depth estimation, the output is scaled by the maximum depth used during pretraining. + """ + + def __init__(self, config): + super().__init__() + + self.head_in_index = config.head_in_index + self.patch_size = config.patch_size + + features = config.fusion_hidden_size + self.conv1 = nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1) + self.conv2 = nn.Conv2d(features // 2, config.head_hidden_size, kernel_size=3, stride=1, padding=1) + self.activation1 = nn.ReLU() + self.conv3 = nn.Conv2d(config.head_hidden_size, 1, kernel_size=1, stride=1, padding=0) + if config.depth_estimation_type == "relative": + self.activation2 = nn.ReLU() + elif config.depth_estimation_type == "metric": + self.activation2 = nn.Sigmoid() + else: + raise ValueError(f"Unknown depth estimation type: {config.depth_estimation_type}") + self.max_depth = config.max_depth + + def forward(self, hidden_states: list[torch.Tensor], patch_height, patch_width) -> torch.Tensor: + hidden_states = hidden_states[self.head_in_index] + + predicted_depth = self.conv1(hidden_states) + predicted_depth = nn.functional.interpolate( + predicted_depth, + (int(patch_height * self.patch_size), int(patch_width * self.patch_size)), + mode="bilinear", + align_corners=True, + ) + predicted_depth = self.conv2(predicted_depth) + predicted_depth = self.activation1(predicted_depth) + predicted_depth = self.conv3(predicted_depth) + predicted_depth = self.activation2(predicted_depth) * self.max_depth + predicted_depth = predicted_depth.squeeze(dim=1) # shape (batch_size, height, width) + + return predicted_depth + + +@auto_docstring( + custom_intro=""" + Depth Anything Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2. + """ +) +class DepthAnythingForDepthEstimation(DepthAnythingPreTrainedModel): + _no_split_modules = ["DPTViTEmbeddings"] + + def __init__(self, config): + super().__init__(config) + + self.backbone = load_backbone(config) + self.neck = DepthAnythingNeck(config) + self.head = DepthAnythingDepthEstimationHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor] | DepthEstimatorOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth depth estimation maps for computing the loss. + + Examples: + ```python + >>> from transformers import AutoImageProcessor, AutoModelForDepthEstimation + >>> import torch + >>> import numpy as np + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-small-hf") + >>> model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-small-hf") + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # interpolate to original size + >>> post_processed_output = image_processor.post_process_depth_estimation( + ... outputs, + ... target_sizes=[(image.height, image.width)], + ... ) + + >>> # visualize the prediction + >>> predicted_depth = post_processed_output[0]["predicted_depth"] + >>> depth = predicted_depth * 255 / predicted_depth.max() + >>> depth = depth.detach().cpu().numpy() + >>> depth = Image.fromarray(depth.astype("uint8")) + ```""" + loss = None + if labels is not None: + raise NotImplementedError("Training is not implemented yet") + + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + + outputs = self.backbone.forward_with_filtered_kwargs( + pixel_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions + ) + hidden_states = outputs.feature_maps + + _, _, height, width = pixel_values.shape + patch_size = self.config.patch_size + patch_height = height // patch_size + patch_width = width // patch_size + + hidden_states = self.neck(hidden_states, patch_height, patch_width) + + predicted_depth = self.head(hidden_states, patch_height, patch_width) + + if not return_dict: + if output_hidden_states: + output = (predicted_depth,) + outputs[1:] + else: + output = (predicted_depth,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return DepthEstimatorOutput( + loss=loss, + predicted_depth=predicted_depth, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=outputs.attentions, + ) + + +__all__ = ["DepthAnythingForDepthEstimation", "DepthAnythingPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/dpr/__init__.py b/third_party/transformers/src/transformers/models/dpr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..386b47bb2ecd07ed85a243e7c1920c8d2bfb4441 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/__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_dpr import * + from .modeling_dpr import * + from .tokenization_dpr import * + from .tokenization_dpr_fast import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/dpr/configuration_dpr.py b/third_party/transformers/src/transformers/models/dpr/configuration_dpr.py new file mode 100644 index 0000000000000000000000000000000000000000..4a5429581d8d28ca12f918a753853c198c1568f5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/configuration_dpr.py @@ -0,0 +1,63 @@ +# Copyright 2010, DPR authors, The Hugging Face 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. +"""DPR model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="facebook/dpr-ctx_encoder-single-nq-base") +@strict +class DPRConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import DPRConfig, DPRContextEncoder + + >>> # Initializing a DPR facebook/dpr-ctx_encoder-single-nq-base style configuration + >>> configuration = DPRConfig() + + >>> # Initializing a model (with random weights) from the facebook/dpr-ctx_encoder-single-nq-base style configuration + >>> model = DPRContextEncoder(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "dpr" + + 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 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + projection_dim: int = 0 + is_decoder: bool = False + add_cross_attention: bool = False + + +__all__ = ["DPRConfig"] diff --git a/third_party/transformers/src/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..5151c0972a7ed72c47d125400b918aba3a0d3c0d --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py @@ -0,0 +1,145 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import collections +from pathlib import Path + +import torch +from torch.serialization import default_restore_location + +from transformers import BertConfig, DPRConfig, DPRContextEncoder, DPRQuestionEncoder, DPRReader + + +CheckpointState = collections.namedtuple( + "CheckpointState", ["model_dict", "optimizer_dict", "scheduler_dict", "offset", "epoch", "encoder_params"] +) + + +def load_states_from_checkpoint(model_file: str) -> CheckpointState: + print(f"Reading saved model from {model_file}") + state_dict = torch.load( + model_file, map_location=lambda s, l: default_restore_location(s, "cpu"), weights_only=True + ) + return CheckpointState(**state_dict) + + +class DPRState: + def __init__(self, src_file: Path): + self.src_file = src_file + + def load_dpr_model(self): + raise NotImplementedError + + @staticmethod + def from_type(comp_type: str, *args, **kwargs) -> "DPRState": + if comp_type.startswith("c"): + return DPRContextEncoderState(*args, **kwargs) + if comp_type.startswith("q"): + return DPRQuestionEncoderState(*args, **kwargs) + if comp_type.startswith("r"): + return DPRReaderState(*args, **kwargs) + else: + raise ValueError("Component type must be either 'ctx_encoder', 'question_encoder' or 'reader'.") + + +class DPRContextEncoderState(DPRState): + def load_dpr_model(self): + model = DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0])) + print(f"Loading DPR biencoder from {self.src_file}") + saved_state = load_states_from_checkpoint(self.src_file) + encoder, prefix = model.ctx_encoder, "ctx_model." + # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3 + state_dict = {"bert_model.embeddings.position_ids": model.ctx_encoder.bert_model.embeddings.position_ids} + for key, value in saved_state.model_dict.items(): + if key.startswith(prefix): + key = key[len(prefix) :] + if not key.startswith("encode_proj."): + key = "bert_model." + key + state_dict[key] = value + encoder.load_state_dict(state_dict) + return model + + +class DPRQuestionEncoderState(DPRState): + def load_dpr_model(self): + model = DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0])) + print(f"Loading DPR biencoder from {self.src_file}") + saved_state = load_states_from_checkpoint(self.src_file) + encoder, prefix = model.question_encoder, "question_model." + # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3 + state_dict = {"bert_model.embeddings.position_ids": model.question_encoder.bert_model.embeddings.position_ids} + for key, value in saved_state.model_dict.items(): + if key.startswith(prefix): + key = key[len(prefix) :] + if not key.startswith("encode_proj."): + key = "bert_model." + key + state_dict[key] = value + encoder.load_state_dict(state_dict) + return model + + +class DPRReaderState(DPRState): + def load_dpr_model(self): + model = DPRReader(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0])) + print(f"Loading DPR reader from {self.src_file}") + saved_state = load_states_from_checkpoint(self.src_file) + # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3 + state_dict = { + "encoder.bert_model.embeddings.position_ids": model.span_predictor.encoder.bert_model.embeddings.position_ids + } + for key, value in saved_state.model_dict.items(): + if key.startswith("encoder.") and not key.startswith("encoder.encode_proj"): + key = "encoder.bert_model." + key[len("encoder.") :] + state_dict[key] = value + model.span_predictor.load_state_dict(state_dict) + return model + + +def convert(comp_type: str, src_file: Path, dest_dir: Path): + dest_dir = Path(dest_dir) + dest_dir.mkdir(exist_ok=True) + + dpr_state = DPRState.from_type(comp_type, src_file=src_file) + model = dpr_state.load_dpr_model() + model.save_pretrained(dest_dir) + model.from_pretrained(dest_dir) # sanity check + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--type", type=str, help="Type of the component to convert: 'ctx_encoder', 'question_encoder' or 'reader'." + ) + parser.add_argument( + "--src", + type=str, + help=( + "Path to the dpr checkpoint file. They can be downloaded from the official DPR repo" + " https://github.com/facebookresearch/DPR. Note that in the official repo, both encoders are stored in the" + " 'retriever' checkpoints." + ), + ) + parser.add_argument("--dest", type=str, default=None, help="Path to the output PyTorch model directory.") + args = parser.parse_args() + + src_file = Path(args.src) + dest_dir = f"converted-{src_file.name}" if args.dest is None else args.dest + dest_dir = Path(dest_dir) + assert src_file.exists() + assert args.type is not None, ( + "Please specify the component type of the DPR model to convert: 'ctx_encoder', 'question_encoder' or 'reader'." + ) + convert(args.type, src_file, dest_dir) diff --git a/third_party/transformers/src/transformers/models/dpr/modeling_dpr.py b/third_party/transformers/src/transformers/models/dpr/modeling_dpr.py new file mode 100644 index 0000000000000000000000000000000000000000..3ed9a1759db6d144a3bb14b9045b797729563784 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/modeling_dpr.py @@ -0,0 +1,576 @@ +# Copyright 2018 DPR Authors, The Hugging Face 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 DPR model for Open Domain Question Answering.""" + +from dataclasses import dataclass + +import torch +from torch import Tensor, nn + +from ...modeling_outputs import BaseModelOutputWithPooling +from ...modeling_utils import PreTrainedModel +from ...utils import ( + ModelOutput, + auto_docstring, + logging, +) +from ..bert.modeling_bert import BertModel +from .configuration_dpr import DPRConfig + + +logger = logging.get_logger(__name__) + + +########## +# Outputs +########## + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`DPRQuestionEncoder`]. + """ +) +class DPRContextEncoderOutput(ModelOutput): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, embeddings_size)`): + The DPR encoder outputs the *pooler_output* that corresponds to the context representation. Last layer + hidden-state of the first token of the sequence (classification token) further processed by a Linear layer. + This output is to be used to embed contexts for nearest neighbors queries with questions embeddings. + """ + + pooler_output: torch.FloatTensor + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`DPRQuestionEncoder`]. + """ +) +class DPRQuestionEncoderOutput(ModelOutput): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, embeddings_size)`): + The DPR encoder outputs the *pooler_output* that corresponds to the question representation. Last layer + hidden-state of the first token of the sequence (classification token) further processed by a Linear layer. + This output is to be used to embed questions for nearest neighbors queries with context embeddings. + """ + + pooler_output: torch.FloatTensor + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`DPRQuestionEncoder`]. + """ +) +class DPRReaderOutput(ModelOutput): + r""" + start_logits (`torch.FloatTensor` of shape `(n_passages, sequence_length)`): + Logits of the start index of the span for each passage. + end_logits (`torch.FloatTensor` of shape `(n_passages, sequence_length)`): + Logits of the end index of the span for each passage. + relevance_logits (`torch.FloatTensor` of shape `(n_passages, )`): + Outputs of the QA classifier of the DPRReader that corresponds to the scores of each passage to answer the + question, compared to all the other passages. + """ + + start_logits: torch.FloatTensor + end_logits: torch.FloatTensor | None = None + relevance_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@auto_docstring +class DPRPreTrainedModel(PreTrainedModel): + _supports_sdpa = True + + +class DPREncoder(DPRPreTrainedModel): + base_model_prefix = "bert_model" + + def __init__(self, config: DPRConfig): + super().__init__(config) + self.bert_model = BertModel(config, add_pooling_layer=False) + if self.bert_model.config.hidden_size <= 0: + raise ValueError("Encoder hidden_size can't be zero") + self.projection_dim = config.projection_dim + if self.projection_dim > 0: + self.encode_proj = nn.Linear(self.bert_model.config.hidden_size, config.projection_dim) + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + input_ids: Tensor, + attention_mask: Tensor | None = None, + token_type_ids: Tensor | None = None, + inputs_embeds: Tensor | None = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ) -> BaseModelOutputWithPooling | tuple[Tensor, ...]: + outputs = self.bert_model( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + pooled_output = sequence_output[:, 0, :] + + if self.projection_dim > 0: + pooled_output = self.encode_proj(pooled_output) + + if not return_dict: + return (sequence_output, pooled_output) + outputs[2:] + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + @property + def embeddings_size(self) -> int: + if self.projection_dim > 0: + return self.encode_proj.out_features + return self.bert_model.config.hidden_size + + +class DPRSpanPredictor(DPRPreTrainedModel): + base_model_prefix = "encoder" + + def __init__(self, config: DPRConfig): + super().__init__(config) + self.encoder = DPREncoder(config) + self.qa_outputs = nn.Linear(self.encoder.embeddings_size, 2) + self.qa_classifier = nn.Linear(self.encoder.embeddings_size, 1) + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + input_ids: Tensor, + attention_mask: Tensor, + inputs_embeds: Tensor | None = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ) -> DPRReaderOutput | tuple[Tensor, ...]: + # notations: N - number of questions in a batch, M - number of passages per questions, L - sequence length + n_passages, sequence_length = input_ids.size() if input_ids is not None else inputs_embeds.size()[:2] + # feed encoder + outputs = self.encoder( + input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + + # compute logits + 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() + relevance_logits = self.qa_classifier(sequence_output[:, 0, :]) + + # resize + start_logits = start_logits.view(n_passages, sequence_length) + end_logits = end_logits.view(n_passages, sequence_length) + relevance_logits = relevance_logits.view(n_passages) + + if not return_dict: + return (start_logits, end_logits, relevance_logits) + outputs[2:] + + return DPRReaderOutput( + start_logits=start_logits, + end_logits=end_logits, + relevance_logits=relevance_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +################## +# PreTrainedModel +################## + + +class DPRPretrainedContextEncoder(DPRPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: DPRConfig + base_model_prefix = "ctx_encoder" + + +class DPRPretrainedQuestionEncoder(DPRPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: DPRConfig + base_model_prefix = "question_encoder" + + +class DPRPretrainedReader(DPRPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: DPRConfig + base_model_prefix = "span_predictor" + + +############### +# Actual Models +############### + + +@auto_docstring( + custom_intro=""" + The bare DPRContextEncoder transformer outputting pooler outputs as context representations. + """ +) +class DPRContextEncoder(DPRPretrainedContextEncoder): + def __init__(self, config: DPRConfig): + super().__init__(config) + self.config = config + self.ctx_encoder = DPREncoder(config) + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + token_type_ids: Tensor | None = None, + inputs_embeds: Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> DPRContextEncoderOutput | tuple[Tensor, ...]: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. To match pretraining, DPR input sequence should be + formatted with [CLS] and [SEP] tokens as follows: + + (a) For sequence pairs (for a pair title+text for example): + + ``` + tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP] + token_type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1 + ``` + + (b) For single sequences (for a question for example): + + ``` + tokens: [CLS] the dog is hairy . [SEP] + token_type_ids: 0 0 0 0 0 0 0 + ``` + + DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right + rather than the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + + Examples: + + ```python + >>> from transformers import DPRContextEncoder, DPRContextEncoderTokenizer + + >>> tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base") + >>> model = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base") + >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="pt")["input_ids"] + >>> embeddings = model(input_ids).pooler_output + ```""" + + 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: + 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 input_ids is None + else (input_ids != self.config.pad_token_id) + ) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + outputs = self.ctx_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if not return_dict: + return outputs[1:] + return DPRContextEncoderOutput( + pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions + ) + + +@auto_docstring( + custom_intro=""" + The bare DPRQuestionEncoder transformer outputting pooler outputs as question representations. + """ +) +class DPRQuestionEncoder(DPRPretrainedQuestionEncoder): + def __init__(self, config: DPRConfig): + super().__init__(config) + self.config = config + self.question_encoder = DPREncoder(config) + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + token_type_ids: Tensor | None = None, + inputs_embeds: Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> DPRQuestionEncoderOutput | tuple[Tensor, ...]: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. To match pretraining, DPR input sequence should be + formatted with [CLS] and [SEP] tokens as follows: + + (a) For sequence pairs (for a pair title+text for example): + + ``` + tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP] + token_type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1 + ``` + + (b) For single sequences (for a question for example): + + ``` + tokens: [CLS] the dog is hairy . [SEP] + token_type_ids: 0 0 0 0 0 0 0 + ``` + + DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right + rather than the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + + Examples: + + ```python + >>> from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer + + >>> tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base") + >>> model = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base") + >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="pt")["input_ids"] + >>> embeddings = model(input_ids).pooler_output + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = ( + torch.ones(input_shape, device=device) + if input_ids is None + else (input_ids != self.config.pad_token_id) + ) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + outputs = self.question_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if not return_dict: + return outputs[1:] + return DPRQuestionEncoderOutput( + pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions + ) + + +@auto_docstring( + custom_intro=""" + The bare DPRReader transformer outputting span predictions. + """ +) +class DPRReader(DPRPretrainedReader): + def __init__(self, config: DPRConfig): + super().__init__(config) + self.config = config + self.span_predictor = DPRSpanPredictor(config) + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + inputs_embeds: Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> DPRReaderOutput | tuple[Tensor, ...]: + r""" + input_ids (`tuple[torch.LongTensor]` of shapes `(n_passages, sequence_length)`): + Indices of input sequence tokens in the vocabulary. It has to be a sequence triplet with 1) the question + and 2) the passages titles and 3) the passages texts To match pretraining, DPR `input_ids` sequence should + be formatted with [CLS] and [SEP] with the format: + + `[CLS] [SEP] [SEP] ` + + DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right + rather than the left. + + Indices can be obtained using [`DPRReaderTokenizer`]. See this class documentation for more details. + + [What are input IDs?](../glossary#input-ids) + inputs_embeds (`torch.FloatTensor` of shape `(n_passages, 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. + + Examples: + + ```python + >>> from transformers import DPRReader, DPRReaderTokenizer + + >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> encoded_inputs = tokenizer( + ... questions=["What is love ?"], + ... titles=["Haddaway"], + ... texts=["'What Is Love' is a song recorded by the artist Haddaway"], + ... return_tensors="pt", + ... ) + >>> outputs = model(**encoded_inputs) + >>> start_logits = outputs.start_logits + >>> end_logits = outputs.end_logits + >>> relevance_logits = outputs.relevance_logits + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = torch.ones(input_shape, device=device) + + return self.span_predictor( + input_ids, + attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +__all__ = [ + "DPRContextEncoder", + "DPRPretrainedContextEncoder", + "DPRPreTrainedModel", + "DPRPretrainedQuestionEncoder", + "DPRPretrainedReader", + "DPRQuestionEncoder", + "DPRReader", +] diff --git a/third_party/transformers/src/transformers/models/dpr/tokenization_dpr.py b/third_party/transformers/src/transformers/models/dpr/tokenization_dpr.py new file mode 100644 index 0000000000000000000000000000000000000000..ff33ec5e0f0742341128924c01770410e4d39fb3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/tokenization_dpr.py @@ -0,0 +1,330 @@ +# Copyright 2018 The HuggingFace Inc. team, The Hugging Face 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 DPR.""" + +import collections + +from ...tokenization_utils_base import BatchEncoding +from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging +from ..bert.tokenization_bert import BertTokenizer + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"} + + +class DPRContextEncoderTokenizer(BertTokenizer): + r""" + Construct a DPRContextEncoder tokenizer. + + [`DPRContextEncoderTokenizer`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation + splitting and wordpiece. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + def __init__(self, *args, do_lower_case=False, **kwargs): + super().__init__(*args, **kwargs) + self.do_lower_case = do_lower_case + + +class DPRQuestionEncoderTokenizer(BertTokenizer): + r""" + Constructs a DPRQuestionEncoder tokenizer. + + [`DPRQuestionEncoderTokenizer`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation + splitting and wordpiece. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + def __init__(self, *args, do_lower_case=False, **kwargs): + super().__init__(*args, **kwargs) + self.do_lower_case = do_lower_case + + +DPRSpanPrediction = collections.namedtuple( + "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"] +) + +DPRReaderOutput = collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"]) + + +CUSTOM_DPR_READER_DOCSTRING = r""" + Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`. + It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers), + using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)` + with the format: + + ``` + [CLS] [SEP] [SEP] + ``` + + Args: + questions (`str` or `list[str]`): + The questions to be encoded. You can specify one question for many passages. In this case, the question + will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in + `titles` or `texts`. + titles (`str` or `list[str]`): + The passages titles to be encoded. This can be a string or a list of strings if there are several passages. + texts (`str` or `list[str]`): + The passages texts to be encoded. This can be a string or a list of strings if there are several passages. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence + if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to + the maximum acceptable input length for the model if that argument is not provided. This will truncate + token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch + of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. This will only truncate the first + sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. This will only truncate the + second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + return_attention_mask (`bool`, *optional*): + Whether or not to return the attention mask. If not set, will return the attention mask according to the + specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + + Returns: + `dict[str, list[list[int]]]`: A dictionary with the following keys: + + - `input_ids`: List of token ids to be fed to a model. + - `attention_mask`: List of indices specifying which tokens should be attended to by the model. + """ + + +@add_start_docstrings(CUSTOM_DPR_READER_DOCSTRING) +class CustomDPRReaderTokenizerMixin: + def __call__( + self, + questions, + titles: str | None = None, + texts: str | None = None, + padding: bool | str = False, + truncation: bool | str = False, + max_length: int | None = None, + return_tensors: str | TensorType | None = None, + return_attention_mask: bool | None = None, + **kwargs, + ) -> BatchEncoding: + if titles is None and texts is None: + return super().__call__( + questions, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_attention_mask=return_attention_mask, + **kwargs, + ) + elif titles is None or texts is None: + text_pair = titles if texts is None else texts + return super().__call__( + questions, + text_pair, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_attention_mask=return_attention_mask, + **kwargs, + ) + titles = titles if not isinstance(titles, str) else [titles] + texts = texts if not isinstance(texts, str) else [texts] + n_passages = len(titles) + questions = questions if not isinstance(questions, str) else [questions] * n_passages + if len(titles) != len(texts): + raise ValueError( + f"There should be as many titles than texts but got {len(titles)} titles and {len(texts)} texts." + ) + encoded_question_and_titles = super().__call__(questions, titles, padding=False, truncation=False)["input_ids"] + encoded_texts = super().__call__(texts, add_special_tokens=False, padding=False, truncation=False)["input_ids"] + encoded_inputs = { + "input_ids": [ + (encoded_question_and_title + encoded_text)[:max_length] + if max_length is not None and truncation + else encoded_question_and_title + encoded_text + for encoded_question_and_title, encoded_text in zip(encoded_question_and_titles, encoded_texts) + ] + } + if return_attention_mask is not False: + attention_mask = [] + for input_ids in encoded_inputs["input_ids"]: + attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids]) + encoded_inputs["attention_mask"] = attention_mask + return self.pad(encoded_inputs, padding=padding, max_length=max_length, return_tensors=return_tensors) + + def decode_best_spans( + self, + reader_input: BatchEncoding, + reader_output: DPRReaderOutput, + num_spans: int = 16, + max_answer_length: int = 64, + num_spans_per_passage: int = 4, + ) -> list[DPRSpanPrediction]: + """ + Get the span predictions for the extractive Q&A model. + + Returns: *List* of *DPRReaderOutput* sorted by descending *(relevance_score, span_score)*. Each + *DPRReaderOutput* is a *Tuple* with: + + - **span_score**: `float` that corresponds to the score given by the reader for this span compared to other + spans in the same passage. It corresponds to the sum of the start and end logits of the span. + - **relevance_score**: `float` that corresponds to the score of the each passage to answer the question, + compared to all the other passages. It corresponds to the output of the QA classifier of the DPRReader. + - **doc_id**: `int` the id of the passage. - **start_index**: `int` the start index of the span + (inclusive). - **end_index**: `int` the end index of the span (inclusive). + + Examples: + + ```python + >>> from transformers import DPRReader, DPRReaderTokenizer + + >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> encoded_inputs = tokenizer( + ... questions=["What is love ?"], + ... titles=["Haddaway"], + ... texts=["'What Is Love' is a song recorded by the artist Haddaway"], + ... return_tensors="pt", + ... ) + >>> outputs = model(**encoded_inputs) + >>> predicted_spans = tokenizer.decode_best_spans(encoded_inputs, outputs) + >>> print(predicted_spans[0].text) # best span + a song + ```""" + input_ids = reader_input["input_ids"] + start_logits, end_logits, relevance_logits = reader_output[:3] + n_passages = len(relevance_logits) + sorted_docs = sorted(range(n_passages), reverse=True, key=relevance_logits.__getitem__) + nbest_spans_predictions: list[DPRReaderOutput] = [] + for doc_id in sorted_docs: + sequence_ids = list(input_ids[doc_id]) + # assuming question & title information is at the beginning of the sequence + passage_offset = sequence_ids.index(self.sep_token_id, 2) + 1 # second sep id + if sequence_ids[-1] == self.pad_token_id: + sequence_len = sequence_ids.index(self.pad_token_id) + else: + sequence_len = len(sequence_ids) + + best_spans = self._get_best_spans( + start_logits=start_logits[doc_id][passage_offset:sequence_len], + end_logits=end_logits[doc_id][passage_offset:sequence_len], + max_answer_length=max_answer_length, + top_spans=num_spans_per_passage, + ) + for start_index, end_index in best_spans: + start_index += passage_offset + end_index += passage_offset + nbest_spans_predictions.append( + DPRSpanPrediction( + span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index], + relevance_score=relevance_logits[doc_id], + doc_id=doc_id, + start_index=start_index, + end_index=end_index, + text=self.decode(sequence_ids[start_index : end_index + 1]), + ) + ) + if len(nbest_spans_predictions) >= num_spans: + break + return nbest_spans_predictions[:num_spans] + + def _get_best_spans( + self, + start_logits: list[int], + end_logits: list[int], + max_answer_length: int, + top_spans: int, + ) -> list[DPRSpanPrediction]: + """ + Finds the best answer span for the extractive Q&A model for one passage. It returns the best span by descending + `span_score` order and keeping max `top_spans` spans. Spans longer that `max_answer_length` are ignored. + """ + scores = [] + for start_index, start_score in enumerate(start_logits): + for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]): + scores.append(((start_index, start_index + answer_length), start_score + end_score)) + scores = sorted(scores, key=lambda x: x[1], reverse=True) + chosen_span_intervals = [] + for (start_index, end_index), score in scores: + if start_index > end_index: + raise ValueError(f"Wrong span indices: [{start_index}:{end_index}]") + length = end_index - start_index + 1 + if length > max_answer_length: + raise ValueError(f"Span is too long: {length} > {max_answer_length}") + if any( + start_index <= prev_start_index <= prev_end_index <= end_index + or prev_start_index <= start_index <= end_index <= prev_end_index + for (prev_start_index, prev_end_index) in chosen_span_intervals + ): + continue + chosen_span_intervals.append((start_index, end_index)) + + if len(chosen_span_intervals) == top_spans: + break + return chosen_span_intervals + + +@add_end_docstrings(CUSTOM_DPR_READER_DOCSTRING) +class DPRReaderTokenizer(CustomDPRReaderTokenizerMixin, BertTokenizer): + r""" + Construct a DPRReader tokenizer. + + [`DPRReaderTokenizer`] is almost identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation + splitting and wordpiece. The difference is that is has three inputs strings: question, titles and texts that are + combined to be fed to the [`DPRReader`] model. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__(self, *args, do_lower_case=False, **kwargs): + super().__init__(*args, **kwargs) + self.do_lower_case = do_lower_case + + +__all__ = ["DPRContextEncoderTokenizer", "DPRQuestionEncoderTokenizer", "DPRReaderOutput", "DPRReaderTokenizer"] diff --git a/third_party/transformers/src/transformers/models/dpr/tokenization_dpr_fast.py b/third_party/transformers/src/transformers/models/dpr/tokenization_dpr_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..96723f63cdfe3904c86dd1de5b3786d0112ae160 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpr/tokenization_dpr_fast.py @@ -0,0 +1,318 @@ +# Copyright 2018 The HuggingFace Inc. team, The Hugging Face 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 DPR.""" + +import collections + +from ...tokenization_utils_base import BatchEncoding +from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging +from ..bert.tokenization_bert import BertTokenizer +from .tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer, DPRReaderTokenizer + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"} + + +class DPRContextEncoderTokenizerFast(BertTokenizer): + r""" + Construct a "fast" DPRContextEncoder tokenizer (backed by HuggingFace's *tokenizers* library). + + [`DPRContextEncoderTokenizerFast`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: + punctuation splitting and wordpiece. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + """ + + vocab_files_names = VOCAB_FILES_NAMES + slow_tokenizer_class = DPRContextEncoderTokenizer + + +class DPRQuestionEncoderTokenizerFast(BertTokenizer): + r""" + Constructs a "fast" DPRQuestionEncoder tokenizer (backed by HuggingFace's *tokenizers* library). + + [`DPRQuestionEncoderTokenizerFast`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: + punctuation splitting and wordpiece. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + """ + + vocab_files_names = VOCAB_FILES_NAMES + slow_tokenizer_class = DPRQuestionEncoderTokenizer + + +DPRSpanPrediction = collections.namedtuple( + "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"] +) + +DPRReaderOutput = collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"]) + + +CUSTOM_DPR_READER_DOCSTRING = r""" + Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`. + It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers), + using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)` + with the format: + + [CLS] [SEP] [SEP] + + Args: + questions (`str` or `list[str]`): + The questions to be encoded. You can specify one question for many passages. In this case, the question + will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in + `titles` or `texts`. + titles (`str` or `list[str]`): + The passages titles to be encoded. This can be a string or a list of strings if there are several passages. + texts (`str` or `list[str]`): + The passages texts to be encoded. This can be a string or a list of strings if there are several passages. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence + if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to + the maximum acceptable input length for the model if that argument is not provided. This will truncate + token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch + of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. This will only truncate the first + sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. This will only truncate the + second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + return_attention_mask (`bool`, *optional*): + Whether or not to return the attention mask. If not set, will return the attention mask according to the + specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + + Return: + `dict[str, list[list[int]]]`: A dictionary with the following keys: + + - `input_ids`: List of token ids to be fed to a model. + - `attention_mask`: List of indices specifying which tokens should be attended to by the model. + """ + + +@add_start_docstrings(CUSTOM_DPR_READER_DOCSTRING) +class CustomDPRReaderTokenizerMixin: + def __call__( + self, + questions, + titles: str | None = None, + texts: str | None = None, + padding: bool | str = False, + truncation: bool | str = False, + max_length: int | None = None, + return_tensors: str | TensorType | None = None, + return_attention_mask: bool | None = None, + **kwargs, + ) -> BatchEncoding: + if titles is None and texts is None: + return super().__call__( + questions, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_attention_mask=return_attention_mask, + **kwargs, + ) + elif titles is None or texts is None: + text_pair = titles if texts is None else texts + return super().__call__( + questions, + text_pair, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_attention_mask=return_attention_mask, + **kwargs, + ) + titles = titles if not isinstance(titles, str) else [titles] + texts = texts if not isinstance(texts, str) else [texts] + n_passages = len(titles) + questions = questions if not isinstance(questions, str) else [questions] * n_passages + assert len(titles) == len(texts), ( + f"There should be as many titles than texts but got {len(titles)} titles and {len(texts)} texts." + ) + encoded_question_and_titles = super().__call__(questions, titles, padding=False, truncation=False)["input_ids"] + encoded_texts = super().__call__(texts, add_special_tokens=False, padding=False, truncation=False)["input_ids"] + encoded_inputs = { + "input_ids": [ + (encoded_question_and_title + encoded_text)[:max_length] + if max_length is not None and truncation + else encoded_question_and_title + encoded_text + for encoded_question_and_title, encoded_text in zip(encoded_question_and_titles, encoded_texts) + ] + } + if return_attention_mask is not False: + attention_mask = [] + for input_ids in encoded_inputs["input_ids"]: + attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids]) + encoded_inputs["attention_mask"] = attention_mask + return self.pad(encoded_inputs, padding=padding, max_length=max_length, return_tensors=return_tensors) + + def decode_best_spans( + self, + reader_input: BatchEncoding, + reader_output: DPRReaderOutput, + num_spans: int = 16, + max_answer_length: int = 64, + num_spans_per_passage: int = 4, + ) -> list[DPRSpanPrediction]: + """ + Get the span predictions for the extractive Q&A model. + + Returns: *List* of *DPRReaderOutput* sorted by descending *(relevance_score, span_score)*. Each + *DPRReaderOutput* is a *Tuple* with: + + - **span_score**: `float` that corresponds to the score given by the reader for this span compared to other + spans in the same passage. It corresponds to the sum of the start and end logits of the span. + - **relevance_score**: `float` that corresponds to the score of the each passage to answer the question, + compared to all the other passages. It corresponds to the output of the QA classifier of the DPRReader. + - **doc_id**: `int` the id of the passage. - ***start_index**: `int` the start index of the span + (inclusive). - **end_index**: `int` the end index of the span (inclusive). + + Examples: + + ```python + >>> from transformers import DPRReader, DPRReaderTokenizer + + >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base") + >>> encoded_inputs = tokenizer( + ... questions=["What is love ?"], + ... titles=["Haddaway"], + ... texts=["'What Is Love' is a song recorded by the artist Haddaway"], + ... return_tensors="pt", + ... ) + >>> outputs = model(**encoded_inputs) + >>> predicted_spans = tokenizer.decode_best_spans(encoded_inputs, outputs) + >>> print(predicted_spans[0].text) # best span + a song + ```""" + input_ids = reader_input["input_ids"] + start_logits, end_logits, relevance_logits = reader_output[:3] + n_passages = len(relevance_logits) + sorted_docs = sorted(range(n_passages), reverse=True, key=relevance_logits.__getitem__) + nbest_spans_predictions: list[DPRReaderOutput] = [] + for doc_id in sorted_docs: + sequence_ids = list(input_ids[doc_id]) + # assuming question & title information is at the beginning of the sequence + passage_offset = sequence_ids.index(self.sep_token_id, 2) + 1 # second sep id + if sequence_ids[-1] == self.pad_token_id: + sequence_len = sequence_ids.index(self.pad_token_id) + else: + sequence_len = len(sequence_ids) + + best_spans = self._get_best_spans( + start_logits=start_logits[doc_id][passage_offset:sequence_len], + end_logits=end_logits[doc_id][passage_offset:sequence_len], + max_answer_length=max_answer_length, + top_spans=num_spans_per_passage, + ) + for start_index, end_index in best_spans: + start_index += passage_offset + end_index += passage_offset + nbest_spans_predictions.append( + DPRSpanPrediction( + span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index], + relevance_score=relevance_logits[doc_id], + doc_id=doc_id, + start_index=start_index, + end_index=end_index, + text=self.decode(sequence_ids[start_index : end_index + 1]), + ) + ) + if len(nbest_spans_predictions) >= num_spans: + break + return nbest_spans_predictions[:num_spans] + + def _get_best_spans( + self, + start_logits: list[int], + end_logits: list[int], + max_answer_length: int, + top_spans: int, + ) -> list[DPRSpanPrediction]: + """ + Finds the best answer span for the extractive Q&A model for one passage. It returns the best span by descending + `span_score` order and keeping max `top_spans` spans. Spans longer that `max_answer_length` are ignored. + """ + scores = [] + for start_index, start_score in enumerate(start_logits): + for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]): + scores.append(((start_index, start_index + answer_length), start_score + end_score)) + scores = sorted(scores, key=lambda x: x[1], reverse=True) + chosen_span_intervals = [] + for (start_index, end_index), score in scores: + assert start_index <= end_index, f"Wrong span indices: [{start_index}:{end_index}]" + length = end_index - start_index + 1 + assert length <= max_answer_length, f"Span is too long: {length} > {max_answer_length}" + if any( + start_index <= prev_start_index <= prev_end_index <= end_index + or prev_start_index <= start_index <= end_index <= prev_end_index + for (prev_start_index, prev_end_index) in chosen_span_intervals + ): + continue + chosen_span_intervals.append((start_index, end_index)) + + if len(chosen_span_intervals) == top_spans: + break + return chosen_span_intervals + + +@add_end_docstrings(CUSTOM_DPR_READER_DOCSTRING) +class DPRReaderTokenizerFast(CustomDPRReaderTokenizerMixin, BertTokenizer): + r""" + Constructs a "fast" DPRReader tokenizer (backed by HuggingFace's *tokenizers* library). + + [`DPRReaderTokenizerFast`] is almost identical to [`BertTokenizer`] and runs end-to-end tokenization: + punctuation splitting and wordpiece. The difference is that is has three inputs strings: question, titles and texts + that are combined to be fed to the [`DPRReader`] model. + + Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters. + + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + slow_tokenizer_class = DPRReaderTokenizer + + +__all__ = ["DPRContextEncoderTokenizerFast", "DPRQuestionEncoderTokenizerFast", "DPRReaderTokenizerFast"] diff --git a/third_party/transformers/src/transformers/models/dpt/__init__.py b/third_party/transformers/src/transformers/models/dpt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d20828a426d14db0b9d3ae0a96749ed6dab2e1e6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/__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_dpt import * + from .image_processing_dpt import * + from .image_processing_pil_dpt import * + from .modeling_dpt 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/dpt/configuration_dpt.py b/third_party/transformers/src/transformers/models/dpt/configuration_dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..24a0d312abaed024470852f24ee8c473e3ffea09 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/configuration_dpt.py @@ -0,0 +1,158 @@ +# 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. +"""DPT 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.configuration_auto import AutoConfig + + +@auto_docstring(checkpoint="Intel/dpt-large") +@strict +class DPTConfig(PreTrainedConfig): + r""" + is_hybrid (`bool`, *optional*, defaults to `False`): + Whether to use a hybrid backbone. Useful in the context of loading DPT-Hybrid models. + backbone_out_indices (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + Indices of the intermediate hidden states to use from backbone. + readout_type (`str`, *optional*, defaults to `"project"`): + The readout type to use when processing the readout token (CLS token) of the intermediate hidden states of + the ViT backbone. Can be one of [`"ignore"`, `"add"`, `"project"`]. + - "ignore" simply ignores the CLS token. + - "add" passes the information from the CLS token to all other tokens by adding the representations. + - "project" passes information to the other tokens by concatenating the readout to all other tokens before + projecting the + representation to the original feature dimension D using a linear layer followed by a GELU non-linearity. + reassemble_factors (`list[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`): + The up/downsampling factors of the reassemble layers. + neck_hidden_sizes (`list[str]`, *optional*, defaults to `[96, 192, 384, 768]`): + The hidden sizes to project to for the feature maps of the backbone. + fusion_hidden_size (`int`, *optional*, defaults to 256): + The number of channels before fusion. + head_in_index (`int`, *optional*, defaults to -1): + The index of the features to use in the heads. + use_batch_norm_in_fusion_residual (`bool`, *optional*, defaults to `False`): + Whether to use batch normalization in the pre-activate residual units of the fusion blocks. + use_bias_in_fusion_residual (`bool`, *optional*, defaults to `True`): + Whether to use bias in the pre-activate residual units of the fusion blocks. + add_projection (`bool`, *optional*, defaults to `False`): + Whether to add a projection layer before the depth estimation head. + use_auxiliary_head (`bool`, *optional*, defaults to `True`): + Whether to use an auxiliary head during training. + auxiliary_loss_weight (`float`, *optional*, defaults to 0.4): + Weight of the cross-entropy loss of the auxiliary head. + semantic_classifier_dropout (`float`, *optional*, defaults to 0.1): + The dropout ratio for the semantic classification head. + backbone_featmap_shape (`list[int]`, *optional*, defaults to `[1, 1024, 24, 24]`): + Used only for the `hybrid` embedding type. The shape of the feature maps of the backbone. + neck_ignore_stages (`list[int]`, *optional*, defaults to `[0, 1]`): + Used only for the `hybrid` embedding type. The stages of the readout layers to ignore. + pooler_output_size (`int`, *optional*): + Dimensionality of the pooler layer. If None, defaults to `hidden_size`. + pooler_act (`str`, *optional*, defaults to `"tanh"`): + The activation function to be used by the pooler. + + Example: + + ```python + >>> from transformers import DPTModel, DPTConfig + + >>> # Initializing a DPT dpt-large style configuration + >>> configuration = DPTConfig() + + >>> # Initializing a model from the dpt-large style configuration + >>> model = DPTModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "dpt" + sub_configs = {"backbone_config": AutoConfig} + + # NOTE: some values are typed as `None` on purpose + # DPT creates one of: backbone or the general model only + # so official checkpoint saved them as `None` + hidden_size: int = 768 + num_hidden_layers: None | int = 12 + num_attention_heads: int | None = 12 + intermediate_size: int | None = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int | None = 0.0 + attention_probs_dropout_prob: float | int | None = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float | None = 1e-12 + image_size: int | list[int] | tuple[int, int] | None = 384 + patch_size: int | list[int] | tuple[int, int] | None = 16 + num_channels: int | None = 3 + is_hybrid: bool = False + qkv_bias: bool | None = True + backbone_out_indices: list[int] | tuple[int, ...] | None = (2, 5, 8, 11) + readout_type: str = "project" + reassemble_factors: list[int | float] | tuple[int | float, ...] = (4, 2, 1, 0.5) + neck_hidden_sizes: list[int] | tuple[int, ...] = (96, 192, 384, 768) + fusion_hidden_size: int = 256 + head_in_index: int = -1 + use_batch_norm_in_fusion_residual: bool | None = False + use_bias_in_fusion_residual: bool | None = None + add_projection: bool = False + use_auxiliary_head: bool | None = True + auxiliary_loss_weight: float = 0.4 + semantic_loss_ignore_index: int = 255 + semantic_classifier_dropout: float | int = 0.1 + backbone_featmap_shape: list[int] | tuple[int, ...] | None = (1, 1024, 24, 24) + neck_ignore_stages: list[int] | tuple[int, ...] = (0, 1) + backbone_config: dict | PreTrainedConfig | None = None + pooler_output_size: int | None = None + pooler_act: str = "tanh" + + def __post_init__(self, **kwargs): + if self.readout_type not in ["ignore", "add", "project"]: + raise ValueError("Readout_type must be one of ['ignore', 'add', 'project']") + + if self.is_hybrid: + if isinstance(self.backbone_config, dict): + self.backbone_config.setdefault("model_type", "bit") + + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="bit", + default_config_kwargs={ + "global_padding": "same", + "layer_type": "bottleneck", + "depths": [3, 4, 9], + "out_features": ["stage1", "stage2", "stage3"], + "embedding_dynamic_padding": True, + }, + **kwargs, + ) + if self.readout_type != "project": + raise ValueError("Readout type must be 'project' when using `DPT-hybrid` mode.") + elif kwargs.get("backbone") is not None or self.backbone_config is not None: + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + **kwargs, + ) + self.backbone_out_indices = None + + self.backbone_featmap_shape = self.backbone_featmap_shape if self.is_hybrid else None + self.neck_ignore_stages = self.neck_ignore_stages if self.is_hybrid else [] + self.pooler_output_size = self.pooler_output_size if self.pooler_output_size else self.hidden_size + super().__post_init__(**kwargs) + + +__all__ = ["DPTConfig"] diff --git a/third_party/transformers/src/transformers/models/dpt/convert_dinov2_depth_to_hf.py b/third_party/transformers/src/transformers/models/dpt/convert_dinov2_depth_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..cb3ce6ec869e348acd89b3f8c9ec30dc2ed2347d --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/convert_dinov2_depth_to_hf.py @@ -0,0 +1,384 @@ +# 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 DINOv2 + DPT checkpoints from the original repository. URL: +https://github.com/facebookresearch/dinov2/tree/main""" + +import argparse +import itertools +import math +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from PIL import Image +from torchvision import transforms + +from transformers import Dinov2Config, DPTConfig, DPTForDepthEstimation, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(model_name): + if "small" in model_name: + # equivalent to stage 3, stage 6, stage 9, stage 12 + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-small", out_indices=[3, 6, 9, 12], apply_layernorm=False, reshape_hidden_states=False + ) + neck_hidden_sizes = [48, 96, 192, 384] + elif "base" in model_name: + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-base", out_indices=[3, 6, 9, 12], apply_layernorm=False, reshape_hidden_states=False + ) + neck_hidden_sizes = [96, 192, 384, 768] + elif "large" in model_name: + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-large", out_indices=[5, 12, 18, 24], apply_layernorm=False, reshape_hidden_states=False + ) + neck_hidden_sizes = [128, 256, 512, 1024] + elif "giant" in model_name: + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-giant", out_indices=[10, 20, 30, 40], apply_layernorm=False, reshape_hidden_states=False + ) + neck_hidden_sizes = [192, 384, 768, 1536] + else: + raise NotImplementedError("To do") + + config = DPTConfig( + backbone_config=backbone_config, + neck_hidden_sizes=neck_hidden_sizes, + use_bias_in_fusion_residual=False, + add_projection=True, + ) + + return config + + +# here we list all DPT keys to be renamed (original name on the left, our name on the right) +def create_rename_keys_dpt(config): + rename_keys = [] + + # fmt: off + # activation postprocessing (projections, readout projections + resize blocks) + for i in range(4): + rename_keys.append((f"decode_head.reassemble_blocks.projects.{i}.conv.weight", f"neck.reassemble_stage.layers.{i}.projection.weight")) + rename_keys.append((f"decode_head.reassemble_blocks.projects.{i}.conv.bias", f"neck.reassemble_stage.layers.{i}.projection.bias")) + + rename_keys.append((f"decode_head.reassemble_blocks.readout_projects.{i}.0.weight", f"neck.reassemble_stage.readout_projects.{i}.0.weight")) + rename_keys.append((f"decode_head.reassemble_blocks.readout_projects.{i}.0.bias", f"neck.reassemble_stage.readout_projects.{i}.0.bias")) + + if i != 2: + rename_keys.append((f"decode_head.reassemble_blocks.resize_layers.{i}.weight", f"neck.reassemble_stage.layers.{i}.resize.weight")) + rename_keys.append((f"decode_head.reassemble_blocks.resize_layers.{i}.bias", f"neck.reassemble_stage.layers.{i}.resize.bias")) + + # fusion layers + for i in range(4): + rename_keys.append((f"decode_head.fusion_blocks.{i}.project.conv.weight", f"neck.fusion_stage.layers.{i}.projection.weight")) + rename_keys.append((f"decode_head.fusion_blocks.{i}.project.conv.bias", f"neck.fusion_stage.layers.{i}.projection.bias")) + if i != 0: + rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit1.conv1.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer1.convolution1.weight")) + rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit1.conv2.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer1.convolution2.weight")) + rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit2.conv1.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer2.convolution1.weight")) + rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit2.conv2.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer2.convolution2.weight")) + + # neck convolutions + for i in range(4): + rename_keys.append((f"decode_head.convs.{i}.conv.weight", f"neck.convs.{i}.weight")) + + # head + rename_keys.append(("decode_head.project.conv.weight", "head.projection.weight")) + rename_keys.append(("decode_head.project.conv.bias", "head.projection.bias")) + + for i in range(0, 5, 2): + rename_keys.append((f"decode_head.conv_depth.head.{i}.weight", f"head.head.{i}.weight")) + rename_keys.append((f"decode_head.conv_depth.head.{i}.bias", f"head.head.{i}.bias")) + # fmt: on + + return rename_keys + + +# here we list all backbone keys to be renamed (original name on the left, our name on the right) +def create_rename_keys_backbone(config): + rename_keys = [] + + # fmt: off + # patch embedding layer + rename_keys.append(("cls_token", "backbone.embeddings.cls_token")) + rename_keys.append(("mask_token", "backbone.embeddings.mask_token")) + rename_keys.append(("pos_embed", "backbone.embeddings.position_embeddings")) + rename_keys.append(("patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight")) + rename_keys.append(("patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias")) + + # Transformer encoder + for i in range(config.backbone_config.num_hidden_layers): + # layernorms + rename_keys.append((f"blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.norm1.weight")) + rename_keys.append((f"blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.norm1.bias")) + rename_keys.append((f"blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.norm2.weight")) + rename_keys.append((f"blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.norm2.bias")) + # MLP + if config.backbone_config.use_swiglu_ffn: + rename_keys.append((f"blocks.{i}.mlp.w12.weight", f"backbone.encoder.layer.{i}.mlp.w12.weight")) + rename_keys.append((f"blocks.{i}.mlp.w12.bias", f"backbone.encoder.layer.{i}.mlp.w12.bias")) + rename_keys.append((f"blocks.{i}.mlp.w3.weight", f"backbone.encoder.layer.{i}.mlp.w3.weight")) + rename_keys.append((f"blocks.{i}.mlp.w3.bias", f"backbone.encoder.layer.{i}.mlp.w3.bias")) + else: + rename_keys.append((f"blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.mlp.fc1.weight")) + rename_keys.append((f"blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.mlp.fc1.bias")) + rename_keys.append((f"blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.mlp.fc2.weight")) + rename_keys.append((f"blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.mlp.fc2.bias")) + # layerscale + rename_keys.append((f"blocks.{i}.ls1.gamma", f"backbone.encoder.layer.{i}.layer_scale1.lambda1")) + rename_keys.append((f"blocks.{i}.ls2.gamma", f"backbone.encoder.layer.{i}.layer_scale2.lambda1")) + # attention projection layer + rename_keys.append((f"blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight")) + rename_keys.append((f"blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias")) + # fmt: on + + rename_keys.append(("norm.weight", "backbone.layernorm.weight")) + rename_keys.append(("norm.bias", "backbone.layernorm.bias")) + + return rename_keys + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + for i in range(config.backbone_config.num_hidden_layers): + # read in weights + bias of input projection layer (in timm, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"blocks.{i}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"blocks.{i}.attn.qkv.bias") + hidden_size = config.backbone_config.hidden_size + # next, add query, keys and values (in that order) to the state dict + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[:hidden_size] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + hidden_size : hidden_size * 2, : + ] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[ + hidden_size : hidden_size * 2 + ] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-hidden_size:] + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "https://dl.fbaipublicfiles.com/dinov2/images/example.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +name_to_url = { + "dpt-dinov2-small-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_nyu_dpt_head.pth", + "dpt-dinov2-small-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_kitti_dpt_head.pth", + "dpt-dinov2-base-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_nyu_dpt_head.pth", + "dpt-dinov2-base-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_kitti_dpt_head.pth", + "dpt-dinov2-large-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_nyu_dpt_head.pth", + "dpt-dinov2-large-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_kitti_dpt_head.pth", + "dpt-dinov2-giant-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_nyu_dpt_head.pth", + "dpt-dinov2-giant-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_kitti_dpt_head.pth", +} + + +def get_original_pixel_values(image): + class CenterPadding: + def __init__(self, multiple): + super().__init__() + self.multiple = multiple + + def _get_pad(self, size): + new_size = math.ceil(size / self.multiple) * self.multiple + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + def __call__(self, img): + pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in img.shape[-2:][::-1])) + output = torch.nn.functional.pad(img, pads) + return output + + def __repr__(self): + return self.__class__.__name__ + "()" + + def make_depth_transform() -> transforms.Compose: + return transforms.Compose( + [ + transforms.ToTensor(), + lambda x: 255.0 * x[:3], # Discard alpha component and scale by 255 + transforms.Normalize( + mean=(123.675, 116.28, 103.53), + std=(58.395, 57.12, 57.375), + ), + CenterPadding(multiple=14), + ] + ) + + transform = make_depth_transform() + original_pixel_values = transform(image).unsqueeze(0) + + return original_pixel_values + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + # define DPT configuration based on URL + checkpoint_url = name_to_url[model_name] + config = get_dpt_config(model_name) + + # load original DPT state_dict from URL + print("URL:", checkpoint_url) + dpt_state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["state_dict"] + # rename keys + rename_keys = create_rename_keys_dpt(config) + for src, dest in rename_keys: + rename_key(dpt_state_dict, src, dest) + + # load original backbone state_dict from URL + if "small" in model_name: + original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14") + elif "base" in model_name: + original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitb14") + elif "large" in model_name: + original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14") + elif "giant" in model_name: + original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitg14") + else: + raise NotImplementedError("To do") + original_model.eval() + backbone_state_dict = original_model.state_dict() + + # rename keys + rename_keys = create_rename_keys_backbone(config) + for src, dest in rename_keys: + rename_key(backbone_state_dict, src, dest) + + # read in qkv matrices + read_in_q_k_v(backbone_state_dict, config) + + for key, val in backbone_state_dict.copy().items(): + val = backbone_state_dict.pop(key) + if "w12" in key: + key = key.replace("w12", "weights_in") + if "w3" in key: + key = key.replace("w3", "weights_out") + backbone_state_dict[key] = val + + # merge state_dicts + state_dict = {**backbone_state_dict, **dpt_state_dict} + + # load HuggingFace model + model = DPTForDepthEstimation(config) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + assert missing_keys == [ + "neck.fusion_stage.layers.0.residual_layer1.convolution1.weight", + "neck.fusion_stage.layers.0.residual_layer1.convolution2.weight", + ] + model.eval() + + # Verify image processor + processor = DPTImageProcessor( + do_resize=False, + do_rescale=False, + do_pad=True, + size_divisor=14, + do_normalize=True, + image_mean=(123.675, 116.28, 103.53), + image_std=(58.395, 57.12, 57.375), + ) + + image = prepare_img() + pixel_values = processor(image, return_tensors="pt").pixel_values.float() + original_pixel_values = get_original_pixel_values(image) + + assert torch.allclose(pixel_values, original_pixel_values) + + # Verify forward pass + with torch.no_grad(): + outputs = model(pixel_values) + + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values of predicted depth:", predicted_depth[0, :3, :3]) + + # assert logits + if verify_logits: + if model_name == "dpt-dinov2-small-nyu": + expected_shape = torch.Size([1, 576, 736]) + expected_slice = torch.tensor( + [[3.3576, 3.4741, 3.4345], [3.4324, 3.5012, 3.2775], [3.2560, 3.3563, 3.2354]] + ) + + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-5) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"facebook/{model_name}") + processor.push_to_hub(repo_id=f"facebook/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="dpt-dinov2-small-nyu", + type=str, + choices=name_to_url.keys(), + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + parser.add_argument( + "--verify_logits", + action="store_true", + required=False, + help="Path to the output PyTorch model directory.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits) diff --git a/third_party/transformers/src/transformers/models/dpt/convert_dpt_beit_to_hf.py b/third_party/transformers/src/transformers/models/dpt/convert_dpt_beit_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..5e913d048ac5e188e68f0d5164353bf19a252864 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/convert_dpt_beit_to_hf.py @@ -0,0 +1,306 @@ +# 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 DPT 3.1 checkpoints from the MiDaS repository. URL: https://github.com/isl-org/MiDaS""" + +import argparse +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from PIL import Image + +from transformers import BeitConfig, DPTConfig, DPTForDepthEstimation, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(model_name): + hidden_size = 768 + num_hidden_layers = 12 + num_attention_heads = 12 + intermediate_size = 3072 + out_features = ["stage3", "stage6", "stage9", "stage12"] # beit-base-384 uses [2, 5, 8, 11] + + if "large" in model_name: + hidden_size = 1024 + num_hidden_layers = 24 + num_attention_heads = 16 + intermediate_size = 4096 + out_features = ["stage6", "stage12", "stage18", "stage24"] # beit-large-512 uses [5, 11, 17, 23] + + if "512" in model_name: + image_size = 512 + elif "384" in model_name: + image_size = 384 + else: + raise ValueError("Model not supported") + + backbone_config = BeitConfig( + image_size=image_size, + num_hidden_layers=num_hidden_layers, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + use_relative_position_bias=True, + reshape_hidden_states=False, + out_features=out_features, + ) + + neck_hidden_sizes = [256, 512, 1024, 1024] if "large" in model_name else [96, 192, 384, 768] + config = DPTConfig(backbone_config=backbone_config, neck_hidden_sizes=neck_hidden_sizes) + + return config, image_size + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def create_rename_keys(config): + rename_keys = [] + + # fmt: off + # stem + rename_keys.append(("pretrained.model.cls_token", "backbone.embeddings.cls_token")) + rename_keys.append(("pretrained.model.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight")) + rename_keys.append(("pretrained.model.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias")) + + # Transformer encoder + for i in range(config.backbone_config.num_hidden_layers): + rename_keys.append((f"pretrained.model.blocks.{i}.gamma_1", f"backbone.encoder.layer.{i}.lambda_1")) + rename_keys.append((f"pretrained.model.blocks.{i}.gamma_2", f"backbone.encoder.layer.{i}.lambda_2")) + rename_keys.append((f"pretrained.model.blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.layernorm_before.weight")) + rename_keys.append((f"pretrained.model.blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.layernorm_before.bias")) + rename_keys.append((f"pretrained.model.blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.layernorm_after.weight")) + rename_keys.append((f"pretrained.model.blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.layernorm_after.bias")) + rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.intermediate.dense.weight")) + rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.intermediate.dense.bias")) + rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.output.dense.weight")) + rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.output.dense.bias")) + rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight")) + rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias")) + rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_bias_table", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table")) + rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_index", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index")) + + # activation postprocessing (readout projections + resize blocks) + for i in range(4): + rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.weight", f"neck.reassemble_stage.readout_projects.{i}.0.weight")) + rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.bias", f"neck.reassemble_stage.readout_projects.{i}.0.bias")) + + rename_keys.append((f"pretrained.act_postprocess{i+1}.3.weight", f"neck.reassemble_stage.layers.{i}.projection.weight")) + rename_keys.append((f"pretrained.act_postprocess{i+1}.3.bias", f"neck.reassemble_stage.layers.{i}.projection.bias")) + + if i != 2: + rename_keys.append((f"pretrained.act_postprocess{i+1}.4.weight", f"neck.reassemble_stage.layers.{i}.resize.weight")) + rename_keys.append((f"pretrained.act_postprocess{i+1}.4.bias", f"neck.reassemble_stage.layers.{i}.resize.bias")) + + # refinenet (tricky here) + mapping = {1:3, 2:2, 3:1, 4:0} + + for i in range(1, 5): + j = mapping[i] + rename_keys.append((f"scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight")) + rename_keys.append((f"scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias")) + + # scratch convolutions + for i in range(4): + rename_keys.append((f"scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight")) + + # head + for i in range(0, 5, 2): + rename_keys.append((f"scratch.output_conv.{i}.weight", f"head.head.{i}.weight")) + rename_keys.append((f"scratch.output_conv.{i}.bias", f"head.head.{i}.bias")) + + return rename_keys + + +def remove_ignore_keys_(state_dict): + ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + hidden_size = config.backbone_config.hidden_size + for i in range(config.backbone_config.num_hidden_layers): + # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"pretrained.model.blocks.{i}.attn.qkv.weight") + q_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.q_bias") + v_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.v_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = q_bias + state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + hidden_size : hidden_size * 2, : + ] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :] + state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = v_bias + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + name_to_url = { + "dpt-beit-large-512": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt", + "dpt-beit-large-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_384.pt", + "dpt-beit-base-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_base_384.pt", + } + + # define DPT configuration based on URL + checkpoint_url = name_to_url[model_name] + config, image_size = get_dpt_config(model_name) + # load original state_dict from URL + state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu") + # remove certain keys + remove_ignore_keys_(state_dict) + # rename keys + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + # read in qkv matrices + read_in_q_k_v(state_dict, config) + + # load HuggingFace model + model = DPTForDepthEstimation(config) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + assert missing_keys == [] + # assert unexpected_keys == ["pretrained.model.fc_norm.weight", "pretrained.model.fc_norm.bias"] + model.eval() + + # Check outputs on an image + # We set `keep_aspect_ratio=False` as our current BEiT does not support arbitrary window sizes + processor = DPTImageProcessor( + size={"height": image_size, "width": image_size}, keep_aspect_ratio=False, ensure_multiple_of=32 + ) + + image = prepare_img() + pixel_values = processor(image, return_tensors="pt").pixel_values + + print("First values of pixel values:", pixel_values[0, 0, :3, :3]) + print("Mean of pixel values:", pixel_values.mean().item()) + print("Shape of pixel values:", pixel_values.shape) + + from PIL import Image + from torchvision import transforms + + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + + transforms = transforms.Compose( + [ + transforms.Resize((image_size, image_size)), + transforms.ToTensor(), + ] + ) + pixel_values = transforms(image).unsqueeze(0) + + # forward pass + with torch.no_grad(): + outputs = model(pixel_values) + + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values of predicted depth:", predicted_depth[0, :3, :3]) + + # assert logits + # TODO there's still a small difference with the original logits + if model_name == "dpt-beit-large-512": + # OK, checked + expected_shape = torch.Size([1, 512, 512]) + expected_slice = torch.tensor( + [[2804.6260, 2792.5708, 2812.9263], [2772.0288, 2780.1118, 2796.2529], [2748.1094, 2766.6558, 2766.9834]] + ) + elif model_name == "dpt-beit-large-384": + # OK, checked + expected_shape = torch.Size([1, 384, 384]) + expected_slice = torch.tensor( + [[1783.2273, 1780.5729, 1792.6453], [1759.9817, 1765.5359, 1778.5002], [1739.1633, 1754.7903, 1757.1990]], + ) + elif model_name == "dpt-beit-base-384": + # OK, checked + expected_shape = torch.Size([1, 384, 384]) + expected_slice = torch.tensor( + [[2898.4482, 2891.3750, 2904.8079], [2858.6685, 2877.2615, 2894.4507], [2842.1235, 2854.1023, 2861.6328]], + ) + + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"nielsr/{model_name}") + processor.push_to_hub(repo_id=f"nielsr/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="dpt-beit-large-512", + type=str, + choices=["dpt-beit-large-512", "dpt-beit-large-384", "dpt-beit-base-384"], + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py b/third_party/transformers/src/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..63df2ba1248597002e0701653f88c7d916472825 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py @@ -0,0 +1,316 @@ +# 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 DPT checkpoints from the original repository. URL: https://github.com/isl-org/DPT""" + +import argparse +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(checkpoint_url): + config = DPTConfig(embedding_type="hybrid") + + if "large" in checkpoint_url: + config.hidden_size = 1024 + config.intermediate_size = 4096 + config.num_hidden_layers = 24 + config.num_attention_heads = 16 + config.backbone_out_indices = [5, 11, 17, 23] + config.neck_hidden_sizes = [256, 512, 1024, 1024] + expected_shape = (1, 384, 384) + + if "nyu" in checkpoint_url or "midas" in checkpoint_url: + config.hidden_size = 768 + config.reassemble_factors = [1, 1, 1, 0.5] + config.neck_hidden_sizes = [256, 512, 768, 768] + config.num_labels = 150 + config.patch_size = 16 + expected_shape = (1, 384, 384) + config.use_batch_norm_in_fusion_residual = False + config.readout_type = "project" + + if "ade" in checkpoint_url: + config.use_batch_norm_in_fusion_residual = True + config.hidden_size = 768 + config.reassemble_stage = [1, 1, 1, 0.5] + config.num_labels = 150 + config.patch_size = 16 + repo_id = "huggingface/label-files" + filename = "ade20k-id2label.json" + id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text()) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + expected_shape = [1, 150, 480, 480] + + return config, expected_shape + + +def remove_ignore_keys_(state_dict): + ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + +def rename_key(name): + if ( + "pretrained.model" in name + and "cls_token" not in name + and "pos_embed" not in name + and "patch_embed" not in name + ): + name = name.replace("pretrained.model", "dpt.encoder") + if "pretrained.model" in name: + name = name.replace("pretrained.model", "dpt.embeddings") + if "patch_embed" in name: + name = name.replace("patch_embed", "") + if "pos_embed" in name: + name = name.replace("pos_embed", "position_embeddings") + if "attn.proj" in name: + name = name.replace("attn.proj", "attention.output.dense") + if "proj" in name and "project" not in name: + name = name.replace("proj", "projection") + if "blocks" in name: + name = name.replace("blocks", "layer") + 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 "norm1" in name and "backbone" not in name: + name = name.replace("norm1", "layernorm_before") + if "norm2" in name and "backbone" not in name: + name = name.replace("norm2", "layernorm_after") + if "scratch.output_conv" in name: + name = name.replace("scratch.output_conv", "head") + if "scratch" in name: + name = name.replace("scratch", "neck") + if "layer1_rn" in name: + name = name.replace("layer1_rn", "convs.0") + if "layer2_rn" in name: + name = name.replace("layer2_rn", "convs.1") + if "layer3_rn" in name: + name = name.replace("layer3_rn", "convs.2") + if "layer4_rn" in name: + name = name.replace("layer4_rn", "convs.3") + if "refinenet" in name: + layer_idx = int(name[len("neck.refinenet") : len("neck.refinenet") + 1]) + # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3 + name = name.replace(f"refinenet{layer_idx}", f"fusion_stage.layers.{abs(layer_idx - 4)}") + if "out_conv" in name: + name = name.replace("out_conv", "projection") + if "resConfUnit1" in name: + name = name.replace("resConfUnit1", "residual_layer1") + if "resConfUnit2" in name: + name = name.replace("resConfUnit2", "residual_layer2") + if "conv1" in name: + name = name.replace("conv1", "convolution1") + if "conv2" in name: + name = name.replace("conv2", "convolution2") + # readout blocks + if "pretrained.act_postprocess1.0.project.0" in name: + name = name.replace("pretrained.act_postprocess1.0.project.0", "neck.reassemble_stage.readout_projects.0.0") + if "pretrained.act_postprocess2.0.project.0" in name: + name = name.replace("pretrained.act_postprocess2.0.project.0", "neck.reassemble_stage.readout_projects.1.0") + if "pretrained.act_postprocess3.0.project.0" in name: + name = name.replace("pretrained.act_postprocess3.0.project.0", "neck.reassemble_stage.readout_projects.2.0") + if "pretrained.act_postprocess4.0.project.0" in name: + name = name.replace("pretrained.act_postprocess4.0.project.0", "neck.reassemble_stage.readout_projects.3.0") + + # resize blocks + if "pretrained.act_postprocess1.3" in name: + name = name.replace("pretrained.act_postprocess1.3", "neck.reassemble_stage.layers.0.projection") + if "pretrained.act_postprocess1.4" in name: + name = name.replace("pretrained.act_postprocess1.4", "neck.reassemble_stage.layers.0.resize") + if "pretrained.act_postprocess2.3" in name: + name = name.replace("pretrained.act_postprocess2.3", "neck.reassemble_stage.layers.1.projection") + if "pretrained.act_postprocess2.4" in name: + name = name.replace("pretrained.act_postprocess2.4", "neck.reassemble_stage.layers.1.resize") + if "pretrained.act_postprocess3.3" in name: + name = name.replace("pretrained.act_postprocess3.3", "neck.reassemble_stage.layers.2.projection") + if "pretrained.act_postprocess4.3" in name: + name = name.replace("pretrained.act_postprocess4.3", "neck.reassemble_stage.layers.3.projection") + if "pretrained.act_postprocess4.4" in name: + name = name.replace("pretrained.act_postprocess4.4", "neck.reassemble_stage.layers.3.resize") + if "pretrained" in name: + name = name.replace("pretrained", "dpt") + if "bn" in name: + name = name.replace("bn", "batch_norm") + if "head" in name: + name = name.replace("head", "head.head") + if "encoder.norm" in name: + name = name.replace("encoder.norm", "layernorm") + if "auxlayer" in name: + name = name.replace("auxlayer", "auxiliary_head.head") + if "backbone" in name: + name = name.replace("backbone", "backbone.bit.encoder") + + if ".." in name: + name = name.replace("..", ".") + + if "stem.conv" in name: + name = name.replace("stem.conv", "bit.embedder.convolution") + if "blocks" in name: + name = name.replace("blocks", "layers") + if "convolution" in name and "backbone" in name: + name = name.replace("convolution", "conv") + if "layer" in name and "backbone" in name: + name = name.replace("layer", "layers") + if "backbone.bit.encoder.bit" in name: + name = name.replace("backbone.bit.encoder.bit", "backbone.bit") + if "embedder.conv" in name: + name = name.replace("embedder.conv", "embedder.convolution") + if "backbone.bit.encoder.stem.norm" in name: + name = name.replace("backbone.bit.encoder.stem.norm", "backbone.bit.embedder.norm") + return name + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + for i in range(config.num_hidden_layers): + # read in weights + bias of input projection layer (in timm, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + config.hidden_size : config.hidden_size * 2, : + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[ + config.hidden_size : config.hidden_size * 2 + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[ + -config.hidden_size :, : + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :] + + +# 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_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name, show_prediction): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + # define DPT configuration based on URL + config, expected_shape = get_dpt_config(checkpoint_url) + # load original state_dict from URL + # state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu") + state_dict = torch.load(checkpoint_url, map_location="cpu", weights_only=True) + # remove certain keys + remove_ignore_keys_(state_dict) + # rename keys + for key in state_dict.copy(): + val = state_dict.pop(key) + state_dict[rename_key(key)] = val + # read in qkv matrices + read_in_q_k_v(state_dict, config) + + # load HuggingFace model + model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config) + model.load_state_dict(state_dict) + model.eval() + + # Check outputs on an image + size = 480 if "ade" in checkpoint_url else 384 + image_processor = DPTImageProcessor(size=size) + + image = prepare_img() + encoding = image_processor(image, return_tensors="pt") + + # forward pass + outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth + + if show_prediction: + prediction = ( + torch.nn.functional.interpolate( + outputs.unsqueeze(1), + size=(image.size[1], image.size[0]), + mode="bicubic", + align_corners=False, + ) + .squeeze() + .cpu() + .numpy() + ) + + Image.fromarray((prediction / prediction.max()) * 255).show() + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + model.push_to_hub("ybelkada/dpt-hybrid-midas") + image_processor.push_to_hub("ybelkada/dpt-hybrid-midas") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--checkpoint_url", + default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt", + type=str, + help="URL of the original DPT checkpoint you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + required=False, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + ) + parser.add_argument( + "--model_name", + default="dpt-large", + type=str, + help="Name of the model, in case you're pushing to the hub.", + ) + parser.add_argument( + "--show_prediction", + action="store_true", + ) + + args = parser.parse_args() + convert_dpt_checkpoint( + args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name, args.show_prediction + ) diff --git a/third_party/transformers/src/transformers/models/dpt/convert_dpt_swinv2_to_hf.py b/third_party/transformers/src/transformers/models/dpt/convert_dpt_swinv2_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..d11651edfaec0663af886a0d8a8b8548ceb18d47 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/convert_dpt_swinv2_to_hf.py @@ -0,0 +1,322 @@ +# 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 DPT 3.1 checkpoints from the MiDaS repository. URL: https://github.com/isl-org/MiDaS""" + +import argparse +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from PIL import Image + +from transformers import DPTConfig, DPTForDepthEstimation, DPTImageProcessor, Swinv2Config +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(model_name): + if "tiny" in model_name: + embed_dim = 96 + depths = (2, 2, 6, 2) + num_heads = (3, 6, 12, 24) + window_size = 16 + # note: for Swinv2-tiny authors used the window_size = 16 variant + # as seen here: https://github.com/isl-org/MiDaS/blob/bdc4ed64c095e026dc0a2f17cabb14d58263decb/midas/backbones/swin2.py#L26 + pretrained_window_sizes = (0, 0, 0, 0) + elif "base" in model_name: + embed_dim = 128 + depths = (2, 2, 18, 2) + num_heads = (4, 8, 16, 32) + window_size = 24 + pretrained_window_sizes = (12, 12, 12, 6) + elif "large" in model_name: + embed_dim = 192 + depths = (2, 2, 18, 2) + num_heads = (6, 12, 24, 48) + window_size = 24 + pretrained_window_sizes = (12, 12, 12, 6) + + if "384" in model_name: + image_size = 384 + elif "256" in model_name: + image_size = 256 + else: + raise ValueError("Model not supported, to do") + + backbone_config = Swinv2Config( + image_size=image_size, + embed_dim=embed_dim, + depths=depths, + window_size=window_size, + pretrained_window_sizes=pretrained_window_sizes, + num_heads=num_heads, + out_features=["stage1", "stage2", "stage3", "stage4"], + ) + + if model_name == "dpt-swinv2-tiny-256": + neck_hidden_sizes = [96, 192, 384, 768] + elif model_name == "dpt-swinv2-base-384": + neck_hidden_sizes = [128, 256, 512, 1024] + elif model_name == "dpt-swinv2-large-384": + neck_hidden_sizes = [192, 384, 768, 1536] + + config = DPTConfig(backbone_config=backbone_config, neck_hidden_sizes=neck_hidden_sizes) + + return config, image_size + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def create_rename_keys(config): + rename_keys = [] + + # fmt: off + # stem + rename_keys.append(("pretrained.model.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight")) + rename_keys.append(("pretrained.model.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias")) + rename_keys.append(("pretrained.model.patch_embed.norm.weight", "backbone.embeddings.norm.weight")) + rename_keys.append(("pretrained.model.patch_embed.norm.bias", "backbone.embeddings.norm.bias")) + + # transformer encoder + for i in range(len(config.backbone_config.depths)): + for j in range(config.backbone_config.depths[i]): + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.logit_scale", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.logit_scale")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.0.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.0.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.0.bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.0.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.2.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.q_bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.query.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.v_bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.value.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.proj.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.proj.bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc2.bias", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm2.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias")) + + # downsample parameters + if i in [0,1,2]: + rename_keys.append((f"pretrained.model.layers.{i}.downsample.reduction.weight", f"backbone.encoder.layers.{i}.downsample.reduction.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.downsample.norm.weight", f"backbone.encoder.layers.{i}.downsample.norm.weight")) + rename_keys.append((f"pretrained.model.layers.{i}.downsample.norm.bias", f"backbone.encoder.layers.{i}.downsample.norm.bias")) + + # note: non-Transformer backbones like Swinv2, LeViT et al don't require activation postprocessing (readout projections + resize blocks) + + # refinenet (tricky here) + mapping = {1:3, 2:2, 3:1, 4:0} + + for i in range(1, 5): + j = mapping[i] + rename_keys.append((f"scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight")) + rename_keys.append((f"scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight")) + rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias")) + + # scratch convolutions + for i in range(4): + rename_keys.append((f"scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight")) + + # head + for i in range(0, 5, 2): + rename_keys.append((f"scratch.output_conv.{i}.weight", f"head.head.{i}.weight")) + rename_keys.append((f"scratch.output_conv.{i}.bias", f"head.head.{i}.bias")) + + return rename_keys + + +def remove_ignore_keys_(state_dict): + ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config, model): + for i in range(len(config.backbone_config.depths)): + for j in range(config.backbone_config.depths[i]): + dim = model.backbone.encoder.layers[i].blocks[j].attention.self.all_head_size + # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"pretrained.model.layers.{i}.blocks.{j}.attn.qkv.weight") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.query.weight"] = in_proj_weight[:dim, :] + state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.key.weight"] = in_proj_weight[ + dim : dim * 2, : + ] + state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.value.weight"] = in_proj_weight[ + -dim:, : + ] + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, verify_logits, push_to_hub): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + name_to_url = { + "dpt-swinv2-tiny-256": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_tiny_256.pt", + "dpt-swinv2-base-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_base_384.pt", + "dpt-swinv2-large-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_large_384.pt", + } + + # define DPT configuration based on URL + checkpoint_url = name_to_url[model_name] + config, image_size = get_dpt_config(model_name) + # load original state_dict from URL + state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu") + + # load HuggingFace model + model = DPTForDepthEstimation(config) + + # remove certain keys + remove_ignore_keys_(state_dict) + # rename keys + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + # read in qkv matrices + read_in_q_k_v(state_dict, config, model) + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + model.eval() + + # Check outputs on an image + processor = DPTImageProcessor(size={"height": image_size, "width": image_size}) + + image = prepare_img() + processor(image, return_tensors="pt") + + if verify_logits: + from torchvision import transforms + + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + transforms = transforms.Compose( + [ + transforms.Resize((image_size, image_size)), + transforms.ToTensor(), + ] + ) + pixel_values = transforms(image).unsqueeze(0) + + # forward pass + with torch.no_grad(): + outputs = model(pixel_values) + + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values of predicted depth:", predicted_depth[0, :3, :3]) + + # assert logits + if model_name == "dpt-swinv2-base-384": + # OK, checked + expected_shape = torch.Size([1, 384, 384]) + expected_slice = torch.tensor( + [ + [1998.5575, 1997.3887, 2009.2981], + [1952.8607, 1979.6488, 2001.0854], + [1953.7697, 1961.7711, 1968.8904], + ], + ) + elif model_name == "dpt-swinv2-tiny-256": + # OK, checked + expected_shape = torch.Size([1, 256, 256]) + expected_slice = torch.tensor( + [[978.9163, 976.5215, 978.5349], [974.1859, 971.7249, 975.8046], [971.3419, 970.3118, 971.6830]], + ) + elif model_name == "dpt-swinv2-large-384": + # OK, checked + expected_shape = torch.Size([1, 384, 384]) + expected_slice = torch.tensor( + [ + [1203.7206, 1200.1495, 1197.8234], + [1196.2484, 1183.5033, 1186.4640], + [1178.8131, 1182.3260, 1174.3975], + ], + ) + + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"Intel/{model_name}") + processor.push_to_hub(repo_id=f"Intel/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="dpt-swinv2-base-384", + type=str, + choices=["dpt-swinv2-tiny-256", "dpt-swinv2-base-384", "dpt-swinv2-large-384"], + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--verify_logits", + action="store_true", + help="Whether to verify logits after conversion.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.verify_logits, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/dpt/convert_dpt_to_pytorch.py b/third_party/transformers/src/transformers/models/dpt/convert_dpt_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..5fea417b4b42133b3771928f2c329cce0a097a86 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/convert_dpt_to_pytorch.py @@ -0,0 +1,276 @@ +# 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 DPT checkpoints from the original repository. URL: https://github.com/isl-org/DPT""" + +import argparse +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(checkpoint_url): + config = DPTConfig() + + if "large" in checkpoint_url: + config.hidden_size = 1024 + config.intermediate_size = 4096 + config.num_hidden_layers = 24 + config.num_attention_heads = 16 + config.backbone_out_indices = [5, 11, 17, 23] + config.neck_hidden_sizes = [256, 512, 1024, 1024] + expected_shape = (1, 384, 384) + + if "ade" in checkpoint_url: + config.use_batch_norm_in_fusion_residual = True + + config.num_labels = 150 + repo_id = "huggingface/label-files" + filename = "ade20k-id2label.json" + id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text()) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + expected_shape = [1, 150, 480, 480] + + return config, expected_shape + + +def remove_ignore_keys_(state_dict): + ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + +def rename_key(name): + if ( + "pretrained.model" in name + and "cls_token" not in name + and "pos_embed" not in name + and "patch_embed" not in name + ): + name = name.replace("pretrained.model", "dpt.encoder") + if "pretrained.model" in name: + name = name.replace("pretrained.model", "dpt.embeddings") + if "patch_embed" in name: + name = name.replace("patch_embed", "patch_embeddings") + if "pos_embed" in name: + name = name.replace("pos_embed", "position_embeddings") + if "attn.proj" in name: + name = name.replace("attn.proj", "attention.output.dense") + if "proj" in name and "project" not in name: + name = name.replace("proj", "projection") + if "blocks" in name: + name = name.replace("blocks", "layer") + 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 "norm1" in name: + name = name.replace("norm1", "layernorm_before") + if "norm2" in name: + name = name.replace("norm2", "layernorm_after") + if "scratch.output_conv" in name: + name = name.replace("scratch.output_conv", "head") + if "scratch" in name: + name = name.replace("scratch", "neck") + if "layer1_rn" in name: + name = name.replace("layer1_rn", "convs.0") + if "layer2_rn" in name: + name = name.replace("layer2_rn", "convs.1") + if "layer3_rn" in name: + name = name.replace("layer3_rn", "convs.2") + if "layer4_rn" in name: + name = name.replace("layer4_rn", "convs.3") + if "refinenet" in name: + layer_idx = int(name[len("neck.refinenet") : len("neck.refinenet") + 1]) + # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3 + name = name.replace(f"refinenet{layer_idx}", f"fusion_stage.layers.{abs(layer_idx - 4)}") + if "out_conv" in name: + name = name.replace("out_conv", "projection") + if "resConfUnit1" in name: + name = name.replace("resConfUnit1", "residual_layer1") + if "resConfUnit2" in name: + name = name.replace("resConfUnit2", "residual_layer2") + if "conv1" in name: + name = name.replace("conv1", "convolution1") + if "conv2" in name: + name = name.replace("conv2", "convolution2") + # readout blocks + if "pretrained.act_postprocess1.0.project.0" in name: + name = name.replace("pretrained.act_postprocess1.0.project.0", "neck.reassemble_stage.readout_projects.0.0") + if "pretrained.act_postprocess2.0.project.0" in name: + name = name.replace("pretrained.act_postprocess2.0.project.0", "neck.reassemble_stage.readout_projects.1.0") + if "pretrained.act_postprocess3.0.project.0" in name: + name = name.replace("pretrained.act_postprocess3.0.project.0", "neck.reassemble_stage.readout_projects.2.0") + if "pretrained.act_postprocess4.0.project.0" in name: + name = name.replace("pretrained.act_postprocess4.0.project.0", "neck.reassemble_stage.readout_projects.3.0") + # resize blocks + if "pretrained.act_postprocess1.3" in name: + name = name.replace("pretrained.act_postprocess1.3", "neck.reassemble_stage.layers.0.projection") + if "pretrained.act_postprocess1.4" in name: + name = name.replace("pretrained.act_postprocess1.4", "neck.reassemble_stage.layers.0.resize") + if "pretrained.act_postprocess2.3" in name: + name = name.replace("pretrained.act_postprocess2.3", "neck.reassemble_stage.layers.1.projection") + if "pretrained.act_postprocess2.4" in name: + name = name.replace("pretrained.act_postprocess2.4", "neck.reassemble_stage.layers.1.resize") + if "pretrained.act_postprocess3.3" in name: + name = name.replace("pretrained.act_postprocess3.3", "neck.reassemble_stage.layers.2.projection") + if "pretrained.act_postprocess4.3" in name: + name = name.replace("pretrained.act_postprocess4.3", "neck.reassemble_stage.layers.3.projection") + if "pretrained.act_postprocess4.4" in name: + name = name.replace("pretrained.act_postprocess4.4", "neck.reassemble_stage.layers.3.resize") + if "pretrained" in name: + name = name.replace("pretrained", "dpt") + if "bn" in name: + name = name.replace("bn", "batch_norm") + if "head" in name: + name = name.replace("head", "head.head") + if "encoder.norm" in name: + name = name.replace("encoder.norm", "layernorm") + if "auxlayer" in name: + name = name.replace("auxlayer", "auxiliary_head.head") + + return name + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + for i in range(config.num_hidden_layers): + # read in weights + bias of input projection layer (in timm, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + config.hidden_size : config.hidden_size * 2, : + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[ + config.hidden_size : config.hidden_size * 2 + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[ + -config.hidden_size :, : + ] + state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :] + + +# 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_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + # define DPT configuration based on URL + config, expected_shape = get_dpt_config(checkpoint_url) + # load original state_dict from URL + state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu") + # remove certain keys + remove_ignore_keys_(state_dict) + # rename keys + for key in state_dict.copy(): + val = state_dict.pop(key) + state_dict[rename_key(key)] = val + # read in qkv matrices + read_in_q_k_v(state_dict, config) + + # load HuggingFace model + model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config) + model.load_state_dict(state_dict) + model.eval() + + # Check outputs on an image + size = 480 if "ade" in checkpoint_url else 384 + image_processor = DPTImageProcessor(size=size) + + image = prepare_img() + encoding = image_processor(image, return_tensors="pt") + + # forward pass + outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth + + # Assert logits + expected_slice = torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]]) + if "ade" in checkpoint_url: + expected_slice = torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]]) + assert outputs.shape == torch.Size(expected_shape) + assert ( + torch.allclose(outputs[0, 0, :3, :3], expected_slice, atol=1e-4) + if "ade" in checkpoint_url + else torch.allclose(outputs[0, :3, :3], expected_slice) + ) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model to hub...") + model.push_to_hub(repo_id=f"nielsr/{model_name}") + image_processor.push_to_hub(repo_id=f"nielsr/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--checkpoint_url", + default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt", + type=str, + help="URL of the original DPT checkpoint you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + required=False, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + ) + parser.add_argument( + "--model_name", + default="dpt-large", + type=str, + required=False, + help="Name of the model, in case you're pushing to the hub.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name) diff --git a/third_party/transformers/src/transformers/models/dpt/image_processing_dpt.py b/third_party/transformers/src/transformers/models/dpt/image_processing_dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..6d157f6385c0e700394aa0cafff3bf9c02488a47 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/image_processing_dpt.py @@ -0,0 +1,415 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/dpt/modular_dpt.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_dpt.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Iterable +from typing import TYPE_CHECKING, Union + +import torch +import torch.nn.functional as F +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_base import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available, requires_backends + + +if TYPE_CHECKING: + from ...modeling_outputs import DepthEstimatorOutput + + +class DPTImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ensure_multiple_of (`int`, *optional*, defaults to 1): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. Can be overridden + by `ensure_multiple_of` in `preprocess`. + keep_aspect_ratio (`bool`, *optional*, defaults to `False`): + If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. Can + be overridden by `keep_aspect_ratio` in `preprocess`. + do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`): + Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. + ADE20k). The background label will be replaced by 255. + """ + + ensure_multiple_of: int + size_divisor: int + keep_aspect_ratio: bool + do_reduce_labels: bool + + +def get_resize_output_image_size( + input_image: "torch.Tensor", + output_size: int | Iterable[int], + keep_aspect_ratio: bool, + multiple: int, +) -> SizeDict: + def constrain_to_multiple_of(val, multiple, min_val=0, max_val=None): + x = round(val / multiple) * multiple + + if max_val is not None and x > max_val: + x = math.floor(val / multiple) * multiple + + if x < min_val: + x = math.ceil(val / multiple) * multiple + + return x + + input_height, input_width = input_image.shape[-2:] + output_height, output_width = output_size + + # determine new height and width + scale_height = output_height / input_height + scale_width = output_width / input_width + + if keep_aspect_ratio: + # scale as little as possible + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + + new_height = constrain_to_multiple_of(scale_height * input_height, multiple=multiple) + new_width = constrain_to_multiple_of(scale_width * input_width, multiple=multiple) + + return SizeDict(height=new_height, width=new_width) + + +@auto_docstring +class DPTImageProcessor(TorchvisionBackend): + """PIL backend for DPT with reduce_label support.""" + + valid_kwargs = DPTImageProcessorKwargs + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 384, "width": 384} + default_to_square = True + + # necessary for modular conversion + crop_size = None + do_resize = True + do_center_crop = None + do_rescale = True + do_normalize = True + do_reduce_labels = None + do_pad = False + rescale_factor = 1 / 255 + ensure_multiple_of = 1 + keep_aspect_ratio = False + + def __init__(self, **kwargs: Unpack[DPTImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[DPTImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to preprocess. + """ + return super().preprocess(images, segmentation_maps, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + return_tensors: str | TensorType | None, + device: Union[str, "torch.device"] | None = None, + **kwargs, + ) -> BatchFeature: + """Handle extra inputs beyond images.""" + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + # Prepare segmentation maps if provided + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + # Process segmentation maps with do_normalize=False and do_rescale=False + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False}) + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + # Convert to int64 and squeeze channel dimension + processed_segmentation_maps = [ + processed_segmentation_map.squeeze(0).to(torch.int64) + for processed_segmentation_map in processed_segmentation_maps + ] + data["labels"] = processed_segmentation_maps + + return BatchFeature(data=data, tensor_type=return_tensors) + + def reduce_label(self, labels: list["torch.Tensor"]) -> list["torch.Tensor"]: + """Reduce label values by 1, replacing 0 with 255.""" + for idx in range(len(labels)): + label = labels[idx] + label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype, device=label.device), label) + label = label - 1 + label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype, device=label.device), label) + labels[idx] = label + return labels + + def _preprocess( + self, + images: list["torch.Tensor"], + do_reduce_labels: bool, + 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, + keep_aspect_ratio: bool, + ensure_multiple_of: int | None, + do_pad: bool, + size_divisor: int | None, + disable_grouping: bool | None, + **kwargs, + ) -> BatchFeature: + """Custom preprocessing for DPT.""" + if do_reduce_labels: + images = self.reduce_label(images) + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_images, + size=size, + resample=resample, + ensure_multiple_of=ensure_multiple_of, + keep_aspect_ratio=keep_aspect_ratio, + ) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + if do_pad: + stacked_images = self.pad_image(stacked_images, size_divisor) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + return processed_images + + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """ + Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. + + Args: + outputs ([`DPTForSemanticSegmentation`]): + Raw outputs of the model. + target_sizes (`list[Tuple]` of length `batch_size`, *optional*): + List of tuples corresponding to the requested final size (height, width) of each prediction. If unset, + predictions will not be resized. + + Returns: + semantic_segmentation: `list[torch.Tensor]` 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. + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for post_process_semantic_segmentation") + + logits = outputs.logits + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + + semantic_segmentation = [] + + for idx in range(len(logits)): + resized_logits = F.interpolate( + logits[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 = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + antialias: bool = True, + ensure_multiple_of: int | None = 1, + keep_aspect_ratio: bool = False, + ) -> "torch.Tensor": + """ + Resize an image to `(size["height"], size["width"])`. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image. + interpolation (`InterpolationMode`, *optional*, defaults to `InterpolationMode.BILINEAR`): + `InterpolationMode` filter to use when resizing the image e.g. `InterpolationMode.BICUBIC`. + antialias (`bool`, *optional*, defaults to `True`): + Whether to use antialiasing when resizing the image + ensure_multiple_of (`int`, *optional*): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value + keep_aspect_ratio (`bool`, *optional*, defaults to `False`): + If `True`, and `do_resize` is `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. + + Returns: + `torch.Tensor`: The resized image. + """ + if not size.height or not size.width: + raise ValueError(f"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}") + + output_size = get_resize_output_image_size( + image, + output_size=(size.height, size.width), + keep_aspect_ratio=keep_aspect_ratio, + multiple=ensure_multiple_of, + ) + return super().resize(image, output_size, resample=resample, antialias=antialias) + + def pad_image( + self, + image: "torch.Tensor", + size_divisor: int = 1, + ) -> "torch.Tensor": + r""" + Center pad a batch of images to be a multiple of `size_divisor`. + + Args: + image (`torch.Tensor`): + Image to pad. Can be a batch of images of dimensions (N, C, H, W) or a single image of dimensions (C, H, W). + size_divisor (`int`): + The width and height of the image will be padded to a multiple of this number. + """ + height, width = image.shape[-2:] + + def _get_pad(size, size_divisor): + new_size = math.ceil(size / size_divisor) * size_divisor + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + pad_top, pad_bottom = _get_pad(height, size_divisor) + pad_left, pad_right = _get_pad(width, size_divisor) + padding = (pad_left, pad_top, pad_right, pad_bottom) + return tvF.pad(image, padding) + + def post_process_depth_estimation( + self, + outputs: "DepthEstimatorOutput", + target_sizes: TensorType | list[tuple[int, int]] | None | None = None, + ) -> list[dict[str, TensorType]]: + """ + Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions and depth PIL images. + Only supports PyTorch. + + Args: + outputs ([`DepthEstimatorOutput`]): + Raw outputs of the model. + target_sizes (`TensorType` 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. + + Returns: + `List[Dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth + predictions. + """ + requires_backends(self, "torch") + + predicted_depth = outputs.predicted_depth + + if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth" + ) + + results = [] + target_sizes = [None] * len(predicted_depth) if target_sizes is None else target_sizes + for depth, target_size in zip(predicted_depth, target_sizes): + if target_size is not None: + depth = torch.nn.functional.interpolate( + depth.unsqueeze(0).unsqueeze(1), size=target_size, mode="bicubic", align_corners=False + ).squeeze() + + results.append({"predicted_depth": depth}) + + return results + + +__all__ = ["DPTImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/dpt/image_processing_pil_dpt.py b/third_party/transformers/src/transformers/models/dpt/image_processing_pil_dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..6f770cac4e5f0619dbae23676fe6aea7b5181159 --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/image_processing_pil_dpt.py @@ -0,0 +1,312 @@ +# 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 DPT.""" + +import math +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, logging +from ...utils.import_utils import requires + + +if TYPE_CHECKING: + from ...modeling_outputs import DepthEstimatorOutput + +logger = logging.get_logger(__name__) + + +# Adapted from transformers.models.dpt.image_processing_dpt.DPTImageProcessorKwargs +class DPTImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ensure_multiple_of (`int`, *optional*, defaults to 1): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. Can be overridden + by `ensure_multiple_of` in `preprocess`. + keep_aspect_ratio (`bool`, *optional*, defaults to `False`): + If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. Can + be overridden by `keep_aspect_ratio` in `preprocess`. + do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`): + Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. + ADE20k). The background label will be replaced by 255. + """ + + ensure_multiple_of: int + size_divisor: int + keep_aspect_ratio: bool + do_reduce_labels: bool + + +# Adapted from transformers.models.dpt.image_processing_dpt.get_resize_output_image_size +def get_resize_output_image_size( + input_image: np.ndarray, + output_size: int | Iterable[int], + keep_aspect_ratio: bool, + multiple: int, +) -> SizeDict: + def constrain_to_multiple_of(val, multiple, min_val=0, max_val=None): + x = round(val / multiple) * multiple + + if max_val is not None and x > max_val: + x = math.floor(val / multiple) * multiple + + if x < min_val: + x = math.ceil(val / multiple) * multiple + + return x + + input_height, input_width = input_image.shape[-2:] + output_height, output_width = output_size + + # determine new height and width + scale_height = output_height / input_height + scale_width = output_width / input_width + + if keep_aspect_ratio: + # scale as little as possible + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + + new_height = constrain_to_multiple_of(scale_height * input_height, multiple=multiple) + new_width = constrain_to_multiple_of(scale_width * input_width, multiple=multiple) + + return SizeDict(height=new_height, width=new_width) + + +@auto_docstring +class DPTImageProcessorPil(PilBackend): + """PIL backend for DPT with custom resize and pad.""" + + valid_kwargs = DPTImageProcessorKwargs + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 384, "width": 384} + default_to_square = True + do_resize = True + do_rescale = True + do_normalize = True + do_pad = False + rescale_factor = 1 / 255 + ensure_multiple_of = 1 + keep_aspect_ratio = False + + def __init__(self, **kwargs: Unpack[DPTImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[DPTImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to preprocess. + """ + return super().preprocess(images, segmentation_maps, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + return_tensors: str | TensorType | None, + device: str | None = None, + **kwargs, + ) -> BatchFeature: + """Handle extra inputs beyond images.""" + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False}) + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + processed_segmentation_maps = [ + processed_segmentation_map.squeeze(0).astype(np.int64) + for processed_segmentation_map in processed_segmentation_maps + ] + data["labels"] = processed_segmentation_maps + + return BatchFeature(data=data, tensor_type=return_tensors) + + def reduce_label(self, image: np.ndarray) -> np.ndarray: + """Reduce label values by 1, replacing 0 with 255.""" + image[image == 0] = 255 + image = image - 1 + image[image == 254] = 255 + return image + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: PILImageResampling | None, + ensure_multiple_of: int = 1, + keep_aspect_ratio: bool = False, + **kwargs, + ) -> np.ndarray: + """Resize with aspect ratio and multiple constraints.""" + if size.height is None or size.width is None: + raise ValueError(f"Size must contain 'height' and 'width' keys. Got {size.keys()}") + output_size = get_resize_output_image_size( + image, + output_size=(size.height, size.width), + keep_aspect_ratio=keep_aspect_ratio, + multiple=ensure_multiple_of, + ) + return super().resize(image, output_size, resample, **kwargs) + + def pad_image(self, image: np.ndarray, size_divisor: int = 1, **kwargs) -> np.ndarray: + """Center pad image to be a multiple of size_divisor.""" + + def _get_pad(size, size_divisor): + new_size = math.ceil(size / size_divisor) * size_divisor + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + height, width = image.shape[-2:] + pad_size_left, pad_size_right = _get_pad(height, size_divisor) + pad_size_top, pad_size_bottom = _get_pad(width, size_divisor) + return np_pad( + image, + padding=((pad_size_left, pad_size_right), (pad_size_top, pad_size_bottom)), + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + do_reduce_labels: bool, + keep_aspect_ratio: bool, + ensure_multiple_of: int, + size_divisor: int | None = None, + **kwargs, + ) -> np.ndarray: + """Custom preprocessing for DPT.""" + processed_images = [] + for image in images: + if do_reduce_labels: + image = self.reduce_label(image) + if do_resize: + image = self.resize( + image, size, resample, ensure_multiple_of=ensure_multiple_of, keep_aspect_ratio=keep_aspect_ratio + ) + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + if do_pad and size_divisor is not None: + image = self.pad_image(image, size_divisor) + processed_images.append(image) + return processed_images + + @requires(backends=("torch",)) + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps.""" + import torch + + logits = outputs.logits + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + semantic_segmentation = [] + for idx in range(len(logits)): + resized_logits = torch.nn.functional.interpolate( + logits[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 = logits.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_depth_estimation( + self, outputs: "DepthEstimatorOutput", target_sizes: TensorType | list[tuple[int, int]] | None = None + ) -> list[dict[str, TensorType]]: + """Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions.""" + import torch + + predicted_depth = outputs.predicted_depth + if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth" + ) + results = [] + target_sizes = [None] * len(predicted_depth) if target_sizes is None else target_sizes + for depth, target_size in zip(predicted_depth, target_sizes): + if target_size is not None: + depth = torch.nn.functional.interpolate( + depth.unsqueeze(0).unsqueeze(1), size=target_size, mode="bicubic", align_corners=False + ).squeeze() + results.append({"predicted_depth": depth}) + return results + + +__all__ = ["DPTImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/dpt/modeling_dpt.py b/third_party/transformers/src/transformers/models/dpt/modeling_dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb8b911d4ae202bd96c96613320883dd0edba4a --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/modeling_dpt.py @@ -0,0 +1,1185 @@ +# Copyright 2022 Intel Labs, OpenMMLab 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 DPT (Dense Prediction Transformers) model. + +This implementation is heavily inspired by OpenMMLab's implementation, found here: +https://github.com/open-mmlab/mmsegmentation/blob/master/mmseg/models/decode_heads/dpt_head.py. + +""" + +import collections.abc +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import load_backbone +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, DepthEstimatorOutput, SemanticSegmenterOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging, torch_int +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_dpt import DPTConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model's outputs that also contains intermediate activations that can be used at later stages. Useful + in the context of Vision models.: + """ +) +class BaseModelOutputWithIntermediateActivations(ModelOutput): + r""" + last_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + intermediate_activations (`tuple(torch.FloatTensor)`, *optional*): + Intermediate activations that can be used to compute hidden states of the model at various layers. + """ + + last_hidden_states: torch.FloatTensor | None = None + intermediate_activations: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model's outputs that also contains a pooling of the last hidden states as well as intermediate + activations that can be used by the model at later stages. + """ +) +class BaseModelOutputWithPoolingAndIntermediateActivations(ModelOutput): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) after further processing + through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns + the classification token after processing through a linear layer and a tanh activation function. The linear + layer weights are trained from the next sentence prediction (classification) objective during pretraining. + intermediate_activations (`tuple(torch.FloatTensor)`, *optional*): + Intermediate activations that can be used to compute hidden states of the model at various layers. + """ + + 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 + intermediate_activations: tuple[torch.FloatTensor, ...] | None = None + + +class DPTViTHybridEmbeddings(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: DPTConfig, feature_size: tuple[int, int] | None = None): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + + self.backbone = load_backbone(config) + feature_dim = self.backbone.channels[-1] + if len(self.backbone.channels) != 3: + raise ValueError(f"Expected backbone to have 3 output features, got {len(self.backbone.channels)}") + self.residual_feature_map_index = [0, 1] # Always take the output of the first and second backbone stage + + if feature_size is None: + feat_map_shape = config.backbone_featmap_shape + feature_size = feat_map_shape[-2:] + feature_dim = feat_map_shape[1] + else: + feature_size = ( + feature_size if isinstance(feature_size, collections.abc.Iterable) else (feature_size, feature_size) + ) + feature_dim = self.backbone.channels[-1] + + self.image_size = image_size + self.patch_size = patch_size[0] + self.num_channels = num_channels + + self.projection = nn.Conv2d(feature_dim, hidden_size, kernel_size=1) + + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size)) + + def _resize_pos_embed(self, posemb, grid_size_height, grid_size_width, start_index=1): + posemb_tok = posemb[:, :start_index] + posemb_grid = posemb[0, start_index:] + + old_grid_size = torch_int(len(posemb_grid) ** 0.5) + + posemb_grid = posemb_grid.reshape(1, old_grid_size, old_grid_size, -1).permute(0, 3, 1, 2) + posemb_grid = nn.functional.interpolate(posemb_grid, size=(grid_size_height, grid_size_width), mode="bilinear") + posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, grid_size_height * grid_size_width, -1) + + posemb = torch.cat([posemb_tok, posemb_grid], dim=1) + + return posemb + + def forward( + self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False + ) -> BaseModelOutputWithIntermediateActivations: + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + if not interpolate_pos_encoding: + if height != self.image_size[0] or width != self.image_size[1]: + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model" + f" ({self.image_size[0]}*{self.image_size[1]})." + ) + + position_embeddings = self._resize_pos_embed( + self.position_embeddings, height // self.patch_size, width // self.patch_size + ) + + backbone_output = self.backbone(pixel_values) + + features = backbone_output.feature_maps[-1] + + # Retrieve also the intermediate activations to use them at later stages + output_hidden_states = [backbone_output.feature_maps[index] for index in self.residual_feature_map_index] + + embeddings = self.projection(features).flatten(2).transpose(1, 2) + + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + position_embeddings + + # Return hidden states and intermediate activations + return BaseModelOutputWithIntermediateActivations( + last_hidden_states=embeddings, + intermediate_activations=output_hidden_states, + ) + + +class DPTViTEmbeddings(nn.Module): + """ + Construct the CLS token, position and patch embeddings. + + """ + + def __init__(self, config): + super().__init__() + + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.patch_embeddings = DPTViTPatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size)) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def _resize_pos_embed(self, posemb, grid_size_height, grid_size_width, start_index=1): + posemb_tok = posemb[:, :start_index] + posemb_grid = posemb[0, start_index:] + + old_grid_size = torch_int(posemb_grid.size(0) ** 0.5) + + posemb_grid = posemb_grid.reshape(1, old_grid_size, old_grid_size, -1).permute(0, 3, 1, 2) + posemb_grid = nn.functional.interpolate(posemb_grid, size=(grid_size_height, grid_size_width), mode="bilinear") + posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, grid_size_height * grid_size_width, -1) + + posemb = torch.cat([posemb_tok, posemb_grid], dim=1) + + return posemb + + def forward(self, pixel_values: torch.Tensor) -> BaseModelOutputWithIntermediateActivations: + batch_size, num_channels, height, width = pixel_values.shape + + # possibly interpolate position encodings to handle varying image sizes + patch_size = self.config.patch_size + position_embeddings = self._resize_pos_embed( + self.position_embeddings, height // patch_size, width // patch_size + ) + + embeddings = self.patch_embeddings(pixel_values) + + batch_size, seq_len, _ = embeddings.size() + + # add the [CLS] token to the embedded patch tokens + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + position_embeddings + + embeddings = self.dropout(embeddings) + + return BaseModelOutputWithIntermediateActivations(last_hidden_states=embeddings) + + +class DPTViTPatchEmbeddings(nn.Module): + """ + Image to Patch Embedding. + + """ + + def __init__(self, config: DPTConfig): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +# Copied from transformers.models.bert.modeling_bert.eager_attention_forward +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float | None = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->DPT +class DPTSelfAttention(nn.Module): + def __init__(self, config: DPTConfig): + 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.dropout_prob = config.attention_probs_dropout_prob + self.scaling = self.attention_head_size**-0.5 + self.is_causal = False + + self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + + def forward( + self, + hidden_states: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size = hidden_states.shape[0] + new_shape = batch_size, -1, self.num_attention_heads, self.attention_head_size + + key_layer = self.key(hidden_states).view(*new_shape).transpose(1, 2) + value_layer = self.value(hidden_states).view(*new_shape).transpose(1, 2) + query_layer = self.query(hidden_states).view(*new_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + context_layer, attention_probs = attention_interface( + self, + query_layer, + key_layer, + value_layer, + None, + is_causal=self.is_causal, + scaling=self.scaling, + dropout=0.0 if not self.training else self.dropout_prob, + **kwargs, + ) + + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.reshape(new_context_layer_shape) + + return context_layer, attention_probs + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViTConfig->DPTConfig, ViTSelfOutput->DPTViTSelfOutput +class DPTViTSelfOutput(nn.Module): + """ + The residual connection is defined in ViTLayer instead of here (as is the case with other models), due to the + layernorm applied before each block. + """ + + def __init__(self, config: DPTConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViTConfig->DPTConfig, ViTSelfAttention->DPTSelfAttention, ViTSelfOutput->DPTViTSelfOutput +class DPTViTAttention(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + self.attention = DPTSelfAttention(config) + self.output = DPTViTSelfOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + self_attn_output, _ = self.attention(hidden_states, **kwargs) + output = self.output(self_attn_output, hidden_states) + return output + + +# Copied from transformers.models.vit.modeling_vit.ViTIntermediate with ViTConfig->DPTConfig, ViTIntermediate->DPTViTIntermediate +class DPTViTIntermediate(nn.Module): + def __init__(self, config: DPTConfig): + 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.vit.modeling_vit.ViTOutput with ViTConfig->DPTConfig, ViTOutput->DPTViTOutput +class DPTViTOutput(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +# Copied from transformers.models.vit.modeling_vit.ViTLayer with ViTConfig->DPTConfig, ViTAttention->DPTViTAttention, ViTIntermediate->DPTViTIntermediate, ViTOutput->DPTViTOutput, ViTLayer->DPTViTLayer +class DPTViTLayer(GradientCheckpointingLayer): + """This corresponds to the Block class in the timm implementation.""" + + def __init__(self, config: DPTConfig): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = DPTViTAttention(config) + self.intermediate = DPTViTIntermediate(config) + self.output = DPTViTOutput(config) + self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + hidden_states_norm = self.layernorm_before(hidden_states) + attention_output = self.attention(hidden_states_norm, **kwargs) + + # first residual connection + hidden_states = attention_output + hidden_states + + # in ViT, layernorm is also applied after self-attention + layer_output = self.layernorm_after(hidden_states) + layer_output = self.intermediate(layer_output) + + # second residual connection is done here + layer_output = self.output(layer_output, hidden_states) + + return layer_output + + +class DPTReassembleStage(nn.Module): + """ + This class reassembles the hidden states of the backbone into image-like feature representations at various + resolutions. + + This happens in 3 stages: + 1. Map the N + 1 tokens to a set of N tokens, by taking into account the readout ([CLS]) token according to + `config.readout_type`. + 2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`. + 3. Resizing the spatial dimensions (height, width). + + Args: + config (`[DPTConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.config = config + self.layers = nn.ModuleList() + if config.is_hybrid: + self._init_reassemble_dpt_hybrid(config) + else: + self._init_reassemble_dpt(config) + + self.neck_ignore_stages = config.neck_ignore_stages + + def _init_reassemble_dpt_hybrid(self, config): + r""" " + For DPT-Hybrid the first 2 reassemble layers are set to `nn.Identity()`, please check the official + implementation: https://github.com/isl-org/DPT/blob/f43ef9e08d70a752195028a51be5e1aff227b913/dpt/vit.py#L438 + for more details. + """ + for i, factor in zip(range(len(config.neck_hidden_sizes)), config.reassemble_factors): + if i <= 1: + self.layers.append(nn.Identity()) + elif i > 1: + self.layers.append(DPTReassembleLayer(config, channels=config.neck_hidden_sizes[i], factor=factor)) + + if config.readout_type != "project": + raise ValueError(f"Readout type {config.readout_type} is not supported for DPT-Hybrid.") + + # When using DPT-Hybrid the readout type is set to "project". The sanity check is done on the config file + self.readout_projects = nn.ModuleList() + hidden_size = _get_backbone_hidden_size(config) + for i in range(len(config.neck_hidden_sizes)): + if i <= 1: + self.readout_projects.append(nn.Sequential(nn.Identity())) + elif i > 1: + self.readout_projects.append( + nn.Sequential(nn.Linear(2 * hidden_size, hidden_size), ACT2FN[config.hidden_act]) + ) + + def _init_reassemble_dpt(self, config): + for i, factor in zip(range(len(config.neck_hidden_sizes)), config.reassemble_factors): + self.layers.append(DPTReassembleLayer(config, channels=config.neck_hidden_sizes[i], factor=factor)) + + if config.readout_type == "project": + self.readout_projects = nn.ModuleList() + hidden_size = _get_backbone_hidden_size(config) + for _ in range(len(config.neck_hidden_sizes)): + self.readout_projects.append( + nn.Sequential(nn.Linear(2 * hidden_size, hidden_size), ACT2FN[config.hidden_act]) + ) + + def forward(self, hidden_states: list[torch.Tensor], patch_height=None, patch_width=None) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length + 1, hidden_size)`): + List of hidden states from the backbone. + """ + out = [] + + for i, hidden_state in enumerate(hidden_states): + if i not in self.neck_ignore_stages: + # reshape to (batch_size, num_channels, height, width) + cls_token, hidden_state = hidden_state[:, 0], hidden_state[:, 1:] + batch_size, sequence_length, num_channels = hidden_state.shape + if patch_height is not None and patch_width is not None: + hidden_state = hidden_state.reshape(batch_size, patch_height, patch_width, num_channels) + else: + size = torch_int(sequence_length**0.5) + hidden_state = hidden_state.reshape(batch_size, size, size, num_channels) + hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() + + feature_shape = hidden_state.shape + if self.config.readout_type == "project": + # reshape to (batch_size, height*width, num_channels) + hidden_state = hidden_state.flatten(2).permute((0, 2, 1)) + readout = cls_token.unsqueeze(1).expand_as(hidden_state) + # concatenate the readout token to the hidden states and project + hidden_state = self.readout_projects[i](torch.cat((hidden_state, readout), -1)) + # reshape back to (batch_size, num_channels, height, width) + hidden_state = hidden_state.permute(0, 2, 1).reshape(feature_shape) + elif self.config.readout_type == "add": + hidden_state = hidden_state.flatten(2) + cls_token.unsqueeze(-1) + hidden_state = hidden_state.reshape(feature_shape) + hidden_state = self.layers[i](hidden_state) + out.append(hidden_state) + + return out + + +def _get_backbone_hidden_size(config): + if config.backbone_config is not None and hasattr(config.backbone_config, "hidden_size"): + return config.backbone_config.hidden_size + else: + return config.hidden_size + + +class DPTReassembleLayer(nn.Module): + def __init__(self, config: DPTConfig, channels: int, factor: int): + super().__init__() + # projection + hidden_size = _get_backbone_hidden_size(config) + self.projection = nn.Conv2d(in_channels=hidden_size, out_channels=channels, kernel_size=1) + + # up/down sampling depending on factor + if factor > 1: + self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0) + elif factor == 1: + self.resize = nn.Identity() + elif factor < 1: + # so should downsample + self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=int(1 / factor), padding=1) + + def forward(self, hidden_state): + hidden_state = self.projection(hidden_state) + hidden_state = self.resize(hidden_state) + return hidden_state + + +class DPTFeatureFusionStage(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + self.layers = nn.ModuleList() + for _ in range(len(config.neck_hidden_sizes)): + self.layers.append(DPTFeatureFusionLayer(config)) + + def forward(self, hidden_states): + # reversing the hidden_states, we start from the last + hidden_states = hidden_states[::-1] + + fused_hidden_states = [] + fused_hidden_state = None + for hidden_state, layer in zip(hidden_states, self.layers): + if fused_hidden_state is None: + # first layer only uses the last hidden_state + fused_hidden_state = layer(hidden_state) + else: + fused_hidden_state = layer(fused_hidden_state, hidden_state) + fused_hidden_states.append(fused_hidden_state) + + return fused_hidden_states + + +class DPTPreActResidualLayer(nn.Module): + """ + ResidualConvUnit, pre-activate residual unit. + + Args: + config (`[DPTConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config: DPTConfig): + super().__init__() + + self.use_batch_norm = config.use_batch_norm_in_fusion_residual + use_bias_in_fusion_residual = ( + config.use_bias_in_fusion_residual + if config.use_bias_in_fusion_residual is not None + else not self.use_batch_norm + ) + + self.activation1 = nn.ReLU() + self.convolution1 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=use_bias_in_fusion_residual, + ) + + self.activation2 = nn.ReLU() + self.convolution2 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=use_bias_in_fusion_residual, + ) + + if self.use_batch_norm: + self.batch_norm1 = nn.BatchNorm2d(config.fusion_hidden_size) + self.batch_norm2 = nn.BatchNorm2d(config.fusion_hidden_size) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + residual = hidden_state + hidden_state = self.activation1(hidden_state) + + hidden_state = self.convolution1(hidden_state) + + if self.use_batch_norm: + hidden_state = self.batch_norm1(hidden_state) + + hidden_state = self.activation2(hidden_state) + hidden_state = self.convolution2(hidden_state) + + if self.use_batch_norm: + hidden_state = self.batch_norm2(hidden_state) + + return hidden_state + residual + + +class DPTFeatureFusionLayer(nn.Module): + """Feature fusion layer, merges feature maps from different stages. + + Args: + config (`[DPTConfig]`): + Model configuration class defining the model architecture. + align_corners (`bool`, *optional*, defaults to `True`): + The align_corner setting for bilinear upsample. + """ + + def __init__(self, config: DPTConfig, align_corners: bool = True): + super().__init__() + + self.align_corners = align_corners + + self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True) + + self.residual_layer1 = DPTPreActResidualLayer(config) + self.residual_layer2 = DPTPreActResidualLayer(config) + + def forward(self, hidden_state: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor: + if residual is not None: + if hidden_state.shape != residual.shape: + residual = nn.functional.interpolate( + residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False + ) + hidden_state = hidden_state + self.residual_layer1(residual) + + hidden_state = self.residual_layer2(hidden_state) + hidden_state = nn.functional.interpolate( + hidden_state, scale_factor=2, mode="bilinear", align_corners=self.align_corners + ) + hidden_state = self.projection(hidden_state) + + return hidden_state + + +@auto_docstring +class DPTPreTrainedModel(PreTrainedModel): + config: DPTConfig + base_model_prefix = "dpt" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": DPTViTLayer, + "attentions": DPTSelfAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, (DPTViTEmbeddings, DPTViTHybridEmbeddings)): + init.zeros_(module.cls_token) + init.zeros_(module.position_embeddings) + + +class DPTViTEncoder(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + self.config = config + self.layer = nn.ModuleList([DPTViTLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward( + self, hidden_states: torch.Tensor, output_hidden_states: bool = False, **kwargs: Unpack[TransformersKwargs] + ) -> BaseModelOutput: + for layer_module in self.layer: + hidden_states = layer_module(hidden_states) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +@auto_docstring +class DPTModel(DPTPreTrainedModel): + def __init__(self, config: DPTConfig, add_pooling_layer: bool = True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + # vit encoder + if config.is_hybrid: + self.embeddings = DPTViTHybridEmbeddings(config) + else: + self.embeddings = DPTViTEmbeddings(config) + self.encoder = DPTViTEncoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.pooler = DPTViTPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + if self.config.is_hybrid: + return self.embeddings + else: + return self.embeddings.patch_embeddings + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPoolingAndIntermediateActivations: + embedding_output: BaseModelOutputWithIntermediateActivations = self.embeddings(pixel_values) + embedding_last_hidden_states = embedding_output.last_hidden_states + + encoder_outputs: BaseModelOutput = self.encoder(embedding_last_hidden_states, **kwargs) + sequence_output = encoder_outputs.last_hidden_state + + sequence_output = self.layernorm(sequence_output) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPoolingAndIntermediateActivations( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + intermediate_activations=embedding_output.intermediate_activations, + ) + + +# Copied from transformers.models.vit.modeling_vit.ViTPooler with ViTConfig->DPTConfig, ViTPooler->DPTViTPooler +class DPTViTPooler(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.pooler_output_size) + self.activation = ACT2FN[config.pooler_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +class DPTNeck(nn.Module): + """ + DPTNeck. A neck is a module that is normally used between the backbone and the head. It takes a list of tensors as + input and produces another list of tensors as output. For DPT, it includes 2 stages: + + * DPTReassembleStage + * DPTFeatureFusionStage. + + Args: + config (dict): config dict. + """ + + def __init__(self, config: DPTConfig): + super().__init__() + self.config = config + + # postprocessing: only required in case of a non-hierarchical backbone (e.g. ViT, BEiT) + if config.backbone_config is not None and config.backbone_config.model_type == "swinv2": + self.reassemble_stage = None + else: + self.reassemble_stage = DPTReassembleStage(config) + + self.convs = nn.ModuleList() + for channel in config.neck_hidden_sizes: + self.convs.append(nn.Conv2d(channel, config.fusion_hidden_size, kernel_size=3, padding=1, bias=False)) + + # fusion + self.fusion_stage = DPTFeatureFusionStage(config) + + def forward( + self, + hidden_states: list[torch.Tensor], + patch_height: int | None = None, + patch_width: int | None = None, + ) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`): + List of hidden states from the backbone. + """ + if not isinstance(hidden_states, (tuple, list)): + raise TypeError("hidden_states should be a tuple or list of tensors") + + if len(hidden_states) != len(self.config.neck_hidden_sizes): + raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.") + + # postprocess hidden states + if self.reassemble_stage is not None: + hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width) + + features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)] + + # fusion blocks + output = self.fusion_stage(features) + + return output + + +class DPTDepthEstimationHead(nn.Module): + """ + Output head consisting of 3 convolutional layers. It progressively halves the feature dimension and upsamples + the predictions to the input resolution after the first convolutional layer (details can be found in the paper's + supplementary material). + """ + + def __init__(self, config: DPTConfig): + super().__init__() + + self.config = config + + self.projection = None + if config.add_projection: + self.projection = nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + + features = config.fusion_hidden_size + self.head = nn.Sequential( + nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1), + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), + nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), + nn.ReLU(), + ) + + def forward(self, hidden_states: list[torch.Tensor]) -> torch.Tensor: + # use last features + hidden_states = hidden_states[self.config.head_in_index] + + if self.projection is not None: + hidden_states = self.projection(hidden_states) + hidden_states = nn.ReLU()(hidden_states) + + predicted_depth = self.head(hidden_states) + predicted_depth = predicted_depth.squeeze(dim=1) + + return predicted_depth + + +@auto_docstring( + custom_intro=""" + DPT Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2. + """ +) +class DPTForDepthEstimation(DPTPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.backbone = None + if config.is_hybrid is False and config.backbone_config is not None: + self.backbone = load_backbone(config) + else: + self.dpt = DPTModel(config, add_pooling_layer=False) + + # Neck + self.neck = DPTNeck(config) + + # Depth estimation head + self.head = DPTDepthEstimationHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> DepthEstimatorOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth depth estimation maps for computing the loss. + + Examples: + ```python + >>> from transformers import AutoImageProcessor, DPTForDepthEstimation + >>> import torch + >>> import numpy as np + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large") + >>> model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # interpolate to original size + >>> post_processed_output = image_processor.post_process_depth_estimation( + ... outputs, + ... target_sizes=[(image.height, image.width)], + ... ) + + >>> # visualize the prediction + >>> predicted_depth = post_processed_output[0]["predicted_depth"] + >>> depth = predicted_depth * 255 / predicted_depth.max() + >>> depth = depth.detach().cpu().numpy() + >>> depth = Image.fromarray(depth.astype("uint8")) + ```""" + loss = None + if labels is not None: + raise NotImplementedError("Training is not implemented yet") + + # Internally the model always needs to output hidden states, we control the output + # per user request on the final output + user_requested_hidden_states = kwargs.get("output_hidden_states") or getattr( + self.config, "output_hidden_states", False + ) + kwargs["output_hidden_states"] = True + + if self.backbone is not None: + outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, **kwargs) + hidden_states = outputs.feature_maps + else: + outputs = self.dpt(pixel_values, **kwargs) + hidden_states = outputs.hidden_states + # only keep certain features based on config.backbone_out_indices + # note that the hidden_states also include the initial embeddings + if not self.config.is_hybrid: + hidden_states = [ + feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices + ] + else: + backbone_hidden_states = outputs.intermediate_activations + backbone_hidden_states.extend( + feature + for idx, feature in enumerate(hidden_states[1:]) + if idx in self.config.backbone_out_indices[2:] + ) + hidden_states = backbone_hidden_states + + patch_height, patch_width = None, None + if self.config.backbone_config is not None and self.config.is_hybrid is False: + _, _, height, width = pixel_values.shape + patch_size = self.config.backbone_config.patch_size + patch_height = height // patch_size + patch_width = width // patch_size + + hidden_states = self.neck(hidden_states, patch_height, patch_width) + predicted_depth = self.head(hidden_states) + + return DepthEstimatorOutput( + loss=loss, + predicted_depth=predicted_depth, + hidden_states=outputs.hidden_states if user_requested_hidden_states else None, + attentions=outputs.attentions, + ) + + +class DPTSemanticSegmentationHead(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + + self.config = config + features = config.fusion_hidden_size + self.head = nn.Sequential( + nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(features), + nn.ReLU(), + nn.Dropout(config.semantic_classifier_dropout), + nn.Conv2d(features, config.num_labels, kernel_size=1), + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), + ) + + def forward(self, hidden_states: list[torch.Tensor]) -> torch.Tensor: + # use last features + hidden_states = hidden_states[self.config.head_in_index] + logits = self.head(hidden_states) + return logits + + +class DPTAuxiliaryHead(nn.Module): + def __init__(self, config: DPTConfig): + super().__init__() + + features = config.fusion_hidden_size + self.head = nn.Sequential( + nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(features), + nn.ReLU(), + nn.Dropout(0.1, False), + nn.Conv2d(features, config.num_labels, kernel_size=1), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.head(hidden_states) + return logits + + +@auto_docstring +class DPTForSemanticSegmentation(DPTPreTrainedModel): + def __init__(self, config: DPTConfig): + super().__init__(config) + + self.dpt = DPTModel(config, add_pooling_layer=False) + + # Neck + self.neck = DPTNeck(config) + + # Segmentation head(s) + self.head = DPTSemanticSegmentationHead(config) + self.auxiliary_head = DPTAuxiliaryHead(config) if config.use_auxiliary_head else None + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> SemanticSegmenterOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy). + + Examples: + ```python + >>> from transformers import AutoImageProcessor, DPTForSemanticSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large-ade") + >>> model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> logits = outputs.logits + ```""" + if labels is not None and self.config.num_labels == 1: + raise ValueError("The number of labels should be greater than one") + + # Internally the model always needs to output hidden states, we control the output + # per user request on the final output + user_requested_hidden_states = kwargs.get("output_hidden_states") or getattr( + self.config, "output_hidden_states", False + ) + kwargs["output_hidden_states"] = True + + outputs: BaseModelOutputWithPoolingAndIntermediateActivations = self.dpt(pixel_values, **kwargs) + hidden_states = outputs.hidden_states + + # only keep certain features based on config.backbone_out_indices + # note that the hidden_states also include the initial embeddings + if not self.config.is_hybrid: + hidden_states = [ + feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices + ] + else: + backbone_hidden_states = outputs.intermediate_activations + backbone_hidden_states.extend( + feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices[2:] + ) + + hidden_states = backbone_hidden_states + + hidden_states = self.neck(hidden_states=hidden_states) + logits = self.head(hidden_states) + + auxiliary_logits = None + if self.auxiliary_head is not None: + auxiliary_logits = self.auxiliary_head(hidden_states[-1]) + + loss = None + if labels is not None: + # upsample logits to the images' original size + upsampled_logits = nn.functional.interpolate( + logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + if auxiliary_logits is not None: + upsampled_auxiliary_logits = nn.functional.interpolate( + auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + # compute weighted loss + loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index) + main_loss = loss_fct(upsampled_logits, labels) + auxiliary_loss = loss_fct(upsampled_auxiliary_logits, labels) + loss = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss + + return SemanticSegmenterOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states if user_requested_hidden_states else None, + attentions=outputs.attentions, + ) + + +__all__ = ["DPTForDepthEstimation", "DPTForSemanticSegmentation", "DPTModel", "DPTPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/dpt/modular_dpt.py b/third_party/transformers/src/transformers/models/dpt/modular_dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..dc024ecf2e4c27e157cc92518596993ed345328a --- /dev/null +++ b/third_party/transformers/src/transformers/models/dpt/modular_dpt.py @@ -0,0 +1,290 @@ +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_base import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs +from ...utils import TensorType, auto_docstring, requires_backends +from ..beit.image_processing_beit import BeitImageProcessor + + +if TYPE_CHECKING: + from ...modeling_outputs import DepthEstimatorOutput + +from torchvision.transforms.v2 import functional as tvF + + +def get_resize_output_image_size( + input_image: "torch.Tensor", + output_size: int | Iterable[int], + keep_aspect_ratio: bool, + multiple: int, +) -> SizeDict: + def constrain_to_multiple_of(val, multiple, min_val=0, max_val=None): + x = round(val / multiple) * multiple + + if max_val is not None and x > max_val: + x = math.floor(val / multiple) * multiple + + if x < min_val: + x = math.ceil(val / multiple) * multiple + + return x + + input_height, input_width = input_image.shape[-2:] + output_height, output_width = output_size + + # determine new height and width + scale_height = output_height / input_height + scale_width = output_width / input_width + + if keep_aspect_ratio: + # scale as little as possible + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + + new_height = constrain_to_multiple_of(scale_height * input_height, multiple=multiple) + new_width = constrain_to_multiple_of(scale_width * input_width, multiple=multiple) + + return SizeDict(height=new_height, width=new_width) + + +class DPTImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ensure_multiple_of (`int`, *optional*, defaults to 1): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. Can be overridden + by `ensure_multiple_of` in `preprocess`. + keep_aspect_ratio (`bool`, *optional*, defaults to `False`): + If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. Can + be overridden by `keep_aspect_ratio` in `preprocess`. + do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`): + Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. + ADE20k). The background label will be replaced by 255. + """ + + ensure_multiple_of: int + size_divisor: int + keep_aspect_ratio: bool + do_reduce_labels: bool + + +@auto_docstring +class DPTImageProcessor(BeitImageProcessor): + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 384, "width": 384} + do_resize = True + do_rescale = True + do_normalize = True + do_pad = False + rescale_factor = 1 / 255 + ensure_multiple_of = 1 + keep_aspect_ratio = False + + # necessary for modular conversion + crop_size = None + do_center_crop = None + do_reduce_labels = None + + valid_kwargs = DPTImageProcessorKwargs + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + antialias: bool = True, + ensure_multiple_of: int | None = 1, + keep_aspect_ratio: bool = False, + ) -> "torch.Tensor": + """ + Resize an image to `(size["height"], size["width"])`. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image. + interpolation (`InterpolationMode`, *optional*, defaults to `InterpolationMode.BILINEAR`): + `InterpolationMode` filter to use when resizing the image e.g. `InterpolationMode.BICUBIC`. + antialias (`bool`, *optional*, defaults to `True`): + Whether to use antialiasing when resizing the image + ensure_multiple_of (`int`, *optional*): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value + keep_aspect_ratio (`bool`, *optional*, defaults to `False`): + If `True`, and `do_resize` is `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. + + Returns: + `torch.Tensor`: The resized image. + """ + if not size.height or not size.width: + raise ValueError(f"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}") + + output_size = get_resize_output_image_size( + image, + output_size=(size.height, size.width), + keep_aspect_ratio=keep_aspect_ratio, + multiple=ensure_multiple_of, + ) + return TorchvisionBackend.resize(self, image, output_size, resample=resample, antialias=antialias) + + def pad_image( + self, + image: "torch.Tensor", + size_divisor: int = 1, + ) -> "torch.Tensor": + r""" + Center pad a batch of images to be a multiple of `size_divisor`. + + Args: + image (`torch.Tensor`): + Image to pad. Can be a batch of images of dimensions (N, C, H, W) or a single image of dimensions (C, H, W). + size_divisor (`int`): + The width and height of the image will be padded to a multiple of this number. + """ + height, width = image.shape[-2:] + + def _get_pad(size, size_divisor): + new_size = math.ceil(size / size_divisor) * size_divisor + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + pad_top, pad_bottom = _get_pad(height, size_divisor) + pad_left, pad_right = _get_pad(width, size_divisor) + padding = (pad_left, pad_top, pad_right, pad_bottom) + return tvF.pad(image, padding) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_reduce_labels: bool, + 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, + keep_aspect_ratio: bool, + ensure_multiple_of: int | None, + do_pad: bool, + size_divisor: int | None, + disable_grouping: bool | None, + **kwargs, + ) -> BatchFeature: + if do_reduce_labels: + images = self.reduce_label(images) + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_images, + size=size, + resample=resample, + ensure_multiple_of=ensure_multiple_of, + keep_aspect_ratio=keep_aspect_ratio, + ) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + if do_pad: + stacked_images = self.pad_image(stacked_images, size_divisor) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + return processed_images + + def post_process_depth_estimation( + self, + outputs: "DepthEstimatorOutput", + target_sizes: TensorType | list[tuple[int, int]] | None | None = None, + ) -> list[dict[str, TensorType]]: + """ + Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions and depth PIL images. + Only supports PyTorch. + + Args: + outputs ([`DepthEstimatorOutput`]): + Raw outputs of the model. + target_sizes (`TensorType` 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. + + Returns: + `List[Dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth + predictions. + """ + requires_backends(self, "torch") + + predicted_depth = outputs.predicted_depth + + if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth" + ) + + results = [] + target_sizes = [None] * len(predicted_depth) if target_sizes is None else target_sizes + for depth, target_size in zip(predicted_depth, target_sizes): + if target_size is not None: + depth = torch.nn.functional.interpolate( + depth.unsqueeze(0).unsqueeze(1), size=target_size, mode="bicubic", align_corners=False + ).squeeze() + + results.append({"predicted_depth": depth}) + + return results + + +__all__ = ["DPTImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/efficientloftr/__init__.py b/third_party/transformers/src/transformers/models/efficientloftr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e9decbfcbc409d40883341b1dbbef8de05aeb087 --- /dev/null +++ b/third_party/transformers/src/transformers/models/efficientloftr/__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_efficientloftr import * + from .image_processing_efficientloftr import * + from .image_processing_pil_efficientloftr import * + from .modeling_efficientloftr 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/efficientloftr/configuration_efficientloftr.py b/third_party/transformers/src/transformers/models/efficientloftr/configuration_efficientloftr.py new file mode 100644 index 0000000000000000000000000000000000000000..0526d95ca37261194cee257f12e88bb09e83b00a --- /dev/null +++ b/third_party/transformers/src/transformers/models/efficientloftr/configuration_efficientloftr.py @@ -0,0 +1,134 @@ +# 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="zju-community/efficientloftr") +@strict +class EfficientLoFTRConfig(PreTrainedConfig): + r""" + stage_num_blocks (`List`, *optional*, defaults to [1, 2, 4, 14]): + The number of blocks in each stages + stage_stride (`List`, *optional*, defaults to [2, 1, 2, 2]): + The stride used in each stage + q_aggregation_kernel_size (`int`, *optional*, defaults to 4): + The kernel size of the aggregation of query states in the fusion network + kv_aggregation_kernel_size (`int`, *optional*, defaults to 4): + The kernel size of the aggregation of key and value states in the fusion network + q_aggregation_stride (`int`, *optional*, defaults to 4): + The stride of the aggregation of query states in the fusion network + kv_aggregation_stride (`int`, *optional*, defaults to 4): + The stride of the aggregation of key and value states in the fusion network + num_attention_layers (`int`, *optional*, defaults to 4): + Number of attention layers in the LocalFeatureTransformer + mlp_activation_function (`str`, *optional*, defaults to `"leaky_relu"`): + Activation function used in the attention mlp layer. + coarse_matching_skip_softmax (`bool`, *optional*, defaults to `False`): + Whether to skip softmax or not at the coarse matching step. + coarse_matching_threshold (`float`, *optional*, defaults to 0.2): + The threshold for the minimum score required for a match. + coarse_matching_temperature (`float`, *optional*, defaults to 0.1): + The temperature to apply to the coarse similarity matrix + coarse_matching_border_removal (`int`, *optional*, defaults to 2): + The size of the border to remove during coarse matching + fine_kernel_size (`int`, *optional*, defaults to 8): + Kernel size used for the fine feature matching + batch_norm_eps (`float`, *optional*, defaults to 1e-05): + The epsilon used by the batch normalization layers + fine_matching_slice_dim (`int`, *optional*, defaults to 8): + The size of the slice used to divide the fine features for the first and second fine matching stages. + fine_matching_regress_temperature (`float`, *optional*, defaults to 10.0): + The temperature to apply to the fine similarity matrix + + Examples: + ```python + >>> from transformers import EfficientLoFTRConfig, EfficientLoFTRForKeypointMatching + + >>> # Initializing a EfficientLoFTR configuration + >>> configuration = EfficientLoFTRConfig() + + >>> # Initializing a model from the EfficientLoFTR configuration + >>> model = EfficientLoFTRForKeypointMatching(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "efficientloftr" + + stage_num_blocks: list[int] | None = None + out_features: list[int] | None = None + stage_stride: list[int] | None = None + hidden_size: int = 256 + activation_function: str = "relu" + q_aggregation_kernel_size: int = 4 + kv_aggregation_kernel_size: int = 4 + q_aggregation_stride: int = 4 + kv_aggregation_stride: int = 4 + num_attention_layers: int = 4 + num_attention_heads: int = 8 + attention_dropout: float | int = 0.0 + attention_bias: bool = False + mlp_activation_function: str = "leaky_relu" + coarse_matching_skip_softmax: bool = False + coarse_matching_threshold: float = 0.2 + coarse_matching_temperature: float = 0.1 + coarse_matching_border_removal: int = 2 + fine_kernel_size: int = 8 + batch_norm_eps: float = 1e-5 + rope_parameters: dict | None = None + fine_matching_slice_dim: int = 8 + fine_matching_regress_temperature: float = 10.0 + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + # Stage level of RepVGG + self.stage_num_blocks = self.stage_num_blocks if self.stage_num_blocks is not None else [1, 2, 4, 14] + self.stage_stride = self.stage_stride if self.stage_stride is not None else [2, 1, 2, 2] + self.out_features = self.out_features if self.out_features is not None else [64, 64, 128, 256] + self.stage_in_channels = [1] + self.out_features[:-1] + + # Block level of RepVGG + self.stage_block_stride = [ + [stride] + [1] * (num_blocks - 1) for stride, num_blocks in zip(self.stage_stride, self.stage_num_blocks) + ] + self.stage_block_out_channels = [ + [self.out_features[stage_idx]] * num_blocks for stage_idx, num_blocks in enumerate(self.stage_num_blocks) + ] + self.stage_block_in_channels = [ + [self.stage_in_channels[stage_idx]] + self.stage_block_out_channels[stage_idx][:-1] + for stage_idx in range(len(self.stage_num_blocks)) + ] + + self.num_key_value_heads = self.num_attention_heads + self.fine_fusion_dims = list(reversed(self.out_features))[:-1] + self.intermediate_size = self.hidden_size * 2 + kwargs.setdefault("partial_rotary_factor", 4.0) # assign default for BC + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size != self.out_features[-1]: + raise ValueError( + f"hidden_size should be equal to the last value in out_features. hidden_size = {self.hidden_size}, out_features = {self.out_features[-1]}" + ) + + +__all__ = ["EfficientLoFTRConfig"] diff --git a/third_party/transformers/src/transformers/models/efficientloftr/image_processing_efficientloftr.py b/third_party/transformers/src/transformers/models/efficientloftr/image_processing_efficientloftr.py new file mode 100644 index 0000000000000000000000000000000000000000..7f2366860df2e5f7897bd1fc365c0ec0e384b358 --- /dev/null +++ b/third_party/transformers/src/transformers/models/efficientloftr/image_processing_efficientloftr.py @@ -0,0 +1,302 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/efficientloftr/modular_efficientloftr.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_efficientloftr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +from typing import TYPE_CHECKING + +from PIL import Image, ImageDraw +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 ( + ImageInput, + ImageType, + PILImageResampling, + SizeDict, + get_image_type, + is_pil_image, + is_valid_image, + to_numpy_array, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from .modeling_efficientloftr import EfficientLoFTRKeypointMatchingOutput + + +class EfficientLoFTRImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`): + Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method. + """ + + do_grayscale: bool + + +def _is_valid_image(image): + return is_pil_image(image) or ( + is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3 + ) + + +def validate_and_format_image_pairs(images: ImageInput): + error_message = ( + "Input images must be a one of the following :", + " - A pair of PIL images.", + " - A pair of 3D arrays.", + " - A list of pairs of PIL images.", + " - A list of pairs of 3D arrays.", + ) + + def _is_valid_image(image): + """images is a PIL Image or a 3D array.""" + return is_pil_image(image) or ( + is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3 + ) + + if isinstance(images, list): + if len(images) == 2 and all((_is_valid_image(image)) for image in images): + return images + if all( + isinstance(image_pair, list) + and len(image_pair) == 2 + and all(_is_valid_image(image) for image in image_pair) + for image_pair in images + ): + return [image for image_pair in images for image in image_pair] + raise ValueError(error_message) + + +def is_grayscale( + image: "torch.Tensor", +): + """Checks if an image is grayscale (all RGB channels are identical).""" + if image.ndim < 3 or image.shape[0 if image.ndim == 3 else 1] == 1: + return True + return torch.all(image[..., 0, :, :] == image[..., 1, :, :]) and torch.all( + image[..., 1, :, :] == image[..., 2, :, :] + ) + + +def convert_to_grayscale( + image: "torch.Tensor", +) -> "torch.Tensor": + """ + Converts an image to grayscale format using the NTSC formula. Only support torch.Tensor. + + This function is supposed to return a 1-channel image, but it returns a 3-channel image with the same value in each + channel, because of an issue that is discussed in : + https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446 + + Args: + image (torch.Tensor): + The image to convert. + """ + if is_grayscale(image): + return image + return tvF.rgb_to_grayscale(image, num_output_channels=3) + + +@auto_docstring +class EfficientLoFTRImageProcessor(TorchvisionBackend): + valid_kwargs = EfficientLoFTRImageProcessorKwargs + resample = PILImageResampling.BILINEAR + size = {"height": 480, "width": 640} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = None + do_grayscale = True + + def __init__(self, **kwargs: Unpack[EfficientLoFTRImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[EfficientLoFTRImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _prepare_images_structure( + self, + images: ImageInput, + **kwargs, + ) -> ImageInput: + # we need to handle image pairs validation and flattening + images = self.fetch_images(images) + return validate_and_format_image_pairs(images) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + do_grayscale: bool = True, + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + processed_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size=size, resample=resample) + processed_images_grouped[shape] = stacked_images + resized_images = reorder_images(processed_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_rescale: + stacked_images = self.rescale(stacked_images, rescale_factor) + if do_grayscale: + stacked_images = convert_to_grayscale(stacked_images) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + # Convert back to pairs format + image_pairs = [processed_images[i : i + 2] for i in range(0, len(processed_images), 2)] + + # Stack each pair into a single tensor to match slow processor format + stacked_pairs = [torch.stack(pair, dim=0) for pair in image_pairs] + + # Return in same format as slow processor + + return BatchFeature(data={"pixel_values": stacked_pairs}, tensor_type=return_tensors) + + def post_process_keypoint_matching( + self, + outputs: "EfficientLoFTRKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, torch.Tensor]]: + """ + Converts the raw output of [`EfficientLoFTRKeypointMatchingOutput`] into lists of keypoints, scores and descriptors + with coordinates absolute to the original image sizes. + Args: + outputs ([`EfficientLoFTRKeypointMatchingOutput`]): + Raw outputs of the model. + target_sizes (`torch.Tensor` or `List[Tuple[Tuple[int, int]]]`, *optional*): + Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`Tuple[int, int]`) containing the + target size `(height, width)` of each image in the batch. This must be the original image size (before + any processing). + threshold (`float`, *optional*, defaults to 0.0): + Threshold to filter out the matches with low scores. + Returns: + `List[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image + of the pair, the matching scores and the matching indices. + """ + if outputs.matches.shape[0] != len(target_sizes): + raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask") + if not all(len(target_size) == 2 for target_size in target_sizes): + raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch") + + if isinstance(target_sizes, list): + image_pair_sizes = torch.tensor(target_sizes, device=outputs.matches.device) + else: + if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2: + raise ValueError( + "Each element of target_sizes must contain the size (h, w) of each image of the batch" + ) + image_pair_sizes = target_sizes + + keypoints = outputs.keypoints.clone() + keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2) + keypoints = keypoints.to(torch.int32) + + results = [] + for keypoints_pair, matches, scores in zip(keypoints, outputs.matches, outputs.matching_scores): + # Filter out matches with low scores + valid_matches = torch.logical_and(scores > threshold, matches > -1) + + matched_keypoints0 = keypoints_pair[0][valid_matches[0]] + matched_keypoints1 = keypoints_pair[1][valid_matches[1]] + matching_scores = scores[0][valid_matches[0]] + + results.append( + { + "keypoints0": matched_keypoints0, + "keypoints1": matched_keypoints1, + "matching_scores": matching_scores, + } + ) + + return results + + def visualize_keypoint_matching( + self, + images, + keypoint_matching_output: list[dict[str, torch.Tensor]], + ) -> list["Image.Image"]: + """ + Plots the image pairs side by side with the detected keypoints as well as the matching between them. + + Args: + images: + Image pairs to plot. Same as `EfficientLoFTRImageProcessor.preprocess`. Expects either a list of 2 + images or a list of list of 2 images list with pixel values ranging from 0 to 255. + keypoint_matching_output (List[Dict[str, torch.Tensor]]]): + A post processed keypoint matching output + + Returns: + `List[PIL.Image.Image]`: A list of PIL images, each containing the image pairs side by side with the detected + keypoints as well as the matching between them. + """ + + images = validate_and_format_image_pairs(images) + images = [to_numpy_array(image) for image in images] + image_pairs = [images[i : i + 2] for i in range(0, len(images), 2)] + + results = [] + for image_pair, pair_output in zip(image_pairs, keypoint_matching_output): + height0, width0 = image_pair[0].shape[:2] + height1, width1 = image_pair[1].shape[:2] + plot_image = torch.zeros((max(height0, height1), width0 + width1, 3), dtype=torch.uint8) + plot_image[:height0, :width0] = torch.from_numpy(image_pair[0]) + plot_image[:height1, width0:] = torch.from_numpy(image_pair[1]) + + plot_image_pil = Image.fromarray(plot_image.numpy()) + draw = ImageDraw.Draw(plot_image_pil) + + keypoints0_x, keypoints0_y = pair_output["keypoints0"].unbind(1) + keypoints1_x, keypoints1_y = pair_output["keypoints1"].unbind(1) + for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip( + keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, pair_output["matching_scores"] + ): + color = self._get_color(matching_score) + draw.line( + (keypoint0_x, keypoint0_y, keypoint1_x + width0, keypoint1_y), + fill=color, + width=3, + ) + draw.ellipse((keypoint0_x - 2, keypoint0_y - 2, keypoint0_x + 2, keypoint0_y + 2), fill="black") + draw.ellipse( + (keypoint1_x + width0 - 2, keypoint1_y - 2, keypoint1_x + width0 + 2, keypoint1_y + 2), + fill="black", + ) + + results.append(plot_image_pil) + return results + + def _get_color(self, score): + """Maps a score to a color.""" + r = int(255 * (1 - score)) + g = int(255 * score) + b = 0 + return r, g, b + + +__all__ = ["EfficientLoFTRImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/efficientloftr/modeling_efficientloftr.py b/third_party/transformers/src/transformers/models/efficientloftr/modeling_efficientloftr.py new file mode 100644 index 0000000000000000000000000000000000000000..e1b5f1fb40a8d9da1b7e06ec0cd5f888c5968e0e --- /dev/null +++ b/third_party/transformers/src/transformers/models/efficientloftr/modeling_efficientloftr.py @@ -0,0 +1,1390 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2CLS, ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BackboneOutput +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS +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, + can_return_tuple, + torch_int, +) +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_efficientloftr import EfficientLoFTRConfig + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of EfficientLoFTR keypoint matching models. Due to the nature of keypoint detection and matching, the number + of keypoints is not fixed and can vary from image to image, which makes batching non-trivial. In the batch of + images, the maximum number of matches is set as the dimension of the matches and matching scores. + """ +) +class EfficientLoFTRKeypointMatchingOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Loss computed during training. + matches (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Index of keypoint matched in the other image. + matching_scores (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Scores of predicted matches. + keypoints (`torch.FloatTensor` of shape `(batch_size, num_keypoints, 2)`): + Absolute (x, y) coordinates of predicted keypoints in a given image. + hidden_states (`tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, 2, num_channels, + num_keypoints)`, returned when `output_hidden_states=True` is passed or when + `config.output_hidden_states=True`) + attentions (`tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, 2, num_heads, num_keypoints, + num_keypoints)`, returned when `output_attentions=True` is passed or when `config.output_attentions=True`) + """ + + loss: torch.FloatTensor | None = None + matches: torch.FloatTensor | None = None + matching_scores: torch.FloatTensor | None = None + keypoints: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@compile_compatible_method_lru_cache(maxsize=32) +def compute_embeddings(inv_freq: torch.Tensor, embed_height: int, embed_width: int, hidden_size: int) -> torch.Tensor: + i_indices = torch.ones(embed_height, embed_width, dtype=inv_freq.dtype, device=inv_freq.device) + j_indices = torch.ones(embed_height, embed_width, dtype=inv_freq.dtype, device=inv_freq.device) + i_indices = i_indices.cumsum(0).unsqueeze(-1) + j_indices = j_indices.cumsum(1).unsqueeze(-1) + + emb = torch.zeros(1, embed_height, embed_width, hidden_size // 2, dtype=inv_freq.dtype, device=inv_freq.device) + emb[:, :, :, 0::2] = i_indices * inv_freq + emb[:, :, :, 1::2] = j_indices * inv_freq + + return emb + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->EfficientLoFTR +class EfficientLoFTRRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + # Ignore copy + def __init__(self, config: EfficientLoFTRConfig, device=None): + super().__init__() + + 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 + # Ignore copy + def compute_default_rope_parameters( + config: EfficientLoFTRConfig | 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 + + # Ignore copy + @torch.no_grad() + def forward( + self, x: torch.Tensor, position_ids: torch.LongTensor | None = None, layer_type=None + ) -> tuple[torch.Tensor, torch.Tensor]: + feats_height, feats_width = x.shape[-2:] + embed_height = (feats_height - self.config.q_aggregation_kernel_size) // self.config.q_aggregation_stride + 1 + embed_width = (feats_width - self.config.q_aggregation_kernel_size) // self.config.q_aggregation_stride + 1 + 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 + emb = compute_embeddings(self.inv_freq, embed_height, embed_width, self.config.hidden_size) + sin = emb.sin() + cos = emb.cos() + + sin = sin.repeat_interleave(2, dim=-1) + cos = cos.repeat_interleave(2, dim=-1) + + sin = sin.to(device=x.device, dtype=x.dtype) + cos = cos.to(device=x.device, dtype=x.dtype) + + return cos, sin + + +# Copied from transformers.models.rt_detr_v2.modeling_rt_detr_v2.RTDetrV2ConvNormLayer with RTDetrV2->EfficientLoFTR +class EfficientLoFTRConvNormLayer(nn.Module): + def __init__(self, config, in_channels, out_channels, kernel_size, stride, padding=None, activation=None): + super().__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding=(kernel_size - 1) // 2 if padding is None else padding, + bias=False, + ) + self.norm = nn.BatchNorm2d(out_channels, config.batch_norm_eps) + self.activation = nn.Identity() if activation is None else ACT2CLS[activation]() + + def forward(self, hidden_state): + hidden_state = self.conv(hidden_state) + hidden_state = self.norm(hidden_state) + hidden_state = self.activation(hidden_state) + return hidden_state + + +class EfficientLoFTRRepVGGBlock(GradientCheckpointingLayer): + """ + RepVGG architecture block introduced by the work "RepVGG: Making VGG-style ConvNets Great Again". + """ + + def __init__(self, config: EfficientLoFTRConfig, stage_idx: int, block_idx: int): + super().__init__() + in_channels = config.stage_block_in_channels[stage_idx][block_idx] + out_channels = config.stage_block_out_channels[stage_idx][block_idx] + stride = config.stage_block_stride[stage_idx][block_idx] + activation = config.activation_function + self.conv1 = EfficientLoFTRConvNormLayer( + config, in_channels, out_channels, kernel_size=3, stride=stride, padding=1 + ) + self.conv2 = EfficientLoFTRConvNormLayer( + config, in_channels, out_channels, kernel_size=1, stride=stride, padding=0 + ) + self.identity = nn.BatchNorm2d(in_channels) if in_channels == out_channels and stride == 1 else None + self.activation = nn.Identity() if activation is None else ACT2FN[activation] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.identity is not None: + identity_out = self.identity(hidden_states) + else: + identity_out = 0 + hidden_states = self.conv1(hidden_states) + self.conv2(hidden_states) + identity_out + hidden_states = self.activation(hidden_states) + return hidden_states + + +class EfficientLoFTRRepVGGStage(nn.Module): + def __init__(self, config: EfficientLoFTRConfig, stage_idx: int): + super().__init__() + self.blocks = nn.ModuleList([]) + for block_idx in range(config.stage_num_blocks[stage_idx]): + self.blocks.append( + EfficientLoFTRRepVGGBlock( + config, + stage_idx, + block_idx, + ) + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + hidden_states = block(hidden_states) + return hidden_states + + +class EfficientLoFTRepVGG(nn.Module): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__() + + self.stages = nn.ModuleList([]) + + for stage_idx in range(len(config.stage_stride)): + stage = EfficientLoFTRRepVGGStage(config, stage_idx) + self.stages.append(stage) + + def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]: + outputs = [] + for stage in self.stages: + hidden_states = stage(hidden_states) + outputs.append(hidden_states) + + # Exclude first stage in outputs + outputs = outputs[1:] + return outputs + + +class EfficientLoFTRAggregationLayer(nn.Module): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__() + + hidden_size = config.hidden_size + + self.q_aggregation = nn.Conv2d( + hidden_size, + hidden_size, + kernel_size=config.q_aggregation_kernel_size, + padding=0, + stride=config.q_aggregation_stride, + bias=False, + groups=hidden_size, + ) + self.kv_aggregation = torch.nn.MaxPool2d( + kernel_size=config.kv_aggregation_kernel_size, stride=config.kv_aggregation_stride + ) + self.norm = nn.LayerNorm(hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + query_states = hidden_states + is_cross_attention = encoder_hidden_states is not None + kv_states = encoder_hidden_states if is_cross_attention else hidden_states + + query_states = self.q_aggregation(query_states) + kv_states = self.kv_aggregation(kv_states) + query_states = query_states.permute(0, 2, 3, 1) + kv_states = kv_states.permute(0, 2, 3, 1) + hidden_states = self.norm(query_states) + encoder_hidden_states = self.norm(kv_states) + return hidden_states, encoder_hidden_states + + +# Copied from transformers.models.cohere.modeling_cohere.rotate_half +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 + + +# Copied from transformers.models.cohere.modeling_cohere.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + 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) + + +# Copied from transformers.models.cohere.modeling_cohere.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.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: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class EfficientLoFTRAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: EfficientLoFTRConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + batch_size, seq_len, dim = hidden_states.shape + input_shape = hidden_states.shape[:-1] + + query_states = self.q_proj(hidden_states).view(batch_size, seq_len, -1, dim) + + current_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + key_states = self.k_proj(current_states).view(batch_size, seq_len, -1, dim) + value_states = self.v_proj(current_states).view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=2) + + query_states = query_states.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(batch_size, seq_len, -1, 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, + query_states, + key_states, + value_states, + attention_mask=None, + 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 EfficientLoFTRMLP(nn.Module): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__() + hidden_size = config.hidden_size + intermediate_size = config.intermediate_size + self.fc1 = nn.Linear(hidden_size * 2, intermediate_size, bias=False) + self.activation = ACT2FN[config.mlp_activation_function] + self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=False) + self.layer_norm = nn.LayerNorm(hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = self.layer_norm(hidden_states) + return hidden_states + + +class EfficientLoFTRAggregatedAttention(nn.Module): + def __init__(self, config: EfficientLoFTRConfig, layer_idx: int): + super().__init__() + + self.q_aggregation_kernel_size = config.q_aggregation_kernel_size + self.aggregation = EfficientLoFTRAggregationLayer(config) + self.attention = EfficientLoFTRAttention(config, layer_idx) + self.mlp = EfficientLoFTRMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + batch_size, embed_dim, _, _ = hidden_states.shape + + # Aggregate features + aggregated_hidden_states, aggregated_encoder_hidden_states = self.aggregation( + hidden_states, encoder_hidden_states + ) + _, aggregated_h, aggregated_w, _ = aggregated_hidden_states.shape + + # Multi-head attention + aggregated_hidden_states = aggregated_hidden_states.reshape(batch_size, -1, embed_dim) + aggregated_encoder_hidden_states = aggregated_encoder_hidden_states.reshape(batch_size, -1, embed_dim) + attn_output, _ = self.attention( + aggregated_hidden_states, + aggregated_encoder_hidden_states, + position_embeddings=position_embeddings, + **kwargs, + ) + + # Upsample features + # (batch_size, seq_len, embed_dim) -> (batch_size, embed_dim, h, w) with seq_len = h * w + attn_output = attn_output.permute(0, 2, 1) + attn_output = attn_output.reshape(batch_size, embed_dim, aggregated_h, aggregated_w) + attn_output = torch.nn.functional.interpolate( + attn_output, scale_factor=self.q_aggregation_kernel_size, mode="bilinear", align_corners=False + ) + intermediate_states = torch.cat([hidden_states, attn_output], dim=1) + intermediate_states = intermediate_states.permute(0, 2, 3, 1) + output_states = self.mlp(intermediate_states) + output_states = output_states.permute(0, 3, 1, 2) + + hidden_states = hidden_states + output_states + + return hidden_states + + +class EfficientLoFTRLocalFeatureTransformerLayer(GradientCheckpointingLayer): + def __init__(self, config: EfficientLoFTRConfig, layer_idx: int): + super().__init__() + + self.self_attention = EfficientLoFTRAggregatedAttention(config, layer_idx) + self.cross_attention = EfficientLoFTRAggregatedAttention(config, layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + batch_size, _, embed_dim, height, width = hidden_states.shape + + hidden_states = hidden_states.reshape(-1, embed_dim, height, width) + hidden_states = self.self_attention(hidden_states, position_embeddings=position_embeddings, **kwargs) + + ### + # Implementation of a bug in the original implementation regarding the cross-attention + # See : https://github.com/zju3dv/MatchAnything/issues/26 + hidden_states = hidden_states.reshape(-1, 2, embed_dim, height, width) + features_0 = hidden_states[:, 0] + features_1 = hidden_states[:, 1] + features_0 = self.cross_attention(features_0, features_1, **kwargs) + features_1 = self.cross_attention(features_1, features_0, **kwargs) + hidden_states = torch.stack((features_0, features_1), dim=1) + ### + + return hidden_states + + +class EfficientLoFTRLocalFeatureTransformer(nn.Module): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__() + self.layers = nn.ModuleList( + [ + EfficientLoFTRLocalFeatureTransformerLayer(config, layer_idx=i) + for i in range(config.num_attention_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings=position_embeddings, **kwargs) + return hidden_states + + +class EfficientLoFTROutConvBlock(nn.Module): + def __init__(self, config: EfficientLoFTRConfig, hidden_size: int, intermediate_size: int): + super().__init__() + + self.out_conv1 = nn.Conv2d(hidden_size, intermediate_size, kernel_size=1, stride=1, padding=0, bias=False) + self.out_conv2 = nn.Conv2d( + intermediate_size, intermediate_size, kernel_size=3, stride=1, padding=1, bias=False + ) + self.batch_norm = nn.BatchNorm2d(intermediate_size) + self.activation = ACT2CLS[config.mlp_activation_function]() + self.out_conv3 = nn.Conv2d(intermediate_size, hidden_size, kernel_size=3, stride=1, padding=1, bias=False) + + def forward(self, hidden_states: torch.Tensor, residual_states: torch.Tensor) -> torch.Tensor: + residual_states = self.out_conv1(residual_states) + residual_states = residual_states + hidden_states + residual_states = self.out_conv2(residual_states) + residual_states = self.batch_norm(residual_states) + residual_states = self.activation(residual_states) + residual_states = self.out_conv3(residual_states) + residual_states = nn.functional.interpolate( + residual_states, scale_factor=2.0, mode="bilinear", align_corners=False + ) + return residual_states + + +class EfficientLoFTRFineFusionLayer(nn.Module): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__() + + self.fine_kernel_size = config.fine_kernel_size + + fine_fusion_dims = config.fine_fusion_dims + self.out_conv = nn.Conv2d( + fine_fusion_dims[0], fine_fusion_dims[0], kernel_size=1, stride=1, padding=0, bias=False + ) + self.out_conv_layers = nn.ModuleList() + for i in range(1, len(fine_fusion_dims)): + out_conv = EfficientLoFTROutConvBlock(config, fine_fusion_dims[i], fine_fusion_dims[i - 1]) + self.out_conv_layers.append(out_conv) + + def forward_pyramid( + self, + hidden_states: torch.Tensor, + residual_states: list[torch.Tensor], + ) -> torch.Tensor: + hidden_states = self.out_conv(hidden_states) + hidden_states = nn.functional.interpolate( + hidden_states, scale_factor=2.0, mode="bilinear", align_corners=False + ) + for i, layer in enumerate(self.out_conv_layers): + hidden_states = layer(hidden_states, residual_states[i]) + + return hidden_states + + def forward( + self, + coarse_features: torch.Tensor, + residual_features: list[torch.Tensor] | tuple[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + For each image pair, compute the fine features of pixels. + In both images, compute a patch of fine features center cropped around each coarse pixel. + In the first image, the feature patch is kernel_size large and long. + In the second image, it is (kernel_size + 2) large and long. + """ + batch_size, _, embed_dim, coarse_height, coarse_width = coarse_features.shape + + coarse_features = coarse_features.reshape(-1, embed_dim, coarse_height, coarse_width) + residual_features = list(reversed(residual_features)) + + # 1. Fine feature extraction + fine_features = self.forward_pyramid(coarse_features, residual_features) + _, fine_embed_dim, fine_height, fine_width = fine_features.shape + + fine_features = fine_features.reshape(batch_size, 2, fine_embed_dim, fine_height, fine_width) + fine_features_0 = fine_features[:, 0] + fine_features_1 = fine_features[:, 1] + + # 2. Unfold all local windows in crops + stride = int(fine_height // coarse_height) + fine_features_0 = nn.functional.unfold( + fine_features_0, kernel_size=self.fine_kernel_size, stride=stride, padding=0 + ) + _, _, seq_len = fine_features_0.shape + fine_features_0 = fine_features_0.reshape(batch_size, -1, self.fine_kernel_size**2, seq_len) + fine_features_0 = fine_features_0.permute(0, 3, 2, 1) + + fine_features_1 = nn.functional.unfold( + fine_features_1, kernel_size=self.fine_kernel_size + 2, stride=stride, padding=1 + ) + fine_features_1 = fine_features_1.reshape(batch_size, -1, (self.fine_kernel_size + 2) ** 2, seq_len) + fine_features_1 = fine_features_1.permute(0, 3, 2, 1) + + return fine_features_0, fine_features_1 + + +@auto_docstring +class EfficientLoFTRPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = EfficientLoFTRConfig + base_model_prefix = "efficientloftr" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _can_record_outputs = { + "hidden_states": EfficientLoFTRRepVGGBlock, + "attentions": EfficientLoFTRAttention, + } + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + """Initialize the weights""" + if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d, nn.BatchNorm2d)): + 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, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + elif isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, EfficientLoFTRRotaryEmbedding): + rope_fn = ( + ROPE_INIT_FUNCTIONS[module.rope_type] + if module.rope_type != "default" + else module.compute_default_rope_parameters + ) + buffer_value, _ = rope_fn(module.config) + init.copy_(module.inv_freq, buffer_value) + init.copy_(module.original_inv_freq, buffer_value) + + # Copied from transformers.models.superpoint.modeling_superpoint.SuperPointPreTrainedModel.extract_one_channel_pixel_values with SuperPoint->EfficientLoFTR + def extract_one_channel_pixel_values(self, pixel_values: torch.FloatTensor) -> torch.FloatTensor: + """ + Assuming pixel_values has shape (batch_size, 3, height, width), and that all channels values are the same, + extract the first channel value to get a tensor of shape (batch_size, 1, height, width) for EfficientLoFTR. This is + a workaround for the issue discussed in : + https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446 + + Args: + pixel_values: torch.FloatTensor of shape (batch_size, 3, height, width) + + Returns: + pixel_values: torch.FloatTensor of shape (batch_size, 1, height, width) + + """ + return pixel_values[:, 0, :, :][:, None, :, :] + + +@auto_docstring( + custom_intro=""" + EfficientLoFTR model taking images as inputs and outputting the features of the images. + """ +) +class EfficientLoFTRModel(EfficientLoFTRPreTrainedModel): + def __init__(self, config: EfficientLoFTRConfig): + super().__init__(config) + + self.config = config + self.backbone = EfficientLoFTRepVGG(config) + self.local_feature_transformer = EfficientLoFTRLocalFeatureTransformer(config) + self.rotary_emb = EfficientLoFTRRotaryEmbedding(config=config) + + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BackboneOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModel + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_78916675_4568141288.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image1 = Image.open(BytesIO(response.read())) + + >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image2 = Image.open(BytesIO(response.read())) + + >>> images = [image1, image2] + + >>> processor = AutoImageProcessor.from_pretrained("zju-community/efficient_loftr") + >>> model = AutoModel.from_pretrained("zju-community/efficient_loftr") + + >>> with torch.no_grad(): + >>> inputs = processor(images, return_tensors="pt") + >>> outputs = model(**inputs) + ```""" + if labels is not None: + raise ValueError("EfficientLoFTR is not trainable, no labels should be provided.") + + if pixel_values.ndim != 5 or pixel_values.size(1) != 2: + raise ValueError("Input must be a 5D tensor of shape (batch_size, 2, num_channels, height, width)") + + batch_size, _, channels, height, width = pixel_values.shape + pixel_values = pixel_values.reshape(batch_size * 2, channels, height, width) + pixel_values = self.extract_one_channel_pixel_values(pixel_values) + + # 1. Local Feature CNN + features = self.backbone(pixel_values) + # Last stage outputs are coarse outputs + coarse_features = features[-1] + # Rest is residual features used in EfficientLoFTRFineFusionLayer + residual_features = features[:-1] + coarse_embed_dim, coarse_height, coarse_width = coarse_features.shape[-3:] + + # 2. Coarse-level LoFTR module + cos, sin = self.rotary_emb(coarse_features) + cos = cos.expand(batch_size * 2, -1, -1, -1).reshape(batch_size * 2, -1, coarse_embed_dim) + sin = sin.expand(batch_size * 2, -1, -1, -1).reshape(batch_size * 2, -1, coarse_embed_dim) + position_embeddings = (cos, sin) + + coarse_features = coarse_features.reshape(batch_size, 2, coarse_embed_dim, coarse_height, coarse_width) + coarse_features = self.local_feature_transformer( + coarse_features, position_embeddings=position_embeddings, **kwargs + ) + features = (coarse_features,) + tuple(residual_features) + + return BackboneOutput(feature_maps=features) + + +def mask_border(tensor: torch.Tensor, border_margin: int, value: bool | float | int) -> torch.Tensor: + """ + Mask a tensor border with a given value + + Args: + tensor (`torch.Tensor` of shape `(batch_size, height_0, width_0, height_1, width_1)`): + The tensor to mask + border_margin (`int`) : + The size of the border + value (`Union[bool, int, float]`): + The value to place in the tensor's borders + + Returns: + tensor (`torch.Tensor` of shape `(batch_size, height_0, width_0, height_1, width_1)`): + The masked tensor + """ + if border_margin <= 0: + return tensor + + tensor[:, :border_margin] = value + tensor[:, :, :border_margin] = value + tensor[:, :, :, :border_margin] = value + tensor[:, :, :, :, :border_margin] = value + tensor[:, -border_margin:] = value + tensor[:, :, -border_margin:] = value + tensor[:, :, :, -border_margin:] = value + tensor[:, :, :, :, -border_margin:] = value + + return tensor + + +def create_meshgrid( + height: int | torch.Tensor, + width: int | torch.Tensor, + normalized_coordinates: bool = False, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """ + Copied from kornia library : kornia/kornia/utils/grid.py:26 + + Generate a coordinate grid for an image. + + When the flag ``normalized_coordinates`` is set to True, the grid is + normalized to be in the range :math:`[-1,1]` to be consistent with the pytorch + function :py:func:`torch.nn.functional.grid_sample`. + + Args: + height (`int`): + The image height (rows). + width (`int`): + The image width (cols). + normalized_coordinates (`bool`): + Whether to normalize coordinates in the range :math:`[-1,1]` in order to be consistent with the + PyTorch function :py:func:`torch.nn.functional.grid_sample`. + device (`torch.device`): + The device on which the grid will be generated. + dtype (`torch.dtype`): + The data type of the generated grid. + + Return: + grid (`torch.Tensor` of shape `(1, height, width, 2)`): + The grid tensor. + + Example: + >>> create_meshgrid(2, 2) + tensor([[[[-1., -1.], + [ 1., -1.]], + + [[-1., 1.], + [ 1., 1.]]]]) + + >>> create_meshgrid(2, 2, normalized_coordinates=False) + tensor([[[[0., 0.], + [1., 0.]], + + [[0., 1.], + [1., 1.]]]]) + + """ + xs = torch.linspace(0, width - 1, width, device=device, dtype=dtype) + ys = torch.linspace(0, height - 1, height, device=device, dtype=dtype) + if normalized_coordinates: + xs = (xs / (width - 1) - 0.5) * 2 + ys = (ys / (height - 1) - 0.5) * 2 + grid = torch.stack(torch.meshgrid(ys, xs, indexing="ij"), dim=-1) + grid = grid.permute(1, 0, 2).unsqueeze(0) + return grid + + +def spatial_expectation2d(input: torch.Tensor, normalized_coordinates: bool = True) -> torch.Tensor: + r""" + Copied from kornia library : kornia/geometry/subpix/dsnt.py:76 + Compute the expectation of coordinate values using spatial probabilities. + + The input heatmap is assumed to represent a valid spatial probability distribution, + which can be achieved using :func:`~kornia.geometry.subpixel.spatial_softmax2d`. + + Args: + input (`torch.Tensor` of shape `(batch_size, embed_dim, height, width)`): + The input tensor representing dense spatial probabilities. + normalized_coordinates (`bool`): + Whether to return the coordinates normalized in the range of :math:`[-1, 1]`. Otherwise, it will return + the coordinates in the range of the input shape. + + Returns: + output (`torch.Tensor` of shape `(batch_size, embed_dim, 2)`) + Expected value of the 2D coordinates. Output order of the coordinates is (x, y). + + Examples: + >>> heatmaps = torch.tensor([[[ + ... [0., 0., 0.], + ... [0., 0., 0.], + ... [0., 1., 0.]]]]) + >>> spatial_expectation2d(heatmaps, False) + tensor([[[1., 2.]]]) + + """ + batch_size, embed_dim, height, width = input.shape + + # Create coordinates grid. + grid = create_meshgrid(height, width, normalized_coordinates, input.device) + grid = grid.to(input.dtype) + + pos_x = grid[..., 0].reshape(-1) + pos_y = grid[..., 1].reshape(-1) + + input_flat = input.view(batch_size, embed_dim, -1) + + # Compute the expectation of the coordinates. + expected_y = torch.sum(pos_y * input_flat, -1, keepdim=True) + expected_x = torch.sum(pos_x * input_flat, -1, keepdim=True) + + output = torch.cat([expected_x, expected_y], -1) + + return output.view(batch_size, embed_dim, 2) + + +@auto_docstring( + custom_intro=""" + EfficientLoFTR model taking images as inputs and outputting the matching of them. + """ +) +class EfficientLoFTRForKeypointMatching(EfficientLoFTRPreTrainedModel): + """EfficientLoFTR dense image matcher + + Given two images, we determine the correspondences by: + 1. Extracting coarse and fine features through a backbone + 2. Transforming coarse features through self and cross attention + 3. Matching coarse features to obtain coarse coordinates of matches + 4. Obtaining full resolution fine features by fusing transformed and backbone coarse features + 5. Refining the coarse matches using fine feature patches centered at each coarse match in a two-stage refinement + + Yifan Wang, Xingyi He, Sida Peng, Dongli Tan and Xiaowei Zhou. + Efficient LoFTR: Semi-Dense Local Feature Matching with Sparse-Like Speed + In CVPR, 2024. https://huggingface.co/papers/2403.04765 + """ + + def __init__(self, config: EfficientLoFTRConfig): + super().__init__(config) + + self.config = config + self.efficientloftr = EfficientLoFTRModel(config) + self.refinement_layer = EfficientLoFTRFineFusionLayer(config) + + self.post_init() + + def _get_matches_from_scores(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Based on a keypoint score matrix, compute the best keypoint matches between the first and second image. + Since each image pair can have different number of matches, the matches are concatenated together for all pair + in the batch and a batch_indices tensor is returned to specify which match belong to which element in the batch. + + Note: + This step can be done as a postprocessing step, because does not involve any model weights/params. + However, we keep it in the modeling code for consistency with other keypoint matching models AND for + easier torch.compile/torch.export (all ops are in torch). + + Args: + scores (`torch.Tensor` of shape `(batch_size, height_0, width_0, height_1, width_1)`): + Scores of keypoints + + Returns: + matched_indices (`torch.Tensor` of shape `(2, num_matches)`): + Indices representing which pixel in the first image matches which pixel in the second image + matching_scores (`torch.Tensor` of shape `(num_matches,)`): + Scores of each match + """ + batch_size, height0, width0, height1, width1 = scores.shape + + scores = scores.view(batch_size, height0 * width0, height1 * width1) + + # For each keypoint, get the best match + max_0 = scores.max(2, keepdim=True).values + max_1 = scores.max(1, keepdim=True).values + + # 1. Thresholding + mask = scores > self.config.coarse_matching_threshold + + # 2. Border removal + mask = mask.reshape(batch_size, height0, width0, height1, width1) + mask = mask_border(mask, self.config.coarse_matching_border_removal, False) + mask = mask.reshape(batch_size, height0 * width0, height1 * width1) + + # 3. Mutual nearest neighbors + mask = mask * (scores == max_0) * (scores == max_1) + + # 4. Fine coarse matches + masked_scores = scores * mask + matching_scores_0, max_indices_0 = masked_scores.max(1) + matching_scores_1, max_indices_1 = masked_scores.max(2) + + matching_indices = torch.cat([max_indices_0, max_indices_1]).reshape(batch_size, 2, -1) + matching_scores = torch.stack([matching_scores_0, matching_scores_1], dim=1) + + # For the keypoints not meeting the threshold score, set the indices to -1 which corresponds to no matches found + matching_indices = torch.where(matching_scores > 0, matching_indices, -1) + + return matching_indices, matching_scores + + def _coarse_matching( + self, coarse_features: torch.Tensor, coarse_scale: float + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + For each image pair, compute the matching confidence between each coarse element (by default (image_height / 8) + * (image_width / 8 elements)) from the first image to the second image. + + Note: + This step can be done as a postprocessing step, because does not involve any model weights/params. + However, we keep it in the modeling code for consistency with other keypoint matching models AND for + easier torch.compile/torch.export (all ops are in torch). + + Args: + coarse_features (`torch.Tensor` of shape `(batch_size, 2, hidden_size, coarse_height, coarse_width)`): + Coarse features + coarse_scale (`float`): Scale between the image size and the coarse size + + Returns: + keypoints (`torch.Tensor` of shape `(batch_size, 2, num_matches, 2)`): + Keypoints coordinates. + matching_scores (`torch.Tensor` of shape `(batch_size, 2, num_matches)`): + The confidence matching score of each keypoint. + matched_indices (`torch.Tensor` of shape `(batch_size, 2, num_matches)`): + Indices which indicates which keypoint in an image matched with which keypoint in the other image. For + both image in the pair. + """ + batch_size, _, embed_dim, height, width = coarse_features.shape + + # (batch_size, 2, embed_dim, height, width) -> (batch_size, 2, height * width, embed_dim) + coarse_features = coarse_features.permute(0, 1, 3, 4, 2) + coarse_features = coarse_features.reshape(batch_size, 2, -1, embed_dim) + + coarse_features = coarse_features / coarse_features.shape[-1] ** 0.5 + coarse_features_0 = coarse_features[:, 0] + coarse_features_1 = coarse_features[:, 1] + + similarity = coarse_features_0 @ coarse_features_1.transpose(-1, -2) + similarity = similarity / self.config.coarse_matching_temperature + + if self.config.coarse_matching_skip_softmax: + confidence = similarity + else: + confidence = nn.functional.softmax(similarity, 1) * nn.functional.softmax(similarity, 2) + + confidence = confidence.view(batch_size, height, width, height, width) + matched_indices, matching_scores = self._get_matches_from_scores(confidence) + + keypoints = torch.stack([matched_indices % width, matched_indices // width], dim=-1) * coarse_scale + + return keypoints, matching_scores, matched_indices + + def _get_first_stage_fine_matching( + self, + fine_confidence: torch.Tensor, + coarse_matched_keypoints: torch.Tensor, + fine_window_size: int, + fine_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + For each coarse pixel, retrieve the highest fine confidence score and index. + The index represents the matching between a pixel position in the fine window in the first image and a pixel + position in the fine window of the second image. + For example, for a fine_window_size of 64 (8 * 8), the index 2474 represents the matching between the index 38 + (2474 // 64) in the fine window of the first image, and the index 42 in the second image. This means that 38 + which corresponds to the position (4, 6) (4 // 8 and 4 % 8) is matched with the position (5, 2). In this example + the coarse matched coordinate will be shifted to the matched fine coordinates in the first and second image. + + Note: + This step can be done as a postprocessing step, because does not involve any model weights/params. + However, we keep it in the modeling code for consistency with other keypoint matching models AND for + easier torch.compile/torch.export (all ops are in torch). + + Args: + fine_confidence (`torch.Tensor` of shape `(num_matches, fine_window_size, fine_window_size)`): + First stage confidence of matching fine features between the first and the second image + coarse_matched_keypoints (`torch.Tensor` of shape `(2, num_matches, 2)`): + Coarse matched keypoint between the first and the second image. + fine_window_size (`int`): + Size of the window used to refine matches + fine_scale (`float`): + Scale between the size of fine features and coarse features + + Returns: + indices (`torch.Tensor` of shape `(2, num_matches, 1)`): + Indices of the fine coordinate matched in the fine window + fine_matches (`torch.Tensor` of shape `(2, num_matches, 2)`): + Coordinates of matched keypoints after the first fine stage + """ + batch_size, num_keypoints, _, _ = fine_confidence.shape + fine_kernel_size = torch_int(fine_window_size**0.5) + + fine_confidence = fine_confidence.reshape(batch_size, num_keypoints, -1) + values, indices = torch.max(fine_confidence, dim=-1) + indices = indices[..., None] + indices_0 = indices // fine_window_size + indices_1 = indices % fine_window_size + + grid = create_meshgrid( + fine_kernel_size, + fine_kernel_size, + normalized_coordinates=False, + device=fine_confidence.device, + dtype=fine_confidence.dtype, + ) + grid = grid - (fine_kernel_size // 2) + 0.5 + grid = grid.reshape(1, 1, -1, 2).expand(batch_size, num_keypoints, -1, -1) + delta_0 = torch.gather(grid, 1, indices_0.unsqueeze(-1).expand(-1, -1, -1, 2)).squeeze(2) + delta_1 = torch.gather(grid, 1, indices_1.unsqueeze(-1).expand(-1, -1, -1, 2)).squeeze(2) + + fine_matches_0 = coarse_matched_keypoints[:, 0] + delta_0 * fine_scale + fine_matches_1 = coarse_matched_keypoints[:, 1] + delta_1 * fine_scale + + indices = torch.stack([indices_0, indices_1], dim=1) + fine_matches = torch.stack([fine_matches_0, fine_matches_1], dim=1) + + return indices, fine_matches + + def _get_second_stage_fine_matching( + self, + indices: torch.Tensor, + fine_matches: torch.Tensor, + fine_confidence: torch.Tensor, + fine_window_size: int, + fine_scale: float, + ) -> torch.Tensor: + """ + For the given position in their respective fine windows, retrieve the 3x3 fine confidences around this position. + After applying softmax to these confidences, compute the 2D spatial expected coordinates. + Shift the first stage fine matching with these expected coordinates. + + Note: + This step can be done as a postprocessing step, because does not involve any model weights/params. + However, we keep it in the modeling code for consistency with other keypoint matching models AND for + easier torch.compile/torch.export (all ops are in torch). + + Args: + indices (`torch.Tensor` of shape `(batch_size, 2, num_keypoints)`): + Indices representing the position of each keypoint in the fine window + fine_matches (`torch.Tensor` of shape `(2, num_matches, 2)`): + Coordinates of matched keypoints after the first fine stage + fine_confidence (`torch.Tensor` of shape `(num_matches, fine_window_size, fine_window_size)`): + Second stage confidence of matching fine features between the first and the second image + fine_window_size (`int`): + Size of the window used to refine matches + fine_scale (`float`): + Scale between the size of fine features and coarse features + + Returns: + fine_matches (`torch.Tensor` of shape `(2, num_matches, 2)`): + Coordinates of matched keypoints after the second fine stage + """ + batch_size, num_keypoints, _, _ = fine_confidence.shape + fine_kernel_size = torch_int(fine_window_size**0.5) + + indices_0 = indices[:, 0] + indices_1 = indices[:, 1] + indices_1_i = indices_1 // fine_kernel_size + indices_1_j = indices_1 % fine_kernel_size + + # matches_indices, indices_0, indices_1_i, indices_1_j of shape (num_matches, 3, 3) + batch_indices = torch.arange(batch_size, device=indices_0.device).reshape(batch_size, 1, 1, 1) + matches_indices = torch.arange(num_keypoints, device=indices_0.device).reshape(1, num_keypoints, 1, 1) + indices_0 = indices_0[..., None] + indices_1_i = indices_1_i[..., None] + indices_1_j = indices_1_j[..., None] + + delta = create_meshgrid(3, 3, normalized_coordinates=True, device=indices_0.device).to(torch.long) + delta = delta[None, ...] + + indices_1_i = indices_1_i + delta[..., 1] + indices_1_j = indices_1_j + delta[..., 0] + + fine_confidence = fine_confidence.reshape( + batch_size, num_keypoints, fine_window_size, fine_kernel_size + 2, fine_kernel_size + 2 + ) + # (batch_size, seq_len, fine_window_size, fine_kernel_size + 2, fine_kernel_size + 2) -> (batch_size, seq_len, 3, 3) + fine_confidence = fine_confidence[batch_indices, matches_indices, indices_0, indices_1_i, indices_1_j] + fine_confidence = fine_confidence.reshape(batch_size, num_keypoints, 9) + fine_confidence = nn.functional.softmax( + fine_confidence / self.config.fine_matching_regress_temperature, dim=-1 + ) + + heatmap = fine_confidence.reshape(batch_size, num_keypoints, 3, 3) + fine_coordinates_normalized = spatial_expectation2d(heatmap, True)[0] + + fine_matches_0 = fine_matches[:, 0] + fine_matches_1 = fine_matches[:, 1] + (fine_coordinates_normalized * (3 // 2) * fine_scale) + + fine_matches = torch.stack([fine_matches_0, fine_matches_1], dim=1) + + return fine_matches + + def _fine_matching( + self, + fine_features_0: torch.Tensor, + fine_features_1: torch.Tensor, + coarse_matched_keypoints: torch.Tensor, + fine_scale: float, + ) -> torch.Tensor: + """ + For each coarse pixel with a corresponding window of fine features, compute the matching confidence between fine + features in the first image and the second image. + + Fine features are sliced in two part : + - The first part used for the first stage are the first fine_hidden_size - config.fine_matching_slicedim (64 - 8 + = 56 by default) features. + - The second part used for the second stage are the last config.fine_matching_slicedim (8 by default) features. + + Each part is used to compute a fine confidence tensor of the following shape : + (batch_size, (coarse_height * coarse_width), fine_window_size, fine_window_size) + They correspond to the score between each fine pixel in the first image and each fine pixel in the second image. + + Args: + fine_features_0 (`torch.Tensor` of shape `(num_matches, fine_kernel_size ** 2, fine_kernel_size ** 2)`): + Fine features from the first image + fine_features_1 (`torch.Tensor` of shape `(num_matches, (fine_kernel_size + 2) ** 2, (fine_kernel_size + 2) + ** 2)`): + Fine features from the second image + coarse_matched_keypoints (`torch.Tensor` of shape `(2, num_matches, 2)`): + Keypoint coordinates found in coarse matching for the first and second image + fine_scale (`int`): + Scale between the size of fine features and coarse features + + Returns: + fine_coordinates (`torch.Tensor` of shape `(2, num_matches, 2)`): + Matched keypoint between the first and the second image. All matched keypoints are concatenated in the + second dimension. + + """ + batch_size, num_keypoints, fine_window_size, fine_embed_dim = fine_features_0.shape + fine_matching_slice_dim = self.config.fine_matching_slice_dim + + fine_kernel_size = torch_int(fine_window_size**0.5) + + # Split fine features into first and second stage features + split_fine_features_0 = torch.split(fine_features_0, fine_embed_dim - fine_matching_slice_dim, -1) + split_fine_features_1 = torch.split(fine_features_1, fine_embed_dim - fine_matching_slice_dim, -1) + + # Retrieve first stage fine features + fine_features_0 = split_fine_features_0[0] + fine_features_1 = split_fine_features_1[0] + + # Normalize first stage fine features + fine_features_0 = fine_features_0 / fine_features_0.shape[-1] ** 0.5 + fine_features_1 = fine_features_1 / fine_features_1.shape[-1] ** 0.5 + + # Compute first stage confidence + fine_confidence = fine_features_0 @ fine_features_1.transpose(-1, -2) + fine_confidence = nn.functional.softmax(fine_confidence, 1) * nn.functional.softmax(fine_confidence, 2) + fine_confidence = fine_confidence.reshape( + batch_size, num_keypoints, fine_window_size, fine_kernel_size + 2, fine_kernel_size + 2 + ) + fine_confidence = fine_confidence[..., 1:-1, 1:-1] + first_stage_fine_confidence = fine_confidence.reshape( + batch_size, num_keypoints, fine_window_size, fine_window_size + ) + + fine_indices, fine_matches = self._get_first_stage_fine_matching( + first_stage_fine_confidence, + coarse_matched_keypoints, + fine_window_size, + fine_scale, + ) + + # Retrieve second stage fine features + fine_features_0 = split_fine_features_0[1] + fine_features_1 = split_fine_features_1[1] + + # Normalize second stage fine features + fine_features_1 = fine_features_1 / fine_matching_slice_dim**0.5 + + # Compute second stage fine confidence + second_stage_fine_confidence = fine_features_0 @ fine_features_1.transpose(-1, -2) + + fine_coordinates = self._get_second_stage_fine_matching( + fine_indices, + fine_matches, + second_stage_fine_confidence, + fine_window_size, + fine_scale, + ) + + return fine_coordinates + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> EfficientLoFTRKeypointMatchingOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModel + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_78916675_4568141288.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image1 = Image.open(BytesIO(response.read())) + + >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image2 = Image.open(BytesIO(response.read())) + + >>> images = [image1, image2] + + >>> processor = AutoImageProcessor.from_pretrained("zju-community/efficient_loftr") + >>> model = AutoModel.from_pretrained("zju-community/efficient_loftr") + + >>> with torch.no_grad(): + >>> inputs = processor(images, return_tensors="pt") + >>> outputs = model(**inputs) + ```""" + if labels is not None: + raise ValueError("SuperGlue is not trainable, no labels should be provided.") + + # 1. Extract coarse and residual features + model_outputs: BackboneOutput = self.efficientloftr(pixel_values, **kwargs) + features = model_outputs.feature_maps + + # 2. Compute coarse-level matching + coarse_features = features[0] + coarse_embed_dim, coarse_height, coarse_width = coarse_features.shape[-3:] + batch_size, _, channels, height, width = pixel_values.shape + coarse_scale = height / coarse_height + coarse_keypoints, coarse_matching_scores, coarse_matched_indices = self._coarse_matching( + coarse_features, coarse_scale + ) + + # 3. Fine-level refinement + residual_features = features[1:] + coarse_features = coarse_features / self.config.hidden_size**0.5 + fine_features_0, fine_features_1 = self.refinement_layer(coarse_features, residual_features) + + # Filter fine features with coarse matches indices + _, _, num_keypoints = coarse_matching_scores.shape + batch_indices = torch.arange(batch_size)[..., None] + fine_features_0 = fine_features_0[batch_indices, coarse_matched_indices[:, 0]] + fine_features_1 = fine_features_1[batch_indices, coarse_matched_indices[:, 1]] + + # 4. Computer fine-level matching + fine_height = torch_int(coarse_height * coarse_scale) + fine_scale = height / fine_height + matching_keypoints = self._fine_matching(fine_features_0, fine_features_1, coarse_keypoints, fine_scale) + + matching_keypoints[:, :, :, 0] = matching_keypoints[:, :, :, 0] / width + matching_keypoints[:, :, :, 1] = matching_keypoints[:, :, :, 1] / height + + loss = None + return EfficientLoFTRKeypointMatchingOutput( + loss=loss, + matches=coarse_matched_indices, + matching_scores=coarse_matching_scores, + keypoints=matching_keypoints, + hidden_states=model_outputs.hidden_states, + attentions=model_outputs.attentions, + ) + + +__all__ = ["EfficientLoFTRPreTrainedModel", "EfficientLoFTRModel", "EfficientLoFTRForKeypointMatching"] diff --git a/third_party/transformers/src/transformers/models/efficientloftr/modular_efficientloftr.py b/third_party/transformers/src/transformers/models/efficientloftr/modular_efficientloftr.py new file mode 100644 index 0000000000000000000000000000000000000000..17e3e399a8df182835d5136b5f1d7281f3f6b233 --- /dev/null +++ b/third_party/transformers/src/transformers/models/efficientloftr/modular_efficientloftr.py @@ -0,0 +1,151 @@ +from typing import TYPE_CHECKING + +from ...processing_utils import ImagesKwargs +from ...utils import TensorType, is_torch_available +from ...utils.import_utils import requires +from ..superglue.image_processing_pil_superglue import SuperGlueImageProcessorPil +from ..superglue.image_processing_superglue import SuperGlueImageProcessor + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from .modeling_efficientloftr import EfficientLoFTRKeypointMatchingOutput + + +class EfficientLoFTRImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`): + Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method. + """ + + do_grayscale: bool + + +class EfficientLoFTRImageProcessor(SuperGlueImageProcessor): + def post_process_keypoint_matching( + self, + outputs: "EfficientLoFTRKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, torch.Tensor]]: + """ + Converts the raw output of [`EfficientLoFTRKeypointMatchingOutput`] into lists of keypoints, scores and descriptors + with coordinates absolute to the original image sizes. + Args: + outputs ([`EfficientLoFTRKeypointMatchingOutput`]): + Raw outputs of the model. + target_sizes (`torch.Tensor` or `List[Tuple[Tuple[int, int]]]`, *optional*): + Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`Tuple[int, int]`) containing the + target size `(height, width)` of each image in the batch. This must be the original image size (before + any processing). + threshold (`float`, *optional*, defaults to 0.0): + Threshold to filter out the matches with low scores. + Returns: + `List[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image + of the pair, the matching scores and the matching indices. + """ + if outputs.matches.shape[0] != len(target_sizes): + raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask") + if not all(len(target_size) == 2 for target_size in target_sizes): + raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch") + + if isinstance(target_sizes, list): + image_pair_sizes = torch.tensor(target_sizes, device=outputs.matches.device) + else: + if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2: + raise ValueError( + "Each element of target_sizes must contain the size (h, w) of each image of the batch" + ) + image_pair_sizes = target_sizes + + keypoints = outputs.keypoints.clone() + keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2) + keypoints = keypoints.to(torch.int32) + + results = [] + for keypoints_pair, matches, scores in zip(keypoints, outputs.matches, outputs.matching_scores): + # Filter out matches with low scores + valid_matches = torch.logical_and(scores > threshold, matches > -1) + + matched_keypoints0 = keypoints_pair[0][valid_matches[0]] + matched_keypoints1 = keypoints_pair[1][valid_matches[1]] + matching_scores = scores[0][valid_matches[0]] + + results.append( + { + "keypoints0": matched_keypoints0, + "keypoints1": matched_keypoints1, + "matching_scores": matching_scores, + } + ) + + return results + + +class EfficientLoFTRImageProcessorPil(SuperGlueImageProcessorPil): + @requires(backends=("torch",)) + def post_process_keypoint_matching( + self, + outputs: "EfficientLoFTRKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, "torch.Tensor"]]: + """ + Converts the raw output of [`EfficientLoFTRKeypointMatchingOutput`] into lists of keypoints, scores and descriptors + with coordinates absolute to the original image sizes. + Args: + outputs ([`EfficientLoFTRKeypointMatchingOutput`]): + Raw outputs of the model. + target_sizes (`torch.Tensor` or `List[Tuple[Tuple[int, int]]]`, *optional*): + Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`Tuple[int, int]`) containing the + target size `(height, width)` of each image in the batch. This must be the original image size (before + any processing). + threshold (`float`, *optional*, defaults to 0.0): + Threshold to filter out the matches with low scores. + Returns: + `List[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image + of the pair, the matching scores and the matching indices. + """ + import torch + + if outputs.matches.shape[0] != len(target_sizes): + raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask") + if not all(len(target_size) == 2 for target_size in target_sizes): + raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch") + + if isinstance(target_sizes, list): + image_pair_sizes = torch.tensor(target_sizes, device=outputs.matches.device) + else: + if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2: + raise ValueError( + "Each element of target_sizes must contain the size (h, w) of each image of the batch" + ) + image_pair_sizes = target_sizes + + keypoints = outputs.keypoints.clone() + keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2) + keypoints = keypoints.to(torch.int32) + + results = [] + for keypoints_pair, matches, scores in zip(keypoints, outputs.matches, outputs.matching_scores): + # Filter out matches with low scores + valid_matches = torch.logical_and(scores > threshold, matches > -1) + + matched_keypoints0 = keypoints_pair[0][valid_matches[0]] + matched_keypoints1 = keypoints_pair[1][valid_matches[1]] + matching_scores = scores[0][valid_matches[0]] + + results.append( + { + "keypoints0": matched_keypoints0, + "keypoints1": matched_keypoints1, + "matching_scores": matching_scores, + } + ) + + return results + + +__all__ = ["EfficientLoFTRImageProcessor", "EfficientLoFTRImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/ernie4_5_moe/__init__.py b/third_party/transformers/src/transformers/models/ernie4_5_moe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb30318fa25ec560d549b08421ea08f02b0dce6f --- /dev/null +++ b/third_party/transformers/src/transformers/models/ernie4_5_moe/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_ernie4_5_moe import * + from .modeling_ernie4_5_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/ernie4_5_moe/configuration_ernie4_5_moe.py b/third_party/transformers/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0c0edbb760c7e577d9db28f0908d165504b1e9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py @@ -0,0 +1,121 @@ +# Copyright (c) 2025 Baidu, Inc. and HuggingFace Inc. team. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ernie 4.5 MoE 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="baidu/ERNIE-4.5-21B-A3B-PT") +@strict +class Ernie4_5_MoeConfig(PreTrainedConfig): + r""" + use_bias (`bool`, *optional*, defaults to `False`): + Whether to use a bias in any of the projections including mlp and attention for example. + moe_k (`int`, *optional*, defaults to 6): + Number of selected experts. + moe_num_experts (`int`, *optional*, defaults to 64): + Number of routed experts. + moe_num_shared_experts (`int`, *optional*, defaults to 2): + The number of experts that are shared for all MoE forwards. + moe_layer_start_index (`int`, *optional*, defaults to 1): + The first index at which MoE layers start to appear. + moe_layer_end_index (`int`, *optional*, defaults to -1): + The last possible index for a MoE layer. + moe_layer_interval (`int`, *optional*, defaults to 1): + The intervals between MoE layers to appear. + moe_norm_min (`float`, *optional*, defaults to 1e-12): + Minimum division value during routing normalization. + + Example: + + ```python + >>> from transformers import Ernie4_5_MoeModel, Ernie4_5_MoEConfig + + >>> # Initializing a Ernie4_5_MoE style configuration + >>> configuration = Ernie4_5_MoEConfig() + + >>> # Initializing a model from the ERNIE-4.5-21B-A3B style configuration + >>> model = Ernie4_5_MoeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "ernie4_5_moe" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_experts": "moe_num_experts", "num_experts_per_tok": "moe_k"} + default_theta = 500000.0 + + # Default tensor parallel plan for base model `Ernie4_5_MoE` + 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", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 103424 + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + hidden_size: int = 2560 + intermediate_size: int = 12288 + num_hidden_layers: int = 28 + num_attention_heads: int = 20 + num_key_value_heads: int | None = 4 + hidden_act: str = "silu" + max_position_embeddings: int = 131072 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + use_cache: bool = True + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + use_bias: bool | None = False + moe_intermediate_size: int = 1536 + moe_k: int | None = 6 + moe_num_experts: int | None = 64 + moe_num_shared_experts: int | None = 2 + moe_layer_start_index: int | None = 1 + moe_layer_end_index: int | None = -1 + moe_layer_interval: int | None = 1 + moe_norm_min: float | None = 1e-12 + output_router_logits: bool | None = False + router_aux_loss_coef: float | None = 0.001 + + def __post_init__(self, **kwargs): + self.moe_layer_end_index = ( + self.num_hidden_layers - 1 if self.moe_layer_end_index == -1 else self.moe_layer_end_index + ) + super().__post_init__(**kwargs) + + +__all__ = ["Ernie4_5_MoeConfig"] diff --git a/third_party/transformers/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/third_party/transformers/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..9c106a90010db57c9316dca096e91f096ddbb1ef --- /dev/null +++ b/third_party/transformers/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -0,0 +1,742 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/ernie4_5_moe/modular_ernie4_5_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_ernie4_5_moe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright (c) 2025 Baidu, Inc. and HuggingFace Inc. team. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.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_ernie4_5_moe import Ernie4_5_MoeConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class Ernie4_5_MoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Ernie4_5_MoeRMSNorm 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 Ernie4_5_MoeMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Ernie4_5_MoeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Ernie4_5_MoeConfig, 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: Ernie4_5_MoeConfig | 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 + + # keeping it in full precision + return cos, sin + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., 0::2] + x2 = x[..., 1::2] + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + # glm rope style (with full dim) and full precision + original_dtype = q.dtype + + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + # Interleave them instead of usual shape + cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1) + sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1) + + q_embed = (q.float() * cos) + (rotate_half(q).float() * sin) + k_embed = (k.float() * cos) + (rotate_half(k).float() * sin) + + return q_embed.to(original_dtype), k_embed.to(original_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 + + +@use_kernelized_func(apply_rotary_pos_emb) +class Ernie4_5_MoeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Ernie4_5_MoeConfig, 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 = 0.0 + self.is_causal = True + + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.use_bias) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.use_bias) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Ernie4_5_MoeStatics(nn.Module): + """ + Stores MoE (Mixture of Experts) statistics + - Bias for the gating + - Additionally, usage per expert in the original codebase + """ + + def __init__(self, config): + super().__init__() + + num_experts_groups = 1 + num_experts = config.moe_num_experts + + self.e_score_correction_bias = nn.Parameter( + torch.zeros(num_experts_groups, num_experts, dtype=torch.float32), + requires_grad=False, + ) + + def forward(self, hidden_states): + # NOTE: This is a workaround to enable TP with a module that only has parameters + # + # Otherwise, it stays as `DTensor` when called in the "super" forward + # 1. All other tensors are local (`torch.Tensor`) + # 2. Isolate does not work on `nn.Module` which only has parameters + return hidden_states + self.e_score_correction_bias.squeeze() + + +@use_experts_implementation +class Ernie4_5_MoeExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config): + super().__init__() + self.num_experts = config.moe_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 = 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 Ernie4_5_MoeTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.weight = nn.Parameter(torch.zeros(config.moe_num_experts, config.hidden_size, dtype=torch.float32)) + self.moe_statics = Ernie4_5_MoeStatics(config) + self.top_k = config.moe_k + self.norm_min = config.moe_norm_min + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + device_type = ( + hidden_states.device.type + if isinstance(hidden_states.device.type, str) and hidden_states.device.type != "mps" + else "cpu" + ) + + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + router_logits = F.linear(hidden_states.float(), self.weight.float()) + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(self.moe_statics(routing_weights), self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) + routing_weights = routing_weights / torch.clamp( + routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min + ) + routing_weights = routing_weights.to(hidden_states.dtype) + return router_logits, selected_experts, routing_weights + + +class Ernie4_5_MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_dim = config.hidden_size + self.num_experts = config.moe_num_experts + self.top_k = config.moe_k + self.gate = Ernie4_5_MoeTopKRouter(config) + self.experts = Ernie4_5_MoeExperts(config) + + self.shared_experts = None + if config.moe_num_shared_experts > 0: + self.shared_experts = Ernie4_5_MoeMLP(config, config.moe_intermediate_size * config.moe_num_shared_experts) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + + _, top_k_index, top_k_weights = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + shared_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, self.hidden_dim) + return final_hidden_states.to(hidden_states.dtype) + + +class Ernie4_5_MoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Ernie4_5_MoeAttention(config, layer_idx) + + if ( + ((layer_idx + 1) % config.moe_layer_interval == 0) + and layer_idx >= config.moe_layer_start_index + and layer_idx <= config.moe_layer_end_index + ): + self.mlp = Ernie4_5_MoeSparseMoeBlock(config) + else: + self.mlp = Ernie4_5_MoeMLP(config) + + self.input_layernorm = Ernie4_5_MoeRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = Ernie4_5_MoeRMSNorm(config.hidden_size, config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class Ernie4_5_MoePreTrainedModel(PreTrainedModel): + config: Ernie4_5_MoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Ernie4_5_MoeDecoderLayer"] + _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(Ernie4_5_MoeTopKRouter, index=0), + "hidden_states": Ernie4_5_MoeDecoderLayer, + "attentions": Ernie4_5_MoeAttention, + } + # Not supporting multi-token prediction (MTP) atm + _keys_to_ignore_on_load_unexpected = ["mtp"] + _keep_in_fp32_modules_strict = ["gate.weight", "moe_statics"] + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Ernie4_5_MoeStatics): + init.zeros_(module.e_score_correction_bias) + elif isinstance(module, Ernie4_5_MoeExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class Ernie4_5_MoeModel(Ernie4_5_MoePreTrainedModel): + def __init__(self, config: Ernie4_5_MoeConfig): + 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( + [Ernie4_5_MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Ernie4_5_MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Ernie4_5_MoeRotaryEmbedding(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) + + 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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +def load_balancing_loss_func( + gate_logits: torch.Tensor | tuple[torch.Tensor] | None, + num_experts: int | None = None, + top_k=2, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor | int: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +@auto_docstring +class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, 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 = Ernie4_5_MoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=config.use_bias) + + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.moe_num_experts + self.num_experts_per_tok = config.moe_k + + # 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]`. + """ + + 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__ = ["Ernie4_5_MoeForCausalLM", "Ernie4_5_MoeModel", "Ernie4_5_MoePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py b/third_party/transformers/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..6a2fa9620d97b56b304be7a5697ba4c9eb0f3004 --- /dev/null +++ b/third_party/transformers/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py @@ -0,0 +1,315 @@ +# Copyright (c) 2025 Baidu, Inc. and HuggingFace Inc. team. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Ernie 4.5 MoE model.""" + +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, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..ernie4_5.modeling_ernie4_5 import Ernie4_5RotaryEmbedding, apply_rotary_pos_emb, rotate_half # noqa: F401 +from ..llama.modeling_llama import LlamaAttention, LlamaRMSNorm +from ..mixtral.modeling_mixtral import ( + MixtralExperts, + MixtralForCausalLM, + MixtralPreTrainedModel, +) +from ..qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeMLP +from .configuration_ernie4_5_moe import Ernie4_5_MoeConfig + + +logger = logging.get_logger(__name__) + + +class Ernie4_5_MoeRMSNorm(LlamaRMSNorm): + pass + + +class Ernie4_5_MoeMLP(Qwen3MoeMLP): + def __init__(self, config, intermediate_size=None): + super().__init__(config, intermediate_size) + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) + + +class Ernie4_5_MoeRotaryEmbedding(Ernie4_5RotaryEmbedding): + def __init__(self, config: Ernie4_5_MoeConfig, device=None): + super().__init__(config, device) + + +class Ernie4_5_MoeAttention(LlamaAttention): + def __init__(self, config: Ernie4_5_MoeConfig, layer_idx: int): + super().__init__(config, layer_idx) + + self.attention_dropout = 0.0 + + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.use_bias) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.use_bias) + + +class Ernie4_5_MoeStatics(nn.Module): + """ + Stores MoE (Mixture of Experts) statistics + - Bias for the gating + - Additionally, usage per expert in the original codebase + """ + + def __init__(self, config): + super().__init__() + + num_experts_groups = 1 + num_experts = config.moe_num_experts + + self.e_score_correction_bias = nn.Parameter( + torch.zeros(num_experts_groups, num_experts, dtype=torch.float32), + requires_grad=False, + ) + + def forward(self, hidden_states): + # NOTE: This is a workaround to enable TP with a module that only has parameters + # + # Otherwise, it stays as `DTensor` when called in the "super" forward + # 1. All other tensors are local (`torch.Tensor`) + # 2. Isolate does not work on `nn.Module` which only has parameters + return hidden_states + self.e_score_correction_bias.squeeze() + + +class Ernie4_5_MoeExperts(MixtralExperts): + def __init__(self, config): + super().__init__() + self.num_experts = config.moe_num_experts + self.intermediate_dim = config.moe_intermediate_size + + +class Ernie4_5_MoeTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.weight = nn.Parameter(torch.zeros(config.moe_num_experts, config.hidden_size, dtype=torch.float32)) + self.moe_statics = Ernie4_5_MoeStatics(config) + self.top_k = config.moe_k + self.norm_min = config.moe_norm_min + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + device_type = ( + hidden_states.device.type + if isinstance(hidden_states.device.type, str) and hidden_states.device.type != "mps" + else "cpu" + ) + + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + router_logits = F.linear(hidden_states.float(), self.weight.float()) + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(self.moe_statics(routing_weights), self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) + routing_weights = routing_weights / torch.clamp( + routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min + ) + routing_weights = routing_weights.to(hidden_states.dtype) + return router_logits, selected_experts, routing_weights + + +class Ernie4_5_MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_dim = config.hidden_size + self.num_experts = config.moe_num_experts + self.top_k = config.moe_k + self.gate = Ernie4_5_MoeTopKRouter(config) + self.experts = Ernie4_5_MoeExperts(config) + + self.shared_experts = None + if config.moe_num_shared_experts > 0: + self.shared_experts = Ernie4_5_MoeMLP(config, config.moe_intermediate_size * config.moe_num_shared_experts) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + + _, top_k_index, top_k_weights = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + shared_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, self.hidden_dim) + return final_hidden_states.to(hidden_states.dtype) + + +class Ernie4_5_MoeDecoderLayer(Qwen3MoeDecoderLayer): + def __init__(self, config, layer_idx): + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + + self.self_attn = Ernie4_5_MoeAttention(config, layer_idx) + + if ( + ((layer_idx + 1) % config.moe_layer_interval == 0) + and layer_idx >= config.moe_layer_start_index + and layer_idx <= config.moe_layer_end_index + ): + self.mlp = Ernie4_5_MoeSparseMoeBlock(config) + else: + self.mlp = Ernie4_5_MoeMLP(config) + + self.input_layernorm = Ernie4_5_MoeRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = Ernie4_5_MoeRMSNorm(config.hidden_size, config.rms_norm_eps) + + +@auto_docstring +class Ernie4_5_MoePreTrainedModel(MixtralPreTrainedModel): + config: Ernie4_5_MoeConfig + _no_split_modules = ["Ernie4_5_MoeDecoderLayer"] + # Not supporting multi-token prediction (MTP) atm + _keys_to_ignore_on_load_unexpected = ["mtp"] + _can_record_outputs = { + "router_logits": OutputRecorder(Ernie4_5_MoeTopKRouter, index=0), + "hidden_states": Ernie4_5_MoeDecoderLayer, + "attentions": Ernie4_5_MoeAttention, + } + _keep_in_fp32_modules_strict = ["gate.weight", "moe_statics"] + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, Ernie4_5_MoeStatics): + init.zeros_(module.e_score_correction_bias) + elif isinstance(module, Ernie4_5_MoeExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class Ernie4_5_MoeModel(Ernie4_5_MoePreTrainedModel): + def __init__(self, config: Ernie4_5_MoeConfig): + 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( + [Ernie4_5_MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Ernie4_5_MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Ernie4_5_MoeRotaryEmbedding(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) + + 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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Ernie4_5_MoeForCausalLM(MixtralForCausalLM): + def __init__(self, config): + PreTrainedModel.__init__(self, config) + self.model = Ernie4_5_MoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=config.use_bias) + + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.moe_num_experts + self.num_experts_per_tok = config.moe_k + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward(self, **super_kwargs): + 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]`. + """ + super().forward(**super_kwargs) + + +__all__ = [ + "Ernie4_5_MoeForCausalLM", + "Ernie4_5_MoeModel", + "Ernie4_5_MoePreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/florence2/__init__.py b/third_party/transformers/src/transformers/models/florence2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..686295d1fbd952e124bab885980ac682261a07fb --- /dev/null +++ b/third_party/transformers/src/transformers/models/florence2/__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_florence2 import * + from .modeling_florence2 import * + from .processing_florence2 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/florence2/configuration_florence2.py b/third_party/transformers/src/transformers/models/florence2/configuration_florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..7a03f35d372c48f976d9d1d1582c3d120f12ac5b --- /dev/null +++ b/third_party/transformers/src/transformers/models/florence2/configuration_florence2.py @@ -0,0 +1,139 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/florence2/modular_florence2.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_florence2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Microsoft and the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +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="florence-community/Florence-2-base") +@strict +class Florence2VisionConfig(PreTrainedConfig): + r""" + depths (`Tuple[int]`, *optional*, defaults to `(1, 1, 9, 1)`): + The depth of the model. + patch_stride (`Tuple[int]`, *optional*, defaults to `(4, 2, 2, 2)`): + The patch stride of the image. + patch_padding (`Tuple[int]`, *optional*, defaults to `(3, 1, 1, 1)`): + The patch padding of the image. + patch_prenorm (`Tuple[bool]`, *optional*, defaults to `(False, True, True, True)`): + Whether to apply layer normalization before the patch embedding layer. + num_groups (`Tuple[int]`, *optional*, defaults to `(4, 8, 16, 32)`): + The number of groups. + window_size (`int`, *optional*, defaults to 12): + The window size of the model. + max_temporal_embeddings (`int`, *optional*, defaults to 100): + The configuration of the visual temporal embedding. + + Example: + + ```python + >>> from transformers import Florence2VisionConfig, Florence2VisionModel + + >>> # Initializing a Florence2 Vision style configuration + >>> configuration = Florence2VisionConfig() + + >>> # Initializing a model (with random weights) + >>> model = Florence2VisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "florence_vision" + + in_channels: int = 3 + depths: list[int] | tuple[int, ...] = (1, 1, 9, 1) + patch_size: list[int] | tuple[int, ...] = (7, 3, 3, 3) + patch_stride: list[int] | tuple[int, ...] = (4, 2, 2, 2) + patch_padding: list[int] | tuple[int, ...] = (3, 1, 1, 1) + patch_prenorm: list[bool] | tuple[bool, ...] = (False, True, True, True) + embed_dim: list[int] | tuple[int, ...] = (128, 256, 512, 1024) + num_heads: list[int] | tuple[int, ...] = (4, 8, 16, 32) + num_groups: list[int] | tuple[int, ...] = (4, 8, 16, 32) + window_size: int = 12 + drop_path_rate: float | int = 0.1 + mlp_ratio: float = 4.0 + qkv_bias: bool = True + activation_function: str = "gelu" + projection_dim: int = 1024 + max_temporal_embeddings: int = 100 + max_position_embeddings: int = 50 + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="florence-community/Florence-2-base") +@strict +class Florence2Config(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import Florence2ForConditionalGeneration, Florence2Config, CLIPVisionConfig, BartConfig + + >>> # Initializing a clip-like vision config + >>> vision_config = CLIPVisionConfig() + + >>> # Initializing a Bart config + >>> text_config = BartConfig() + + >>> # Initializing a Florence-2 configuration + >>> configuration = Florence2Config(vision_config, text_config) + + >>> # Initializing a model from the florence-2 configuration + >>> model = Florence2ForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "florence2" + sub_configs = { + "text_config": AutoConfig, + "vision_config": Florence2VisionConfig, + } + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + image_token_id: int = 51289 + is_encoder_decoder: bool = True + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "bart") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["bart"]() + + if isinstance(self.vision_config, dict): + self.vision_config = Florence2VisionConfig(**self.vision_config) + elif self.vision_config is None: + logger.info("vision_config is None. Initializing the Florence2VisionConfig with default values.") + self.vision_config = Florence2VisionConfig() + + super().__post_init__(**kwargs) + + +__all__ = ["Florence2Config", "Florence2VisionConfig"] diff --git a/third_party/transformers/src/transformers/models/florence2/modeling_florence2.py b/third_party/transformers/src/transformers/models/florence2/modeling_florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..fd941b85ce66323041c673526f923e62dee94c28 --- /dev/null +++ b/third_party/transformers/src/transformers/models/florence2/modeling_florence2.py @@ -0,0 +1,1014 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/florence2/modular_florence2.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_florence2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Microsoft and the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...generation import GenerationMixin +from ...modeling_outputs import BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + is_torch_available, + 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_florence2 import Florence2Config, Florence2VisionConfig + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +class Florence2VisionDropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float | None = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return f"p={self.drop_prob}" + + +class Florence2VisionLearnedAbsolutePositionEmbedding2D(nn.Module): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, config: Florence2Config): + super().__init__() + num_pos = config.vision_config.max_position_embeddings + embedding_dim = config.vision_config.embed_dim[-1] + self.row_embeddings = nn.Embedding(num_pos, embedding_dim // 2) + self.column_embeddings = nn.Embedding(num_pos, embedding_dim - (embedding_dim // 2)) + + def forward(self, pixel_values, pixel_mask=None): + height, width = pixel_values.shape[-2:] + width_values = torch.arange(width, device=pixel_values.device) + height_values = torch.arange(height, device=pixel_values.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(pixel_values.shape[0], 1, 1, 1) + return pos + + +class Florence2VisionPositionalEmbeddingCosine1D(nn.Module): + """ + This module generates 1D cosine positional embeddings using precomputed sinusoidal functions. + """ + + def __init__(self, config: Florence2Config): + super().__init__() + self.embed_dim = config.vision_config.embed_dim[-1] + self.max_seq_len = config.vision_config.max_temporal_embeddings + pos_idx_to_embed = torch.empty((self.max_seq_len, self.embed_dim)) + sine, cosine = self.get_sinusoid_embeddings( + max_positions=self.max_seq_len, + embed_dim=self.embed_dim, + ) + pos_idx_to_embed[:, 0::2] = sine + pos_idx_to_embed[:, 1::2] = cosine + # Save the positional embeddings in a constant buffer. + self.register_buffer("pos_idx_to_embed", pos_idx_to_embed) + + @staticmethod + def get_sinusoid_embeddings(max_positions: int, embed_dim: int): + half_dim = embed_dim // 2 + emb = math.log(10000) / half_dim + emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb) + emb = torch.arange(max_positions, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0) + return torch.sin(emb), torch.cos(emb) + + def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor: + len_seq = seq_embeds.size(1) + if len_seq > self.max_seq_len: + raise ValueError(f"Maximum sequence length {self.max_seq_len}, got {len_seq}") + pos_embeds = self.pos_idx_to_embed[0:len_seq, :] + return pos_embeds + + +class Florence2VisionMLP(nn.Module): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.activation_function] + self.fc1 = nn.Linear(config.embed_dim[stage_idx], int(config.embed_dim[stage_idx] * config.mlp_ratio)) + self.fc2 = nn.Linear(int(config.embed_dim[stage_idx] * config.mlp_ratio), config.embed_dim[stage_idx]) + + 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 Florence2VisionConvEmbed(nn.Module): + """Image to Patch Embedding""" + + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.stage_idx = stage_idx + self.patch_size = config.patch_size[stage_idx] + self.in_channels = config.in_channels if stage_idx == 0 else config.embed_dim[stage_idx - 1] + self.embed_dim = config.embed_dim[stage_idx] + self.stride = config.patch_stride[stage_idx] + self.padding = config.patch_padding[stage_idx] + self.pre_norm = config.patch_prenorm[stage_idx] + + self.conv = nn.Conv2d( + self.in_channels, + self.embed_dim, + kernel_size=self.patch_size, + stride=self.stride, + padding=self.padding, + ) + + dim_norm = self.in_channels if self.pre_norm else self.embed_dim + self.norm = nn.LayerNorm(dim_norm) + + def forward(self, hidden_states: torch.Tensor): + if self.norm and self.pre_norm: + hidden_states = hidden_states.permute(0, 2, 3, 1) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 3, 1, 2) + + hidden_states = self.conv(hidden_states) + + if self.norm and not self.pre_norm: + hidden_states = hidden_states.permute(0, 2, 3, 1) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 3, 1, 2) + return hidden_states + + +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 Florence2VisionChannelAttention(nn.Module): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.dim = config.embed_dim[stage_idx] + self.groups = config.num_groups[stage_idx] + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=config.qkv_bias) + self.proj = nn.Linear(self.dim, self.dim) + self.is_causal = False + + def forward(self, hidden_states: torch.Tensor): + batch_size, num_tokens, hidden_size = hidden_states.shape + + # Reshape for grouped channel attention + qkv = self.qkv(hidden_states).reshape(batch_size, num_tokens, 3, self.groups, hidden_size // self.groups) + qkv = qkv.permute(2, 0, 3, 4, 1) + query, key, value = qkv.unbind(0) + + scale = num_tokens**-0.5 + # Channel-to-channel attention within groups: + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + hidden_states, _ = attention_interface( + self, + query, + key, + value, + attention_mask=None, + scaling=scale, + ) + hidden_states = hidden_states.permute(0, 3, 2, 1) + hidden_states = hidden_states.reshape(batch_size, num_tokens, hidden_size) + + # Final projection + hidden_states = self.proj(hidden_states) + return hidden_states + + +class Florence2VisionChannelBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + drop_path_rate: float, + ): + super().__init__() + + self.config = config + dim_in = config.embed_dim[stage_idx] + + self.conv1 = nn.Conv2d( + dim_in, + dim_in, + kernel_size=3, + padding=1, + groups=dim_in, + ) + self.norm1 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.channel_attn = Florence2VisionChannelAttention(config=config, stage_idx=stage_idx) + self.drop_path1 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + self.conv2 = nn.Conv2d( + dim_in, + dim_in, + kernel_size=3, + padding=1, + groups=dim_in, + ) + self.norm2 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.ffn = Florence2VisionMLP(config=config, stage_idx=stage_idx) + self.drop_path2 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + def forward(self, hidden_states: torch.Tensor): + batch_size, embed_dim, height, width = hidden_states.shape + + # First channel block: Depthwise Conv + Channel Attention + hidden_states = self.conv1(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # Channel group attention self-attention mechanism + hidden_states = self.norm1(hidden_states) + hidden_states = self.channel_attn(hidden_states) + hidden_states = residual + self.drop_path1(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + # Second channel block: Depthwise Conv + FFN + hidden_states = self.conv2(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # FFN + hidden_states = self.norm2(hidden_states) + hidden_states = self.ffn(hidden_states) + hidden_states = residual + self.drop_path2(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + return hidden_states + + +class Florence2VisionWindowAttention(nn.Module): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.dim = config.embed_dim[stage_idx] + self.window_size = config.window_size + self.num_heads = config.num_heads[stage_idx] + head_dim = self.dim // self.num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=config.qkv_bias) + self.proj = nn.Linear(self.dim, self.dim) + self.is_causal = False + + def forward(self, hidden_states: torch.Tensor): + batch_size, height, width, embed_dim = hidden_states.shape + + # Pad the input if necessary + pad_left = pad_top = 0 + pad_right = (self.window_size - width % self.window_size) % self.window_size + pad_bottom = (self.window_size - height % self.window_size) % self.window_size + hidden_states = F.pad(hidden_states, (0, 0, pad_left, pad_right, pad_top, pad_bottom)) + _, padded_height, padded_width, _ = hidden_states.shape + + # Partition input into non-overlapping windows (for local spatial attention in DaViT) + hidden_states = hidden_states.view( + batch_size, + padded_height // self.window_size, + self.window_size, + padded_width // self.window_size, + self.window_size, + embed_dim, + ) + windowed_hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous() + windowed_hidden_states = windowed_hidden_states.view(-1, self.window_size * self.window_size, embed_dim) + + # Generate Q, K, V for each window + num_windows_per_batch, num_tokens_per_window, embed_dim = windowed_hidden_states.shape + qkv = self.qkv(windowed_hidden_states).reshape( + num_windows_per_batch, num_tokens_per_window, 3, self.num_heads, embed_dim // self.num_heads + ) + qkv = qkv.permute(2, 0, 3, 1, 4) + query, key, value = qkv.unbind(0) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + windowed_hidden_states, _ = attention_interface( + self, + query, + key, + value, + attention_mask=None, + scaling=self.scale, + ) + windowed_hidden_states = windowed_hidden_states.view(num_windows_per_batch, num_tokens_per_window, embed_dim) + windowed_hidden_states = self.proj(windowed_hidden_states) + + # Merge windows back to original spatial layout + windowed_hidden_states = windowed_hidden_states.view(-1, self.window_size, self.window_size, embed_dim) + hidden_states = windowed_hidden_states.view( + -1, + padded_height // self.window_size, + padded_width // self.window_size, + self.window_size, + self.window_size, + embed_dim, + ) + hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous() + hidden_states = hidden_states.view(-1, padded_height, padded_width, embed_dim) + hidden_states = hidden_states[:, :height, :width, :].contiguous() + hidden_states = hidden_states.view(batch_size, height * width, embed_dim) + + return hidden_states + + +class Florence2VisionSpatialBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + drop_path_rate: float, + ): + super().__init__() + + self.conv1 = nn.Conv2d( + config.embed_dim[stage_idx], + config.embed_dim[stage_idx], + kernel_size=3, + padding=1, + groups=config.embed_dim[stage_idx], + ) + self.norm1 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.window_attn = Florence2VisionWindowAttention(config=config, stage_idx=stage_idx) + self.drop_path1 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + self.conv2 = nn.Conv2d( + config.embed_dim[stage_idx], + config.embed_dim[stage_idx], + kernel_size=3, + padding=1, + groups=config.embed_dim[stage_idx], + ) + self.norm2 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.ffn = Florence2VisionMLP(config=config, stage_idx=stage_idx) + self.drop_path2 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + def forward(self, hidden_states: torch.Tensor): + batch_size, embed_dim, height, width = hidden_states.shape + + # First spatial mixing block: Conv + Window Attention + hidden_states = self.conv1(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # Spatial Window-based self-attention mechanism + hidden_states = self.norm1(hidden_states) + hidden_states = hidden_states.view(batch_size, height, width, embed_dim) + hidden_states = self.window_attn(hidden_states) + hidden_states = residual + self.drop_path1(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + # Second spatial mixing block: Conv + FFN + hidden_states = self.conv2(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # FFN + hidden_states = self.norm2(hidden_states) + hidden_states = self.ffn(hidden_states) + hidden_states = residual + self.drop_path2(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + return hidden_states + + +class Florence2VisionBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + spatial_drop_path_rate: float, + channel_drop_path_rate: float, + ): + super().__init__() + self.spatial_block = Florence2VisionSpatialBlock( + config=config, + stage_idx=stage_idx, + drop_path_rate=spatial_drop_path_rate, + ) + self.channel_block = Florence2VisionChannelBlock( + config=config, + stage_idx=stage_idx, + drop_path_rate=channel_drop_path_rate, + ) + + def forward(self, hidden_states: torch.Tensor): + hidden_states = self.spatial_block(hidden_states) + hidden_states = self.channel_block(hidden_states) + return hidden_states + + +@auto_docstring +class Florence2VisionPreTrainedModel(PreTrainedModel): + config_class = Florence2VisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _can_record_outputs = { + "hidden_states": Florence2VisionBlock, + "attentions": [Florence2VisionChannelAttention, Florence2VisionWindowAttention], + } + + +@auto_docstring +class Florence2VisionBackbone(Florence2VisionPreTrainedModel): + def __init__(self, config: Florence2VisionConfig): + super().__init__(config) + self.config = config + + self.embed_dim = config.embed_dim + self.num_heads = config.num_heads + self.num_groups = config.num_groups + self.num_stages = len(self.embed_dim) + + if not (self.num_stages == len(self.num_heads) == len(self.num_groups)): + raise ValueError( + f"Expected self.num_stages ({self.num_stages}) == " + f"len(self.num_heads) ({len(self.num_heads)}) == " + f"len(self.num_groups) ({len(self.num_groups)})" + ) + + dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths) * 2, device="cpu")] + depth_offset = 0 + + convs = [] + blocks = [] + for stage_idx in range(self.num_stages): + conv_embed = Florence2VisionConvEmbed( + config=config, + stage_idx=stage_idx, + ) + convs.append(conv_embed) + + block = nn.ModuleList( + Florence2VisionBlock( + config=config, + stage_idx=stage_idx, + spatial_drop_path_rate=dpr[depth_offset + block_idx * 2], + channel_drop_path_rate=dpr[depth_offset + block_idx * 2 + 1], + ) + for block_idx in range(config.depths[stage_idx]) + ) + blocks.append(block) + depth_offset += config.depths[stage_idx] * 2 + + self.convs = nn.ModuleList(convs) + self.blocks = nn.ModuleList(blocks) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + target_dtype = self.convs[0].conv.weight.dtype + hidden_states = hidden_states.to(dtype=target_dtype) + for conv, block in zip(self.convs, self.blocks): + hidden_states = conv(hidden_states) + for layer in block: + hidden_states = layer(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + ) + + +class Florence2MultiModalProjector(nn.Module): + def __init__(self, config: Florence2Config): + super().__init__() + self.vision_embedding_dim = config.vision_config.embed_dim[-1] + self.vision_projection_dim = config.vision_config.projection_dim + self.image_projection = nn.Linear(self.vision_embedding_dim, self.vision_projection_dim, bias=False) + self.image_proj_norm = nn.LayerNorm(self.vision_projection_dim) + self.image_position_embed = Florence2VisionLearnedAbsolutePositionEmbedding2D(config=config) + self.visual_temporal_embed = Florence2VisionPositionalEmbeddingCosine1D(config=config) + + def forward(self, image_features): + position_features = image_features + self.image_position_embed(image_features) + position_features = position_features.flatten(2).transpose(1, 2) + temporal_features = self.visual_temporal_embed(position_features[:, :1, :]) + temporal_features = temporal_features.unsqueeze(1) + visual_token_features = position_features + temporal_features + visual_token_features = visual_token_features.unsqueeze(1) + spatial_image_features = visual_token_features.mean(dim=2) + temporal_image_features = visual_token_features.mean(dim=1) + image_features = torch.cat([spatial_image_features, temporal_image_features], dim=1) + image_features = self.image_projection(image_features) + image_features = self.image_proj_norm(image_features) + return image_features + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Florence-2 base model's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class Florence2Seq2SeqModelOutput(Seq2SeqModelOutput): + r""" + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_image_tokens, 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 + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Florence-2 model's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class Florence2Seq2SeqLMOutput(Seq2SeqLMOutput): + 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). + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_image_tokens, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + image_hidden_states: tuple[torch.FloatTensor, ...] | None = None + + +@auto_docstring +class Florence2PreTrainedModel(PreTrainedModel): + config: Florence2Config + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _skip_keys_device_placement = "past_key_values" + + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _supports_flex_attn = True + + _supports_attention_backend = False + config_class = Florence2Config + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Florence2VisionPositionalEmbeddingCosine1D): + pos_idx_to_embed = torch.empty((module.max_seq_len, module.embed_dim)) + sine, cosine = module.get_sinusoid_embeddings( + max_positions=module.max_seq_len, + embed_dim=module.embed_dim, + ) + pos_idx_to_embed[:, 0::2] = sine + pos_idx_to_embed[:, 1::2] = cosine + init.copy_(module.pos_idx_to_embed, pos_idx_to_embed) + + +@auto_docstring( + custom_intro=""" + Florence-2 is a vision model for captioning, detection, and segmentation. + """ +) +class Florence2Model(Florence2PreTrainedModel): + def __init__(self, config: Florence2Config): + super().__init__(config) + self.vision_tower = Florence2VisionBackbone(config=config.vision_config) + + self.multi_modal_projector = Florence2MultiModalProjector(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.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. + """ + image_outputs = self.vision_tower(pixel_values, **kwargs) + image_outputs.pooler_output = self.multi_modal_projector(image_outputs.last_hidden_state) + + return image_outputs + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + + n_image_tokens = special_image_mask.sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", + ) + return special_image_mask + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs, + ) -> tuple | Florence2Seq2SeqModelOutput: + if encoder_outputs is None: + 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).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + encoder_outputs = self.language_model.encoder( + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + if decoder_input_ids is None: + decoder_start_token_id = self.config.text_config.decoder_start_token_id + decoder_input_ids = torch.ones((inputs_embeds.size()[0], 1), dtype=torch.long, device=inputs_embeds.device) + decoder_input_ids *= decoder_start_token_id + + decoder_outputs = self.language_model.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs[0], + encoder_attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Florence2Seq2SeqModelOutput( + 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, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + def get_encoder(self, modality=None): + if modality is None: + return self.language_model.get_encoder() + else: + return super().get_encoder(modality=modality) + + +def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): + """ + Shift input ids one token to the right. + """ + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +@auto_docstring( + custom_intro=""" + Florence-2 is a vision model for captioning, detection, and segmentation. + """ +) +class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixin): + _tied_weights_keys = { + "lm_head.weight": "model.language_model.shared.weight", + } + + def __init__(self, config: Florence2Config): + super().__init__(config) + self.model = Florence2Model(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.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Florence2Seq2SeqLMOutput: + 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 PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, Florence2ForConditionalGeneration + + >>> model = Florence2ForConditionalGeneration.from_pretrained("florence-community/Florence-2-large") + >>> processor = AutoProcessor.from_pretrained("florence-community/Florence-2-large") + + >>> prompt = "" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(text=prompt, images=image, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_length=100) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "A green car parked in front of a yellow building." + ```""" + 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.text_config.pad_token_id, self.config.text_config.decoder_start_token_id + ) + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + 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 Florence2Seq2SeqLMOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + 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 + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + return self.model.get_placeholder_mask( + input_ids=input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + + def _prepare_encoder_decoder_kwargs_for_generation( + self, + inputs_tensor: torch.Tensor, + model_kwargs, + model_input_name: str | None, + generation_config, + ) -> dict[str, Any]: + # override to handle merging image and text embeddings before passing to language encoder + inputs_embeds = model_kwargs.pop("inputs_embeds", None) + pixel_values = model_kwargs.pop("pixel_values", None) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(inputs_tensor) + + if pixel_values is not None: + image_features = self.get_image_features(pixel_values).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + inputs_tensor, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + model_kwargs["inputs_embeds"] = inputs_embeds + model_kwargs = super()._prepare_encoder_decoder_kwargs_for_generation( + None, model_kwargs, model_input_name, generation_config + ) + model_kwargs.pop("inputs_embeds", None) + return model_kwargs + + +__all__ = [ + "Florence2Model", + "Florence2ForConditionalGeneration", + "Florence2PreTrainedModel", + "Florence2VisionBackbone", + "Florence2VisionPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/florence2/modular_florence2.py b/third_party/transformers/src/transformers/models/florence2/modular_florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..843fe79f2bb4ce72709044757931595cb85cf9e0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/florence2/modular_florence2.py @@ -0,0 +1,1676 @@ +# Copyright 2025 Microsoft and the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +import re +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch.nn as nn +import torch.nn.functional as F +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...configuration_utils import PreTrainedConfig +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...modeling_outputs import BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import MultiModalData, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import CONFIG_MAPPING, AutoConfig +from ..bart.modeling_bart import eager_attention_forward, shift_tokens_right +from ..beit.modeling_beit import BeitDropPath +from ..llama4.modeling_llama4 import Llama4VisionMLP +from ..llava.modeling_llava import LlavaForConditionalGeneration, LlavaModel, LlavaPreTrainedModel +from ..llava.processing_llava import LlavaProcessorKwargs + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="florence-community/Florence-2-base") +@strict +class Florence2VisionConfig(PreTrainedConfig): + r""" + depths (`Tuple[int]`, *optional*, defaults to `(1, 1, 9, 1)`): + The depth of the model. + patch_stride (`Tuple[int]`, *optional*, defaults to `(4, 2, 2, 2)`): + The patch stride of the image. + patch_padding (`Tuple[int]`, *optional*, defaults to `(3, 1, 1, 1)`): + The patch padding of the image. + patch_prenorm (`Tuple[bool]`, *optional*, defaults to `(False, True, True, True)`): + Whether to apply layer normalization before the patch embedding layer. + num_groups (`Tuple[int]`, *optional*, defaults to `(4, 8, 16, 32)`): + The number of groups. + window_size (`int`, *optional*, defaults to 12): + The window size of the model. + max_temporal_embeddings (`int`, *optional*, defaults to 100): + The configuration of the visual temporal embedding. + + Example: + + ```python + >>> from transformers import Florence2VisionConfig, Florence2VisionModel + + >>> # Initializing a Florence2 Vision style configuration + >>> configuration = Florence2VisionConfig() + + >>> # Initializing a model (with random weights) + >>> model = Florence2VisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "florence_vision" + + in_channels: int = 3 + depths: list[int] | tuple[int, ...] = (1, 1, 9, 1) + patch_size: list[int] | tuple[int, ...] = (7, 3, 3, 3) + patch_stride: list[int] | tuple[int, ...] = (4, 2, 2, 2) + patch_padding: list[int] | tuple[int, ...] = (3, 1, 1, 1) + patch_prenorm: list[bool] | tuple[bool, ...] = (False, True, True, True) + embed_dim: list[int] | tuple[int, ...] = (128, 256, 512, 1024) + num_heads: list[int] | tuple[int, ...] = (4, 8, 16, 32) + num_groups: list[int] | tuple[int, ...] = (4, 8, 16, 32) + window_size: int = 12 + drop_path_rate: float | int = 0.1 + mlp_ratio: float = 4.0 + qkv_bias: bool = True + activation_function: str = "gelu" + projection_dim: int = 1024 + max_temporal_embeddings: int = 100 + max_position_embeddings: int = 50 + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="florence-community/Florence-2-base") +@strict +class Florence2Config(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import Florence2ForConditionalGeneration, Florence2Config, CLIPVisionConfig, BartConfig + + >>> # Initializing a clip-like vision config + >>> vision_config = CLIPVisionConfig() + + >>> # Initializing a Bart config + >>> text_config = BartConfig() + + >>> # Initializing a Florence-2 configuration + >>> configuration = Florence2Config(vision_config, text_config) + + >>> # Initializing a model from the florence-2 configuration + >>> model = Florence2ForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "florence2" + sub_configs = { + "text_config": AutoConfig, + "vision_config": Florence2VisionConfig, + } + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + image_token_id: int = 51289 + is_encoder_decoder: bool = True + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "bart") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["bart"]() + + if isinstance(self.vision_config, dict): + self.vision_config = Florence2VisionConfig(**self.vision_config) + elif self.vision_config is None: + logger.info("vision_config is None. Initializing the Florence2VisionConfig with default values.") + self.vision_config = Florence2VisionConfig() + + super().__post_init__(**kwargs) + + +class Florence2ProcessorKwargs(LlavaProcessorKwargs): + pass + + +@auto_docstring +class Florence2Processor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + num_additional_image_tokens: int = 0, + post_processor_config: dict | None = None, + **kwargs, + ): + r""" + num_additional_image_tokens (`int`, *optional*, defaults to 0): + Number of additional tokens added to the image embeddings, such as CLS (+1). If the backbone has no CLS or other + extra tokens appended, no need to set this arg. + post_processor_config (`dict`, *optional*, defaults to 0): + Task-specific parsing rules for [`Florence2PostProcessor`], e.g. regex patterns, + thresholds, or banned tokens. + """ + self.tasks_answer_post_processing_type = { + "": "pure_text", + "": "ocr", + "": "pure_text", + "": "pure_text", + "": "pure_text", + "": "description_with_bboxes", + "": "description_with_bboxes", + "": "phrase_grounding", + "": "polygons", + "": "polygons", + "": "description_with_bboxes_or_polygons", + "": "pure_text", + "": "pure_text", + "": "pure_text", + "": "bboxes", + } + + self.task_prompts_without_inputs = { + "": "What is the text in the image?", + "": "What is the text in the image, with regions?", + "": "What does the image describe?", + "": "Describe in detail what is shown in the image.", + "": "Describe with a paragraph what is shown in the image.", + "": "Locate the objects with category name in the image.", + "": "Locate the objects in the image, with their descriptions.", + "": "Locate the region proposals in the image.", + } + + self.task_prompts_with_input = { + "": "Locate the phrases in the caption: {input}", + "": "Locate {input} in the image with mask", + "": "What is the polygon mask of region {input}", + "": "Locate {input} in the image.", + "": "What is the region {input}?", + "": "What does the region {input} describe?", + "": "What text is in the region {input}?", + } + + self.num_image_tokens = image_processor.image_seq_length + self.num_additional_image_tokens = num_additional_image_tokens + self.post_processor_config = post_processor_config + self.post_processor = Florence2PostProcessor(config=post_processor_config, tokenizer=tokenizer) + self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id + + super().__init__(image_processor, tokenizer, **kwargs) + + def _construct_prompts(self, text: str | list[str]) -> list[str]: + """ + Construct prompts by replacing task tokens with corresponding prompt strings. + """ + if isinstance(text, str): + text = [text] + + prompts = [] + for prompt in text: + # Check for tasks without inputs + for task_token, task_prompt in self.task_prompts_without_inputs.items(): + if task_token in prompt: + if prompt != task_token: + raise ValueError(f"Task token {task_token} should be the only content in the prompt.") + prompt = task_prompt + break + # Check for tasks with inputs + for task_token, task_prompt in self.task_prompts_with_input.items(): + if task_token in prompt: + input_text = prompt.replace(task_token, "").strip() + prompt = task_prompt.format(input=input_text) + break + prompts.append(prompt) + return prompts + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[Florence2ProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + if images is None and text is None: + raise ValueError("You have to specify at least one of `images` or `text`.") + + output_kwargs = self._merge_kwargs( + Florence2ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + image_inputs = {} + if images is not None: + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + + if text is None: + logger.warning_once("You are using Florence-2 without a text prefix.") + text = [""] * (1 if not isinstance(images, list) else len(images)) + elif isinstance(text, str): + text = [text] + + if not isinstance(text, list) or not all(isinstance(token, str) for token in text): + raise ValueError("`text` must be a string or list of strings.") + + if isinstance(images, list) and len(images) != len(text): + raise ValueError(f"Number of images ({len(images)}) must match number of texts ({len(text)}).") + + prompt_strings = self._construct_prompts(text) + + # Add image tokens and special tokens if images are provided + if image_inputs.get("pixel_values") is not None: + # Replace the image token with the expanded image token sequence + expanded_image_prompts = [] + for sample in prompt_strings: + sample = ( + self.image_token * self.num_image_tokens + + self.tokenizer.bos_token + + sample + + self.tokenizer.eos_token + ) + expanded_image_prompts.append(sample) + prompt_strings = expanded_image_prompts + + # Construct and tokenize prompts + output_kwargs["text_kwargs"].pop("add_special_tokens", None) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer( + prompt_strings, **output_kwargs["text_kwargs"], add_special_tokens=False, return_tensors=None + ) + self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**image_inputs, **text_inputs}, tensor_type=return_tensors) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + + vision_data = {} + if image_sizes is not None: + num_image_tokens = [self.num_image_tokens] * len(image_sizes) + num_image_patches = [1] * len(image_sizes) + + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=False, **kwargs): + """ + Post-processes the output of `FuyuForConditionalGeneration` to only return the text output. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + containing the token ids of the generated sequences. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `list[str]`: The decoded text output. + """ + return self.batch_decode(generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs) + + def post_process_generation(self, text=None, sequence=None, task=None, image_size=None) -> dict[str, Any]: + """ + Post-process generation outputs based on the task. + + Args: + text (`str`, *optional*): + Generated text. + sequence (`Union[List[int], torch.Tensor]`, *optional*): + Generated token sequence. + task (`str`, *optional*): + The task for post-processing. + image_size (`Tuple[int, int]`, *optional*): + Image size for dequantization. + + Returns: + `Dict[str, Any]`: Post-processed results keyed by task. + """ + if task is None: + raise ValueError("`task` must be provided for post-processing.") + + post_proc_type = self.tasks_answer_post_processing_type.get(task, "pure_text") + parsed = self.post_processor( + text=text, + sequence=sequence, + image_size=image_size, + parse_tasks=[post_proc_type], + )[post_proc_type] + + if post_proc_type == "pure_text": + final_answer = parsed.replace("", "").replace("", "").strip() + elif post_proc_type in ["description_with_bboxes", "bboxes"]: + bboxes = [inst["bbox"] for inst in parsed] + labels = [inst["cat_name"] for inst in parsed] + final_answer = {"bboxes": bboxes, "labels": labels} + if parsed and "score" in parsed[0]: + final_answer["scores"] = [inst["score"] for inst in parsed] + elif post_proc_type == "ocr": + quad_boxes = [inst["quad_box"] for inst in parsed] + labels = [inst["text"] for inst in parsed] + final_answer = {"quad_boxes": quad_boxes, "labels": labels} + elif post_proc_type == "phrase_grounding": + bboxes = [] + labels = [] + for inst in parsed: + for bbox in inst["bbox"]: + bboxes.append(bbox) + labels.append(inst["cat_name"]) + final_answer = {"bboxes": bboxes, "labels": labels} + elif post_proc_type in ["description_with_polygons", "polygons"]: + polygons = [inst["polygons"] for inst in parsed] + labels = [inst["cat_name"] for inst in parsed] + final_answer = {"polygons": polygons, "labels": labels} + elif post_proc_type == "description_with_bboxes_or_polygons": + bboxes = [] + bboxes_labels = [] + polygons = [] + polygons_labels = [] + for inst in parsed: + label = inst["cat_name"] + if "polygons" in inst: + polygons.append(inst["polygons"]) + polygons_labels.append(label) + else: + bboxes.append(inst["bbox"]) + bboxes_labels.append(label) + final_answer = { + "bboxes": bboxes, + "bboxes_labels": bboxes_labels, + "polygons": polygons, + "polygons_labels": polygons_labels, + } + else: + raise ValueError(f"Unknown post-processing type: {post_proc_type}") + + return {task: final_answer} + + +class Florence2PostProcessor: + """ + Post-processor for Florence-2 model outputs. Parses generated text into structured results for various tasks + like object detection, OCR, phrase grounding, etc. + + Args: + tokenizer (`PreTrainedTokenizer`): + The tokenizer used for decoding model outputs. + """ + + def __init__(self, config, tokenizer): + self.tokenizer = tokenizer + self.parse_task_config = config or {} + self.banned_grounding_tokens = set( + self.parse_task_config.get("phrase_grounding", {}).get("banned_grounding_tokens", []) + ) + self.all_special_tokens = set(self.tokenizer.all_special_tokens) + self.quantize_bins = (1000, 1000) + + def quantize(self, locations: "torch.Tensor", size: tuple[int, int]) -> "torch.Tensor": + """ + Quantize locations. + + Args: + locations (`torch.Tensor`): + Tensor of shape (N, 4) for boxes or (N, 2) for points/coordinates. + size (`tuple[int, int]`): + Original image size (width, height). + + Returns: + `torch.Tensor`: Quantized locations as integers. + """ + bins_w, bins_h = self.quantize_bins + size_w, size_h = size + per_bin_w = size_w / bins_w + per_bin_h = size_h / bins_h + + if locations.shape[-1] == 4: # Bounding boxes: [xmin, ymin, xmax, ymax] + xmin, ymin, xmax, ymax = locations.split(1, dim=-1) + q_xmin = (xmin / per_bin_w).floor().clamp(0, bins_w - 1) + q_ymin = (ymin / per_bin_h).floor().clamp(0, bins_h - 1) + q_xmax = (xmax / per_bin_w).floor().clamp(0, bins_w - 1) + q_ymax = (ymax / per_bin_h).floor().clamp(0, bins_h - 1) + return torch.cat([q_xmin, q_ymin, q_xmax, q_ymax], dim=-1).int() + + elif locations.shape[-1] == 2: # Points/coordinates: [x, y] + x, y = locations.split(1, dim=-1) + q_x = (x / per_bin_w).floor().clamp(0, bins_w - 1) + q_y = (y / per_bin_h).floor().clamp(0, bins_h - 1) + return torch.cat([q_x, q_y], dim=-1).int() + + else: + raise ValueError(f"Unsupported location shape: last dim must be 2 or 4, got {locations.shape[-1]}.") + + def dequantize(self, locations: "torch.Tensor", size: tuple[int, int]) -> "torch.Tensor": + """ + Dequantize locations back to original scale. + + Args: + locations (`torch.Tensor`): + Quantized tensor of shape (N, 4) for boxes or (N, 2) for points/coordinates. + size (`tuple[int, int]`): + Original image size (width, height). + + Returns: + `torch.Tensor`: Dequantized locations as floats. + """ + bins_w, bins_h = self.quantize_bins + size_w, size_h = size + per_bin_w = size_w / bins_w + per_bin_h = size_h / bins_h + + # Add 0.5 to use the center position of the bin as the coordinate. + if locations.shape[-1] == 4: # Bounding boxes + xmin, ymin, xmax, ymax = locations.split(1, dim=-1) + dq_xmin = (xmin + 0.5) * per_bin_w + dq_ymin = (ymin + 0.5) * per_bin_h + dq_xmax = (xmax + 0.5) * per_bin_w + dq_ymax = (ymax + 0.5) * per_bin_h + return torch.cat([dq_xmin, dq_ymin, dq_xmax, dq_ymax], dim=-1).int() + + elif locations.shape[-1] == 2: # Points/coordinates + x, y = locations.split(1, dim=-1) + dq_x = (x + 0.5) * per_bin_w + dq_y = (y + 0.5) * per_bin_h + return torch.cat([dq_x, dq_y], dim=-1).int() + + else: + raise ValueError(f"Unsupported location shape: last dim must be 2 or 4, got {locations.shape[-1]}.") + + def decode_with_spans(self, token_ids: list[int]) -> tuple[str, list[tuple[int, int]]]: + """ + Decode token IDs to text and compute character spans. + + Args: + token_ids (`list[int]`): + list of token IDs to decode. + + Returns: + `tuple[str, list[tuple[int, int]]]`: Decoded text and list of spans (start, end) for each token. + """ + filtered_tokens = self.tokenizer.convert_ids_to_tokens(token_ids, skip_special_tokens=False) + text = "" + spans = [] + for token in filtered_tokens: + if token in self.all_special_tokens: + sub_text = token + else: + sub_text = self.tokenizer.convert_tokens_to_string([token]) + span = (len(text), len(text) + len(sub_text)) + text += sub_text + spans.append(span) + return text, spans + + def parse_ocr_from_text_and_spans( + self, text: str, pattern: str | None, image_size: tuple[int, int], area_threshold: float = 0.0 + ) -> list[dict[str, Any]]: + """ + Parse OCR results with quadrilateral boxes. + + Args: + text (`str`): + The generated text. + pattern (`str`): + Regex pattern for matching. + image_size (`tuple[int, int]`): + Image size (width, height). + area_threshold (`float`, *optional*, defaults to 0.0): + Minimum area threshold for filtering boxes. + + Returns: + `list[dict[str, Any]]`: list of instances with 'quad_box' and 'text'. + """ + text = text.replace("", "").replace("", "").replace("", "") + if pattern is None: + pattern = r"(.+?)" + + matches = re.findall(pattern, text) + instances = [] + width, height = image_size + + for content, *quad_str in matches: + quad_bins = [int(i) for i in quad_str] + quad_box = self.dequantize(torch.tensor(quad_bins).reshape(-1, 2), size=image_size).flatten().tolist() + + if area_threshold > 0: + x_coords = quad_box[0::2] + y_coords = quad_box[1::2] + # Apply the Shoelace formula + area = 0.5 * abs( + sum(x_coords[i] * y_coords[i + 1] - x_coords[i + 1] * y_coords[i] for i in range(4 - 1)) + ) + + if area < (width * height) * area_threshold: + continue + + instances.append({"quad_box": quad_box, "text": content.strip()}) + return instances + + def parse_phrase_grounding_from_text_and_spans( + self, text: str, image_size: tuple[int, int] + ) -> list[dict[str, Any]]: + """ + Parse phrase grounding results. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + + Returns: + `list[dict[str, Any]]`: list of instances with 'bbox' and 'cat_name'. + """ + text = text.replace("", "").replace("", "").replace("", "") + phrase_pattern = r"([^<]+(?:){4,})" + phrases = re.findall(phrase_pattern, text) + text_pattern = r"^\s*(.*?)(?=||||||" + + instances = [] + for phrase_text in phrases: + phrase_text = phrase_text.replace("", "", 1).replace("", "", 1) + if not phrase_text: + continue + match = re.search(text_pattern, phrase_text) + if not match: + continue + phrase = match.group().strip() + if phrase in self.banned_grounding_tokens: + continue + boxes_matches = list(re.finditer(box_pattern, phrase_text)) + if not boxes_matches: + continue + bbox_bins = [[int(m.group(j)) for j in range(1, 5)] for m in boxes_matches] + bboxes = self.dequantize(torch.tensor(bbox_bins), size=image_size).tolist() + phrase = phrase.encode("ascii", "ignore").decode("ascii") + instances.append({"bbox": bboxes, "cat_name": phrase}) + return instances + + def _find_matched_token_indices(self, cur_span: tuple[int, int], token_spans: list[tuple[int, int]]) -> list[int]: + return [i for i, span in enumerate(token_spans) if not (span[1] <= cur_span[0] or span[0] >= cur_span[1])] + + def parse_description_with_bboxes_from_text_and_spans( + self, + text: str, + image_size: tuple[int, int], + allow_empty_phrase: bool = False, + ) -> list[dict[str, Any]]: + """ + Parse descriptions with bounding boxes. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + allow_empty_phrase (`bool`, *optional*, defaults to `False`): + Allow phrases without text. + + Returns: + `list[dict[str, Any]]`: list of instances with 'bbox', 'cat_name', and optional 'score'. + """ + text = text.replace("", "").replace("", "").replace("", "") + + if allow_empty_phrase: + pattern = r"(?:(?:){4,})" + else: + pattern = r"([^<]+(?:){4,})" + phrases = re.findall(pattern, text) + + text_pattern = r"^\s*(.*?)(?=||||||" + + instances = [] + for phrase_text in phrases: + phrase_text = phrase_text.replace("", "", 1).replace("", "", 1) + if not phrase_text and not allow_empty_phrase: + continue + match = re.search(text_pattern, phrase_text) + if not match: + continue + phrase = match.group().strip() + boxes_matches = list(re.finditer(box_pattern, phrase_text)) + if not boxes_matches: + continue + bbox_bins = [[int(m.group(j)) for j in range(1, 5)] for m in boxes_matches] + bboxes = self.dequantize(torch.tensor(bbox_bins), size=image_size).tolist() + + phrase = phrase.encode("ascii", "ignore").decode("ascii") + for bbox in bboxes: + instance = {"bbox": bbox, "cat_name": phrase} + instances.append(instance) + + return instances + + def parse_description_with_polygons_from_text_and_spans( + self, + text: str, + image_size: tuple[int, int], + allow_empty_phrase: bool = False, + polygon_sep_token: str = "", + polygon_start_token: str = "", + polygon_end_token: str = "", + with_box_at_start: bool = False, + ) -> list[dict[str, Any]]: + """ + Parse descriptions with polygons. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + allow_empty_phrase (`bool`, *optional*, defaults to `False`): + Allow phrases without text. + polygon_sep_token (`str`, *optional*, defaults to ""): + Token separating polygons. + polygon_start_token (`str`, *optional*, defaults to ""): + Start token for polygons. + polygon_end_token (`str`, *optional*, defaults to ""): + End token for polygons. + with_box_at_start (`bool`, *optional*, defaults to `False`): + Whether a bounding box is at the start of polygons. + + Returns: + `list[dict[str, Any]]`: list of instances with 'polygons', 'cat_name', and optional 'bbox'. + """ + text = text.replace("", "").replace("", "").replace("", "") + + if allow_empty_phrase: + pattern = rf"(?:(?:|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})" + else: + pattern = rf"([^<]+(?:|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})" + phrases = re.findall(pattern, text) + phrase_pattern = r"^\s*(.*?)(?=||||||)" + poly_instance_pattern = rf"{re.escape(polygon_start_token)}(.*?){re.escape(polygon_end_token)}" + box_pattern = rf"((?:)+)(?:{re.escape(polygon_sep_token)}|$)" + + instances = [] + for phrase_text in phrases: + phrase_text_strip = re.sub(r"^", "", phrase_text, count=1) + if not phrase_text_strip and not allow_empty_phrase: + continue + match = re.search(phrase_pattern, phrase_text_strip) + if not match: + continue + phrase = match.group().strip() + + if polygon_start_token in phrase_text and polygon_end_token in phrase_text: + poly_instances = [m.group(1) for m in re.finditer(poly_instance_pattern, phrase_text)] + else: + poly_instances = [phrase_text] + + for poly_inst in poly_instances: + poly_matches = list(re.finditer(box_pattern, poly_inst)) + if len(poly_matches) == 0: + continue + bbox = [] + polygons = [] + for poly_match in poly_matches: + poly_str = poly_match.group(1) + poly_bins = [int(m.group(1)) for m in re.finditer(r"", poly_str)] + if with_box_at_start and not bbox: + if len(poly_bins) > 4: + bbox = poly_bins[:4] + poly_bins = poly_bins[4:] + else: + bbox = [0, 0, 0, 0] + if len(poly_bins) % 2 == 1: + poly_bins = poly_bins[:-1] + poly_coords = ( + self.dequantize(torch.tensor(poly_bins).reshape(-1, 2), size=image_size).flatten().tolist() + ) + polygons.append(poly_coords) + + instance = {"cat_name": phrase, "polygons": polygons} + if bbox: + instance["bbox"] = self.dequantize(torch.tensor([bbox]), size=image_size)[0].tolist() + instances.append(instance) + return instances + + def __call__(self, text=None, sequence=None, image_size=None, parse_tasks=None) -> dict[str, Any]: + """ + Process model output and parse into task-specific results. + + Args: + text (`Optional[str]`, *optional*): + Generated text. Either this or `sequence` must be provided. + sequence (`Optional[Union[list[int], torch.Tensor]]`, *optional*): + Token sequence. Either this or `text` must be provided. + image_size (`Optional[tuple[int, int]]`, *optional*): + Image size (width, height) required for dequantization. + parse_tasks (`Optional[Union[str, list[str]]]`, *optional*): + Specific tasks to parse. If None, parse all supported tasks. + + Returns: + `dict[str, Any]`: Parsed results for each task, including the raw 'text'. + """ + if parse_tasks is not None: + parse_tasks = [parse_tasks] if isinstance(parse_tasks, str) else parse_tasks + for task in parse_tasks: + if task not in self.parse_task_config.keys(): + raise ValueError(f"Unsupported parse task: {task}") + + if (text is None and sequence is None) or (text is not None and sequence is not None): + raise ValueError("Exactly one of 'text' or 'sequence' must be provided.") + + if sequence is not None: + if isinstance(sequence, torch.Tensor): + sequence = sequence.tolist() + sequence = sequence[1:] if sequence[0] == self.tokenizer.bos_token_id else sequence # Skip BOS if present + text, _ = self.decode_with_spans(sequence) + + parsed_dict = {"text": text} + + tasks_to_parse = parse_tasks or self.parse_task_config.keys() + for task in tasks_to_parse: + config = self.parse_task_config[task] + pattern = config.get("PATTERN") + + if task == "ocr": + parsed_dict["ocr"] = self.parse_ocr_from_text_and_spans( + text, pattern=pattern, image_size=image_size, area_threshold=config.get("AREA_THRESHOLD", 0.0) + ) + elif task == "phrase_grounding": + parsed_dict["phrase_grounding"] = self.parse_phrase_grounding_from_text_and_spans( + text, image_size=image_size + ) + elif task == "pure_text": + parsed_dict["pure_text"] = text + elif task == "description_with_bboxes": + parsed_dict["description_with_bboxes"] = self.parse_description_with_bboxes_from_text_and_spans( + text, image_size=image_size + ) + elif task == "description_with_polygons": + parsed_dict["description_with_polygons"] = self.parse_description_with_polygons_from_text_and_spans( + text, image_size=image_size + ) + elif task == "polygons": + parsed_dict["polygons"] = self.parse_description_with_polygons_from_text_and_spans( + text, image_size=image_size, allow_empty_phrase=True + ) + elif task == "bboxes": + parsed_dict["bboxes"] = self.parse_description_with_bboxes_from_text_and_spans( + text, image_size=image_size, allow_empty_phrase=True + ) + elif task == "description_with_bboxes_or_polygons": + if "" in text: + instances = self.parse_description_with_polygons_from_text_and_spans(text, image_size=image_size) + else: + instances = self.parse_description_with_bboxes_from_text_and_spans(text, image_size=image_size) + parsed_dict["description_with_bboxes_or_polygons"] = instances + else: + raise ValueError(f"task {task} is not supported") + + return parsed_dict + + +class Florence2VisionDropPath(BeitDropPath): + pass + + +class Florence2VisionLearnedAbsolutePositionEmbedding2D(nn.Module): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, config: Florence2Config): + super().__init__() + num_pos = config.vision_config.max_position_embeddings + embedding_dim = config.vision_config.embed_dim[-1] + self.row_embeddings = nn.Embedding(num_pos, embedding_dim // 2) + self.column_embeddings = nn.Embedding(num_pos, embedding_dim - (embedding_dim // 2)) + + def forward(self, pixel_values, pixel_mask=None): + height, width = pixel_values.shape[-2:] + width_values = torch.arange(width, device=pixel_values.device) + height_values = torch.arange(height, device=pixel_values.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(pixel_values.shape[0], 1, 1, 1) + return pos + + +class Florence2VisionPositionalEmbeddingCosine1D(nn.Module): + """ + This module generates 1D cosine positional embeddings using precomputed sinusoidal functions. + """ + + def __init__(self, config: Florence2Config): + super().__init__() + self.embed_dim = config.vision_config.embed_dim[-1] + self.max_seq_len = config.vision_config.max_temporal_embeddings + pos_idx_to_embed = torch.empty((self.max_seq_len, self.embed_dim)) + sine, cosine = self.get_sinusoid_embeddings( + max_positions=self.max_seq_len, + embed_dim=self.embed_dim, + ) + pos_idx_to_embed[:, 0::2] = sine + pos_idx_to_embed[:, 1::2] = cosine + # Save the positional embeddings in a constant buffer. + self.register_buffer("pos_idx_to_embed", pos_idx_to_embed) + + @staticmethod + def get_sinusoid_embeddings(max_positions: int, embed_dim: int): + half_dim = embed_dim // 2 + emb = math.log(10000) / half_dim + emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb) + emb = torch.arange(max_positions, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0) + return torch.sin(emb), torch.cos(emb) + + def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor: + len_seq = seq_embeds.size(1) + if len_seq > self.max_seq_len: + raise ValueError(f"Maximum sequence length {self.max_seq_len}, got {len_seq}") + pos_embeds = self.pos_idx_to_embed[0:len_seq, :] + return pos_embeds + + +class Florence2VisionMLP(Llama4VisionMLP): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__(config) + self.fc1 = nn.Linear(config.embed_dim[stage_idx], int(config.embed_dim[stage_idx] * config.mlp_ratio)) + self.activation_fn = ACT2FN[config.activation_function] + self.fc2 = nn.Linear(int(config.embed_dim[stage_idx] * config.mlp_ratio), config.embed_dim[stage_idx]) + + +class Florence2VisionConvEmbed(nn.Module): + """Image to Patch Embedding""" + + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.stage_idx = stage_idx + self.patch_size = config.patch_size[stage_idx] + self.in_channels = config.in_channels if stage_idx == 0 else config.embed_dim[stage_idx - 1] + self.embed_dim = config.embed_dim[stage_idx] + self.stride = config.patch_stride[stage_idx] + self.padding = config.patch_padding[stage_idx] + self.pre_norm = config.patch_prenorm[stage_idx] + + self.conv = nn.Conv2d( + self.in_channels, + self.embed_dim, + kernel_size=self.patch_size, + stride=self.stride, + padding=self.padding, + ) + + dim_norm = self.in_channels if self.pre_norm else self.embed_dim + self.norm = nn.LayerNorm(dim_norm) + + def forward(self, hidden_states: torch.Tensor): + if self.norm and self.pre_norm: + hidden_states = hidden_states.permute(0, 2, 3, 1) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 3, 1, 2) + + hidden_states = self.conv(hidden_states) + + if self.norm and not self.pre_norm: + hidden_states = hidden_states.permute(0, 2, 3, 1) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 3, 1, 2) + return hidden_states + + +class Florence2VisionChannelAttention(nn.Module): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.dim = config.embed_dim[stage_idx] + self.groups = config.num_groups[stage_idx] + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=config.qkv_bias) + self.proj = nn.Linear(self.dim, self.dim) + self.is_causal = False + + def forward(self, hidden_states: torch.Tensor): + batch_size, num_tokens, hidden_size = hidden_states.shape + + # Reshape for grouped channel attention + qkv = self.qkv(hidden_states).reshape(batch_size, num_tokens, 3, self.groups, hidden_size // self.groups) + qkv = qkv.permute(2, 0, 3, 4, 1) + query, key, value = qkv.unbind(0) + + scale = num_tokens**-0.5 + # Channel-to-channel attention within groups: + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + hidden_states, _ = attention_interface( + self, + query, + key, + value, + attention_mask=None, + scaling=scale, + ) + hidden_states = hidden_states.permute(0, 3, 2, 1) + hidden_states = hidden_states.reshape(batch_size, num_tokens, hidden_size) + + # Final projection + hidden_states = self.proj(hidden_states) + return hidden_states + + +class Florence2VisionChannelBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + drop_path_rate: float, + ): + super().__init__() + + self.config = config + dim_in = config.embed_dim[stage_idx] + + self.conv1 = nn.Conv2d( + dim_in, + dim_in, + kernel_size=3, + padding=1, + groups=dim_in, + ) + self.norm1 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.channel_attn = Florence2VisionChannelAttention(config=config, stage_idx=stage_idx) + self.drop_path1 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + self.conv2 = nn.Conv2d( + dim_in, + dim_in, + kernel_size=3, + padding=1, + groups=dim_in, + ) + self.norm2 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.ffn = Florence2VisionMLP(config=config, stage_idx=stage_idx) + self.drop_path2 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + def forward(self, hidden_states: torch.Tensor): + batch_size, embed_dim, height, width = hidden_states.shape + + # First channel block: Depthwise Conv + Channel Attention + hidden_states = self.conv1(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # Channel group attention self-attention mechanism + hidden_states = self.norm1(hidden_states) + hidden_states = self.channel_attn(hidden_states) + hidden_states = residual + self.drop_path1(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + # Second channel block: Depthwise Conv + FFN + hidden_states = self.conv2(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # FFN + hidden_states = self.norm2(hidden_states) + hidden_states = self.ffn(hidden_states) + hidden_states = residual + self.drop_path2(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + return hidden_states + + +class Florence2VisionWindowAttention(nn.Module): + def __init__(self, config: Florence2VisionConfig, stage_idx: int): + super().__init__() + self.config = config + self.dim = config.embed_dim[stage_idx] + self.window_size = config.window_size + self.num_heads = config.num_heads[stage_idx] + head_dim = self.dim // self.num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=config.qkv_bias) + self.proj = nn.Linear(self.dim, self.dim) + self.is_causal = False + + def forward(self, hidden_states: torch.Tensor): + batch_size, height, width, embed_dim = hidden_states.shape + + # Pad the input if necessary + pad_left = pad_top = 0 + pad_right = (self.window_size - width % self.window_size) % self.window_size + pad_bottom = (self.window_size - height % self.window_size) % self.window_size + hidden_states = F.pad(hidden_states, (0, 0, pad_left, pad_right, pad_top, pad_bottom)) + _, padded_height, padded_width, _ = hidden_states.shape + + # Partition input into non-overlapping windows (for local spatial attention in DaViT) + hidden_states = hidden_states.view( + batch_size, + padded_height // self.window_size, + self.window_size, + padded_width // self.window_size, + self.window_size, + embed_dim, + ) + windowed_hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous() + windowed_hidden_states = windowed_hidden_states.view(-1, self.window_size * self.window_size, embed_dim) + + # Generate Q, K, V for each window + num_windows_per_batch, num_tokens_per_window, embed_dim = windowed_hidden_states.shape + qkv = self.qkv(windowed_hidden_states).reshape( + num_windows_per_batch, num_tokens_per_window, 3, self.num_heads, embed_dim // self.num_heads + ) + qkv = qkv.permute(2, 0, 3, 1, 4) + query, key, value = qkv.unbind(0) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + windowed_hidden_states, _ = attention_interface( + self, + query, + key, + value, + attention_mask=None, + scaling=self.scale, + ) + windowed_hidden_states = windowed_hidden_states.view(num_windows_per_batch, num_tokens_per_window, embed_dim) + windowed_hidden_states = self.proj(windowed_hidden_states) + + # Merge windows back to original spatial layout + windowed_hidden_states = windowed_hidden_states.view(-1, self.window_size, self.window_size, embed_dim) + hidden_states = windowed_hidden_states.view( + -1, + padded_height // self.window_size, + padded_width // self.window_size, + self.window_size, + self.window_size, + embed_dim, + ) + hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous() + hidden_states = hidden_states.view(-1, padded_height, padded_width, embed_dim) + hidden_states = hidden_states[:, :height, :width, :].contiguous() + hidden_states = hidden_states.view(batch_size, height * width, embed_dim) + + return hidden_states + + +class Florence2VisionSpatialBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + drop_path_rate: float, + ): + super().__init__() + + self.conv1 = nn.Conv2d( + config.embed_dim[stage_idx], + config.embed_dim[stage_idx], + kernel_size=3, + padding=1, + groups=config.embed_dim[stage_idx], + ) + self.norm1 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.window_attn = Florence2VisionWindowAttention(config=config, stage_idx=stage_idx) + self.drop_path1 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + self.conv2 = nn.Conv2d( + config.embed_dim[stage_idx], + config.embed_dim[stage_idx], + kernel_size=3, + padding=1, + groups=config.embed_dim[stage_idx], + ) + self.norm2 = nn.LayerNorm(config.embed_dim[stage_idx]) + self.ffn = Florence2VisionMLP(config=config, stage_idx=stage_idx) + self.drop_path2 = Florence2VisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + + def forward(self, hidden_states: torch.Tensor): + batch_size, embed_dim, height, width = hidden_states.shape + + # First spatial mixing block: Conv + Window Attention + hidden_states = self.conv1(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # Spatial Window-based self-attention mechanism + hidden_states = self.norm1(hidden_states) + hidden_states = hidden_states.view(batch_size, height, width, embed_dim) + hidden_states = self.window_attn(hidden_states) + hidden_states = residual + self.drop_path1(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + # Second spatial mixing block: Conv + FFN + hidden_states = self.conv2(hidden_states) + hidden_states + hidden_states = hidden_states.flatten(2).transpose(1, 2) + residual = hidden_states + + # FFN + hidden_states = self.norm2(hidden_states) + hidden_states = self.ffn(hidden_states) + hidden_states = residual + self.drop_path2(hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(batch_size, embed_dim, height, width) + + return hidden_states + + +class Florence2VisionBlock(nn.Module): + def __init__( + self, + config: Florence2VisionConfig, + stage_idx: int, + spatial_drop_path_rate: float, + channel_drop_path_rate: float, + ): + super().__init__() + self.spatial_block = Florence2VisionSpatialBlock( + config=config, + stage_idx=stage_idx, + drop_path_rate=spatial_drop_path_rate, + ) + self.channel_block = Florence2VisionChannelBlock( + config=config, + stage_idx=stage_idx, + drop_path_rate=channel_drop_path_rate, + ) + + def forward(self, hidden_states: torch.Tensor): + hidden_states = self.spatial_block(hidden_states) + hidden_states = self.channel_block(hidden_states) + return hidden_states + + +@auto_docstring +class Florence2VisionPreTrainedModel(PreTrainedModel): + config_class = Florence2VisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _can_record_outputs = { + "hidden_states": Florence2VisionBlock, + "attentions": [Florence2VisionChannelAttention, Florence2VisionWindowAttention], + } + + +@auto_docstring +class Florence2VisionBackbone(Florence2VisionPreTrainedModel): + def __init__(self, config: Florence2VisionConfig): + super().__init__(config) + self.config = config + + self.embed_dim = config.embed_dim + self.num_heads = config.num_heads + self.num_groups = config.num_groups + self.num_stages = len(self.embed_dim) + + if not (self.num_stages == len(self.num_heads) == len(self.num_groups)): + raise ValueError( + f"Expected self.num_stages ({self.num_stages}) == " + f"len(self.num_heads) ({len(self.num_heads)}) == " + f"len(self.num_groups) ({len(self.num_groups)})" + ) + + dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths) * 2, device="cpu")] + depth_offset = 0 + + convs = [] + blocks = [] + for stage_idx in range(self.num_stages): + conv_embed = Florence2VisionConvEmbed( + config=config, + stage_idx=stage_idx, + ) + convs.append(conv_embed) + + block = nn.ModuleList( + Florence2VisionBlock( + config=config, + stage_idx=stage_idx, + spatial_drop_path_rate=dpr[depth_offset + block_idx * 2], + channel_drop_path_rate=dpr[depth_offset + block_idx * 2 + 1], + ) + for block_idx in range(config.depths[stage_idx]) + ) + blocks.append(block) + depth_offset += config.depths[stage_idx] * 2 + + self.convs = nn.ModuleList(convs) + self.blocks = nn.ModuleList(blocks) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + target_dtype = self.convs[0].conv.weight.dtype + hidden_states = hidden_states.to(dtype=target_dtype) + for conv, block in zip(self.convs, self.blocks): + hidden_states = conv(hidden_states) + for layer in block: + hidden_states = layer(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + ) + + +class Florence2MultiModalProjector(nn.Module): + def __init__(self, config: Florence2Config): + super().__init__() + self.vision_embedding_dim = config.vision_config.embed_dim[-1] + self.vision_projection_dim = config.vision_config.projection_dim + self.image_projection = nn.Linear(self.vision_embedding_dim, self.vision_projection_dim, bias=False) + self.image_proj_norm = nn.LayerNorm(self.vision_projection_dim) + self.image_position_embed = Florence2VisionLearnedAbsolutePositionEmbedding2D(config=config) + self.visual_temporal_embed = Florence2VisionPositionalEmbeddingCosine1D(config=config) + + def forward(self, image_features): + position_features = image_features + self.image_position_embed(image_features) + position_features = position_features.flatten(2).transpose(1, 2) + temporal_features = self.visual_temporal_embed(position_features[:, :1, :]) + temporal_features = temporal_features.unsqueeze(1) + visual_token_features = position_features + temporal_features + visual_token_features = visual_token_features.unsqueeze(1) + spatial_image_features = visual_token_features.mean(dim=2) + temporal_image_features = visual_token_features.mean(dim=1) + image_features = torch.cat([spatial_image_features, temporal_image_features], dim=1) + image_features = self.image_projection(image_features) + image_features = self.image_proj_norm(image_features) + return image_features + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Florence-2 base model's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class Florence2Seq2SeqModelOutput(Seq2SeqModelOutput): + r""" + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_image_tokens, 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 + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Florence-2 model's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class Florence2Seq2SeqLMOutput(Seq2SeqLMOutput): + 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). + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_image_tokens, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + image_hidden_states: tuple[torch.FloatTensor, ...] | None = None + + +@auto_docstring +class Florence2PreTrainedModel(LlavaPreTrainedModel): + config_class = Florence2Config + base_model_prefix = "model" + + _supports_attention_backend = False + + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, Florence2VisionPositionalEmbeddingCosine1D): + pos_idx_to_embed = torch.empty((module.max_seq_len, module.embed_dim)) + sine, cosine = module.get_sinusoid_embeddings( + max_positions=module.max_seq_len, + embed_dim=module.embed_dim, + ) + pos_idx_to_embed[:, 0::2] = sine + pos_idx_to_embed[:, 1::2] = cosine + init.copy_(module.pos_idx_to_embed, pos_idx_to_embed) + + +@auto_docstring( + custom_intro=""" + Florence-2 is a vision model for captioning, detection, and segmentation. + """ +) +class Florence2Model(LlavaModel): + def __init__(self, config: Florence2Config): + super().__init__(config) + self.vision_tower = Florence2VisionBackbone(config=config.vision_config) + + def get_encoder(self, modality=None): + if modality is None: + return self.language_model.get_encoder() + else: + return super().get_encoder(modality=modality) + + @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.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. + """ + image_outputs = self.vision_tower(pixel_values, **kwargs) + image_outputs.pooler_output = self.multi_modal_projector(image_outputs.last_hidden_state) + + return image_outputs + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs, + ) -> tuple | Florence2Seq2SeqModelOutput: + if encoder_outputs is None: + 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).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + encoder_outputs = self.language_model.encoder( + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + if decoder_input_ids is None: + decoder_start_token_id = self.config.text_config.decoder_start_token_id + decoder_input_ids = torch.ones((inputs_embeds.size()[0], 1), dtype=torch.long, device=inputs_embeds.device) + decoder_input_ids *= decoder_start_token_id + + decoder_outputs = self.language_model.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs[0], + encoder_attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Florence2Seq2SeqModelOutput( + 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, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + +@auto_docstring( + custom_intro=""" + Florence-2 is a vision model for captioning, detection, and segmentation. + """ +) +class Florence2ForConditionalGeneration(LlavaForConditionalGeneration): + _tied_weights_keys = { + "lm_head.weight": "model.language_model.shared.weight", + } + + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Florence2Seq2SeqLMOutput: + 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 PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, Florence2ForConditionalGeneration + + >>> model = Florence2ForConditionalGeneration.from_pretrained("florence-community/Florence-2-large") + >>> processor = AutoProcessor.from_pretrained("florence-community/Florence-2-large") + + >>> prompt = "" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(text=prompt, images=image, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_length=100) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "A green car parked in front of a yellow building." + ```""" + 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.text_config.pad_token_id, self.config.text_config.decoder_start_token_id + ) + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + 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 Florence2Seq2SeqLMOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + return self.model.get_placeholder_mask( + input_ids=input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + + def _prepare_encoder_decoder_kwargs_for_generation( + self, + inputs_tensor: torch.Tensor, + model_kwargs, + model_input_name: str | None, + generation_config, + ) -> dict[str, Any]: + # override to handle merging image and text embeddings before passing to language encoder + inputs_embeds = model_kwargs.pop("inputs_embeds", None) + pixel_values = model_kwargs.pop("pixel_values", None) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(inputs_tensor) + + if pixel_values is not None: + image_features = self.get_image_features(pixel_values).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + inputs_tensor, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + model_kwargs["inputs_embeds"] = inputs_embeds + model_kwargs = super()._prepare_encoder_decoder_kwargs_for_generation( + None, model_kwargs, model_input_name, generation_config + ) + model_kwargs.pop("inputs_embeds", None) + return model_kwargs + + +__all__ = [ + "Florence2Config", + "Florence2Processor", + "Florence2VisionConfig", + "Florence2Model", + "Florence2ForConditionalGeneration", + "Florence2PreTrainedModel", + "Florence2VisionBackbone", + "Florence2VisionPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/florence2/processing_florence2.py b/third_party/transformers/src/transformers/models/florence2/processing_florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..94fb8bed3abc46a51553db9db983684686a96227 --- /dev/null +++ b/third_party/transformers/src/transformers/models/florence2/processing_florence2.py @@ -0,0 +1,760 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/florence2/modular_florence2.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_florence2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Microsoft and the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import re +from typing import Any + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, is_torch_available, logging + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class Florence2ProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": {"padding": False, "return_mm_token_type_ids": False}, + } + + +@auto_docstring +class Florence2Processor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + num_additional_image_tokens: int = 0, + post_processor_config: dict | None = None, + **kwargs, + ): + r""" + num_additional_image_tokens (`int`, *optional*, defaults to 0): + Number of additional tokens added to the image embeddings, such as CLS (+1). If the backbone has no CLS or other + extra tokens appended, no need to set this arg. + post_processor_config (`dict`, *optional*, defaults to 0): + Task-specific parsing rules for [`Florence2PostProcessor`], e.g. regex patterns, + thresholds, or banned tokens. + """ + self.tasks_answer_post_processing_type = { + "": "pure_text", + "": "ocr", + "": "pure_text", + "": "pure_text", + "": "pure_text", + "": "description_with_bboxes", + "": "description_with_bboxes", + "": "phrase_grounding", + "": "polygons", + "": "polygons", + "": "description_with_bboxes_or_polygons", + "": "pure_text", + "": "pure_text", + "": "pure_text", + "": "bboxes", + } + + self.task_prompts_without_inputs = { + "": "What is the text in the image?", + "": "What is the text in the image, with regions?", + "": "What does the image describe?", + "": "Describe in detail what is shown in the image.", + "": "Describe with a paragraph what is shown in the image.", + "": "Locate the objects with category name in the image.", + "": "Locate the objects in the image, with their descriptions.", + "": "Locate the region proposals in the image.", + } + + self.task_prompts_with_input = { + "": "Locate the phrases in the caption: {input}", + "": "Locate {input} in the image with mask", + "": "What is the polygon mask of region {input}", + "": "Locate {input} in the image.", + "": "What is the region {input}?", + "": "What does the region {input} describe?", + "": "What text is in the region {input}?", + } + + self.num_image_tokens = image_processor.image_seq_length + self.num_additional_image_tokens = num_additional_image_tokens + self.post_processor_config = post_processor_config + self.post_processor = Florence2PostProcessor(config=post_processor_config, tokenizer=tokenizer) + self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id + + super().__init__(image_processor, tokenizer, **kwargs) + + def _construct_prompts(self, text: str | list[str]) -> list[str]: + """ + Construct prompts by replacing task tokens with corresponding prompt strings. + """ + if isinstance(text, str): + text = [text] + + prompts = [] + for prompt in text: + # Check for tasks without inputs + for task_token, task_prompt in self.task_prompts_without_inputs.items(): + if task_token in prompt: + if prompt != task_token: + raise ValueError(f"Task token {task_token} should be the only content in the prompt.") + prompt = task_prompt + break + # Check for tasks with inputs + for task_token, task_prompt in self.task_prompts_with_input.items(): + if task_token in prompt: + input_text = prompt.replace(task_token, "").strip() + prompt = task_prompt.format(input=input_text) + break + prompts.append(prompt) + return prompts + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[Florence2ProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + if images is None and text is None: + raise ValueError("You have to specify at least one of `images` or `text`.") + + output_kwargs = self._merge_kwargs( + Florence2ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + image_inputs = {} + if images is not None: + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + + if text is None: + logger.warning_once("You are using Florence-2 without a text prefix.") + text = [""] * (1 if not isinstance(images, list) else len(images)) + elif isinstance(text, str): + text = [text] + + if not isinstance(text, list) or not all(isinstance(token, str) for token in text): + raise ValueError("`text` must be a string or list of strings.") + + if isinstance(images, list) and len(images) != len(text): + raise ValueError(f"Number of images ({len(images)}) must match number of texts ({len(text)}).") + + prompt_strings = self._construct_prompts(text) + + # Add image tokens and special tokens if images are provided + if image_inputs.get("pixel_values") is not None: + # Replace the image token with the expanded image token sequence + expanded_image_prompts = [] + for sample in prompt_strings: + sample = ( + self.image_token * self.num_image_tokens + + self.tokenizer.bos_token + + sample + + self.tokenizer.eos_token + ) + expanded_image_prompts.append(sample) + prompt_strings = expanded_image_prompts + + # Construct and tokenize prompts + output_kwargs["text_kwargs"].pop("add_special_tokens", None) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer( + prompt_strings, **output_kwargs["text_kwargs"], add_special_tokens=False, return_tensors=None + ) + self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**image_inputs, **text_inputs}, tensor_type=return_tensors) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + + vision_data = {} + if image_sizes is not None: + num_image_tokens = [self.num_image_tokens] * len(image_sizes) + num_image_patches = [1] * len(image_sizes) + + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=False, **kwargs): + """ + Post-processes the output of `FuyuForConditionalGeneration` to only return the text output. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + containing the token ids of the generated sequences. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `list[str]`: The decoded text output. + """ + return self.batch_decode(generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs) + + def post_process_generation(self, text=None, sequence=None, task=None, image_size=None) -> dict[str, Any]: + """ + Post-process generation outputs based on the task. + + Args: + text (`str`, *optional*): + Generated text. + sequence (`Union[List[int], torch.Tensor]`, *optional*): + Generated token sequence. + task (`str`, *optional*): + The task for post-processing. + image_size (`Tuple[int, int]`, *optional*): + Image size for dequantization. + + Returns: + `Dict[str, Any]`: Post-processed results keyed by task. + """ + if task is None: + raise ValueError("`task` must be provided for post-processing.") + + post_proc_type = self.tasks_answer_post_processing_type.get(task, "pure_text") + parsed = self.post_processor( + text=text, + sequence=sequence, + image_size=image_size, + parse_tasks=[post_proc_type], + )[post_proc_type] + + if post_proc_type == "pure_text": + final_answer = parsed.replace("", "").replace("", "").strip() + elif post_proc_type in ["description_with_bboxes", "bboxes"]: + bboxes = [inst["bbox"] for inst in parsed] + labels = [inst["cat_name"] for inst in parsed] + final_answer = {"bboxes": bboxes, "labels": labels} + if parsed and "score" in parsed[0]: + final_answer["scores"] = [inst["score"] for inst in parsed] + elif post_proc_type == "ocr": + quad_boxes = [inst["quad_box"] for inst in parsed] + labels = [inst["text"] for inst in parsed] + final_answer = {"quad_boxes": quad_boxes, "labels": labels} + elif post_proc_type == "phrase_grounding": + bboxes = [] + labels = [] + for inst in parsed: + for bbox in inst["bbox"]: + bboxes.append(bbox) + labels.append(inst["cat_name"]) + final_answer = {"bboxes": bboxes, "labels": labels} + elif post_proc_type in ["description_with_polygons", "polygons"]: + polygons = [inst["polygons"] for inst in parsed] + labels = [inst["cat_name"] for inst in parsed] + final_answer = {"polygons": polygons, "labels": labels} + elif post_proc_type == "description_with_bboxes_or_polygons": + bboxes = [] + bboxes_labels = [] + polygons = [] + polygons_labels = [] + for inst in parsed: + label = inst["cat_name"] + if "polygons" in inst: + polygons.append(inst["polygons"]) + polygons_labels.append(label) + else: + bboxes.append(inst["bbox"]) + bboxes_labels.append(label) + final_answer = { + "bboxes": bboxes, + "bboxes_labels": bboxes_labels, + "polygons": polygons, + "polygons_labels": polygons_labels, + } + else: + raise ValueError(f"Unknown post-processing type: {post_proc_type}") + + return {task: final_answer} + + +class Florence2PostProcessor: + """ + Post-processor for Florence-2 model outputs. Parses generated text into structured results for various tasks + like object detection, OCR, phrase grounding, etc. + + Args: + tokenizer (`PreTrainedTokenizer`): + The tokenizer used for decoding model outputs. + """ + + def __init__(self, config, tokenizer): + self.tokenizer = tokenizer + self.parse_task_config = config or {} + self.banned_grounding_tokens = set( + self.parse_task_config.get("phrase_grounding", {}).get("banned_grounding_tokens", []) + ) + self.all_special_tokens = set(self.tokenizer.all_special_tokens) + self.quantize_bins = (1000, 1000) + + def quantize(self, locations: "torch.Tensor", size: tuple[int, int]) -> "torch.Tensor": + """ + Quantize locations. + + Args: + locations (`torch.Tensor`): + Tensor of shape (N, 4) for boxes or (N, 2) for points/coordinates. + size (`tuple[int, int]`): + Original image size (width, height). + + Returns: + `torch.Tensor`: Quantized locations as integers. + """ + bins_w, bins_h = self.quantize_bins + size_w, size_h = size + per_bin_w = size_w / bins_w + per_bin_h = size_h / bins_h + + if locations.shape[-1] == 4: # Bounding boxes: [xmin, ymin, xmax, ymax] + xmin, ymin, xmax, ymax = locations.split(1, dim=-1) + q_xmin = (xmin / per_bin_w).floor().clamp(0, bins_w - 1) + q_ymin = (ymin / per_bin_h).floor().clamp(0, bins_h - 1) + q_xmax = (xmax / per_bin_w).floor().clamp(0, bins_w - 1) + q_ymax = (ymax / per_bin_h).floor().clamp(0, bins_h - 1) + return torch.cat([q_xmin, q_ymin, q_xmax, q_ymax], dim=-1).int() + + elif locations.shape[-1] == 2: # Points/coordinates: [x, y] + x, y = locations.split(1, dim=-1) + q_x = (x / per_bin_w).floor().clamp(0, bins_w - 1) + q_y = (y / per_bin_h).floor().clamp(0, bins_h - 1) + return torch.cat([q_x, q_y], dim=-1).int() + + else: + raise ValueError(f"Unsupported location shape: last dim must be 2 or 4, got {locations.shape[-1]}.") + + def dequantize(self, locations: "torch.Tensor", size: tuple[int, int]) -> "torch.Tensor": + """ + Dequantize locations back to original scale. + + Args: + locations (`torch.Tensor`): + Quantized tensor of shape (N, 4) for boxes or (N, 2) for points/coordinates. + size (`tuple[int, int]`): + Original image size (width, height). + + Returns: + `torch.Tensor`: Dequantized locations as floats. + """ + bins_w, bins_h = self.quantize_bins + size_w, size_h = size + per_bin_w = size_w / bins_w + per_bin_h = size_h / bins_h + + # Add 0.5 to use the center position of the bin as the coordinate. + if locations.shape[-1] == 4: # Bounding boxes + xmin, ymin, xmax, ymax = locations.split(1, dim=-1) + dq_xmin = (xmin + 0.5) * per_bin_w + dq_ymin = (ymin + 0.5) * per_bin_h + dq_xmax = (xmax + 0.5) * per_bin_w + dq_ymax = (ymax + 0.5) * per_bin_h + return torch.cat([dq_xmin, dq_ymin, dq_xmax, dq_ymax], dim=-1).int() + + elif locations.shape[-1] == 2: # Points/coordinates + x, y = locations.split(1, dim=-1) + dq_x = (x + 0.5) * per_bin_w + dq_y = (y + 0.5) * per_bin_h + return torch.cat([dq_x, dq_y], dim=-1).int() + + else: + raise ValueError(f"Unsupported location shape: last dim must be 2 or 4, got {locations.shape[-1]}.") + + def decode_with_spans(self, token_ids: list[int]) -> tuple[str, list[tuple[int, int]]]: + """ + Decode token IDs to text and compute character spans. + + Args: + token_ids (`list[int]`): + list of token IDs to decode. + + Returns: + `tuple[str, list[tuple[int, int]]]`: Decoded text and list of spans (start, end) for each token. + """ + filtered_tokens = self.tokenizer.convert_ids_to_tokens(token_ids, skip_special_tokens=False) + text = "" + spans = [] + for token in filtered_tokens: + if token in self.all_special_tokens: + sub_text = token + else: + sub_text = self.tokenizer.convert_tokens_to_string([token]) + span = (len(text), len(text) + len(sub_text)) + text += sub_text + spans.append(span) + return text, spans + + def parse_ocr_from_text_and_spans( + self, text: str, pattern: str | None, image_size: tuple[int, int], area_threshold: float = 0.0 + ) -> list[dict[str, Any]]: + """ + Parse OCR results with quadrilateral boxes. + + Args: + text (`str`): + The generated text. + pattern (`str`): + Regex pattern for matching. + image_size (`tuple[int, int]`): + Image size (width, height). + area_threshold (`float`, *optional*, defaults to 0.0): + Minimum area threshold for filtering boxes. + + Returns: + `list[dict[str, Any]]`: list of instances with 'quad_box' and 'text'. + """ + text = text.replace("", "").replace("", "").replace("", "") + if pattern is None: + pattern = r"(.+?)" + + matches = re.findall(pattern, text) + instances = [] + width, height = image_size + + for content, *quad_str in matches: + quad_bins = [int(i) for i in quad_str] + quad_box = self.dequantize(torch.tensor(quad_bins).reshape(-1, 2), size=image_size).flatten().tolist() + + if area_threshold > 0: + x_coords = quad_box[0::2] + y_coords = quad_box[1::2] + # Apply the Shoelace formula + area = 0.5 * abs( + sum(x_coords[i] * y_coords[i + 1] - x_coords[i + 1] * y_coords[i] for i in range(4 - 1)) + ) + + if area < (width * height) * area_threshold: + continue + + instances.append({"quad_box": quad_box, "text": content.strip()}) + return instances + + def parse_phrase_grounding_from_text_and_spans( + self, text: str, image_size: tuple[int, int] + ) -> list[dict[str, Any]]: + """ + Parse phrase grounding results. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + + Returns: + `list[dict[str, Any]]`: list of instances with 'bbox' and 'cat_name'. + """ + text = text.replace("", "").replace("", "").replace("", "") + phrase_pattern = r"([^<]+(?:){4,})" + phrases = re.findall(phrase_pattern, text) + text_pattern = r"^\s*(.*?)(?=||||||" + + instances = [] + for phrase_text in phrases: + phrase_text = phrase_text.replace("", "", 1).replace("", "", 1) + if not phrase_text: + continue + match = re.search(text_pattern, phrase_text) + if not match: + continue + phrase = match.group().strip() + if phrase in self.banned_grounding_tokens: + continue + boxes_matches = list(re.finditer(box_pattern, phrase_text)) + if not boxes_matches: + continue + bbox_bins = [[int(m.group(j)) for j in range(1, 5)] for m in boxes_matches] + bboxes = self.dequantize(torch.tensor(bbox_bins), size=image_size).tolist() + phrase = phrase.encode("ascii", "ignore").decode("ascii") + instances.append({"bbox": bboxes, "cat_name": phrase}) + return instances + + def _find_matched_token_indices(self, cur_span: tuple[int, int], token_spans: list[tuple[int, int]]) -> list[int]: + return [i for i, span in enumerate(token_spans) if not (span[1] <= cur_span[0] or span[0] >= cur_span[1])] + + def parse_description_with_bboxes_from_text_and_spans( + self, + text: str, + image_size: tuple[int, int], + allow_empty_phrase: bool = False, + ) -> list[dict[str, Any]]: + """ + Parse descriptions with bounding boxes. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + allow_empty_phrase (`bool`, *optional*, defaults to `False`): + Allow phrases without text. + + Returns: + `list[dict[str, Any]]`: list of instances with 'bbox', 'cat_name', and optional 'score'. + """ + text = text.replace("", "").replace("", "").replace("", "") + + if allow_empty_phrase: + pattern = r"(?:(?:){4,})" + else: + pattern = r"([^<]+(?:){4,})" + phrases = re.findall(pattern, text) + + text_pattern = r"^\s*(.*?)(?=||||||" + + instances = [] + for phrase_text in phrases: + phrase_text = phrase_text.replace("", "", 1).replace("", "", 1) + if not phrase_text and not allow_empty_phrase: + continue + match = re.search(text_pattern, phrase_text) + if not match: + continue + phrase = match.group().strip() + boxes_matches = list(re.finditer(box_pattern, phrase_text)) + if not boxes_matches: + continue + bbox_bins = [[int(m.group(j)) for j in range(1, 5)] for m in boxes_matches] + bboxes = self.dequantize(torch.tensor(bbox_bins), size=image_size).tolist() + + phrase = phrase.encode("ascii", "ignore").decode("ascii") + for bbox in bboxes: + instance = {"bbox": bbox, "cat_name": phrase} + instances.append(instance) + + return instances + + def parse_description_with_polygons_from_text_and_spans( + self, + text: str, + image_size: tuple[int, int], + allow_empty_phrase: bool = False, + polygon_sep_token: str = "", + polygon_start_token: str = "", + polygon_end_token: str = "", + with_box_at_start: bool = False, + ) -> list[dict[str, Any]]: + """ + Parse descriptions with polygons. + + Args: + text (`str`): + The generated text. + image_size (`tuple[int, int]`): + Image size (width, height). + allow_empty_phrase (`bool`, *optional*, defaults to `False`): + Allow phrases without text. + polygon_sep_token (`str`, *optional*, defaults to ""): + Token separating polygons. + polygon_start_token (`str`, *optional*, defaults to ""): + Start token for polygons. + polygon_end_token (`str`, *optional*, defaults to ""): + End token for polygons. + with_box_at_start (`bool`, *optional*, defaults to `False`): + Whether a bounding box is at the start of polygons. + + Returns: + `list[dict[str, Any]]`: list of instances with 'polygons', 'cat_name', and optional 'bbox'. + """ + text = text.replace("", "").replace("", "").replace("", "") + + if allow_empty_phrase: + pattern = rf"(?:(?:|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})" + else: + pattern = rf"([^<]+(?:|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})" + phrases = re.findall(pattern, text) + phrase_pattern = r"^\s*(.*?)(?=||||||)" + poly_instance_pattern = rf"{re.escape(polygon_start_token)}(.*?){re.escape(polygon_end_token)}" + box_pattern = rf"((?:)+)(?:{re.escape(polygon_sep_token)}|$)" + + instances = [] + for phrase_text in phrases: + phrase_text_strip = re.sub(r"^", "", phrase_text, count=1) + if not phrase_text_strip and not allow_empty_phrase: + continue + match = re.search(phrase_pattern, phrase_text_strip) + if not match: + continue + phrase = match.group().strip() + + if polygon_start_token in phrase_text and polygon_end_token in phrase_text: + poly_instances = [m.group(1) for m in re.finditer(poly_instance_pattern, phrase_text)] + else: + poly_instances = [phrase_text] + + for poly_inst in poly_instances: + poly_matches = list(re.finditer(box_pattern, poly_inst)) + if len(poly_matches) == 0: + continue + bbox = [] + polygons = [] + for poly_match in poly_matches: + poly_str = poly_match.group(1) + poly_bins = [int(m.group(1)) for m in re.finditer(r"", poly_str)] + if with_box_at_start and not bbox: + if len(poly_bins) > 4: + bbox = poly_bins[:4] + poly_bins = poly_bins[4:] + else: + bbox = [0, 0, 0, 0] + if len(poly_bins) % 2 == 1: + poly_bins = poly_bins[:-1] + poly_coords = ( + self.dequantize(torch.tensor(poly_bins).reshape(-1, 2), size=image_size).flatten().tolist() + ) + polygons.append(poly_coords) + + instance = {"cat_name": phrase, "polygons": polygons} + if bbox: + instance["bbox"] = self.dequantize(torch.tensor([bbox]), size=image_size)[0].tolist() + instances.append(instance) + return instances + + def __call__(self, text=None, sequence=None, image_size=None, parse_tasks=None) -> dict[str, Any]: + """ + Process model output and parse into task-specific results. + + Args: + text (`Optional[str]`, *optional*): + Generated text. Either this or `sequence` must be provided. + sequence (`Optional[Union[list[int], torch.Tensor]]`, *optional*): + Token sequence. Either this or `text` must be provided. + image_size (`Optional[tuple[int, int]]`, *optional*): + Image size (width, height) required for dequantization. + parse_tasks (`Optional[Union[str, list[str]]]`, *optional*): + Specific tasks to parse. If None, parse all supported tasks. + + Returns: + `dict[str, Any]`: Parsed results for each task, including the raw 'text'. + """ + if parse_tasks is not None: + parse_tasks = [parse_tasks] if isinstance(parse_tasks, str) else parse_tasks + for task in parse_tasks: + if task not in self.parse_task_config.keys(): + raise ValueError(f"Unsupported parse task: {task}") + + if (text is None and sequence is None) or (text is not None and sequence is not None): + raise ValueError("Exactly one of 'text' or 'sequence' must be provided.") + + if sequence is not None: + if isinstance(sequence, torch.Tensor): + sequence = sequence.tolist() + sequence = sequence[1:] if sequence[0] == self.tokenizer.bos_token_id else sequence # Skip BOS if present + text, _ = self.decode_with_spans(sequence) + + parsed_dict = {"text": text} + + tasks_to_parse = parse_tasks or self.parse_task_config.keys() + for task in tasks_to_parse: + config = self.parse_task_config[task] + pattern = config.get("PATTERN") + + if task == "ocr": + parsed_dict["ocr"] = self.parse_ocr_from_text_and_spans( + text, pattern=pattern, image_size=image_size, area_threshold=config.get("AREA_THRESHOLD", 0.0) + ) + elif task == "phrase_grounding": + parsed_dict["phrase_grounding"] = self.parse_phrase_grounding_from_text_and_spans( + text, image_size=image_size + ) + elif task == "pure_text": + parsed_dict["pure_text"] = text + elif task == "description_with_bboxes": + parsed_dict["description_with_bboxes"] = self.parse_description_with_bboxes_from_text_and_spans( + text, image_size=image_size + ) + elif task == "description_with_polygons": + parsed_dict["description_with_polygons"] = self.parse_description_with_polygons_from_text_and_spans( + text, image_size=image_size + ) + elif task == "polygons": + parsed_dict["polygons"] = self.parse_description_with_polygons_from_text_and_spans( + text, image_size=image_size, allow_empty_phrase=True + ) + elif task == "bboxes": + parsed_dict["bboxes"] = self.parse_description_with_bboxes_from_text_and_spans( + text, image_size=image_size, allow_empty_phrase=True + ) + elif task == "description_with_bboxes_or_polygons": + if "" in text: + instances = self.parse_description_with_polygons_from_text_and_spans(text, image_size=image_size) + else: + instances = self.parse_description_with_bboxes_from_text_and_spans(text, image_size=image_size) + parsed_dict["description_with_bboxes_or_polygons"] = instances + else: + raise ValueError(f"task {task} is not supported") + + return parsed_dict + + +__all__ = ["Florence2Processor"] diff --git a/third_party/transformers/src/transformers/models/gptj/__init__.py b/third_party/transformers/src/transformers/models/gptj/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a814910a88850c60031508df776718566ef30d60 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gptj/__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_gptj import * + from .modeling_gptj 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/gptj/configuration_gptj.py b/third_party/transformers/src/transformers/models/gptj/configuration_gptj.py new file mode 100644 index 0000000000000000000000000000000000000000..81bdf28c1a81930e4abee58676245a2dd82a7ca5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gptj/configuration_gptj.py @@ -0,0 +1,72 @@ +# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPT-J model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="EleutherAI/gpt-j-6B") +@strict +class GPTJConfig(PreTrainedConfig): + r""" + rotary_dim (`int`, *optional*, defaults to 64): + Number of dimensions in the embedding that Rotary Position Embedding is applied to. + + Example: + + ```python + >>> from transformers import GPTJModel, GPTJConfig + + >>> # Initializing a GPT-J 6B configuration + >>> configuration = GPTJConfig() + + >>> # Initializing a model from the configuration + >>> model = GPTJModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gptj" + attribute_map = { + "max_position_embeddings": "n_positions", + "hidden_size": "n_embd", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + vocab_size: int = 50400 + n_positions: int = 2048 + n_embd: int = 4096 + n_layer: int = 28 + n_head: int = 16 + rotary_dim: int = 64 + n_inner: int | None = None + activation_function: str = "gelu_new" + resid_pdrop: float | int = 0.0 + embd_pdrop: float | int = 0.0 + attn_pdrop: float | int = 0.0 + layer_norm_epsilon: float = 1e-5 + initializer_range: float = 0.02 + use_cache: bool = True + bos_token_id: int | None = 50256 + eos_token_id: int | list[int] | None = 50256 + pad_token_id: int | None = None + tie_word_embeddings: bool = False + + +__all__ = ["GPTJConfig"] diff --git a/third_party/transformers/src/transformers/models/gptj/modeling_gptj.py b/third_party/transformers/src/transformers/models/gptj/modeling_gptj.py new file mode 100644 index 0000000000000000000000000000000000000000..c92026f2f209c687378ed6bcf0951df36421a7f3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gptj/modeling_gptj.py @@ -0,0 +1,856 @@ +# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch GPT-J 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 ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_gptj import GPTJConfig + + +if is_flash_attn_available(): + from ...modeling_flash_attention_utils import _flash_attention_forward + + +logger = logging.get_logger(__name__) + + +def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor: + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim)) + sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(num_pos, dtype=torch.int64).float(), inv_freq).float() + return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1) + + +def get_embed_positions(embed_positions, position_ids): + return embed_positions.to(position_ids.device).repeat(position_ids.shape[0], 1, 1) + + +def rotate_every_two(x: torch.Tensor) -> torch.Tensor: + x1 = x[:, :, :, ::2] + x2 = x[:, :, :, 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)') + + +def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor: + sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3) + cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3) + return (tensor * cos) + (rotate_every_two(tensor) * sin) + + +class GPTJAttention(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + self.config = config + self.max_positions = config.max_position_embeddings + + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + self.is_causal = True + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.embed_dim = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_attention_heads + if self.head_dim * self.num_attention_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and" + f" `num_attention_heads`: {self.num_attention_heads})." + ) + self.scale_attn = math.sqrt(self.head_dim) + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) + self.rotary_dim = config.rotary_dim + self.pos_embd_dim = self.rotary_dim or self.embed_dim + self.register_buffer( + "embed_positions", create_sinusoidal_positions(self.max_positions, self.pos_embd_dim), persistent=False + ) + + def _split_heads(self, tensor, num_attention_heads, attn_head_size, rotary): + """ + Splits hidden dim into attn_head_size and num_attention_heads + """ + new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size) + tensor = tensor.view(new_shape) + if rotary: + return tensor + if len(tensor.shape) == 5: + return tensor.permute(0, 1, 3, 2, 4) # (batch, blocks, head, block_length, head_features) + elif len(tensor.shape) == 4: + return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features) + else: + raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}") + + def _merge_heads(self, tensor, num_attention_heads, attn_head_size): + """ + Merges attn_head_size dim and num_attn_heads dim into hidden dim + """ + if len(tensor.shape) == 5: + tensor = tensor.permute(0, 1, 3, 2, 4).contiguous() + elif len(tensor.shape) == 4: + tensor = tensor.permute(0, 2, 1, 3).contiguous() + else: + raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}") + new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,) + return tensor.view(new_shape) + + def _attn( + self, + query, + key, + value, + attention_mask=None, + ): + # Keep the attention weights computation in fp32 to avoid overflow issues + query = query.to(torch.float32) + key = key.to(torch.float32) + + attn_weights = torch.matmul(query, key.transpose(-1, -2)) + attn_weights = attn_weights / self.scale_attn + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = attn_weights.to(value.dtype) + attn_weights = self.attn_dropout(attn_weights) + + attn_output = torch.matmul(attn_weights, value) + + return attn_output, attn_weights + + def _get_embed_positions(self, position_ids): + embed_positions = self.embed_positions + if embed_positions.device != position_ids.device: + embed_positions = embed_positions.to(position_ids.device) + self.embed_positions = embed_positions + return embed_positions.repeat(position_ids.shape[0], 1, 1) + + def forward( + self, + hidden_states: torch.FloatTensor, + layer_past: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> ( + tuple[torch.Tensor, tuple[torch.Tensor]] + | tuple[torch.Tensor, tuple[torch.Tensor], tuple[torch.Tensor, ...]] + | None + ): + query = self.q_proj(hidden_states) + key = self.k_proj(hidden_states) + value = self.v_proj(hidden_states) + + query = self._split_heads(query, self.num_attention_heads, self.head_dim, True) + key = self._split_heads(key, self.num_attention_heads, self.head_dim, True) + value = self._split_heads(value, self.num_attention_heads, self.head_dim, False) + + embed_positions = self._get_embed_positions(position_ids) + + repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1]) + sincos = torch.gather(embed_positions, 1, repeated_position_ids).to(key.dtype) + sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1) + + if self.rotary_dim is not None: + k_rot = key[:, :, :, : self.rotary_dim] + k_pass = key[:, :, :, self.rotary_dim :] + + q_rot = query[:, :, :, : self.rotary_dim] + q_pass = query[:, :, :, self.rotary_dim :] + + k_rot = apply_rotary_pos_emb(k_rot, sin, cos) + q_rot = apply_rotary_pos_emb(q_rot, sin, cos) + + key = torch.cat([k_rot, k_pass], dim=-1) + query = torch.cat([q_rot, q_pass], dim=-1) + else: + key = apply_rotary_pos_emb(key, sin, cos) + query = apply_rotary_pos_emb(query, sin, cos) + + key = key.permute(0, 2, 1, 3) + query = query.permute(0, 2, 1, 3) + + if layer_past is not None: + key, value = layer_past.update(key, value, self.layer_idx) + + # compute self-attention: V x Softmax(QK^T) + attn_output, attn_weights = self._attn(query, key, value, attention_mask) + + attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim) + attn_output = self.out_proj(attn_output) + attn_output = self.resid_dropout(attn_output) + + return attn_output, attn_weights + + +class GPTJFlashAttention2(GPTJAttention): + """ + GPTJ flash attention module. This module inherits from `GPTJAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask() + + def forward( + self, + hidden_states: torch.FloatTensor, + layer_past: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> ( + tuple[torch.Tensor, tuple[torch.Tensor]] + | tuple[torch.Tensor, tuple[torch.Tensor], tuple[torch.Tensor, ...]] + | None + ): + query = self.q_proj(hidden_states) + key = self.k_proj(hidden_states) + value = self.v_proj(hidden_states) + + query = self._split_heads(query, self.num_attention_heads, self.head_dim, True) + key = self._split_heads(key, self.num_attention_heads, self.head_dim, True) + value = self._split_heads(value, self.num_attention_heads, self.head_dim, False) + + embed_positions = self._get_embed_positions(position_ids) + + repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1]) + sincos = torch.gather(embed_positions, 1, repeated_position_ids).to(key.dtype) + sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1) + + if self.rotary_dim is not None: + k_rot = key[:, :, :, : self.rotary_dim] + k_pass = key[:, :, :, self.rotary_dim :] + + q_rot = query[:, :, :, : self.rotary_dim] + q_pass = query[:, :, :, self.rotary_dim :] + + k_rot = apply_rotary_pos_emb(k_rot, sin, cos) + q_rot = apply_rotary_pos_emb(q_rot, sin, cos) + + key = torch.cat([k_rot, k_pass], dim=-1) + query = torch.cat([q_rot, q_pass], dim=-1) + else: + key = apply_rotary_pos_emb(key, sin, cos) + query = apply_rotary_pos_emb(query, sin, cos) + + # tanspose to have the desired shape + # before transpose: batch_size x seq_length x num_attention_heads x head_dim + # after transpose: batch_size x num_attention_heads x seq_length x head_dim + key = key.permute(0, 2, 1, 3) + query = query.permute(0, 2, 1, 3) + # value: batch_size x num_attention_heads x seq_length x head_dim + + if layer_past is not None: + key, value = layer_past.update(key, value, self.layer_idx) + + # The Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we need to keep the original shape for query and key, and reshape value + # to have the correct shape. + key = key.permute(0, 2, 1, 3).contiguous() + query = query.permute(0, 2, 1, 3).contiguous() + value = value.permute(0, 2, 1, 3).contiguous() + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + input_dtype = query.dtype + device_type = query.device.type if query.device.type != "mps" else "cpu" + if input_dtype == torch.float32: + if torch.is_autocast_enabled(device_type): + target_dtype = torch.get_autocast_dtype(device_type) + # Handle the case where the model is quantized + elif hasattr(self.config, "_is_quantized"): + target_dtype = self.config.dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query = query.to(target_dtype) + key = key.to(target_dtype) + value = value.to(target_dtype) + + attention_dropout = self.config.attn_pdrop if self.training else 0.0 # attn_pdrop in gptj + + query_length = query.shape[1] + + # Compute attention + attn_weights = _flash_attention_forward( + query, + key, + value, + attention_mask, + query_length, + dropout=attention_dropout, + is_causal=self.is_causal, + use_top_left_mask=self._flash_attn_uses_top_left_mask, + ) + + # Reshape outputs + attn_output = attn_weights.reshape( + attn_weights.shape[0], attn_weights.shape[1], attn_weights.shape[2] * attn_weights.shape[3] + ) + attn_output = self.out_proj(attn_output) + attn_output = self.resid_dropout(attn_output) + return attn_output, attn_weights + + +GPTJ_ATTENTION_CLASSES = { + "eager": GPTJAttention, + "flash_attention_2": GPTJFlashAttention2, +} + + +class GPTJMLP(nn.Module): + def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim + super().__init__() + embed_dim = config.n_embd + + self.fc_in = nn.Linear(embed_dim, intermediate_size) + self.fc_out = nn.Linear(intermediate_size, embed_dim) + + self.act = ACT2FN[config.activation_function] + self.dropout = nn.Dropout(config.resid_pdrop) + + def forward(self, hidden_states: torch.FloatTensor | None) -> torch.FloatTensor: + hidden_states = self.fc_in(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.fc_out(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class GPTJBlock(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd + self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.attn = GPTJ_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + self.mlp = GPTJMLP(inner_dim, config) + + def forward( + self, + hidden_states: torch.FloatTensor | None, + layer_past: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, tuple[torch.FloatTensor, ...]] | None: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_outputs, attn_weights = self.attn( + hidden_states=hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + ) + feed_forward_hidden_states = self.mlp(hidden_states) + hidden_states = attn_outputs + feed_forward_hidden_states + residual + + return hidden_states, attn_weights + + +@auto_docstring +class GPTJPreTrainedModel(PreTrainedModel): + config: GPTJConfig + base_model_prefix = "transformer" + supports_gradient_checkpointing = True + _no_split_modules = ["GPTJBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _can_compile_fullgraph = True + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, GPTJAttention): + init.copy_(module.embed_positions, create_sinusoidal_positions(module.max_positions, module.pos_embd_dim)) + + +@auto_docstring +class GPTJModel(GPTJPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embed_dim = config.n_embd + self.vocab_size = config.vocab_size + self.wte = nn.Embedding(config.vocab_size, self.embed_dim) + self.drop = nn.Dropout(config.embd_pdrop) + self.h = nn.ModuleList([GPTJBlock(config, layer_idx=i) for i in range(config.n_layer)]) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.wte + + def set_input_embeddings(self, new_embeddings): + self.wte = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPast: + r""" + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert *input_ids* indices into associated vectors than the + model's internal embedding lookup matrix. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + seq_length = inputs_embeds.shape[1] + if position_ids is None: + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + + if token_type_ids is not None: + token_type_ids = token_type_ids.view(-1, seq_length) + token_type_embeds = self.wte(token_type_ids) + hidden_states = hidden_states + token_type_embeds + + hidden_states = self.drop(hidden_states) + output_shape = (-1, seq_length, hidden_states.size(-1)) + + all_self_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + for i, block in enumerate(self.h): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + outputs = block( + hidden_states, + layer_past=past_key_values, + attention_mask=causal_mask, + position_ids=position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + hidden_states = outputs[0] + if output_attentions: + all_self_attentions = all_self_attentions + (outputs[1],) + + hidden_states = self.ln_f(hidden_states) + + hidden_states = hidden_states.view(output_shape) + # Add last hidden state + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions] if v is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring( + custom_intro=""" + The GPT-J Model transformer with a language modeling head on top. + """ +) +class GPTJForCausalLM(GPTJPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} + + def __init__(self, config): + super().__init__(config) + self.transformer = GPTJModel(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithPast: + r""" + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert *input_ids* indices into associated vectors than the + model's internal embedding lookup matrix. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The GPT-J Model transformer with a sequence classification head on top (linear layer). + + [`GPTJForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT, GPT-2, GPT-Neo) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """ +) +class GPTJForSequenceClassification(GPTJPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = GPTJModel(config) + self.score = nn.Linear(config.n_embd, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutputWithPast: + r""" + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert *input_ids* indices into associated vectors than the + model's internal embedding lookup matrix. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + labels = labels.to(pooled_logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class GPTJForQuestionAnswering(GPTJPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = GPTJModel(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.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, + ) -> tuple | QuestionAnsweringModelOutput: + r""" + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert *input_ids* indices into associated vectors than the + model's internal embedding lookup matrix. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.transformer( + 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).to(start_logits.device) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1).to(end_logits.device) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "GPTJForCausalLM", + "GPTJForQuestionAnswering", + "GPTJForSequenceClassification", + "GPTJModel", + "GPTJPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/granite/__init__.py b/third_party/transformers/src/transformers/models/granite/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..08b74d14ca871a7752b61fb71f08b7a7886d80f7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/granite/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2024 EleutherAI and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_granite import * + from .modeling_granite import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/granite/configuration_granite.py b/third_party/transformers/src/transformers/models/granite/configuration_granite.py new file mode 100644 index 0000000000000000000000000000000000000000..e026cbbe5ff3e094e3f03036f181e0437e5a6476 --- /dev/null +++ b/third_party/transformers/src/transformers/models/granite/configuration_granite.py @@ -0,0 +1,95 @@ +# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Granite model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="ibm-granite/granite-3.0-8b-base") +@strict +class GraniteConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import GraniteModel, GraniteConfig + + >>> # Initializing a Granite granite-3b style configuration + >>> configuration = GraniteConfig() + + >>> # Initializing a model from the granite-7b style configuration + >>> model = GraniteModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "granite" + keys_to_ignore_at_inference = ["past_key_values"] + # Default tensor parallel plan for base model `GraniteModel` + 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 = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + mlp_bias: bool = False + embedding_multiplier: float | int = 1.0 + logits_scaling: float | int = 1.0 + residual_multiplier: float | int = 1.0 + attention_multiplier: float | int = 1.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + +__all__ = ["GraniteConfig"] diff --git a/third_party/transformers/src/transformers/models/granite/modeling_granite.py b/third_party/transformers/src/transformers/models/granite/modeling_granite.py new file mode 100644 index 0000000000000000000000000000000000000000..934345fe67235441896e6e4e09b473c18b151301 --- /dev/null +++ b/third_party/transformers/src/transformers/models/granite/modeling_granite.py @@ -0,0 +1,519 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/granite/modular_granite.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_granite.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, 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_granite import GraniteConfig + + +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 GraniteAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GraniteConfig, 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 = config.attention_multiplier + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +@use_kernel_forward_from_hub("RMSNorm") +class GraniteRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + GraniteRMSNorm 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 GraniteMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class GraniteDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GraniteConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = GraniteAttention(config=config, layer_idx=layer_idx) + + self.mlp = GraniteMLP(config) + self.input_layernorm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.residual_multiplier = config.residual_multiplier + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + 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. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states * self.residual_multiplier + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states * self.residual_multiplier + + return hidden_states + + +@auto_docstring +class GranitePreTrainedModel(PreTrainedModel): + config: GraniteConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GraniteDecoderLayer"] + _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": GraniteDecoderLayer, + "attentions": GraniteAttention, + } + + +class GraniteRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: GraniteConfig, 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: GraniteConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@auto_docstring +class GraniteModel(GranitePreTrainedModel): + def __init__(self, config: GraniteConfig): + 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( + [GraniteDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = GraniteRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.embedding_multiplier = config.embedding_multiplier + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> 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) + + inputs_embeds = inputs_embeds * self.embedding_multiplier + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_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 BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class GraniteForCausalLM(GranitePreTrainedModel, 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 = GraniteModel(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, GraniteForCausalLM + + >>> model = GraniteForCausalLM.from_pretrained("meta-granite/Granite-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-granite/Granite-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 + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + logits = logits / self.config.logits_scaling # main diff with 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__ = ["GraniteForCausalLM", "GraniteModel", "GranitePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/granite/modular_granite.py b/third_party/transformers/src/transformers/models/granite/modular_granite.py new file mode 100644 index 0000000000000000000000000000000000000000..3b62306582eefea2b5dfb0eb7b82301500d5a147 --- /dev/null +++ b/third_party/transformers/src/transformers/models/granite/modular_granite.py @@ -0,0 +1,223 @@ +# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from torch import nn + +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaModel, + LlamaPreTrainedModel, +) +from .configuration_granite import GraniteConfig + + +logger = logging.get_logger(__name__) + + +class GraniteAttention(LlamaAttention): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GraniteConfig, layer_idx: int | None = None): + super().__init__(config, layer_idx) + self.scaling = config.attention_multiplier + + +class GraniteDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: GraniteConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.residual_multiplier = config.residual_multiplier + self.self_attn = GraniteAttention(config=config, layer_idx=layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + 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. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states * self.residual_multiplier + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states * self.residual_multiplier + + return hidden_states + + +class GranitePreTrainedModel(LlamaPreTrainedModel): + _can_record_outputs = { + "hidden_states": GraniteDecoderLayer, + "attentions": GraniteAttention, + } + + +class GraniteModel(LlamaModel): + def __init__(self, config: GraniteConfig): + super().__init__(config) + self.embedding_multiplier = config.embedding_multiplier + self.layers = nn.ModuleList( + [GraniteDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + @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) + + inputs_embeds = inputs_embeds * self.embedding_multiplier + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_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 BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class GraniteForCausalLM(LlamaForCausalLM): + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + 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.config.logits_scaling # main diff with 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__ = ["GraniteForCausalLM", "GraniteModel", "GranitePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/informer/__init__.py b/third_party/transformers/src/transformers/models/informer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cd7a901eab235ecd14d2127a9eb78022c720c0f3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/informer/__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_informer import * + from .modeling_informer 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/informer/configuration_informer.py b/third_party/transformers/src/transformers/models/informer/configuration_informer.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b4fe6a329fbc65312f3011eb09e483cea66edf --- /dev/null +++ b/third_party/transformers/src/transformers/models/informer/configuration_informer.py @@ -0,0 +1,176 @@ +# 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. +"""Informer model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="huggingface/informer-tourism-monthly") +@strict +class InformerConfig(PreTrainedConfig): + r""" + prediction_length (`int`): + The prediction length for the decoder. In other words, the prediction horizon of the model. This value is + typically dictated by the dataset and we recommend to set it appropriately. + context_length (`int`, *optional*, defaults to `prediction_length`): + The context length for the encoder. If `None`, the context length will be the same as the + `prediction_length`. + distribution_output (`string`, *optional*, defaults to `"student_t"`): + The distribution emission head for the model. Could be either "student_t", "normal" or "negative_binomial". + loss (`string`, *optional*, defaults to `"nll"`): + The loss function for the model corresponding to the `distribution_output` head. For parametric + distributions it is the negative log likelihood (nll) - which currently is the only supported one. + input_size (`int`, *optional*, defaults to 1): + The size of the target variable which by default is 1 for univariate targets. Would be > 1 in case of + multivariate targets. + lags_sequence (`list[int]`, *optional*, defaults to `[1, 2, 3, 4, 5, 6, 7]`): + The lags of the input time series as covariates often dictated by the frequency of the data. Default is + `[1, 2, 3, 4, 5, 6, 7]` but we recommend to change it based on the dataset appropriately. + scaling (`string` or `bool`, *optional* defaults to `"mean"`): + Whether to scale the input targets via "mean" scaler, "std" scaler or no scaler if `None`. If `True`, the + scaler is set to "mean". + num_dynamic_real_features (`int`, *optional*, defaults to 0): + The number of dynamic real valued features. + num_static_real_features (`int`, *optional*, defaults to 0): + The number of static real valued features. + num_static_categorical_features (`int`, *optional*, defaults to 0): + The number of static categorical features. + num_time_features (`int`, *optional*, defaults to 0): + The number of time features in the input time series. + cardinality (`list[int]`, *optional*): + The cardinality (number of different values) for each of the static categorical features. Should be a list + of integers, having the same length as `num_static_categorical_features`. Cannot be `None` if + `num_static_categorical_features` is > 0. + embedding_dimension (`list[int]`, *optional*): + The dimension of the embedding for each of the static categorical features. Should be a list of integers, + having the same length as `num_static_categorical_features`. Cannot be `None` if + `num_static_categorical_features` is > 0. + activation_dropout (`float`, *optional*, defaults to 0.1): + The dropout probability used between the two layers of the feed-forward networks. + num_parallel_samples (`int`, *optional*, defaults to 100): + The number of samples to generate in parallel for each time step of inference. + attention_type (`str`, *optional*, defaults to "prob"): + Attention used in encoder. This can be set to "prob" (Informer's ProbAttention) or "full" (vanilla + transformer's canonical self-attention). + sampling_factor (`int`, *optional*, defaults to 5): + ProbSparse sampling factor (only makes affect when `attention_type`="prob"). It is used to control the + reduced query matrix (Q_reduce) input length. + distil (`bool`, *optional*, defaults to `True`): + Whether to use distilling in encoder. + + Example: + + ```python + >>> from transformers import InformerConfig, InformerModel + + >>> # Initializing an Informer configuration with 12 time steps for prediction + >>> configuration = InformerConfig(prediction_length=12) + + >>> # Randomly initializing a model (with random weights) from the configuration + >>> model = InformerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "informer" + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + "num_hidden_layers": "encoder_layers", + "initializer_range": "init_std", + } + + prediction_length: int | None = None + context_length: int | None = None + distribution_output: str = "student_t" + loss: str = "nll" + input_size: int = 1 + lags_sequence: list[int] | None = None + scaling: str | bool | None = "mean" + num_dynamic_real_features: int = 0 + num_static_real_features: int = 0 + num_static_categorical_features: int = 0 + num_time_features: int = 0 + cardinality: list[int] | None = None + embedding_dimension: list[int] | None = None + d_model: int = 64 + encoder_ffn_dim: int = 32 + decoder_ffn_dim: int = 32 + encoder_attention_heads: int = 2 + decoder_attention_heads: int = 2 + encoder_layers: int = 2 + decoder_layers: int = 2 + is_encoder_decoder: bool = True + activation_function: str = "gelu" + dropout: float | int = 0.05 + encoder_layerdrop: float | int = 0.1 + decoder_layerdrop: float | int = 0.1 + attention_dropout: float | int = 0.1 + activation_dropout: float | int = 0.1 + num_parallel_samples: int = 100 + init_std: float = 0.02 + use_cache: bool = True + attention_type: str = "prob" + sampling_factor: int = 5 + distil: bool = True + + def __post_init__(self, **kwargs): + self.context_length = self.context_length or self.prediction_length + self.lags_sequence = self.lags_sequence if self.lags_sequence is not None else [1, 2, 3, 4, 5, 6, 7] + + if not (self.cardinality and self.num_static_categorical_features > 0): + self.cardinality = [0] + + if not (self.embedding_dimension and self.num_static_categorical_features > 0): + self.embedding_dimension = [min(50, (cat + 1) // 2) for cat in self.cardinality] + + self.feature_size = self.input_size * len(self.lags_sequence) + self._number_of_features + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if ( + self.cardinality + and self.num_static_categorical_features > 0 + and len(self.cardinality) != self.num_static_categorical_features + ): + raise ValueError( + "The cardinality should be a list of the same length as `num_static_categorical_features`" + ) + + if ( + self.embedding_dimension + and self.num_static_categorical_features > 0 + and len(self.embedding_dimension) != self.num_static_categorical_features + ): + raise ValueError( + "The embedding dimension should be a list of the same length as `num_static_categorical_features`" + ) + + @property + def _number_of_features(self) -> int: + return ( + sum(self.embedding_dimension) + + self.num_dynamic_real_features + + self.num_time_features + + self.num_static_real_features + + self.input_size * 2 # the log1p(abs(loc)) and log(scale) features + ) + + +__all__ = ["InformerConfig"] diff --git a/third_party/transformers/src/transformers/models/informer/modeling_informer.py b/third_party/transformers/src/transformers/models/informer/modeling_informer.py new file mode 100644 index 0000000000000000000000000000000000000000..e458516d47bb4819e148adf11e082a9bb346d0c3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/informer/modeling_informer.py @@ -0,0 +1,1773 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/informer/modular_informer.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_informer.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2023 Amazon and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable + +import numpy as np +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + SampleTSPredictionOutput, + Seq2SeqTSModelOutput, + Seq2SeqTSPredictionOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_informer import InformerConfig + + +logger = logging.get_logger(__name__) + + +class InformerFeatureEmbedder(nn.Module): + """ + Embed a sequence of categorical features. + + Args: + cardinalities (`list[int]`): + List of cardinalities of the categorical features. + embedding_dims (`list[int]`): + List of embedding dimensions of the categorical features. + """ + + def __init__(self, cardinalities: list[int], embedding_dims: list[int]) -> None: + super().__init__() + + self.num_features = len(cardinalities) + self.embedders = nn.ModuleList([nn.Embedding(c, d) for c, d in zip(cardinalities, embedding_dims)]) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + if self.num_features > 1: + # we slice the last dimension, giving an array of length + # self.num_features with shape (N,T) or (N) + cat_feature_slices = torch.chunk(features, self.num_features, dim=-1) + else: + cat_feature_slices = [features] + + return torch.cat( + [ + embed(cat_feature_slice.squeeze(-1)) + for embed, cat_feature_slice in zip(self.embedders, cat_feature_slices) + ], + dim=-1, + ) + + +class InformerStdScaler(nn.Module): + """ + Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by + subtracting from the mean and dividing by the standard deviation. + """ + + def __init__(self, config: InformerConfig): + super().__init__() + self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1 + self.keepdim = config.keepdim if hasattr(config, "keepdim") else True + self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5 + + def forward( + self, data: torch.Tensor, observed_indicator: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Parameters: + data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): + input for Batch norm calculation + observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`): + Calculating the scale on the observed indicator. + Returns: + tuple of `torch.Tensor` of shapes + (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, + `(batch_size, 1, num_input_channels)`) + """ + denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim) + denominator = denominator.clamp_min(1.0) + loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator + + variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator + scale = torch.sqrt(variance + self.minimum_scale) + return (data - loc) / scale, loc, scale + + +class InformerMeanScaler(nn.Module): + """ + Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data + accordingly. + """ + + def __init__(self, config: InformerConfig): + super().__init__() + self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1 + self.keepdim = config.keepdim if hasattr(config, "keepdim") else True + self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10 + self.default_scale = config.default_scale if hasattr(config, "default_scale") else None + + def forward( + self, data: torch.Tensor, observed_indicator: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Parameters: + data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): + input for Batch norm calculation + observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`): + Calculating the scale on the observed indicator. + Returns: + tuple of `torch.Tensor` of shapes + (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, + `(batch_size, 1, num_input_channels)`) + """ + ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True) + num_observed = observed_indicator.sum(self.dim, keepdim=True) + + scale = ts_sum / torch.clamp(num_observed, min=1) + + # If `default_scale` is provided, we use it, otherwise we use the scale + # of the batch. + if self.default_scale is None: + batch_sum = ts_sum.sum(dim=0) + batch_observations = torch.clamp(num_observed.sum(0), min=1) + default_scale = torch.squeeze(batch_sum / batch_observations) + else: + default_scale = self.default_scale * torch.ones_like(scale) + + # apply default scale where there are no observations + scale = torch.where(num_observed > 0, scale, default_scale) + + # ensure the scale is at least `self.minimum_scale` + scale = torch.clamp(scale, min=self.minimum_scale) + scaled_data = data / scale + + if not self.keepdim: + scale = scale.squeeze(dim=self.dim) + + return scaled_data, torch.zeros_like(scale), scale + + +class InformerNOPScaler(nn.Module): + """ + Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data. + """ + + def __init__(self, config: InformerConfig): + super().__init__() + self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1 + self.keepdim = config.keepdim if hasattr(config, "keepdim") else True + + def forward( + self, data: torch.Tensor, observed_indicator: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Parameters: + data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): + input for Batch norm calculation + Returns: + tuple of `torch.Tensor` of shapes + (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, + `(batch_size, 1, num_input_channels)`) + """ + scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim) + loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim) + return data, loc, scale + + +class InformerSinusoidalPositionalEmbedding(nn.Embedding): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None) -> None: + super().__init__(num_positions, embedding_dim, _freeze=True) + + def create_weight(self): + """ + Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in + the 2nd half of the vector. [dim // 2:] + """ + n_pos, dim = self.weight.shape + position_enc = np.array( + [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)] + ) + out = torch.empty(n_pos, dim, dtype=self.weight.dtype, requires_grad=False) + sentinel = dim // 2 if dim % 2 == 0 else (dim // 2) + 1 + out[:, 0:sentinel] = torch.FloatTensor(np.sin(position_enc[:, 0::2])) + out[:, sentinel:] = torch.FloatTensor(np.cos(position_enc[:, 1::2])) + return out + + @torch.no_grad() + def forward( + self, input_ids_shape: torch.Size, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None + ) -> torch.Tensor: + """`input_ids_shape` is expected to be [bsz x seqlen].""" + if position_ids is None: + bsz, seq_len = input_ids_shape[:2] + position_ids = torch.arange( + past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device + ) + return super().forward(position_ids) + + +class InformerValueEmbedding(nn.Module): + def __init__(self, feature_size, d_model): + super().__init__() + self.value_projection = nn.Linear(in_features=feature_size, out_features=d_model, bias=False) + + def forward(self, x): + return self.value_projection(x) + + +@auto_docstring +class InformerPreTrainedModel(PreTrainedModel): + config: InformerConfig + base_model_prefix = "model" + main_input_name = "past_values" + input_modalities = ("time",) + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + super()._init_weights(module) + if isinstance(module, InformerSinusoidalPositionalEmbedding): + init.copy_(module.weight, module.create_weight()) + + +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 InformerAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + is_causal: bool = False, + config: InformerConfig | None = None, + layer_idx: int | None = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + self.layer_idx = layer_idx + if layer_idx is None and self.is_decoder: + logger.warning_once( + f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and " + "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2) + + is_updated = False + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + + current_states = key_value_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = curr_past_key_values.layers[self.layer_idx].keys + value_states = curr_past_key_values.layers[self.layer_idx].values + else: + key_states = self.k_proj(current_states) + value_states = self.v_proj(current_states) + key_states = key_states.view(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class InformerProbSparseAttention(nn.Module): + """Probabilistic Attention mechanism to select the "active" + queries rather than the "lazy" queries and provides a sparse Transformer thus mitigating the quadratic compute and + memory requirements of vanilla attention""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float | int = 0.0, + is_decoder: bool = False, + sampling_factor: int = 5, + bias: bool = True, + layer_idx: int | None = None, + ): + super().__init__() + self.factor = sampling_factor + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.layer_idx = layer_idx + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + bsz, tgt_len, _ = hidden_states.size() + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # 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(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.reshape(*proj_shape) + value_states = value_states.reshape(*proj_shape) + + key_states_time_length = key_states.size(1) # L_K + log_key_states_time_length = np.ceil(np.log1p(key_states_time_length)).astype("int").item() # log_L_K + + query_states_time_length = query_states.size(1) # L_Q + log_query_states_time_length = np.ceil(np.log1p(query_states_time_length)).astype("int").item() # log_L_Q + + u_part = min(self.factor * query_states_time_length * log_key_states_time_length, key_states_time_length) + u = min(self.factor * log_query_states_time_length, query_states_time_length) + + if key_states_time_length > 0: + index_sample = torch.randint(0, key_states_time_length, (u_part,)) + k_sample = key_states[:, index_sample, :] + else: + k_sample = key_states + + queries_keys_sample = torch.bmm(query_states, k_sample.transpose(1, 2)) # Q_K_sampled + + # find the Top_k query with sparsity measurement + if u > 0: + sparsity_measurement = queries_keys_sample.max(dim=-1)[0] - torch.div( + queries_keys_sample.sum(dim=-1), key_states_time_length + ) # M + top_u_sparsity_measurement = sparsity_measurement.topk(u, sorted=False)[1] # M_top + + # calculate q_reduce: query_states[:, top_u_sparsity_measurement] + dim_for_slice = torch.arange(query_states.size(0)).unsqueeze(-1) + q_reduce = query_states[dim_for_slice, top_u_sparsity_measurement] + else: + q_reduce = query_states + top_u_sparsity_measurement = None + + # Use q_reduce to calculate attention weights + attn_weights = torch.bmm(q_reduce, key_states.transpose(1, 2)) + + src_len = key_states.size(1) + if attn_weights.size() != (bsz * self.num_heads, u, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, u, 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()}" + ) + prob_mask = attention_mask.expand(bsz, self.num_heads, tgt_len, src_len).reshape( + bsz * self.num_heads, tgt_len, src_len + ) + + if top_u_sparsity_measurement is not None: + dim_for_slice = torch.arange(prob_mask.size(0)).unsqueeze(-1) + prob_mask = prob_mask[dim_for_slice, top_u_sparsity_measurement, :] + + attn_weights = attn_weights.view(bsz, self.num_heads, u, src_len) + prob_mask.view( + bsz, self.num_heads, u, src_len + ) + attn_weights = attn_weights.view(bsz * self.num_heads, u, src_len) + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + # this operation is a bit awkward, but it's required to + # make sure that attn_weights keeps its gradient. + # In order to do so, attn_weights have to be reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, u, src_len) + attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, u, src_len) + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + attn_output = torch.bmm(attn_probs, value_states) + + # calculate context for updating the attn_output, based on: + # https://github.com/zhouhaoyi/Informer2020/blob/ac59c7447135473fb2aafeafe94395f884d5c7a5/models/attn.py#L74 + if self.is_decoder: + # cast to float32 before operation to avoid overflow + context = value_states.cumsum(dim=-2, dtype=torch.float32).to(value_states.dtype) + else: + v_mean_dim_time = value_states.mean(dim=-2) + context = ( + v_mean_dim_time.unsqueeze(dim=1) + .expand(bsz * self.num_heads, query_states_time_length, v_mean_dim_time.size(-1)) + .clone() + ) + + if top_u_sparsity_measurement is not None: + # update context: copy the attention output to the context at top_u_sparsity_measurement index + dim_for_slice = torch.arange(context.size(0)).unsqueeze(-1) + context[dim_for_slice, top_u_sparsity_measurement, :] = attn_output + attn_output = context + + if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights_reshaped + + +# source: https://github.com/zhouhaoyi/Informer2020/blob/main/models/encoder.py +class InformerConvLayer(GradientCheckpointingLayer): + def __init__(self, c_in): + super().__init__() + self.downConv = nn.Conv1d( + in_channels=c_in, + out_channels=c_in, + kernel_size=3, + padding=1, + padding_mode="circular", + ) + self.norm = nn.BatchNorm1d(c_in) + self.activation = nn.ELU() + self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1) + + def forward(self, x): + x = self.downConv(x.permute(0, 2, 1)) + x = self.norm(x) + x = self.activation(x) + x = self.maxPool(x) + x = x.transpose(1, 2) + return x + + +class InformerEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: InformerConfig): + super().__init__() + self.embed_dim = config.d_model + 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) + + if config.attention_type == "prob": + self.self_attn = InformerProbSparseAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + sampling_factor=config.sampling_factor, + ) + else: + self.self_attn = InformerAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + config=config, + ) + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states, _ = self.self_attn( + hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + residual = hidden_states + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + if hidden_states.dtype == torch.float16 and not torch.isfinite(hidden_states).all(): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + return hidden_states + + +class InformerDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: InformerConfig, layer_idx: int | None = None): + super().__init__() + self.embed_dim = config.d_model + 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 = InformerAttention( + self.embed_dim, + config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + config=config, + layer_idx=layer_idx, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim) + self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + if config.attention_type == "prob": + self.self_attn = InformerProbSparseAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + sampling_factor=config.sampling_factor, + is_decoder=True, + layer_idx=layer_idx, + ) + else: + self.self_attn = InformerAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + config=config, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = True, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Cross-Attention Block + if encoder_hidden_states is not None: + residual = hidden_states + + hidden_states, _ = self.encoder_attn( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + # Fully Connected + residual = hidden_states + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states + + +class InformerEncoder(InformerPreTrainedModel): + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`InformerEncoderLayer`]. + + Args: + config: InformerConfig + """ + + _can_record_outputs = { + "hidden_states": InformerEncoderLayer, + "attentions": OutputRecorder(target_class=None, index=1, class_name="self_attn"), + } + + def __init__(self, config: InformerConfig): + super().__init__(config) + + self.dropout = config.dropout + self.layerdrop = config.encoder_layerdrop + if config.prediction_length is None: + raise ValueError("The `prediction_length` config needs to be specified.") + + self.value_embedding = InformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model) + self.embed_positions = InformerSinusoidalPositionalEmbedding( + config.context_length + config.prediction_length, config.d_model + ) + self.layers = nn.ModuleList([InformerEncoderLayer(config) for _ in range(config.encoder_layers)]) + self.layernorm_embedding = nn.LayerNorm(config.d_model) + self.gradient_checkpointing = False + + if config.distil: + self.conv_layers = nn.ModuleList( + [InformerConvLayer(config.d_model) for _ in range(config.encoder_layers - 1)] + ) + self.conv_layers.append(None) + else: + self.conv_layers = [None] * config.encoder_layers + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + hidden_states = self.value_embedding(inputs_embeds) + embed_pos = self.embed_positions(inputs_embeds.size()) + + hidden_states = self.layernorm_embedding(hidden_states + embed_pos) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + for idx, (encoder_layer, conv_layer) in enumerate(zip(self.layers, self.conv_layers)): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if not to_drop: + hidden_states = encoder_layer( + hidden_states, + attention_mask, + **kwargs, + ) + if conv_layer is not None: + hidden_states = conv_layer(hidden_states) + + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +class InformerDecoder(InformerPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a + [`InformerDecoderLayer`] + + Args: + config: InformerConfig + """ + + _can_record_outputs = { + "hidden_states": InformerDecoderLayer, + "attentions": OutputRecorder(target_class=None, index=1, class_name="self_attn"), + "cross_attentions": OutputRecorder(InformerAttention, index=1, layer_name="encoder_attn"), + } + + def __init__(self, config: InformerConfig): + super().__init__(config) + self.dropout = config.dropout + self.layerdrop = config.decoder_layerdrop + if config.prediction_length is None: + raise ValueError("The `prediction_length` config needs to be specified.") + + self.value_embedding = InformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model) + self.embed_positions = InformerSinusoidalPositionalEmbedding( + config.context_length + config.prediction_length, config.d_model + ) + self.layers = nn.ModuleList([InformerDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)]) + self.layernorm_embedding = nn.LayerNorm(config.d_model) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: + r""" + Args: + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + 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. + """ + + # initialize `past_key_values` + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + + self_attn_cache = ( + past_key_values.self_attention_cache + if isinstance(past_key_values, EncoderDecoderCache) + else past_key_values + ) + + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=self_attn_cache, + ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + hidden_states = self.value_embedding(inputs_embeds) + embed_pos = self.embed_positions(inputs_embeds.size(), past_key_values_length=self.config.context_length) + hidden_states = self.layernorm_embedding(hidden_states + embed_pos) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + hidden_states = decoder_layer( + hidden_states, + attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class InformerModel(InformerPreTrainedModel): + def __init__(self, config: InformerConfig): + super().__init__(config) + + if config.scaling == "mean" or config.scaling is True: + self.scaler = InformerMeanScaler(config) + elif config.scaling == "std": + self.scaler = InformerStdScaler(config) + else: + self.scaler = InformerNOPScaler(config) + + if config.num_static_categorical_features > 0: + self.embedder = InformerFeatureEmbedder( + cardinalities=config.cardinality, + embedding_dims=config.embedding_dimension, + ) + + # transformer encoder-decoder and mask initializer + self.encoder = InformerEncoder(config) + self.decoder = InformerDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + @property + def _past_length(self) -> int: + return self.config.context_length + max(self.config.lags_sequence) + + def get_lagged_subsequences( + self, sequence: torch.Tensor, subsequences_length: int, shift: int = 0 + ) -> torch.Tensor: + """ + Returns lagged subsequences of a given sequence. Returns a tensor of shape (N, S, C, I), + where S = subsequences_length and I = len(indices), containing lagged subsequences. Specifically, lagged[i, + j, :, k] = sequence[i, -indices[k]-S+j, :]. + + Args: + sequence: Tensor + The sequence from which lagged subsequences should be extracted. Shape: (N, T, C). + subsequences_length : int + Length of the subsequences to be extracted. + shift: int + Shift the lags by this amount back. + """ + sequence_length = sequence.shape[1] + indices = [lag - shift for lag in self.config.lags_sequence] + + if max(indices) + subsequences_length > sequence_length: + raise ValueError( + f"lags cannot go further than history length, found lag {max(indices)} " + f"while history length is only {sequence_length}" + ) + + lagged_values = [] + for lag_index in indices: + begin_index = -lag_index - subsequences_length + end_index = -lag_index if lag_index > 0 else None + lagged_values.append(sequence[:, begin_index:end_index, ...]) + return torch.stack(lagged_values, dim=-1) + + def create_network_inputs( + self, + past_values: torch.Tensor, + past_time_features: torch.Tensor, + static_categorical_features: torch.Tensor | None = None, + static_real_features: torch.Tensor | None = None, + past_observed_mask: torch.Tensor | None = None, + future_values: torch.Tensor | None = None, + future_time_features: torch.Tensor | None = None, + ): + # time feature + time_feat = ( + torch.cat( + ( + past_time_features[:, self._past_length - self.config.context_length :, ...], + future_time_features, + ), + dim=1, + ) + if future_values is not None + else past_time_features[:, self._past_length - self.config.context_length :, ...] + ) + + # target + if past_observed_mask is None: + past_observed_mask = torch.ones_like(past_values) + + context = past_values[:, -self.config.context_length :] + observed_context = past_observed_mask[:, -self.config.context_length :] + _, loc, scale = self.scaler(context, observed_context) + + inputs = ( + (torch.cat((past_values, future_values), dim=1) - loc) / scale + if future_values is not None + else (past_values - loc) / scale + ) + + # static features + if loc.ndim == 3: + squeezed_loc = loc.squeeze(1) + squeezed_scale = scale.squeeze(1) + else: + squeezed_loc = loc + squeezed_scale = scale + log_abs_loc = squeezed_loc.abs().log1p() + log_scale = squeezed_scale.log() + static_feat = torch.cat((log_abs_loc, log_scale), dim=1) + + if static_real_features is not None: + static_feat = torch.cat((static_real_features, static_feat), dim=1) + if static_categorical_features is not None: + embedded_cat = self.embedder(static_categorical_features) + static_feat = torch.cat((embedded_cat, static_feat), dim=1) + expanded_static_feat = static_feat.unsqueeze(1).expand(-1, time_feat.shape[1], -1) + + # all features + features = torch.cat((expanded_static_feat, time_feat), dim=-1) + + # lagged features + subsequences_length = ( + self.config.context_length + self.config.prediction_length + if future_values is not None + else self.config.context_length + ) + lagged_sequence = self.get_lagged_subsequences(sequence=inputs, subsequences_length=subsequences_length) + lags_shape = lagged_sequence.shape + reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1) + + if reshaped_lagged_sequence.shape[1] != time_feat.shape[1]: + raise ValueError( + f"input length {reshaped_lagged_sequence.shape[1]} and time feature lengths {time_feat.shape[1]} does not match" + ) + + # transformer inputs + transformer_inputs = torch.cat((reshaped_lagged_sequence, features), dim=-1) + + return transformer_inputs, loc, scale, static_feat + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + past_values: torch.Tensor, + past_time_features: torch.Tensor, + past_observed_mask: torch.Tensor, + static_categorical_features: torch.Tensor | None = None, + static_real_features: torch.Tensor | None = None, + future_values: torch.Tensor | None = None, + future_time_features: torch.Tensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqTSModelOutput | tuple: + r""" + past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`): + Past values of the time series, that serve as context in order to predict the future. The sequence size of + this tensor must be larger than the `context_length` of the model, since the model will use the larger size + to construct lag features, i.e. additional values from the past which are added in order to serve as "extra + context". + + The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if no + `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest + look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length of + the past. + + The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as + `static_categorical_features`, `static_real_features`, `past_time_features` and lags). + + Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`): + Required time features, which the model internally will add to `past_values`. These could be things like + "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These + could also be so-called "age" features, which basically help the model know "at which point in life" a + time-series is. Age features have small values for distant past time steps and increase monotonically the + more we approach the current time step. Holiday features are also a good example of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in + `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*): + Optional static categorical features for which the model will learn an embedding, which it will add to the + values of the time series. + + Static categorical features are features which have the same value for all time steps (static over time). + + A typical example of a static categorical feature is a time series ID. + static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*): + Optional static real features which the model will add to the values of the time series. + + Static real features are features which have the same value for all time steps (static over time). + + A typical example of a static real feature is promotion information. + future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)` or `(batch_size, prediction_length, input_size)`, *optional*): + Future values of the time series, that serve as labels for the model. The `future_values` is what the + Transformer needs during training to learn to output, given the `past_values`. + + The sequence length here is equal to `prediction_length`. + + See the demo notebook and code snippets for details. + + Optionally, during training any missing values need to be replaced with zeros and indicated via the + `future_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`): + Required time features for the prediction window, which the model internally will add to `future_values`. + These could be things like "month of year", "day of the month", etc. encoded as vectors (for instance as + Fourier features). These could also be so-called "age" features, which basically help the model know "at + which point in life" a time-series is. Age features have small values for distant past time steps and + increase monotonically the more we approach the current time step. Holiday features are also a good example + of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.num_time_features` + `config.num_dynamic_real_features`. + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of + hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. + + Examples: + + ```python + >>> from huggingface_hub import hf_hub_download + >>> import torch + >>> from transformers import InformerModel + + >>> file = hf_hub_download( + ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset" + ... ) + >>> batch = torch.load(file) + + >>> model = InformerModel.from_pretrained("huggingface/informer-tourism-monthly") + + >>> # during training, one provides both past and future values + >>> # as well as possible additional features + >>> outputs = model( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_values=batch["future_values"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> last_hidden_state = outputs.last_hidden_state + ```""" + transformer_inputs, loc, scale, static_feat = self.create_network_inputs( + past_values=past_values, + past_time_features=past_time_features, + past_observed_mask=past_observed_mask, + static_categorical_features=static_categorical_features, + static_real_features=static_real_features, + future_values=future_values, + future_time_features=future_time_features, + ) + + if encoder_outputs is None: + enc_input = transformer_inputs[:, : self.config.context_length, ...] + encoder_outputs = self.encoder( + inputs_embeds=enc_input, + **kwargs, + ) + # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True + 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, + ) + + # Avoid empty tensors and instead create a zeroes tensor which + # will be treated the same in torch, i.e. matmul with empty == all 0s + if self.config.context_length >= transformer_inputs.shape[1]: + bsz, _, dim = transformer_inputs.shape + dec_input = torch.zeros( + size=(bsz, 1, dim), device=transformer_inputs.device, dtype=transformer_inputs.dtype + ) + else: + dec_input = transformer_inputs[:, self.config.context_length :, ...] + + decoder_outputs = self.decoder( + inputs_embeds=dec_input, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs[0], + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + return Seq2SeqTSModelOutput( + 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, + loc=loc, + scale=scale, + static_features=static_feat, + ) + + +def weighted_average(input_tensor: torch.Tensor, weights: torch.Tensor | None = None, dim=None) -> torch.Tensor: + """ + Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero, + meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`. + + Args: + input_tensor (`torch.FloatTensor`): + Input tensor, of which the average must be computed. + weights (`torch.FloatTensor`, *optional*): + Weights tensor, of the same shape as `input_tensor`. + dim (`int`, *optional*): + The dim along which to average `input_tensor`. + + Returns: + `torch.FloatTensor`: The tensor with values averaged along the specified `dim`. + """ + if weights is not None: + weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor)) + sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0) + return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights + else: + return input_tensor.mean(dim=dim) + + +def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor: + """ + Computes the negative log likelihood loss from input distribution with respect to target. + """ + return -input.log_prob(target) + + +@auto_docstring +class InformerForPrediction(InformerPreTrainedModel): + def __init__(self, config: InformerConfig): + super().__init__(config) + + self.model = InformerModel(config) + if config.distribution_output == "student_t": + self.distribution_output = StudentTOutput(dim=config.input_size) + elif config.distribution_output == "normal": + self.distribution_output = NormalOutput(dim=config.input_size) + elif config.distribution_output == "negative_binomial": + self.distribution_output = NegativeBinomialOutput(dim=config.input_size) + else: + raise ValueError(f"Unknown distribution output {config.distribution_output}") + + self.parameter_projection = self.distribution_output.get_parameter_projection(self.model.config.d_model) + self.target_shape = self.distribution_output.event_shape + + if config.loss == "nll": + self.loss = nll + else: + raise ValueError(f"Unknown loss function {config.loss}") + + # Initialize weights of distribution_output and apply final processing + self.post_init() + + def output_params(self, dec_output): + return self.parameter_projection(dec_output) + + @torch.jit.ignore + def output_distribution(self, params, loc=None, scale=None, trailing_n=None) -> torch.distributions.Distribution: + sliced_params = params + if trailing_n is not None: + sliced_params = [p[:, -trailing_n:] for p in params] + return self.distribution_output.distribution(sliced_params, loc=loc, scale=scale) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + past_values: torch.Tensor, + past_time_features: torch.Tensor, + past_observed_mask: torch.Tensor, + static_categorical_features: torch.Tensor | None = None, + static_real_features: torch.Tensor | None = None, + future_values: torch.Tensor | None = None, + future_time_features: torch.Tensor | None = None, + future_observed_mask: torch.Tensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqTSModelOutput | tuple: + r""" + past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`): + Past values of the time series, that serve as context in order to predict the future. The sequence size of + this tensor must be larger than the `context_length` of the model, since the model will use the larger size + to construct lag features, i.e. additional values from the past which are added in order to serve as "extra + context". + + The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if no + `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest + look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length of + the past. + + The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as + `static_categorical_features`, `static_real_features`, `past_time_features` and lags). + + Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`): + Required time features, which the model internally will add to `past_values`. These could be things like + "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These + could also be so-called "age" features, which basically help the model know "at which point in life" a + time-series is. Age features have small values for distant past time steps and increase monotonically the + more we approach the current time step. Holiday features are also a good example of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in + `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*): + Optional static categorical features for which the model will learn an embedding, which it will add to the + values of the time series. + + Static categorical features are features which have the same value for all time steps (static over time). + + A typical example of a static categorical feature is a time series ID. + static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*): + Optional static real features which the model will add to the values of the time series. + + Static real features are features which have the same value for all time steps (static over time). + + A typical example of a static real feature is promotion information. + future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)` or `(batch_size, prediction_length, input_size)`, *optional*): + Future values of the time series, that serve as labels for the model. The `future_values` is what the + Transformer needs during training to learn to output, given the `past_values`. + + The sequence length here is equal to `prediction_length`. + + See the demo notebook and code snippets for details. + + Optionally, during training any missing values need to be replaced with zeros and indicated via the + `future_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`): + Required time features for the prediction window, which the model internally will add to `future_values`. + These could be things like "month of year", "day of the month", etc. encoded as vectors (for instance as + Fourier features). These could also be so-called "age" features, which basically help the model know "at + which point in life" a time-series is. Age features have small values for distant past time steps and + increase monotonically the more we approach the current time step. Holiday features are also a good example + of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.num_time_features` + `config.num_dynamic_real_features`. + future_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `future_values` were observed and which were missing. Mask values selected + in `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + + This mask is used to filter out missing values for the final loss calculation. + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of + hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. + + Examples: + + ```python + >>> from huggingface_hub import hf_hub_download + >>> import torch + >>> from transformers import InformerForPrediction + + >>> file = hf_hub_download( + ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset" + ... ) + >>> batch = torch.load(file) + + >>> model = InformerForPrediction.from_pretrained( + ... "huggingface/informer-tourism-monthly" + ... ) + + >>> # during training, one provides both past and future values + >>> # as well as possible additional features + >>> outputs = model( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_values=batch["future_values"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> loss = outputs.loss + >>> loss.backward() + + >>> # during inference, one only provides past values + >>> # as well as possible additional features + >>> # the model autoregressively generates future values + >>> outputs = model.generate( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> mean_prediction = outputs.sequences.mean(dim=1) + ```""" + if future_values is not None: + use_cache = False + + outputs: Seq2SeqTSModelOutput = self.model( + past_values=past_values, + past_time_features=past_time_features, + past_observed_mask=past_observed_mask, + static_categorical_features=static_categorical_features, + static_real_features=static_real_features, + future_values=future_values, + future_time_features=future_time_features, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + prediction_loss = None + params = None + if future_values is not None: + params = self.output_params(outputs.last_hidden_state) + distribution = self.output_distribution(params, loc=outputs.loc, scale=outputs.scale) + + loss = self.loss(distribution, future_values) + + if future_observed_mask is None: + future_observed_mask = torch.ones_like(future_values) + + if len(self.target_shape) == 0: + loss_weights = future_observed_mask + else: + loss_weights, _ = future_observed_mask.min(dim=-1, keepdim=False) + + prediction_loss = weighted_average(loss, weights=loss_weights) + + return Seq2SeqTSPredictionOutput( + loss=prediction_loss, + params=params, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + loc=outputs.loc, + scale=outputs.scale, + static_features=outputs.static_features, + ) + + @torch.no_grad() + def generate( + self, + past_values: torch.Tensor, + past_time_features: torch.Tensor, + future_time_features: torch.Tensor, + past_observed_mask: torch.Tensor | None = None, + static_categorical_features: torch.Tensor | None = None, + static_real_features: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + ) -> SampleTSPredictionOutput: + r""" + Greedily generate sequences of sample predictions from a model with a probability distribution head. + + Parameters: + past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`): + Past values of the time series, that serve as context in order to predict the future. The sequence size + of this tensor must be larger than the `context_length` of the model, since the model will use the + larger size to construct lag features, i.e. additional values from the past which are added in order to + serve as "extra context". + + The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if + no `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest + look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length + of the past. + + The `past_values` is what the Transformer encoder gets as input (with optional additional features, + such as `static_categorical_features`, `static_real_features`, `past_time_features` and lags). + + Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number + of variates in the time series per time step. + past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`): + Required time features, which the model internally will add to `past_values`. These could be things + like "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). + These could also be so-called "age" features, which basically help the model know "at which point in + life" a time-series is. Age features have small values for distant past time steps and increase + monotonically the more we approach the current time step. Holiday features are also a good example of + time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, + where the position encodings are learned from scratch internally as parameters of the model, the Time + Series Transformer requires to provide additional time features. The Time Series Transformer only + learns additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these + features must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`): + Required time features for the prediction window, which the model internally will add to sampled + predictions. These could be things like "month of year", "day of the month", etc. encoded as vectors + (for instance as Fourier features). These could also be so-called "age" features, which basically help + the model know "at which point in life" a time-series is. Age features have small values for distant + past time steps and increase monotonically the more we approach the current time step. Holiday features + are also a good example of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, + where the position encodings are learned from scratch internally as parameters of the model, the Time + Series Transformer requires to provide additional time features. The Time Series Transformer only + learns additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these + features must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected + in `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + + static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*): + Optional static categorical features for which the model will learn an embedding, which it will add to + the values of the time series. + + Static categorical features are features which have the same value for all time steps (static over + time). + + A typical example of a static categorical feature is a time series ID. + static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*): + Optional static real features which the model will add to the values of the time series. + + Static real features are features which have the same value for all time steps (static over time). + + A typical example of a static real feature is promotion information. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. + + Return: + [`SampleTSPredictionOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of + samples, prediction_length)` or `(batch_size, number of samples, prediction_length, input_size)` for + multivariate predictions. + """ + outputs = self( + static_categorical_features=static_categorical_features, + static_real_features=static_real_features, + past_time_features=past_time_features, + past_values=past_values, + past_observed_mask=past_observed_mask, + future_time_features=future_time_features, + future_values=None, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + use_cache=True, + ) + + decoder = self.model.get_decoder() + enc_last_hidden = outputs.encoder_last_hidden_state + loc = outputs.loc + scale = outputs.scale + static_feat = outputs.static_features + + num_parallel_samples = self.config.num_parallel_samples + repeated_loc = loc.repeat_interleave(repeats=num_parallel_samples, dim=0) + repeated_scale = scale.repeat_interleave(repeats=num_parallel_samples, dim=0) + + repeated_past_values = ( + past_values.repeat_interleave(repeats=num_parallel_samples, dim=0) - repeated_loc + ) / repeated_scale + + expanded_static_feat = static_feat.unsqueeze(1).expand(-1, future_time_features.shape[1], -1) + features = torch.cat((expanded_static_feat, future_time_features), dim=-1) + repeated_features = features.repeat_interleave(repeats=num_parallel_samples, dim=0) + + repeated_enc_last_hidden = enc_last_hidden.repeat_interleave(repeats=num_parallel_samples, dim=0) + + future_samples = [] + + # greedy decoding + for k in range(self.config.prediction_length): + lagged_sequence = self.model.get_lagged_subsequences( + sequence=repeated_past_values, + subsequences_length=1 + k, + shift=1, + ) + + lags_shape = lagged_sequence.shape + reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1) + + decoder_input = torch.cat((reshaped_lagged_sequence, repeated_features[:, : k + 1]), dim=-1) + + dec_output = decoder(inputs_embeds=decoder_input, encoder_hidden_states=repeated_enc_last_hidden) + dec_last_hidden = dec_output.last_hidden_state + + params = self.parameter_projection(dec_last_hidden[:, -1:]) + distr = self.output_distribution(params, loc=repeated_loc, scale=repeated_scale) + next_sample = distr.sample() + + repeated_past_values = torch.cat( + (repeated_past_values, (next_sample - repeated_loc) / repeated_scale), dim=1 + ) + future_samples.append(next_sample) + + concat_future_samples = torch.cat(future_samples, dim=1) + + return SampleTSPredictionOutput( + sequences=concat_future_samples.reshape( + (-1, num_parallel_samples, self.config.prediction_length) + self.target_shape, + ) + ) + + +__all__ = ["InformerForPrediction", "InformerModel", "InformerPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/informer/modular_informer.py b/third_party/transformers/src/transformers/models/informer/modular_informer.py new file mode 100644 index 0000000000000000000000000000000000000000..4c6558ecbd656e336868b1da94bf7e7e79de16ac --- /dev/null +++ b/third_party/transformers/src/transformers/models/informer/modular_informer.py @@ -0,0 +1,795 @@ +# Copyright 2023 Amazon 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 Informer model.""" + +import numpy as np +import torch +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, EncoderDecoderCache +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..bart.modeling_bart import BartAttention +from ..time_series_transformer.modeling_time_series_transformer import ( + TimeSeriesFeatureEmbedder, + TimeSeriesMeanScaler, + TimeSeriesNOPScaler, + TimeSeriesSinusoidalPositionalEmbedding, + TimeSeriesStdScaler, + TimeSeriesTransformerDecoder, + TimeSeriesTransformerDecoderLayer, + TimeSeriesTransformerEncoder, + TimeSeriesTransformerEncoderLayer, + TimeSeriesTransformerForPrediction, + TimeSeriesTransformerModel, + TimeSeriesValueEmbedding, +) +from .configuration_informer import InformerConfig + + +def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor: + """ + Computes the negative log likelihood loss from input distribution with respect to target. + """ + return -input.log_prob(target) + + +class InformerFeatureEmbedder(TimeSeriesFeatureEmbedder): + pass + + +class InformerStdScaler(TimeSeriesStdScaler): + pass + + +class InformerMeanScaler(TimeSeriesMeanScaler): + pass + + +class InformerNOPScaler(TimeSeriesNOPScaler): + pass + + +class InformerSinusoidalPositionalEmbedding(TimeSeriesSinusoidalPositionalEmbedding): + pass + + +class InformerValueEmbedding(TimeSeriesValueEmbedding): + pass + + +@auto_docstring +class InformerPreTrainedModel(PreTrainedModel): + config: InformerConfig + base_model_prefix = "model" + main_input_name = "past_values" + input_modalities = ("time",) + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + super()._init_weights(module) + if isinstance(module, InformerSinusoidalPositionalEmbedding): + init.copy_(module.weight, module.create_weight()) + + +class InformerAttention(BartAttention): + pass + + +class InformerProbSparseAttention(nn.Module): + """Probabilistic Attention mechanism to select the "active" + queries rather than the "lazy" queries and provides a sparse Transformer thus mitigating the quadratic compute and + memory requirements of vanilla attention""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float | int = 0.0, + is_decoder: bool = False, + sampling_factor: int = 5, + bias: bool = True, + layer_idx: int | None = None, + ): + super().__init__() + self.factor = sampling_factor + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.layer_idx = layer_idx + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + bsz, tgt_len, _ = hidden_states.size() + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # 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(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.reshape(*proj_shape) + value_states = value_states.reshape(*proj_shape) + + key_states_time_length = key_states.size(1) # L_K + log_key_states_time_length = np.ceil(np.log1p(key_states_time_length)).astype("int").item() # log_L_K + + query_states_time_length = query_states.size(1) # L_Q + log_query_states_time_length = np.ceil(np.log1p(query_states_time_length)).astype("int").item() # log_L_Q + + u_part = min(self.factor * query_states_time_length * log_key_states_time_length, key_states_time_length) + u = min(self.factor * log_query_states_time_length, query_states_time_length) + + if key_states_time_length > 0: + index_sample = torch.randint(0, key_states_time_length, (u_part,)) + k_sample = key_states[:, index_sample, :] + else: + k_sample = key_states + + queries_keys_sample = torch.bmm(query_states, k_sample.transpose(1, 2)) # Q_K_sampled + + # find the Top_k query with sparsity measurement + if u > 0: + sparsity_measurement = queries_keys_sample.max(dim=-1)[0] - torch.div( + queries_keys_sample.sum(dim=-1), key_states_time_length + ) # M + top_u_sparsity_measurement = sparsity_measurement.topk(u, sorted=False)[1] # M_top + + # calculate q_reduce: query_states[:, top_u_sparsity_measurement] + dim_for_slice = torch.arange(query_states.size(0)).unsqueeze(-1) + q_reduce = query_states[dim_for_slice, top_u_sparsity_measurement] + else: + q_reduce = query_states + top_u_sparsity_measurement = None + + # Use q_reduce to calculate attention weights + attn_weights = torch.bmm(q_reduce, key_states.transpose(1, 2)) + + src_len = key_states.size(1) + if attn_weights.size() != (bsz * self.num_heads, u, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, u, 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()}" + ) + prob_mask = attention_mask.expand(bsz, self.num_heads, tgt_len, src_len).reshape( + bsz * self.num_heads, tgt_len, src_len + ) + + if top_u_sparsity_measurement is not None: + dim_for_slice = torch.arange(prob_mask.size(0)).unsqueeze(-1) + prob_mask = prob_mask[dim_for_slice, top_u_sparsity_measurement, :] + + attn_weights = attn_weights.view(bsz, self.num_heads, u, src_len) + prob_mask.view( + bsz, self.num_heads, u, src_len + ) + attn_weights = attn_weights.view(bsz * self.num_heads, u, src_len) + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + # this operation is a bit awkward, but it's required to + # make sure that attn_weights keeps its gradient. + # In order to do so, attn_weights have to be reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, u, src_len) + attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, u, src_len) + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + attn_output = torch.bmm(attn_probs, value_states) + + # calculate context for updating the attn_output, based on: + # https://github.com/zhouhaoyi/Informer2020/blob/ac59c7447135473fb2aafeafe94395f884d5c7a5/models/attn.py#L74 + if self.is_decoder: + # cast to float32 before operation to avoid overflow + context = value_states.cumsum(dim=-2, dtype=torch.float32).to(value_states.dtype) + else: + v_mean_dim_time = value_states.mean(dim=-2) + context = ( + v_mean_dim_time.unsqueeze(dim=1) + .expand(bsz * self.num_heads, query_states_time_length, v_mean_dim_time.size(-1)) + .clone() + ) + + if top_u_sparsity_measurement is not None: + # update context: copy the attention output to the context at top_u_sparsity_measurement index + dim_for_slice = torch.arange(context.size(0)).unsqueeze(-1) + context[dim_for_slice, top_u_sparsity_measurement, :] = attn_output + attn_output = context + + if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights_reshaped + + +# source: https://github.com/zhouhaoyi/Informer2020/blob/main/models/encoder.py +class InformerConvLayer(GradientCheckpointingLayer): + def __init__(self, c_in): + super().__init__() + self.downConv = nn.Conv1d( + in_channels=c_in, + out_channels=c_in, + kernel_size=3, + padding=1, + padding_mode="circular", + ) + self.norm = nn.BatchNorm1d(c_in) + self.activation = nn.ELU() + self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1) + + def forward(self, x): + x = self.downConv(x.permute(0, 2, 1)) + x = self.norm(x) + x = self.activation(x) + x = self.maxPool(x) + x = x.transpose(1, 2) + return x + + +class InformerEncoderLayer(TimeSeriesTransformerEncoderLayer): + def __init__(self, config: InformerConfig): + super().__init__(config) + + del self.self_attn + + if config.attention_type == "prob": + self.self_attn = InformerProbSparseAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + sampling_factor=config.sampling_factor, + ) + else: + self.self_attn = InformerAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + config=config, + ) + + +class InformerDecoderLayer(TimeSeriesTransformerDecoderLayer): + def __init__(self, config: InformerConfig, layer_idx: int | None = None): + super().__init__(config) + + del self.self_attn + + if config.attention_type == "prob": + self.self_attn = InformerProbSparseAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + sampling_factor=config.sampling_factor, + is_decoder=True, + layer_idx=layer_idx, + ) + else: + self.self_attn = InformerAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + config=config, + layer_idx=layer_idx, + ) + + +class InformerEncoder(TimeSeriesTransformerEncoder): + _can_record_outputs = { + "hidden_states": InformerEncoderLayer, + "attentions": OutputRecorder(target_class=None, index=1, class_name="self_attn"), + } + + def __init__(self, config: InformerConfig): + super().__init__(config) + + self.dropout = config.dropout + self.layerdrop = config.encoder_layerdrop + self.gradient_checkpointing = False + if config.prediction_length is None: + raise ValueError("The `prediction_length` config needs to be specified.") + + self.value_embedding = InformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model) + self.embed_positions = InformerSinusoidalPositionalEmbedding( + config.context_length + config.prediction_length, config.d_model + ) + self.layers = nn.ModuleList([InformerEncoderLayer(config) for _ in range(config.encoder_layers)]) + self.layernorm_embedding = nn.LayerNorm(config.d_model) + + if config.distil: + self.conv_layers = nn.ModuleList( + [InformerConvLayer(config.d_model) for _ in range(config.encoder_layers - 1)] + ) + self.conv_layers.append(None) + else: + self.conv_layers = [None] * config.encoder_layers + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + hidden_states = self.value_embedding(inputs_embeds) + embed_pos = self.embed_positions(inputs_embeds.size()) + + hidden_states = self.layernorm_embedding(hidden_states + embed_pos) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + for idx, (encoder_layer, conv_layer) in enumerate(zip(self.layers, self.conv_layers)): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if not to_drop: + hidden_states = encoder_layer( + hidden_states, + attention_mask, + **kwargs, + ) + if conv_layer is not None: + hidden_states = conv_layer(hidden_states) + + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +class InformerDecoder(TimeSeriesTransformerDecoder): + _can_record_outputs = { + "hidden_states": InformerDecoderLayer, + "attentions": OutputRecorder(target_class=None, index=1, class_name="self_attn"), + "cross_attentions": OutputRecorder(InformerAttention, index=1, layer_name="encoder_attn"), + } + + def __init__(self, config: InformerConfig): + super().__init__(config) + self.dropout = config.dropout + self.layerdrop = config.decoder_layerdrop + if config.prediction_length is None: + raise ValueError("The `prediction_length` config needs to be specified.") + + self.value_embedding = InformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model) + self.embed_positions = InformerSinusoidalPositionalEmbedding( + config.context_length + config.prediction_length, config.d_model + ) + self.layers = nn.ModuleList([InformerDecoderLayer(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() + + +class InformerModel(TimeSeriesTransformerModel): + def __init__(self, config: InformerConfig): + PreTrainedModel.__init__(self, config) + + if config.scaling == "mean" or config.scaling is True: + self.scaler = InformerMeanScaler(config) + elif config.scaling == "std": + self.scaler = InformerStdScaler(config) + else: + self.scaler = InformerNOPScaler(config) + + if config.num_static_categorical_features > 0: + self.embedder = InformerFeatureEmbedder( + cardinalities=config.cardinality, + embedding_dims=config.embedding_dimension, + ) + + # transformer encoder-decoder and mask initializer + self.encoder = InformerEncoder(config) + self.decoder = InformerDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def forward(self, **super_kwargs): + r""" + past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`): + Past values of the time series, that serve as context in order to predict the future. The sequence size of + this tensor must be larger than the `context_length` of the model, since the model will use the larger size + to construct lag features, i.e. additional values from the past which are added in order to serve as "extra + context". + + The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if no + `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest + look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length of + the past. + + The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as + `static_categorical_features`, `static_real_features`, `past_time_features` and lags). + + Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`): + Required time features, which the model internally will add to `past_values`. These could be things like + "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These + could also be so-called "age" features, which basically help the model know "at which point in life" a + time-series is. Age features have small values for distant past time steps and increase monotonically the + more we approach the current time step. Holiday features are also a good example of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in + `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*): + Optional static categorical features for which the model will learn an embedding, which it will add to the + values of the time series. + + Static categorical features are features which have the same value for all time steps (static over time). + + A typical example of a static categorical feature is a time series ID. + static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*): + Optional static real features which the model will add to the values of the time series. + + Static real features are features which have the same value for all time steps (static over time). + + A typical example of a static real feature is promotion information. + future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)` or `(batch_size, prediction_length, input_size)`, *optional*): + Future values of the time series, that serve as labels for the model. The `future_values` is what the + Transformer needs during training to learn to output, given the `past_values`. + + The sequence length here is equal to `prediction_length`. + + See the demo notebook and code snippets for details. + + Optionally, during training any missing values need to be replaced with zeros and indicated via the + `future_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`): + Required time features for the prediction window, which the model internally will add to `future_values`. + These could be things like "month of year", "day of the month", etc. encoded as vectors (for instance as + Fourier features). These could also be so-called "age" features, which basically help the model know "at + which point in life" a time-series is. Age features have small values for distant past time steps and + increase monotonically the more we approach the current time step. Holiday features are also a good example + of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.num_time_features` + `config.num_dynamic_real_features`. + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of + hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. + + Examples: + + ```python + >>> from huggingface_hub import hf_hub_download + >>> import torch + >>> from transformers import InformerModel + + >>> file = hf_hub_download( + ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset" + ... ) + >>> batch = torch.load(file) + + >>> model = InformerModel.from_pretrained("huggingface/informer-tourism-monthly") + + >>> # during training, one provides both past and future values + >>> # as well as possible additional features + >>> outputs = model( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_values=batch["future_values"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> last_hidden_state = outputs.last_hidden_state + ```""" + super().forward(**super_kwargs) + + +class InformerForPrediction(TimeSeriesTransformerForPrediction): + def __init__(self, config: InformerConfig): + PreTrainedModel.__init__(self, config) + + self.model = InformerModel(config) + if config.distribution_output == "student_t": + self.distribution_output = StudentTOutput(dim=config.input_size) + elif config.distribution_output == "normal": + self.distribution_output = NormalOutput(dim=config.input_size) + elif config.distribution_output == "negative_binomial": + self.distribution_output = NegativeBinomialOutput(dim=config.input_size) + else: + raise ValueError(f"Unknown distribution output {config.distribution_output}") + + self.parameter_projection = self.distribution_output.get_parameter_projection(self.model.config.d_model) + self.target_shape = self.distribution_output.event_shape + + if config.loss == "nll": + self.loss = nll + else: + raise ValueError(f"Unknown loss function {config.loss}") + + # Initialize weights of distribution_output and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward(self, **super_kwargs): + r""" + past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`): + Past values of the time series, that serve as context in order to predict the future. The sequence size of + this tensor must be larger than the `context_length` of the model, since the model will use the larger size + to construct lag features, i.e. additional values from the past which are added in order to serve as "extra + context". + + The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if no + `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest + look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length of + the past. + + The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as + `static_categorical_features`, `static_real_features`, `past_time_features` and lags). + + Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`): + Required time features, which the model internally will add to `past_values`. These could be things like + "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These + could also be so-called "age" features, which basically help the model know "at which point in life" a + time-series is. Age features have small values for distant past time steps and increase monotonically the + more we approach the current time step. Holiday features are also a good example of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`. + past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in + `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*): + Optional static categorical features for which the model will learn an embedding, which it will add to the + values of the time series. + + Static categorical features are features which have the same value for all time steps (static over time). + + A typical example of a static categorical feature is a time series ID. + static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*): + Optional static real features which the model will add to the values of the time series. + + Static real features are features which have the same value for all time steps (static over time). + + A typical example of a static real feature is promotion information. + future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)` or `(batch_size, prediction_length, input_size)`, *optional*): + Future values of the time series, that serve as labels for the model. The `future_values` is what the + Transformer needs during training to learn to output, given the `past_values`. + + The sequence length here is equal to `prediction_length`. + + See the demo notebook and code snippets for details. + + Optionally, during training any missing values need to be replaced with zeros and indicated via the + `future_observed_mask`. + + For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number of + variates in the time series per time step. + future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`): + Required time features for the prediction window, which the model internally will add to `future_values`. + These could be things like "month of year", "day of the month", etc. encoded as vectors (for instance as + Fourier features). These could also be so-called "age" features, which basically help the model know "at + which point in life" a time-series is. Age features have small values for distant past time steps and + increase monotonically the more we approach the current time step. Holiday features are also a good example + of time features. + + These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where + the position encodings are learned from scratch internally as parameters of the model, the Time Series + Transformer requires to provide additional time features. The Time Series Transformer only learns + additional embeddings for `static_categorical_features`. + + Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these features + must but known at prediction time. + + The `num_features` here is equal to `config.num_time_features` + `config.num_dynamic_real_features`. + future_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*): + Boolean mask to indicate which `future_values` were observed and which were missing. Mask values selected + in `[0, 1]`: + + - 1 for values that are **observed**, + - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). + + This mask is used to filter out missing values for the final loss calculation. + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of + hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. + + Examples: + + ```python + >>> from huggingface_hub import hf_hub_download + >>> import torch + >>> from transformers import InformerForPrediction + + >>> file = hf_hub_download( + ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset" + ... ) + >>> batch = torch.load(file) + + >>> model = InformerForPrediction.from_pretrained( + ... "huggingface/informer-tourism-monthly" + ... ) + + >>> # during training, one provides both past and future values + >>> # as well as possible additional features + >>> outputs = model( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_values=batch["future_values"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> loss = outputs.loss + >>> loss.backward() + + >>> # during inference, one only provides past values + >>> # as well as possible additional features + >>> # the model autoregressively generates future values + >>> outputs = model.generate( + ... past_values=batch["past_values"], + ... past_time_features=batch["past_time_features"], + ... past_observed_mask=batch["past_observed_mask"], + ... static_categorical_features=batch["static_categorical_features"], + ... static_real_features=batch["static_real_features"], + ... future_time_features=batch["future_time_features"], + ... ) + + >>> mean_prediction = outputs.sequences.mean(dim=1) + ```""" + super().forward(**super_kwargs) + + +__all__ = ["InformerForPrediction", "InformerModel", "InformerPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/jetmoe/__init__.py b/third_party/transformers/src/transformers/models/jetmoe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7058590acc8859f934330ced9556d6ca66b50a51 --- /dev/null +++ b/third_party/transformers/src/transformers/models/jetmoe/__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_jetmoe import * + from .modeling_jetmoe 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/jetmoe/configuration_jetmoe.py b/third_party/transformers/src/transformers/models/jetmoe/configuration_jetmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..a977eb5abbdaa2bbef6b680a7f64c8210504b6b5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/jetmoe/configuration_jetmoe.py @@ -0,0 +1,82 @@ +# Copyright 2024 JetMoe 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. +"""JetMoe 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="jetmoe/jetmoe-8b") +@strict +class JetMoeConfig(PreTrainedConfig): + r""" + kv_channels (`int`, *optional*, defaults to 128): + Defines the number of channels for the key and value tensors. + num_local_experts (`int`, *optional*, defaults to 8): + Defines the number of experts in the MoE and MoA. + + ```python + >>> from transformers import JetMoeModel, JetMoeConfig + + >>> # Initializing a JetMoe 4B style configuration + >>> configuration = JetMoeConfig() + + >>> # Initializing a model from the JetMoe 4B style configuration + >>> model = JetMoeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "jetmoe" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"head_dim": "kv_channels"} + + vocab_size: int = 32000 + hidden_size: int = 2048 + num_hidden_layers: int = 12 + num_key_value_heads: int = 16 + kv_channels: int = 128 + intermediate_size: int = 5632 + max_position_embeddings: int = 4096 + activation_function: str = "silu" + num_local_experts: int = 8 + num_experts_per_tok: int = 2 + output_router_logits: bool = False + aux_loss_coef: float = 0.01 + use_cache: bool = True + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + pad_token_id: int | None = None + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + rms_norm_eps: float = 1e-6 + initializer_range: float = 0.01 + attention_dropout: float | int = 0.0 + + def __post_init__(self, **kwargs): + self.num_attention_heads = self.num_key_value_heads * self.num_experts_per_tok + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.num_experts_per_tok > self.num_local_experts: + raise ValueError("`num_experts_per_tok` must be less than or equal to `num_local_experts`") + + +__all__ = ["JetMoeConfig"] diff --git a/third_party/transformers/src/transformers/models/jetmoe/modeling_jetmoe.py b/third_party/transformers/src/transformers/models/jetmoe/modeling_jetmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..d3ee0bb1487575adaac82277a2964fd39fdb46fd --- /dev/null +++ b/third_party/transformers/src/transformers/models/jetmoe/modeling_jetmoe.py @@ -0,0 +1,830 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/jetmoe/modular_jetmoe.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_jetmoe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 JetMoe 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 collections.abc import Callable +from typing import Optional + +import torch +from torch import nn +from torch.nn import functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...masking_utils import create_causal_mask +from ...modeling_layers import GenericForSequenceClassification, 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, logging +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_jetmoe import JetMoeConfig + + +logger = logging.get_logger(__name__) + + +@use_kernel_forward_from_hub("RMSNorm") +class JetMoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + JetMoeRMSNorm 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 JetMoeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: JetMoeConfig, 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: JetMoeConfig | 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 JetMoeParallelExperts(nn.Module): + def __init__(self, num_experts: int, input_size: int, output_size: int) -> None: + """ + Initialize the JetMoeParallelExperts module. + The experts weights are stored in [num_experts, output_size, input_size] format. Such that it's compatible with + many MoE libraries, such as [Megablock](https://github.com/databricks/megablocks) and + [ScatterMoE](https://github.com/shawntan/scattermoe), as well as the + [MoE kernel](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py) + used in vllm. + + Args: + num_experts (int): + Number of experts. + input_size (int): + Size of the input. + output_size (int): + Size of the output. + """ + super().__init__() + self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size)) + self.num_experts = num_experts + self.input_size = input_size + self.output_size = output_size + + def forward(self, inputs, expert_size): + """ + Forward pass of the JetMoeParallelExperts module. + + Args: + inputs (Tensor): + Input tensor. + expert_size: + Expert size information. + + Returns: + Tensor: Output tensor. + """ + input_list = inputs.split(expert_size, dim=0) + output_list = [] + for i in range(self.num_experts): + output_list.append(F.linear(input_list[i], self.weight[i])) + results = torch.cat(output_list, dim=0) + return results + + +class JetMoeTopKGating(nn.Module): + def __init__(self, input_size: int, num_experts: int, top_k: int): + """ + Initialize the top-k gating mechanism. + + Args: + input_size (`int`): + Size of the input. + num_experts (`int`): + Number of experts. + top_k (`int`): + Number of top experts to select. + """ + super().__init__() + + self.num_experts = num_experts + self.input_size = input_size + self.top_k = top_k + + self.layer = nn.Linear(input_size, num_experts, bias=False) + + def forward(self, hidden_states): + # compute the top_k routing decision + logits = self.layer(hidden_states).float() # [batch_size x seq_len, num_experts] + top_k_logits, top_k_indices = logits.topk(self.top_k, dim=1) # [num_tokens, top_k] + top_k_gates = torch.softmax(top_k_logits, dim=1).type_as(hidden_states) # [num_tokens, top_k] + + # compute number of input given to each expert + zeros = torch.zeros( + [top_k_gates.size(0), self.num_experts], dtype=top_k_gates.dtype, device=top_k_gates.device + ) # [num_tokens, num_experts] + gates = zeros.scatter(1, top_k_indices, 1) # [num_tokens, num_experts] + expert_size = gates.long().sum(0) # [num_experts,] + # (This cause torch.compile to fail with `torch._dynamo.exc.Unsupported: Backend compiler failed with a fake tensor exception at`) + # (and `DataDependentOutputException`) + expert_size = expert_size.tolist() + + # sort and group input tokens according to expert assignment + top_k_experts = top_k_indices.flatten() # [num_tokens * top_k] + _, index_sorted_experts = top_k_experts.sort(0) # [num_tokens * top_k] + batch_index = index_sorted_experts.div(self.top_k, rounding_mode="trunc") # [num_tokens * top_k] + + # gather the gate values for grouped input tokens + top_k_gates = top_k_gates.flatten() # [num_tokens * top_k] + batch_gates = top_k_gates[index_sorted_experts] # [num_tokens * top_k] + + return index_sorted_experts, batch_index, batch_gates, expert_size, logits + + +class JetMoeMoE(nn.Module): + """ + A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts. + + Args: + config: + Configuration object with model hyperparameters. + """ + + def __init__(self, config: JetMoeConfig): + super().__init__() + + self.input_size = config.hidden_size + self.hidden_size = config.intermediate_size + self.activation = ACT2FN[config.activation_function] + self.bias = torch.nn.Parameter(torch.empty(self.input_size)) + self.input_linear = JetMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2) + self.output_linear = JetMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size) + + self.router = JetMoeTopKGating( + input_size=self.input_size, + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + ) + + def forward(self, layer_input): + """ + Forward pass of the mixture of experts layer. + + Args: + layer_input (Tensor): + Input tensor. + + Returns: + Tensor: + Output tensor. + Tensor: + Router logits. + """ + bsz, length, emb_size = layer_input.size() + layer_input = layer_input.reshape(-1, emb_size) + _, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input) + + expert_inputs = layer_input[batch_index] + hidden_states = self.input_linear(expert_inputs, expert_size) + chunked_hidden_states = hidden_states.chunk(2, dim=-1) + hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1] + expert_outputs = self.output_linear(hidden_states, expert_size) + + expert_outputs = expert_outputs * batch_gates[:, None] + + zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device) + layer_output = zeros.index_add(0, batch_index, expert_outputs) + layer_output = layer_output.view(bsz, length, self.input_size) + layer_output = layer_output + self.bias + return layer_output + + +class JetMoeMoA(nn.Module): + """ + A Sparsely gated mixture of attention layer with pairs of query- and output-projections as experts. + + Args: + config: + Configuration object with model hyperparameters. + """ + + def __init__(self, config: JetMoeConfig): + super().__init__() + + self.num_experts = config.num_local_experts + self.input_size = config.hidden_size + self.hidden_size = config.kv_channels * config.num_key_value_heads + self.top_k = config.num_experts_per_tok + self.bias = torch.nn.Parameter(torch.empty(self.input_size)) + + self.input_linear = JetMoeParallelExperts(self.num_experts, self.input_size, self.hidden_size) + self.output_linear = JetMoeParallelExperts(self.num_experts, self.hidden_size, self.input_size) + + self.router = JetMoeTopKGating( + input_size=self.input_size, + num_experts=self.num_experts, + top_k=self.top_k, + ) + + def map(self, layer_input): + """ + Map inputs to attention experts according to routing decision and compute query projection inside each experts. + """ + + # Compute gating topology + bsz, length, emb_size = layer_input.size() + layer_input = layer_input.reshape(-1, emb_size) # [bsz * length, emb_size] + index_sorted_experts, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input) + topo_info = (index_sorted_experts, batch_index, batch_gates, expert_size) + + # Group inputs according to topology and compute query projection + expert_inputs = layer_input[batch_index] # [bsz * length * top_k, emb_size] + expert_outputs = self.input_linear(expert_inputs, expert_size) # [bsz * length * top_k, hidden_size] + + # Ungroup queries back to original order + zeros = torch.zeros( + (bsz * length * self.top_k, self.hidden_size), dtype=expert_outputs.dtype, device=expert_outputs.device + ) + layer_output = zeros.index_add(0, index_sorted_experts, expert_outputs) + layer_output = layer_output.view(bsz, length, self.top_k, -1) # [bsz, length, top_k, hidden_size] + return layer_output, router_logits, topo_info + + def reduce(self, layer_input, topo_info): + """ + Compute output projection inside each attention experts and merge the outputs of different experts. + """ + bsz, length, k, hidden_size = layer_input.size() + layer_input = layer_input.reshape(-1, hidden_size) # [bsz * length * k, hidden_size] + index_sorted_experts, batch_index, batch_gates, expert_size = topo_info + + # Group inputs according to topology and compute output projection + expert_inputs = layer_input[index_sorted_experts] # [bsz * length * top_k, hidden_size] + expert_outputs = self.output_linear(expert_inputs, expert_size) # [bsz * length * top_k, emb_size] + + # Apply gates to attention expert outputs + expert_outputs = expert_outputs * batch_gates[:, None] + + # Ungroup and merge outputs to original order + zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device) + layer_output = zeros.index_add(0, batch_index, expert_outputs) + layer_output = layer_output.view(bsz, length, self.input_size) + layer_output = layer_output + self.bias + return layer_output + + def forward(self, layer_input): + raise NotImplementedError("This module doesn't support call and forward.") + + +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 + + +class JetMoeAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. + """ + + def __init__(self, config: JetMoeConfig, layer_idx: int | None = None): + """ + Initialize the JetMoeAttention module. + + Args: + config: + Configuration object with model hyperparameters. + layer_idx: + Index of the layer in the model. + """ + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.is_causal = True + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.num_key_value_groups = 1 # We ignore this by setting it to 1 as we have different repeat patterns + self.top_k = config.num_experts_per_tok + self.attention_dropout = config.attention_dropout + self.kv_projection_size = config.kv_channels * config.num_key_value_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_heads = config.num_attention_heads + self.head_dim = config.kv_channels + self.scaling = self.head_dim**-0.5 + self.experts = JetMoeMoA(config) + + self.kv_proj = torch.nn.Linear(config.hidden_size, self.kv_projection_size * 2, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> 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, router_logits, topo_info = self.experts.map(hidden_states) + key_states, value_states = self.kv_proj(hidden_states).chunk(2, dim=-1) + + 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 + ) + + # This is different from other models where we repeat k/v heads + # instead of repeat interleaving them + key_states = key_states.repeat(1, self.top_k, 1, 1) + value_states = value_states.repeat(1, self.top_k, 1, 1) + + 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.view(*input_shape, self.top_k, -1) + attn_output = self.experts.reduce(attn_output, topo_info) + attn_output = attn_output.view(*input_shape, -1) + return attn_output, attn_weights, router_logits + + +class JetMoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: JetMoeConfig, layer_idx: int | None = None): + super().__init__() + self.hidden_size = config.hidden_size + self.mlp = JetMoeMoE(config) + self.input_layernorm = JetMoeRMSNorm(config.hidden_size) + self.post_attention_layernorm = JetMoeRMSNorm(config.hidden_size) + self.self_attention = JetMoeAttention(config, layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _, _ = self.self_attention( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class JetMoePreTrainedModel(PreTrainedModel): + config: JetMoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = False + _no_split_modules = ["JetMoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False # TopK gating fails fullgraph compilation at "expert_size = expert_size.tolist()" + _supports_attention_backend = True + _can_record_outputs = { + "router_logits": [OutputRecorder(JetMoeAttention, index=2), OutputRecorder(JetMoeTopKGating, index=4)], + "hidden_states": JetMoeDecoderLayer, + "attentions": OutputRecorder(JetMoeAttention, index=1), + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, JetMoeParallelExperts): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, JetMoeMoA | JetMoeMoE): + init.zeros_(module.bias) + + +@auto_docstring +class JetMoeModel(JetMoePreTrainedModel): + def __init__(self, config: JetMoeConfig): + 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( + [JetMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = JetMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = JetMoeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self._attn_implementation = config._attn_implementation + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + 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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + past_key_values=past_key_values, + use_cache=use_cache, + position_ids=position_ids, + **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 + + +class JetMoeForCausalLM(JetMoePreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = JetMoeModel(config) + self.vocab_size = config.vocab_size + self.aux_loss_coef = config.aux_loss_coef + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.tie_word_embeddings = config.tie_word_embeddings + 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, + logits_to_keep: int | torch.Tensor = 0, + output_router_logits: bool | None = False, + **kwargs, + ) -> MoeCausalLMOutputWithPast: + 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, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.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 JetMoeForSequenceClassification(GenericForSequenceClassification, JetMoePreTrainedModel): ... + + +__all__ = ["JetMoeForCausalLM", "JetMoeModel", "JetMoePreTrainedModel", "JetMoeForSequenceClassification"] diff --git a/third_party/transformers/src/transformers/models/jetmoe/modular_jetmoe.py b/third_party/transformers/src/transformers/models/jetmoe/modular_jetmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..6daa402e8df972e390cbf7685ec60bf0ba40910d --- /dev/null +++ b/third_party/transformers/src/transformers/models/jetmoe/modular_jetmoe.py @@ -0,0 +1,597 @@ +# Copyright 2024 JetMoe AI and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch JetMoe model.""" + +from collections.abc import Callable + +import torch +from torch import nn +from torch.nn import functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_layers import ( + GenericForSequenceClassification, +) +from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..llama.modeling_llama import LlamaDecoderLayer +from ..mixtral.modeling_mixtral import ( + MixtralModel, + MixtralPreTrainedModel, + MixtralRMSNorm, + MixtralRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, + load_balancing_loss_func, +) +from .configuration_jetmoe import JetMoeConfig + + +logger = logging.get_logger(__name__) + + +class JetMoeRMSNorm(MixtralRMSNorm): + pass + + +class JetMoeRotaryEmbedding(MixtralRotaryEmbedding): + pass + + +class JetMoeParallelExperts(nn.Module): + def __init__(self, num_experts: int, input_size: int, output_size: int) -> None: + """ + Initialize the JetMoeParallelExperts module. + The experts weights are stored in [num_experts, output_size, input_size] format. Such that it's compatible with + many MoE libraries, such as [Megablock](https://github.com/databricks/megablocks) and + [ScatterMoE](https://github.com/shawntan/scattermoe), as well as the + [MoE kernel](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py) + used in vllm. + + Args: + num_experts (int): + Number of experts. + input_size (int): + Size of the input. + output_size (int): + Size of the output. + """ + super().__init__() + self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size)) + self.num_experts = num_experts + self.input_size = input_size + self.output_size = output_size + + def forward(self, inputs, expert_size): + """ + Forward pass of the JetMoeParallelExperts module. + + Args: + inputs (Tensor): + Input tensor. + expert_size: + Expert size information. + + Returns: + Tensor: Output tensor. + """ + input_list = inputs.split(expert_size, dim=0) + output_list = [] + for i in range(self.num_experts): + output_list.append(F.linear(input_list[i], self.weight[i])) + results = torch.cat(output_list, dim=0) + return results + + +class JetMoeTopKGating(nn.Module): + def __init__(self, input_size: int, num_experts: int, top_k: int): + """ + Initialize the top-k gating mechanism. + + Args: + input_size (`int`): + Size of the input. + num_experts (`int`): + Number of experts. + top_k (`int`): + Number of top experts to select. + """ + super().__init__() + + self.num_experts = num_experts + self.input_size = input_size + self.top_k = top_k + + self.layer = nn.Linear(input_size, num_experts, bias=False) + + def forward(self, hidden_states): + # compute the top_k routing decision + logits = self.layer(hidden_states).float() # [batch_size x seq_len, num_experts] + top_k_logits, top_k_indices = logits.topk(self.top_k, dim=1) # [num_tokens, top_k] + top_k_gates = torch.softmax(top_k_logits, dim=1).type_as(hidden_states) # [num_tokens, top_k] + + # compute number of input given to each expert + zeros = torch.zeros( + [top_k_gates.size(0), self.num_experts], dtype=top_k_gates.dtype, device=top_k_gates.device + ) # [num_tokens, num_experts] + gates = zeros.scatter(1, top_k_indices, 1) # [num_tokens, num_experts] + expert_size = gates.long().sum(0) # [num_experts,] + # (This cause torch.compile to fail with `torch._dynamo.exc.Unsupported: Backend compiler failed with a fake tensor exception at`) + # (and `DataDependentOutputException`) + expert_size = expert_size.tolist() + + # sort and group input tokens according to expert assignment + top_k_experts = top_k_indices.flatten() # [num_tokens * top_k] + _, index_sorted_experts = top_k_experts.sort(0) # [num_tokens * top_k] + batch_index = index_sorted_experts.div(self.top_k, rounding_mode="trunc") # [num_tokens * top_k] + + # gather the gate values for grouped input tokens + top_k_gates = top_k_gates.flatten() # [num_tokens * top_k] + batch_gates = top_k_gates[index_sorted_experts] # [num_tokens * top_k] + + return index_sorted_experts, batch_index, batch_gates, expert_size, logits + + +class JetMoeMoE(nn.Module): + """ + A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts. + + Args: + config: + Configuration object with model hyperparameters. + """ + + def __init__(self, config: JetMoeConfig): + super().__init__() + + self.input_size = config.hidden_size + self.hidden_size = config.intermediate_size + self.activation = ACT2FN[config.activation_function] + self.bias = torch.nn.Parameter(torch.empty(self.input_size)) + self.input_linear = JetMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2) + self.output_linear = JetMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size) + + self.router = JetMoeTopKGating( + input_size=self.input_size, + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + ) + + def forward(self, layer_input): + """ + Forward pass of the mixture of experts layer. + + Args: + layer_input (Tensor): + Input tensor. + + Returns: + Tensor: + Output tensor. + Tensor: + Router logits. + """ + bsz, length, emb_size = layer_input.size() + layer_input = layer_input.reshape(-1, emb_size) + _, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input) + + expert_inputs = layer_input[batch_index] + hidden_states = self.input_linear(expert_inputs, expert_size) + chunked_hidden_states = hidden_states.chunk(2, dim=-1) + hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1] + expert_outputs = self.output_linear(hidden_states, expert_size) + + expert_outputs = expert_outputs * batch_gates[:, None] + + zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device) + layer_output = zeros.index_add(0, batch_index, expert_outputs) + layer_output = layer_output.view(bsz, length, self.input_size) + layer_output = layer_output + self.bias + return layer_output + + +class JetMoeMoA(nn.Module): + """ + A Sparsely gated mixture of attention layer with pairs of query- and output-projections as experts. + + Args: + config: + Configuration object with model hyperparameters. + """ + + def __init__(self, config: JetMoeConfig): + super().__init__() + + self.num_experts = config.num_local_experts + self.input_size = config.hidden_size + self.hidden_size = config.kv_channels * config.num_key_value_heads + self.top_k = config.num_experts_per_tok + self.bias = torch.nn.Parameter(torch.empty(self.input_size)) + + self.input_linear = JetMoeParallelExperts(self.num_experts, self.input_size, self.hidden_size) + self.output_linear = JetMoeParallelExperts(self.num_experts, self.hidden_size, self.input_size) + + self.router = JetMoeTopKGating( + input_size=self.input_size, + num_experts=self.num_experts, + top_k=self.top_k, + ) + + def map(self, layer_input): + """ + Map inputs to attention experts according to routing decision and compute query projection inside each experts. + """ + + # Compute gating topology + bsz, length, emb_size = layer_input.size() + layer_input = layer_input.reshape(-1, emb_size) # [bsz * length, emb_size] + index_sorted_experts, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input) + topo_info = (index_sorted_experts, batch_index, batch_gates, expert_size) + + # Group inputs according to topology and compute query projection + expert_inputs = layer_input[batch_index] # [bsz * length * top_k, emb_size] + expert_outputs = self.input_linear(expert_inputs, expert_size) # [bsz * length * top_k, hidden_size] + + # Ungroup queries back to original order + zeros = torch.zeros( + (bsz * length * self.top_k, self.hidden_size), dtype=expert_outputs.dtype, device=expert_outputs.device + ) + layer_output = zeros.index_add(0, index_sorted_experts, expert_outputs) + layer_output = layer_output.view(bsz, length, self.top_k, -1) # [bsz, length, top_k, hidden_size] + return layer_output, router_logits, topo_info + + def reduce(self, layer_input, topo_info): + """ + Compute output projection inside each attention experts and merge the outputs of different experts. + """ + bsz, length, k, hidden_size = layer_input.size() + layer_input = layer_input.reshape(-1, hidden_size) # [bsz * length * k, hidden_size] + index_sorted_experts, batch_index, batch_gates, expert_size = topo_info + + # Group inputs according to topology and compute output projection + expert_inputs = layer_input[index_sorted_experts] # [bsz * length * top_k, hidden_size] + expert_outputs = self.output_linear(expert_inputs, expert_size) # [bsz * length * top_k, emb_size] + + # Apply gates to attention expert outputs + expert_outputs = expert_outputs * batch_gates[:, None] + + # Ungroup and merge outputs to original order + zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device) + layer_output = zeros.index_add(0, batch_index, expert_outputs) + layer_output = layer_output.view(bsz, length, self.input_size) + layer_output = layer_output + self.bias + return layer_output + + def forward(self, layer_input): + raise NotImplementedError("This module doesn't support call and forward.") + + +class JetMoeAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. + """ + + def __init__(self, config: JetMoeConfig, layer_idx: int | None = None): + """ + Initialize the JetMoeAttention module. + + Args: + config: + Configuration object with model hyperparameters. + layer_idx: + Index of the layer in the model. + """ + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.is_causal = True + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.num_key_value_groups = 1 # We ignore this by setting it to 1 as we have different repeat patterns + self.top_k = config.num_experts_per_tok + self.attention_dropout = config.attention_dropout + self.kv_projection_size = config.kv_channels * config.num_key_value_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_heads = config.num_attention_heads + self.head_dim = config.kv_channels + self.scaling = self.head_dim**-0.5 + self.experts = JetMoeMoA(config) + + self.kv_proj = torch.nn.Linear(config.hidden_size, self.kv_projection_size * 2, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> 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, router_logits, topo_info = self.experts.map(hidden_states) + key_states, value_states = self.kv_proj(hidden_states).chunk(2, dim=-1) + + 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 + ) + + # This is different from other models where we repeat k/v heads + # instead of repeat interleaving them + key_states = key_states.repeat(1, self.top_k, 1, 1) + value_states = value_states.repeat(1, self.top_k, 1, 1) + + 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.view(*input_shape, self.top_k, -1) + attn_output = self.experts.reduce(attn_output, topo_info) + attn_output = attn_output.view(*input_shape, -1) + return attn_output, attn_weights, router_logits + + +class JetMoeDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: JetMoeConfig, layer_idx: int | None = None): + super().__init__(config, layer_idx) + self.input_layernorm = JetMoeRMSNorm(config.hidden_size) + self.self_attention = JetMoeAttention(config, layer_idx) + self.post_attention_layernorm = JetMoeRMSNorm(config.hidden_size) + self.mlp = JetMoeMoE(config) + del self.self_attn + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _, _ = self.self_attention( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class JetMoePreTrainedModel(MixtralPreTrainedModel): + _can_record_outputs = { + "router_logits": [OutputRecorder(JetMoeAttention, index=2), OutputRecorder(JetMoeTopKGating, index=4)], + "hidden_states": JetMoeDecoderLayer, + "attentions": OutputRecorder(JetMoeAttention, index=1), + } + config: JetMoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = False + _no_split_modules = ["JetMoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _can_compile_fullgraph = False # TopK gating fails fullgraph compilation at "expert_size = expert_size.tolist()" + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + PreTrainedModel._init_weights(self, module) + if isinstance(module, JetMoeParallelExperts): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, JetMoeMoA | JetMoeMoE): + init.zeros_(module.bias) + + +@auto_docstring +class JetMoeModel(MixtralModel): + def __init__(self, config: JetMoeConfig): + 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( + [JetMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = JetMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + 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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + past_key_values=past_key_values, + use_cache=use_cache, + position_ids=position_ids, + **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 JetMoeForCausalLM(JetMoePreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = JetMoeModel(config) + self.vocab_size = config.vocab_size + self.aux_loss_coef = config.aux_loss_coef + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.tie_word_embeddings = config.tie_word_embeddings + 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, + logits_to_keep: int | torch.Tensor = 0, + output_router_logits: bool | None = False, + **kwargs, + ) -> MoeCausalLMOutputWithPast: + 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, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.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 JetMoeForSequenceClassification(GenericForSequenceClassification, JetMoePreTrainedModel): ... + + +__all__ = ["JetMoeForCausalLM", "JetMoeModel", "JetMoePreTrainedModel", "JetMoeForSequenceClassification"] diff --git a/third_party/transformers/src/transformers/models/lasr/__init__.py b/third_party/transformers/src/transformers/models/lasr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f4c7c98261a1c11dcb0046bb3964999484ad2c0e --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/__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_lasr import * + from .feature_extraction_lasr import * + from .modeling_lasr import * + from .tokenization_lasr 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/lasr/configuration_lasr.py b/third_party/transformers/src/transformers/models/lasr/configuration_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..d55ea6449d3724ac56545df2989addd4f835ef7e --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/configuration_lasr.py @@ -0,0 +1,149 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lasr/modular_lasr.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_lasr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team and Google LLC. 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="google/medasr") +@strict +class LasrEncoderConfig(PreTrainedConfig): + r""" + convolution_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in convolutions of the conformer's convolution module. + conv_kernel_size (`int`, *optional*, defaults to 32): + The kernel size of the convolution layers in the Conformer block. + subsampling_conv_channels (`int`, *optional*, defaults to 256): + The number of channels in the subsampling convolution layers. + subsampling_conv_kernel_size (`int`, *optional*, defaults to 5): + The kernel size of the subsampling convolution layers. + subsampling_conv_stride (`int`, *optional*, defaults to 2): + The stride of the subsampling convolution layers. + dropout_positions (`float`, *optional*, defaults to 0.0): + The dropout ratio for the positions in the input sequence. + feed_forward_residual_weights (`tuple[float, float]`, *optional*, defaults to `[1.5, 0.5]`): + The residual weights for the feed forward layers. + conv_residual_weights (`tuple[float, float]`, *optional*, defaults to `[2.0, 1.0]`): + The residual weights for the convolution layers. + batch_norm_momentum (`float`, *optional*, defaults to 0.01): + The momentum for the batch normalization layers + + Example: + ```python + >>> from transformers import LasrEncoderModel, LasrEncoderConfig + + >>> # Initializing a `LasrEncoder` configuration + >>> configuration = LasrEncoderConfig() + + >>> # Initializing a model from the configuration + >>> model = LasrEncoderModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + + This configuration class is based on the LasrEncoder architecture from Google Health AI. You can find more details + and pre-trained models at [google/medasr](https://huggingface.co/google/medasr). + """ + + model_type = "lasr_encoder" + keys_to_ignore_at_inference = ["past_key_values"] + + hidden_size: int = 512 + num_hidden_layers: int = 17 + num_attention_heads: int = 8 + intermediate_size: int = 2048 + hidden_act: str = "silu" + attention_bias: bool = False + convolution_bias: bool = False + conv_kernel_size: int = 32 + subsampling_conv_channels: int = 256 + num_mel_bins: int = 128 + subsampling_conv_kernel_size: int = 5 + subsampling_conv_stride: int = 2 + dropout: float | int = 0.1 + dropout_positions: float | int = 0.0 + layerdrop: float | int = 0.1 + activation_dropout: float | int = 0.1 + attention_dropout: float | int = 0.1 + max_position_embeddings: int = 10000 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-6 + feed_forward_residual_weights: list[float] | tuple[float, ...] = (1.5, 0.5) + conv_residual_weights: list[float] | tuple[float, ...] = (2.0, 1.0) + batch_norm_momentum: float = 0.01 + rope_parameters: dict | None = None + + def __post_init__(self, **kwargs): + self.num_key_value_heads = self.num_attention_heads + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="google/medasr") +@strict +class LasrCTCConfig(PreTrainedConfig): + r""" + ctc_loss_reduction (`str`, *optional*, defaults to `"mean"`): + Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an + instance of [`LasrForCTC`]. + ctc_zero_infinity (`bool`, *optional*, defaults to `True`): + Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly + occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance + of [`LasrForCTC`]. + + Example: + ```python + >>> from transformers import LasrForCTC, LasrCTCConfig + >>> # Initializing a Lasr configuration + >>> configuration = LasrCTCConfig() + >>> # Initializing a model from the configuration + >>> model = LasrForCTC(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + This configuration class is based on the Lasr CTC architecture from Google Health AI. You can find more details + and pre-trained models at [google/medasr](https://huggingface.co/google/medasr). + """ + + model_type = "lasr_ctc" + sub_configs = {"encoder_config": LasrEncoderConfig} + + vocab_size: int = 512 + ctc_loss_reduction: str = "mean" + ctc_zero_infinity: bool = True + encoder_config: dict | PreTrainedConfig | None = None + pad_token_id: int = 0 + + def __post_init__(self, **kwargs): + if isinstance(self.encoder_config, dict): + self.encoder_config = LasrEncoderConfig(**self.encoder_config) + elif self.encoder_config is None: + self.encoder_config = LasrEncoderConfig() + self.initializer_range = self.encoder_config.initializer_range + super().__post_init__(**kwargs) + + @property + def inputs_to_logits_ratio(self): + return self.encoder_config.subsampling_conv_stride**2 + + +__all__ = ["LasrEncoderConfig", "LasrCTCConfig"] diff --git a/third_party/transformers/src/transformers/models/lasr/feature_extraction_lasr.py b/third_party/transformers/src/transformers/models/lasr/feature_extraction_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..7cf1822ee40d8a112929ed6e66f87f825ecfb033 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/feature_extraction_lasr.py @@ -0,0 +1,275 @@ +# Copyright 2025 The HuggingFace Inc. team and Google LLC. 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 numpy as np +import torch + +from ...audio_utils import hertz_to_mel +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import TensorType, logging +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + + +# TODO: @eustlb, we should be able to remove this and use mel_filter_bank from audio_utils +def linear_to_mel_weight_matrix( + num_mel_bins: int, + num_spectrogram_bins: int, + sample_rate: float, + lower_edge_hertz: float, + upper_edge_hertz: float, + dtype, +) -> np.ndarray: + """NumPy-port of the JAX mel weight matrix logic.""" + # We use float64 for precision, matching the JAX implementation. + internal_dtype = np.float64 + + # HTK excludes the spectrogram DC bin. + bands_to_zero = 1 + nyquist_hertz = sample_rate / 2.0 + linear_frequencies = np.linspace(0.0, nyquist_hertz, num_spectrogram_bins, dtype=internal_dtype)[bands_to_zero:] + spectrogram_bins_mel = hertz_to_mel(linear_frequencies, mel_scale="kaldi")[:, np.newaxis] + + edges = np.linspace( + hertz_to_mel(lower_edge_hertz, mel_scale="kaldi"), + hertz_to_mel(upper_edge_hertz, mel_scale="kaldi"), + num_mel_bins + 2, + dtype=internal_dtype, + ) + + lower_edge_mel, center_mel, upper_edge_mel = ( + edges[:-2][np.newaxis, :], + edges[1:-1][np.newaxis, :], + edges[2:][np.newaxis, :], + ) + + lower_slopes = (spectrogram_bins_mel - lower_edge_mel) / (center_mel - lower_edge_mel) + upper_slopes = (upper_edge_mel - spectrogram_bins_mel) / (upper_edge_mel - center_mel) + mel_weights_matrix = np.maximum(0.0, np.minimum(lower_slopes, upper_slopes)) + return np.pad(mel_weights_matrix, [[bands_to_zero, 0], [0, 0]]).astype(dtype) + + +@requires(backends=("torch",)) +class LasrFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs a LASR feature extractor. + + This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains + most of the main methods. Users should refer to this superclass for more information regarding those methods. + + This class extracts mel-filter bank features from raw speech using a custom numpy implementation of the `Short Time + Fourier Transform` which should match pytorch's `torch.stft` equivalent. + + Args: + feature_size (`int`, *optional*, defaults to 128): + The feature dimension of the extracted features. + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + hop_length (`int`, *optional*, defaults to 160): + Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients. + n_fft (`int`, *optional*, defaults to 512): + Size of the Fourier transform. + win_length (`int`, *optional*, defaults to 400): + The window length for the STFT computation. + padding_value (`float`, *optional*, defaults to 0.0): + Padding value used to pad the audio. Should correspond to silences. + """ + + model_input_names = ["input_features", "attention_mask"] + + def __init__( + self, + feature_size=128, + sampling_rate=16000, + hop_length=160, + n_fft=512, + win_length=400, + padding_value=0.0, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + + self.hop_length = hop_length + self.n_fft = n_fft + self.win_length = win_length + self.mel_filters = torch.from_numpy( + linear_to_mel_weight_matrix( + num_mel_bins=feature_size, + num_spectrogram_bins=n_fft // 2 + 1, + sample_rate=sampling_rate, + lower_edge_hertz=125.0, + upper_edge_hertz=7500.0, + dtype=np.float64, + ) + ) + + def _torch_extract_fbank_features(self, waveform, device="cpu"): + # spectrogram + window = torch.hann_window(self.win_length, periodic=False, device=device, dtype=torch.float64) + waveform = waveform.to(torch.float64) + + # TODO: @eustlb, to be standardized + # here we cannot use directly torch.stft because every fft frame is padded with zeros + # due to unfold then rfft, while torch.stft unfolds with the number of fft points + frames = waveform.unfold(-1, self.win_length, self.hop_length) + stft = torch.fft.rfft(window * frames, n=self.n_fft) + power_spec = torch.abs(stft) ** 2 + + # log mel spectrogram + mel_filters = self.mel_filters.to(device) + mel_spec = torch.clamp(power_spec @ mel_filters, min=1e-5) + mel_spec = torch.log(mel_spec) + + return mel_spec + + def __call__( + self, + raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]], + truncation: bool = False, + pad_to_multiple_of: int | None = None, + return_tensors: str | TensorType | None = None, + return_attention_mask: bool | None = None, + padding: str | None = "longest", + max_length: int | None = None, + sampling_rate: int | None = None, + do_normalize: bool | None = None, + device: str | None = "cpu", + return_token_timestamps: bool | None = None, + **kwargs, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for + the STFT computation if available, otherwise a slower NumPy based one. + + Args: + raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`): + The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not + stereo, i.e. single float per timestep. + truncation (`bool`, *optional*, default to `True`): + Activates truncation to cut input sequences longer than *max_length* to *max_length*. + pad_to_multiple_of (`int`, *optional*, defaults to None): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific feature_extractor's default. + + [What are attention masks?](../glossary#attention-mask) + + + + For Parakeet models, `attention_mask` should always be passed for batched inference, to avoid subtle + bugs. + + + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + sampling_rate (`int`, *optional*): + The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition + pipeline. + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values / vectors. + do_normalize (`bool`, *optional*, defaults to `False`): + Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly + improve the performance of the model. + device (`str`, *optional*, defaults to `'cpu'`): + Specifies the device for computation of the log-mel spectrogram of audio signals in the + `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda") + return_token_timestamps (`bool`, *optional*, defaults to `None`): + Deprecated. Use `return_attention_mask` instead from which the number of frames can be inferred. + + Whether or not to return the number of frames of the input raw_speech. + These num_frames can be used by the model to compute word level timestamps. + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a" + f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input" + f" was sampled with {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + # Convert to torch tensor + if isinstance(raw_speech, np.ndarray): + raw_speech = torch.tensor(raw_speech) + elif isinstance(raw_speech, (list, tuple)): + if isinstance(raw_speech[0], (list, np.ndarray)): + raw_speech = [torch.tensor(speech) for speech in raw_speech] + else: # list[float] + raw_speech = torch.tensor(raw_speech) + + is_batched_torch = isinstance(raw_speech, torch.Tensor) and len(raw_speech.shape) > 1 + if is_batched_torch and len(raw_speech.shape) > 2: + logger.warning( + f"Only mono-channel audio is supported for input to {self.__class__.__name__}. " + "We will take the mean of the channels to convert to mono." + ) + raw_speech = raw_speech.mean(-1) + + is_batched_sequence = isinstance(raw_speech, (list, tuple)) + if is_batched_sequence: + for speech in raw_speech: + if len(speech.shape) > 1: + logger.warning( + f"Only mono-channel audio is supported for input to {self.__class__.__name__}. " + "We will take the mean of the channels to convert to mono." + ) + speech = speech.mean(-1) + + if is_batched_torch or is_batched_sequence: + raw_speech = [speech[:, None].to(torch.float32) for speech in raw_speech] + else: + raw_speech = [raw_speech[:, None].to(torch.float32)] + + batched_speech = BatchFeature({"input_features": raw_speech}) + padded_inputs = self.pad( + batched_speech, + padding=padding, + max_length=max_length, + truncation=truncation, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + return_tensors="pt", + ) + input_features = padded_inputs.input_features.squeeze(-1) + input_features = self._torch_extract_fbank_features(input_features, device) + data = { + "input_features": input_features.to(torch.float32), + } + + if return_attention_mask: + attention_mask = padded_inputs.attention_mask[:, self.win_length - 1 :: self.hop_length] + data["attention_mask"] = attention_mask.to(torch.bool) + + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["LasrFeatureExtractor"] diff --git a/third_party/transformers/src/transformers/models/lasr/modeling_lasr.py b/third_party/transformers/src/transformers/models/lasr/modeling_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..7ecea9099410ff1ced1ba72eeb050a0fef6e6620 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/modeling_lasr.py @@ -0,0 +1,728 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lasr/modular_lasr.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_lasr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team and Google LLC. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, CausalLMOutput +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 ModelOutput, 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_lasr import LasrCTCConfig, LasrEncoderConfig + + +class LasrEncoderSubsampling(nn.Module): + def __init__(self, config: LasrEncoderConfig): + super().__init__() + self.dense_0 = nn.Linear(config.num_mel_bins, config.hidden_size) + self.conv_0 = nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size=config.subsampling_conv_kernel_size, + stride=config.subsampling_conv_stride, + ) + self.conv_1 = nn.Conv1d( + config.hidden_size, + config.subsampling_conv_channels, + kernel_size=config.subsampling_conv_kernel_size, + stride=config.subsampling_conv_stride, + ) + self.dense_1 = nn.Linear(config.subsampling_conv_channels, config.hidden_size) + self.act_fn = nn.ReLU() + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + hidden_states = self.act_fn(self.dense_0(input_features)) + hidden_states = hidden_states.transpose(1, 2) + hidden_states = self.act_fn(self.conv_0(hidden_states)) + hidden_states = self.act_fn(self.conv_1(hidden_states)) + hidden_states = hidden_states.transpose(1, 2) + return self.dense_1(hidden_states) + + +class LasrEncoderRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: LasrEncoderConfig, 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: LasrEncoderConfig | 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 LasrEncoderAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LasrEncoderConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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 LasrEncoderConvolutionModule(nn.Module): + def __init__(self, config: LasrEncoderConfig, module_config=None): + """ + Args: + config (LasrEncoderConfig): Configuration for the model. + module_config (dict): Configuration for the module (e.g., encoder or decoder). + """ + super().__init__() + channels = config.hidden_size + # kernel_size should be an odd number for 'SAME' padding + if module_config is None: + # e.g. using `LasrEncoderEncoderConfig` in src/transformers/models/lasr_encoder/configuration_lasr_encoder.py + kernel_size = config.conv_kernel_size + self.activation = ACT2FN[getattr(config, "hidden_act", "silu")] + else: + kernel_size = module_config["kernel_size"] + self.activation = ACT2FN[module_config.get("activation", "silu")] + self.padding = "same" + self.pointwise_conv1 = nn.Conv1d( + channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=config.convolution_bias + ) + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=self.padding, + groups=channels, + bias=config.convolution_bias, + ) + self.norm = nn.BatchNorm1d(config.hidden_size, momentum=config.batch_norm_momentum) + self.pointwise_conv2 = nn.Conv1d( + channels, channels, kernel_size=1, stride=1, padding=0, bias=config.convolution_bias + ) + + def forward(self, hidden_states, attention_mask=None): + """ + Compute convolution module. + + Args: + hidden_states (`torch.Tensor` of shape `(batch, time, channels)`): Input tensor. + attention_mask (`torch.Tensor` of shape `(batch, 1, time, time)`): Attention mask. + + Returns: + `torch.Tensor`: Output tensor of shape `(batch, time, channels)`. + + """ + # exchange the temporal dimension and the feature dimension + hidden_states = hidden_states.transpose(1, 2) + + # GLU mechanism, (batch_size, 2*channel, dim) + hidden_states = self.pointwise_conv1(hidden_states) + # (batch_size, channel, dim) + hidden_states = nn.functional.glu(hidden_states, dim=1) + + # Apply padding mask before convolution + if attention_mask is not None: + if attention_mask.dtype == torch.bool: + all_masked_rows = torch.all(~attention_mask, dim=2) + else: + all_masked_rows = torch.all(~(attention_mask == 0.0), dim=2) + hidden_states = hidden_states.masked_fill(all_masked_rows, 0.0) + + # 1D Depthwise Conv + hidden_states = self.depthwise_conv(hidden_states) + hidden_states = self.norm(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.pointwise_conv2(hidden_states) + + return hidden_states.transpose(1, 2) + + +class LasrEncoderFeedForward(nn.Module): + def __init__(self, config: LasrEncoderConfig): + super().__init__() + self.linear1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias) + self.activation = ACT2FN[config.hidden_act] + self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias) + self.activation_dropout = config.activation_dropout + + def forward(self, hidden_states): + hidden_states = self.activation(self.linear1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.linear2(hidden_states) + return hidden_states + + +class LasrEncoderBlock(GradientCheckpointingLayer): + def __init__(self, config: LasrEncoderConfig, layer_idx: int): + super().__init__() + self.gradient_checkpointing = False + + self.feed_forward1 = LasrEncoderFeedForward(config) + self.self_attn = LasrEncoderAttention(config, layer_idx) + self.conv = LasrEncoderConvolutionModule(config) + self.feed_forward2 = LasrEncoderFeedForward(config) + + self.norm_feed_forward1 = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_self_att = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_conv = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_feed_forward2 = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_out = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + + self.feed_forward_residual_weights = config.feed_forward_residual_weights + self.conv_residual_weights = config.conv_residual_weights + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states)) + hidden_states = ( + self.feed_forward_residual_weights[0] * residual + self.feed_forward_residual_weights[1] * hidden_states + ) + + normalized_hidden_states = self.norm_self_att(hidden_states) + attn_output, _ = self.self_attn( + hidden_states=normalized_hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + attn_output + + conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask) + hidden_states = self.conv_residual_weights[0] * hidden_states + self.conv_residual_weights[1] * conv_output + + residual = hidden_states + hidden_states = self.feed_forward2(self.norm_feed_forward2(hidden_states)) + hidden_states = ( + self.feed_forward_residual_weights[0] * residual + self.feed_forward_residual_weights[1] * hidden_states + ) + + hidden_states = self.norm_out(hidden_states) + + return hidden_states + + +@auto_docstring +class LasrPreTrainedModel(PreTrainedModel): + config: LasrCTCConfig + base_model_prefix = "model" + main_input_name = "input_features" + input_modalities = "audio" + supports_gradient_checkpointing = True + _no_split_modules = ["LasrEncoderBlock"] + _supports_flat_attention_mask = True + _supports_sdpa = True + # padding is incompatible with flex attention as the resulting mask cannot be used to apply padding + _supports_flex_attn = False + + # TODO: @eustlb, add support when flash attention supports custom attention bias + _supports_flash_attn = False + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": LasrEncoderBlock, + "attentions": LasrEncoderAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + + def _get_subsampling_output_length(self, input_lengths: torch.Tensor): + encoder_config = self.config.encoder_config if isinstance(self.config, LasrCTCConfig) else self.config + kernel_size = encoder_config.subsampling_conv_kernel_size + stride = encoder_config.subsampling_conv_stride + + num_layers = 2 + for _ in range(num_layers): + input_lengths = (input_lengths - kernel_size) // stride + 1 + + return input_lengths + + def _get_output_attention_mask(self, attention_mask: torch.Tensor, target_length: int | None = None): + """ + Convert the input attention mask to its subsampled form. `target_length` sets the desired output length, useful + when the attention mask length differs from `sum(-1).max()` (i.e., when the longest sequence in the batch is padded) + """ + output_lengths = self._get_subsampling_output_length(attention_mask.sum(-1)) + # Use target_length if provided, otherwise use max length in batch + max_length = target_length if target_length is not None else output_lengths.max() + attention_mask = torch.arange(max_length, device=attention_mask.device) < output_lengths[:, None] + return attention_mask + + +@auto_docstring( + custom_intro=""" + The LasrEncoder model, based on the Conformer architecture](https://arxiv.org/abs/2005.08100). + """ +) +class LasrEncoder(LasrPreTrainedModel): + config: LasrEncoderConfig + base_model_prefix = "encoder" + + def __init__(self, config: LasrEncoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + + self.dropout = config.dropout + self.dropout_positions = config.dropout_positions + self.layerdrop = config.layerdrop + + self.subsampler = LasrEncoderSubsampling(config) + self.rotary_emb = LasrEncoderRotaryEmbedding(config) + self.layers = nn.ModuleList( + [LasrEncoderBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.out_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, bias=False) + + self.post_init() + + @auto_docstring + @merge_with_config_defaults + @capture_outputs + @can_return_tuple + def forward( + self, + input_features: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + Example: + + ```python + >>> from transformers import AutoProcessor, LasrEncoder + >>> from datasets import load_dataset, Audio + + >>> model_id = TODO + >>> processor = AutoProcessor.from_pretrained(model_id) + >>> encoder = ParakeetEncoder.from_pretrained(model_id) + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) + + >>> inputs = processor(ds[0]["audio"]["array"]) + >>> encoder_outputs = encoder(**inputs) + + >>> print(encoder_outputs.last_hidden_state.shape) + ``` + """ + + hidden_states = self.subsampler(input_features) + cos, sin = self.rotary_emb( + hidden_states, torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + cos = nn.functional.dropout(cos, p=self.dropout_positions, training=self.training) + sin = nn.functional.dropout(sin, p=self.dropout_positions, training=self.training) + + if attention_mask is not None: + attention_mask = self._get_output_attention_mask(attention_mask, target_length=hidden_states.shape[1]) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + for encoder_layer in self.layers: + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if not to_drop: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=(cos, sin), + **kwargs, + ) + + hidden_states = self.out_norm(hidden_states) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +@dataclass +class LasrGenerateOutput(ModelOutput): + """ + Outputs of Lasr models. + + Args: + sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): + Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`. + """ + + sequences: torch.LongTensor + logits: tuple[torch.FloatTensor] | None = None + attentions: tuple[tuple[torch.FloatTensor]] | None = None + hidden_states: tuple[tuple[torch.FloatTensor]] | None = None + + +@auto_docstring( + custom_intro=""" + Lasr Encoder with a Connectionist Temporal Classification (CTC) head. + """ +) +class LasrForCTC(LasrPreTrainedModel): + config: LasrCTCConfig + + def __init__(self, config: LasrCTCConfig): + super().__init__(config) + self.encoder = LasrEncoder(config.encoder_config) + # Conv rather than linear to be consistent with NeMO decoding layer + self.ctc_head = nn.Conv1d(config.encoder_config.hidden_size, config.vocab_size, kernel_size=1) + + self.post_init() + + @auto_docstring + @can_return_tuple + def forward( + self, + input_features: torch.Tensor, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutput: + r""" + Example: + + ```python + >>> from transformers import AutoProcessor, LasrForCTC + >>> from datasets import load_dataset, Audio + + >>> model_id = "nvidia/lasr-ctc-1.1b" + >>> processor = AutoProcessor.from_pretrained(model_id) + >>> model = LasrForCTC.from_pretrained(model_id) + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) + + >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"]) + >>> outputs = model(**inputs) + + >>> print(outputs.loss) + ```""" + + encoder_outputs = self.encoder( + input_features=input_features, + attention_mask=attention_mask, + **kwargs, + ) + + hidden_states = encoder_outputs.last_hidden_state + logits = self.ctc_head(hidden_states.transpose(1, 2)).transpose(1, 2) + + loss = None + if labels is not None: + # retrieve loss input_lengths from attention_mask + attention_mask = ( + attention_mask if attention_mask is not None else torch.ones_like(input_features, dtype=torch.long) + ) + input_lengths = self._get_subsampling_output_length(attention_mask.sum(-1)) + + # assuming that padded tokens are filled with -100 + # when not being attended to + labels_mask = labels != self.config.pad_token_id + target_lengths = labels_mask.sum(-1) + flattened_targets = labels.masked_select(labels_mask) + + # ctc_loss doesn't support fp16 + log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1) + + with torch.backends.cudnn.flags(enabled=False): + loss = nn.functional.ctc_loss( + log_probs, + flattened_targets, + input_lengths, + target_lengths, + blank=self.config.pad_token_id, + reduction=self.config.ctc_loss_reduction, + zero_infinity=self.config.ctc_zero_infinity, + ) + + return CausalLMOutput( + loss=loss, + logits=logits, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + @torch.no_grad() + def generate( + self, + input_features: torch.Tensor, + attention_mask: torch.Tensor | None = None, + return_dict_in_generate: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> LasrGenerateOutput | torch.LongTensor: + r""" + Example: + + ```python + >>> from transformers import AutoProcessor, LasrForCTC + >>> from datasets import load_dataset, Audio + + >>> model_id = TODO + >>> processor = AutoProcessor.from_pretrained(model_id) + >>> model = LasrForCTC.from_pretrained(model_id) + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) + + >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"]) + >>> predicted_ids = model.generate(**inputs) + >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) + + >>> print(transcription) + ``` + """ + kwargs["return_dict"] = True + outputs: CausalLMOutput = self.forward( + input_features=input_features, + attention_mask=attention_mask, + **kwargs, + ) + + # greedy decoding + sequences = outputs.logits.argmax(dim=-1) + + # mask out padded tokens + if attention_mask is not None: + attention_mask = self._get_output_attention_mask(attention_mask, target_length=sequences.shape[1]) + sequences[~attention_mask] = self.config.pad_token_id + + if return_dict_in_generate: + return LasrGenerateOutput( + sequences=sequences, + logits=outputs.logits, + attentions=outputs.attentions, + hidden_states=outputs.hidden_states, + ) + + return sequences + + +__all__ = ["LasrForCTC", "LasrEncoder", "LasrPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/lasr/modular_lasr.py b/third_party/transformers/src/transformers/models/lasr/modular_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..f016f16cff45702e32c7709dfd55c1da9dfec120 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/modular_lasr.py @@ -0,0 +1,523 @@ +# Copyright 2025 The HuggingFace Inc. team and Google LLC. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +from collections.abc import Callable + +import torch +from huggingface_hub.dataclasses import strict +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import Unigram +from torch import nn + +from ...masking_utils import create_bidirectional_mask +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..llama.modeling_llama import LlamaAttention, LlamaRotaryEmbedding, apply_rotary_pos_emb, eager_attention_forward +from ..parakeet.configuration_parakeet import ParakeetCTCConfig, ParakeetEncoderConfig +from ..parakeet.modeling_parakeet import ( + ParakeetEncoderBlock, + ParakeetEncoderConvolutionModule, + ParakeetForCTC, + ParakeetPreTrainedModel, +) +from ..parakeet.processing_parakeet import ParakeetProcessor +from ..t5.tokenization_t5 import T5Tokenizer + + +class LasrTokenizer(T5Tokenizer, TokenizersBackend): + def __init__( + self, + eos_token="", + unk_token="", + pad_token="", + _spm_precompiled_charsmap=None, + extra_ids=100, + additional_special_tokens=None, + vocab=None, + vocab_file=None, + **kwargs, + ): + self._extra_ids = extra_ids + + # Handle extra_ids and additional_special_tokens + if additional_special_tokens is not None: + extra_tokens = [x for x in additional_special_tokens if "" for i in range(extra_ids)] + elif extra_ids > 0 and extra_ids != len(extra_tokens): + raise ValueError( + f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are" + " provided to LasrTokenizer. In this case the additional_special_tokens must include the extra_ids" + " tokens" + ) + else: + extra_tokens = [f"" for i in range(extra_ids)] + additional_special_tokens = extra_tokens + + # LASR vocab structure: =0, =1, =2, then regular vocab, then extra_ids in reverse + if vocab is not None: + self._vocab_scores = vocab + else: + self._vocab_scores = [ + (str(pad_token), 0.0), + (str(eos_token), 0.0), + (str(unk_token), 0.0), + ("▁", -2.0), # Space token + ] + for i in range(extra_ids - 1, -1, -1): + self._vocab_scores.append((f"", 0.0)) + self._tokenizer = Tokenizer( + Unigram( + self._vocab_scores, + unk_id=3, + byte_fallback=False, + ) + ) + + if _spm_precompiled_charsmap is not None: + self._tokenizer.normalizer = normalizers.Precompiled(_spm_precompiled_charsmap) + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.WhitespaceSplit(), + pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True), + ] + ) + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True) + + TokenizersBackend.__init__( + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + extra_ids=extra_ids, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.eos_token_id), + ], + ) + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + group_tokens: bool = True, + **kwargs, + ) -> str: + if isinstance(token_ids, int): + token_ids = [token_ids] + if group_tokens: + token_ids = [token_group[0] for token_group in itertools.groupby(token_ids)] + + # for CTC we filter out the blank token, which is the pad token + token_ids = [token for token in token_ids if token != self.pad_token_id] + + return TokenizersBackend._decode( + self, + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + +class LasrProcessor(ParakeetProcessor): + pass + + +@auto_docstring(checkpoint="google/medasr") +@strict +class LasrEncoderConfig(ParakeetEncoderConfig): + r""" + convolution_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in convolutions of the conformer's convolution module. + conv_kernel_size (`int`, *optional*, defaults to 32): + The kernel size of the convolution layers in the Conformer block. + subsampling_conv_channels (`int`, *optional*, defaults to 256): + The number of channels in the subsampling convolution layers. + subsampling_conv_kernel_size (`int`, *optional*, defaults to 5): + The kernel size of the subsampling convolution layers. + subsampling_conv_stride (`int`, *optional*, defaults to 2): + The stride of the subsampling convolution layers. + dropout_positions (`float`, *optional*, defaults to 0.0): + The dropout ratio for the positions in the input sequence. + feed_forward_residual_weights (`tuple[float, float]`, *optional*, defaults to `[1.5, 0.5]`): + The residual weights for the feed forward layers. + conv_residual_weights (`tuple[float, float]`, *optional*, defaults to `[2.0, 1.0]`): + The residual weights for the convolution layers. + batch_norm_momentum (`float`, *optional*, defaults to 0.01): + The momentum for the batch normalization layers + + Example: + ```python + >>> from transformers import LasrEncoderModel, LasrEncoderConfig + + >>> # Initializing a `LasrEncoder` configuration + >>> configuration = LasrEncoderConfig() + + >>> # Initializing a model from the configuration + >>> model = LasrEncoderModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + + This configuration class is based on the LasrEncoder architecture from Google Health AI. You can find more details + and pre-trained models at [google/medasr](https://huggingface.co/google/medasr). + """ + + hidden_size: int = 512 + num_hidden_layers: int = 17 + intermediate_size: int = 2048 + attention_bias: bool = False + convolution_bias: bool = False + conv_kernel_size: int = 32 + subsampling_conv_kernel_size: int = 5 + num_mel_bins: int = 128 + max_position_embeddings: int = 10000 + layer_norm_eps: float = 1e-6 + feed_forward_residual_weights: list[float] | tuple[float, ...] = (1.5, 0.5) + conv_residual_weights: list[float] | tuple[float, ...] = (2.0, 1.0) + batch_norm_momentum: float = 0.01 + rope_parameters: dict | None = None + + subsampling_factor = AttributeError() + scale_input = AttributeError() + + +@auto_docstring(checkpoint="google/medasr") +@strict +class LasrCTCConfig(ParakeetCTCConfig): + r""" + ctc_loss_reduction (`str`, *optional*, defaults to `"mean"`): + Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an + instance of [`LasrForCTC`]. + ctc_zero_infinity (`bool`, *optional*, defaults to `True`): + Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly + occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance + of [`LasrForCTC`]. + + Example: + ```python + >>> from transformers import LasrForCTC, LasrCTCConfig + >>> # Initializing a Lasr configuration + >>> configuration = LasrCTCConfig() + >>> # Initializing a model from the configuration + >>> model = LasrForCTC(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + This configuration class is based on the Lasr CTC architecture from Google Health AI. You can find more details + and pre-trained models at [google/medasr](https://huggingface.co/google/medasr). + """ + + vocab_size: int = 512 + pad_token_id: int = 0 + + @property + def inputs_to_logits_ratio(self): + return self.encoder_config.subsampling_conv_stride**2 + + +class LasrEncoderSubsampling(nn.Module): + def __init__(self, config: LasrEncoderConfig): + super().__init__() + self.dense_0 = nn.Linear(config.num_mel_bins, config.hidden_size) + self.conv_0 = nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size=config.subsampling_conv_kernel_size, + stride=config.subsampling_conv_stride, + ) + self.conv_1 = nn.Conv1d( + config.hidden_size, + config.subsampling_conv_channels, + kernel_size=config.subsampling_conv_kernel_size, + stride=config.subsampling_conv_stride, + ) + self.dense_1 = nn.Linear(config.subsampling_conv_channels, config.hidden_size) + self.act_fn = nn.ReLU() + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + hidden_states = self.act_fn(self.dense_0(input_features)) + hidden_states = hidden_states.transpose(1, 2) + hidden_states = self.act_fn(self.conv_0(hidden_states)) + hidden_states = self.act_fn(self.conv_1(hidden_states)) + hidden_states = hidden_states.transpose(1, 2) + return self.dense_1(hidden_states) + + +class LasrEncoderRotaryEmbedding(LlamaRotaryEmbedding): ... + + +class LasrEncoderAttention(LlamaAttention): + def __init__(self, config: LasrEncoderConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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 LasrEncoderConvolutionModule(ParakeetEncoderConvolutionModule): + def __init__(self, config: LasrEncoderConfig, module_config=None): + super().__init__(config, module_config) + self.padding = "same" + self.norm = nn.BatchNorm1d(config.hidden_size, momentum=config.batch_norm_momentum) + + +class LasrEncoderBlock(ParakeetEncoderBlock): + def __init__(self, config: LasrEncoderConfig, layer_idx: int): + super().__init__(config, layer_idx) + + self.feed_forward_residual_weights = config.feed_forward_residual_weights + self.conv_residual_weights = config.conv_residual_weights + + self.norm_feed_forward1 = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_self_att = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_conv = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_feed_forward2 = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + self.norm_out = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states)) + hidden_states = ( + self.feed_forward_residual_weights[0] * residual + self.feed_forward_residual_weights[1] * hidden_states + ) + + normalized_hidden_states = self.norm_self_att(hidden_states) + attn_output, _ = self.self_attn( + hidden_states=normalized_hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + attn_output + + conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask) + hidden_states = self.conv_residual_weights[0] * hidden_states + self.conv_residual_weights[1] * conv_output + + residual = hidden_states + hidden_states = self.feed_forward2(self.norm_feed_forward2(hidden_states)) + hidden_states = ( + self.feed_forward_residual_weights[0] * residual + self.feed_forward_residual_weights[1] * hidden_states + ) + + hidden_states = self.norm_out(hidden_states) + + return hidden_states + + +class LasrPreTrainedModel(ParakeetPreTrainedModel): + # padding is incompatible with flex attention as the resulting mask cannot be used to apply padding + _supports_flex_attn = False + + def _init_weights(self, module): + PreTrainedModel._init_weights(module) + + def _get_subsampling_output_length(self, input_lengths: torch.Tensor): + encoder_config = self.config.encoder_config if isinstance(self.config, LasrCTCConfig) else self.config + kernel_size = encoder_config.subsampling_conv_kernel_size + stride = encoder_config.subsampling_conv_stride + + num_layers = 2 + for _ in range(num_layers): + input_lengths = (input_lengths - kernel_size) // stride + 1 + + return input_lengths + + +@auto_docstring( + custom_intro=""" + The LasrEncoder model, based on the Conformer architecture](https://arxiv.org/abs/2005.08100). + """ +) +class LasrEncoder(LasrPreTrainedModel): + config: LasrEncoderConfig + base_model_prefix = "encoder" + + def __init__(self, config: LasrEncoderConfig): + super().__init__(config) + self.gradient_checkpointing = False + + self.dropout = config.dropout + self.dropout_positions = config.dropout_positions + self.layerdrop = config.layerdrop + + self.subsampler = LasrEncoderSubsampling(config) + self.rotary_emb = LasrEncoderRotaryEmbedding(config) + self.layers = nn.ModuleList( + [LasrEncoderBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.out_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, bias=False) + + self.post_init() + + @auto_docstring + @merge_with_config_defaults + @capture_outputs + @can_return_tuple + def forward( + self, + input_features: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + Example: + + ```python + >>> from transformers import AutoProcessor, LasrEncoder + >>> from datasets import load_dataset, Audio + + >>> model_id = TODO + >>> processor = AutoProcessor.from_pretrained(model_id) + >>> encoder = ParakeetEncoder.from_pretrained(model_id) + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) + + >>> inputs = processor(ds[0]["audio"]["array"]) + >>> encoder_outputs = encoder(**inputs) + + >>> print(encoder_outputs.last_hidden_state.shape) + ``` + """ + + hidden_states = self.subsampler(input_features) + cos, sin = self.rotary_emb( + hidden_states, torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + cos = nn.functional.dropout(cos, p=self.dropout_positions, training=self.training) + sin = nn.functional.dropout(sin, p=self.dropout_positions, training=self.training) + + if attention_mask is not None: + attention_mask = self._get_output_attention_mask(attention_mask, target_length=hidden_states.shape[1]) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + for encoder_layer in self.layers: + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if not to_drop: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=(cos, sin), + **kwargs, + ) + + hidden_states = self.out_norm(hidden_states) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +class LasrForCTC(ParakeetForCTC): + def generate(**super_kwargs): + r""" + Example: + + ```python + >>> from transformers import AutoProcessor, LasrForCTC + >>> from datasets import load_dataset, Audio + + >>> model_id = TODO + >>> processor = AutoProcessor.from_pretrained(model_id) + >>> model = LasrForCTC.from_pretrained(model_id) + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) + + >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"]) + >>> predicted_ids = model.generate(**inputs) + >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) + + >>> print(transcription) + ``` + """ + return super().generate(**super_kwargs) + + +__all__ = [ + "LasrForCTC", + "LasrEncoder", + "LasrPreTrainedModel", + "LasrProcessor", + "LasrEncoderConfig", + "LasrCTCConfig", + "LasrTokenizer", +] diff --git a/third_party/transformers/src/transformers/models/lasr/processing_lasr.py b/third_party/transformers/src/transformers/models/lasr/processing_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..c1acaebaae078d568d8c8a79f2155bb8ee064634 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/processing_lasr.py @@ -0,0 +1,100 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lasr/modular_lasr.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_lasr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team and Google LLC. 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 ...audio_utils import AudioInput, make_list_of_audio +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +class LasrProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "audio_kwargs": { + "sampling_rate": 16000, + "padding": "longest", + "return_attention_mask": True, + }, + "text_kwargs": { + "padding": True, + "padding_side": "right", + "add_special_tokens": False, + }, + "common_kwargs": {"return_tensors": "pt"}, + } + + +@auto_docstring +class LasrProcessor(ProcessorMixin): + def __init__(self, feature_extractor, tokenizer): + super().__init__(feature_extractor, tokenizer) + + @auto_docstring + def __call__( + self, + audio: AudioInput, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + sampling_rate: int | None = None, + **kwargs: Unpack[LasrProcessorKwargs], + ): + r""" + sampling_rate (`int`, *optional*): + The sampling rate of the input audio in Hz. This should match the sampling rate expected by the feature + extractor (defaults to 16000 Hz). If provided, it will be validated against the processor's expected + sampling rate, and an error will be raised if they don't match. If not provided, a warning will be + issued and the default sampling rate will be assumed. + """ + audio = make_list_of_audio(audio) + + output_kwargs = self._merge_kwargs( + LasrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if sampling_rate is None: + logger.warning_once( + f"You've provided audio without specifying the sampling rate. It will be assumed to be {output_kwargs['audio_kwargs']['sampling_rate']}, which can result in silent errors." + ) + elif sampling_rate != output_kwargs["audio_kwargs"]["sampling_rate"]: + raise ValueError( + f"The sampling rate of the audio ({sampling_rate}) does not match the sampling rate of the processor ({output_kwargs['audio_kwargs']['sampling_rate']}). Please provide resampled the audio to the expected sampling rate." + ) + + if audio is not None: + inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"]) + if text is not None: + encodings = self.tokenizer(text, **output_kwargs["text_kwargs"]) + + if text is None: + return inputs + else: + inputs["labels"] = encodings["input_ids"] + return inputs + + @property + def model_input_names(self): + feature_extractor_input_names = self.feature_extractor.model_input_names + return feature_extractor_input_names + ["labels"] + + +__all__ = ["LasrProcessor"] diff --git a/third_party/transformers/src/transformers/models/lasr/tokenization_lasr.py b/third_party/transformers/src/transformers/models/lasr/tokenization_lasr.py new file mode 100644 index 0000000000000000000000000000000000000000..f9f59acb79607e691725628af86b19d5131ecd21 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lasr/tokenization_lasr.py @@ -0,0 +1,185 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lasr/modular_lasr.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_lasr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team and Google LLC. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +import re + +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import Unigram + +from ...tokenization_utils_tokenizers import TokenizersBackend + + +VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"} + + +class LasrTokenizer(TokenizersBackend): + """ + Construct a LASR tokenizer (backed by HuggingFace's *tokenizers* library). Based on + [Unigram](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=unigram#models). + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that + contains the vocabulary necessary to instantiate a tokenizer. + 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`. + + + + 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. + extra_ids (`int`, *optional*, defaults to 100): + Add a number of extra ids added to the vocabulary for use as sentinels. These tokens are accessible as + "" where "{%d}" is a number between 0 and extra_ids-1. These tokens can be retrieved by + calling get_sentinel_tokens method and token ids can be by calling get_sentinel_token_ids method + additional_special_tokens (`list[str]`, *optional*): + Additional special tokens used by the tokenizer. + vocab (`str`, `dict` or `list`, *optional*): + Custom vocabulary dict. If not provided, a minimal vocabulary is created using the special tokens. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = Unigram + + def __init__( + self, + eos_token="", + unk_token="", + pad_token="", + _spm_precompiled_charsmap=None, + extra_ids=100, + additional_special_tokens=None, + vocab=None, + vocab_file=None, + **kwargs, + ): + self._extra_ids = extra_ids + + # Handle extra_ids and additional_special_tokens + if additional_special_tokens is not None: + extra_tokens = [x for x in additional_special_tokens if "" for i in range(extra_ids)] + elif extra_ids > 0 and extra_ids != len(extra_tokens): + raise ValueError( + f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are" + " provided to LasrTokenizer. In this case the additional_special_tokens must include the extra_ids" + " tokens" + ) + else: + extra_tokens = [f"" for i in range(extra_ids)] + additional_special_tokens = extra_tokens + + # LASR vocab structure: =0, =1, =2, then regular vocab, then extra_ids in reverse + if vocab is not None: + self._vocab_scores = vocab + else: + self._vocab_scores = [ + (str(pad_token), 0.0), + (str(eos_token), 0.0), + (str(unk_token), 0.0), + ("▁", -2.0), # Space token + ] + for i in range(extra_ids - 1, -1, -1): + self._vocab_scores.append((f"", 0.0)) + self._tokenizer = Tokenizer( + Unigram( + self._vocab_scores, + unk_id=3, + byte_fallback=False, + ) + ) + + if _spm_precompiled_charsmap is not None: + self._tokenizer.normalizer = normalizers.Precompiled(_spm_precompiled_charsmap) + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.WhitespaceSplit(), + pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True), + ] + ) + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True) + + super().__init__( + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + extra_ids=extra_ids, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.eos_token_id), + ], + ) + + def get_sentinel_tokens(self): + """Get the list of sentinel tokens (extra_id tokens) from additional_special_tokens.""" + return list( + set(filter(lambda x: bool(re.search(r"", x)) is not None, self.additional_special_tokens)) + ) + + def get_sentinel_token_ids(self): + """Get the token IDs for sentinel tokens.""" + return [self.convert_tokens_to_ids(token) for token in self.get_sentinel_tokens()] + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + group_tokens: bool = True, + **kwargs, + ) -> str: + if isinstance(token_ids, int): + token_ids = [token_ids] + if group_tokens: + token_ids = [token_group[0] for token_group in itertools.groupby(token_ids)] + + # for CTC we filter out the blank token, which is the pad token + token_ids = [token for token in token_ids if token != self.pad_token_id] + + return super()._decode( + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + +__all__ = ["LasrTokenizer"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/__init__.py b/third_party/transformers/src/transformers/models/layoutlmv3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f1922f0b4c37b7b568d51ea088bb9015605b52 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/__init__.py @@ -0,0 +1,32 @@ +# 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_layoutlmv3 import * + from .feature_extraction_layoutlmv3 import * + from .image_processing_layoutlmv3 import * + from .image_processing_pil_layoutlmv3 import * + from .modeling_layoutlmv3 import * + from .processing_layoutlmv3 import * + from .tokenization_layoutlmv3 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/layoutlmv3/configuration_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/configuration_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..f2db7d4bc1d8b5881e4d61118b3cfe964de3e935 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/configuration_layoutlmv3.py @@ -0,0 +1,101 @@ +# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""LayoutLMv3 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/layoutlmv3-base") +@strict +class LayoutLMv3Config(PreTrainedConfig): + r""" + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum value that the 2D position embedding might ever be used with. Typically set this to something + large just in case (e.g., 1024). + coordinate_size (`int`, *optional*, defaults to `128`): + Dimension of the coordinate embeddings. + shape_size (`int`, *optional*, defaults to `128`): + Dimension of the width and height embeddings. + has_relative_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a relative attention bias in the self-attention mechanism. + rel_pos_bins (`int`, *optional*, defaults to 32): + The number of relative position bins to be used in the self-attention mechanism. + max_rel_pos (`int`, *optional*, defaults to 128): + The maximum number of relative positions to be used in the self-attention mechanism. + rel_2d_pos_bins (`int`, *optional*, defaults to 64): + The number of 2D relative position bins in the self-attention mechanism. + max_rel_2d_pos (`int`, *optional*, defaults to 256): + The maximum number of relative 2D positions in the self-attention mechanism. + has_spatial_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a spatial attention bias in the self-attention mechanism. + text_embed (`bool`, *optional*, defaults to `True`): + Whether or not to add text embeddings. + visual_embed (`bool`, *optional*, defaults to `True`): + Whether or not to add patch embeddings. + input_size (`int`, *optional*, defaults to `224`): + The size (resolution) of the images. + + Example: + + ```python + >>> from transformers import LayoutLMv3Config, LayoutLMv3Model + + >>> # Initializing a LayoutLMv3 microsoft/layoutlmv3-base style configuration + >>> configuration = LayoutLMv3Config() + + >>> # Initializing a model (with random weights) from the microsoft/layoutlmv3-base style configuration + >>> model = LayoutLMv3Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "layoutlmv3" + + 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 = 2 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + pad_token_id: int | None = 1 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 2 + max_2d_position_embeddings: int = 1024 + coordinate_size: int = 128 + shape_size: int = 128 + has_relative_attention_bias: bool = True + rel_pos_bins: int = 32 + max_rel_pos: int = 128 + rel_2d_pos_bins: int = 64 + max_rel_2d_pos: int = 256 + has_spatial_attention_bias: bool = True + text_embed: bool = True + visual_embed: bool = True + input_size: int = 224 + num_channels: int = 3 + patch_size: int | list[int] | tuple[int, int] = 16 + classifier_dropout: float | int | None = None + + +__all__ = ["LayoutLMv3Config"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..8d7933b3aeed015b553dc1ecbbc608d102f8ec50 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_layoutlmv3.py @@ -0,0 +1,213 @@ +# Copyright 2022 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for LayoutLMv3.""" + +import numpy as np +import torch + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import ChannelDimension, group_images_by_shape, reorder_images, to_pil_image +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, + is_pytesseract_available, + logging, + requires_backends, +) + + +if is_pytesseract_available(): + import pytesseract + +from torchvision.transforms.v2 import functional as tvF + + +logger = logging.get_logger(__name__) + + +def normalize_box(box, width, height): + return [ + int(1000 * (box[0] / width)), + int(1000 * (box[1] / height)), + int(1000 * (box[2] / width)), + int(1000 * (box[3] / height)), + ] + + +def apply_tesseract( + image: "np.ndarray | torch.Tensor", + lang: str | None, + tesseract_config: str | None = None, + input_data_format: str | ChannelDimension | None = None, +): + """Applies Tesseract OCR on a document image, and returns recognized words + normalized bounding boxes.""" + requires_backends(apply_tesseract, ["pytesseract"]) + + # Convert torch tensor to numpy if needed + if hasattr(image, "cpu"): + image = image.cpu().numpy() + elif not isinstance(image, np.ndarray): + image = np.array(image) + + tesseract_config = tesseract_config if tesseract_config is not None else "" + + # apply OCR + pil_image = to_pil_image(image, input_data_format=input_data_format) + image_width, image_height = pil_image.size + data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config) + words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"] + + # filter empty words and corresponding coordinates + irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()] + words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices] + left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices] + top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices] + width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices] + height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices] + + # turn coordinates into (left, top, left+width, top+height) format + actual_boxes = [] + for x, y, w, h in zip(left, top, width, height): + actual_box = [x, y, x + w, y + h] + actual_boxes.append(actual_box) + + # finally, normalize the bounding boxes + normalized_boxes = [] + for box in actual_boxes: + normalized_boxes.append(normalize_box(box, image_width, image_height)) + + assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes" + + return words, normalized_boxes + + +class LayoutLMv3ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + apply_ocr (`bool`, *optional*, defaults to `True`): + Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by + the `apply_ocr` parameter in the `preprocess` method. + ocr_lang (`str`, *optional*): + The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is + used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method. + tesseract_config (`str`, *optional*): + Any additional custom configuration flags that are forwarded to the `config` parameter when calling + Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the + `preprocess` method. + """ + + apply_ocr: bool + ocr_lang: str | None + tesseract_config: str | None + + +@auto_docstring +class LayoutLMv3ImageProcessor(TorchvisionBackend): + valid_kwargs = LayoutLMv3ImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + do_resize = True + do_rescale = True + do_normalize = True + apply_ocr = True + ocr_lang = None + tesseract_config = "" + + def __init__(self, **kwargs: Unpack[LayoutLMv3ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LayoutLMv3ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_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, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + apply_ocr: bool = True, + ocr_lang: str | None = None, + tesseract_config: str | None = None, + **kwargs, + ) -> BatchFeature: + # Tesseract OCR to get words + normalized bounding boxes + if apply_ocr: + requires_backends(self, "pytesseract") + words_batch = [] + boxes_batch = [] + for image in images: + if image.is_cuda: + logger.warning_once( + "apply_ocr can only be performed on cpu. Tensors will be transferred to cpu before processing." + ) + words, boxes = apply_tesseract( + image.cpu(), ocr_lang, tesseract_config, input_data_format=ChannelDimension.FIRST + ) + words_batch.append(words) + boxes_batch.append(boxes) + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, resample=resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + data = BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + if apply_ocr: + data["words"] = words_batch + data["boxes"] = boxes_batch + + return data + + +__all__ = ["LayoutLMv3ImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_pil_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_pil_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..34b75d236c84411ebcbf4b99d112d35cc49c0e4a --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/image_processing_pil_layoutlmv3.py @@ -0,0 +1,195 @@ +# Copyright 2022 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for LayoutLMv3.""" + +from typing import TYPE_CHECKING + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import to_pil_image +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, requires_backends + + +if TYPE_CHECKING: + import torch + +try: + import pytesseract +except ImportError: + pytesseract = None + + +# Adapted from transformers.models.layoutlmv3.image_processing_layoutlmv3.LayoutLMv3ImageProcessorKwargs +class LayoutLMv3ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + apply_ocr (`bool`, *optional*, defaults to `True`): + Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by + the `apply_ocr` parameter in the `preprocess` method. + ocr_lang (`str`, *optional*): + The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is + used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method. + tesseract_config (`str`, *optional*): + Any additional custom configuration flags that are forwarded to the `config` parameter when calling + Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the + `preprocess` method. + """ + + apply_ocr: bool + ocr_lang: str | None + tesseract_config: str | None + + +# Adapted from transformers.models.layoutlmv3.image_processing_layoutlmv3.normalize_box +def normalize_box(box, width, height): + return [ + int(1000 * (box[0] / width)), + int(1000 * (box[1] / height)), + int(1000 * (box[2] / width)), + int(1000 * (box[3] / height)), + ] + + +# Adapted from transformers.models.layoutlmv3.image_processing_layoutlmv3.apply_tesseract +def apply_tesseract( + image: "np.ndarray | torch.Tensor", + lang: str | None, + tesseract_config: str | None = None, + input_data_format: str | ChannelDimension | None = None, +): + """Applies Tesseract OCR on a document image, and returns recognized words + normalized bounding boxes.""" + requires_backends(apply_tesseract, ["pytesseract"]) + + # Convert torch tensor to numpy if needed + if hasattr(image, "cpu"): + image = image.cpu().numpy() + elif not isinstance(image, np.ndarray): + image = np.array(image) + + tesseract_config = tesseract_config if tesseract_config is not None else "" + + # apply OCR + pil_image = to_pil_image(image, input_data_format=input_data_format) + image_width, image_height = pil_image.size + data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config) + words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"] + + # filter empty words and corresponding coordinates + irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()] + words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices] + left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices] + top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices] + width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices] + height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices] + + # turn coordinates into (left, top, left+width, top+height) format + actual_boxes = [] + for x, y, w, h in zip(left, top, width, height): + actual_box = [x, y, x + w, y + h] + actual_boxes.append(actual_box) + + # finally, normalize the bounding boxes + normalized_boxes = [] + for box in actual_boxes: + normalized_boxes.append(normalize_box(box, image_width, image_height)) + + assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes" + + return words, normalized_boxes + + +@auto_docstring +class LayoutLMv3ImageProcessorPil(PilBackend): + valid_kwargs = LayoutLMv3ImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + do_resize = True + do_rescale = True + do_normalize = True + apply_ocr = True + ocr_lang = None + tesseract_config = "" + + def __init__(self, **kwargs: Unpack[LayoutLMv3ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LayoutLMv3ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + return_tensors: str | TensorType | None, + apply_ocr: bool = True, + ocr_lang: str | None = None, + tesseract_config: str | None = None, + **kwargs, + ) -> BatchFeature: + # Tesseract OCR to get words + normalized bounding boxes + if apply_ocr: + requires_backends(self, "pytesseract") + words_batch = [] + boxes_batch = [] + for image in images: + words, boxes = apply_tesseract( + image, ocr_lang, tesseract_config, input_data_format=ChannelDimension.FIRST + ) + words_batch.append(words) + boxes_batch.append(boxes) + + processed_images = [] + for image in images: + if do_resize: + image = self.resize(image=image, size=size, resample=resample) + if do_center_crop: + image = self.center_crop(image=image, crop_size=crop_size) + if do_rescale: + image = self.rescale(image=image, scale=rescale_factor) + if do_normalize: + image = self.normalize(image=image, mean=image_mean, std=image_std) + processed_images.append(image) + + data = BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + if apply_ocr: + data["words"] = words_batch + data["boxes"] = boxes_batch + + return data + + +__all__ = ["LayoutLMv3ImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..74c5eadd4fb9d29875f1ef507ca0d9abae2ab9db --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py @@ -0,0 +1,1159 @@ +# Copyright 2022 Microsoft Research 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 LayoutLMv3 model.""" + +import collections +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +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, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ( + auto_docstring, + can_return_tuple, + logging, + torch_int, +) +from ...utils.generic import TransformersKwargs, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_layoutlmv3 import LayoutLMv3Config + + +logger = logging.get_logger(__name__) + + +class LayoutLMv3PatchEmbeddings(nn.Module): + """LayoutLMv3 image (patch) embeddings. This class also automatically interpolates the position embeddings for varying + image sizes.""" + + def __init__(self, config): + super().__init__() + + image_size = ( + config.input_size + if isinstance(config.input_size, collections.abc.Iterable) + else (config.input_size, config.input_size) + ) + patch_size = ( + config.patch_size + if isinstance(config.patch_size, collections.abc.Iterable) + else (config.patch_size, config.patch_size) + ) + self.patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1]) + self.proj = nn.Conv2d(config.num_channels, config.hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values, position_embedding=None): + embeddings = self.proj(pixel_values) + + if position_embedding is not None: + # interpolate the position embedding to the corresponding size + position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1) + position_embedding = position_embedding.permute(0, 3, 1, 2) + patch_height, patch_width = embeddings.shape[2], embeddings.shape[3] + position_embedding = F.interpolate(position_embedding, size=(patch_height, patch_width), mode="bicubic") + embeddings = embeddings + position_embedding + + embeddings = embeddings.flatten(2).transpose(1, 2) + return embeddings + + +class LayoutLMv3TextEmbeddings(nn.Module): + """ + LayoutLMv3 text embeddings. Same as `RobertaEmbeddings` but with added spatial (layout) embeddings. + """ + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + self.padding_idx = config.pad_token_id + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx + ) + + self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + + def calculate_spatial_position_embeddings(self, bbox): + 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(torch.clip(bbox[:, :, 3] - bbox[:, :, 1], 0, 1023)) + w_position_embeddings = self.w_position_embeddings(torch.clip(bbox[:, :, 2] - bbox[:, :, 0], 0, 1023)) + + # below is the difference between LayoutLMEmbeddingsV2 (torch.cat) and LayoutLMEmbeddingsV1 (add) + spatial_position_embeddings = torch.cat( + [ + left_position_embeddings, + upper_position_embeddings, + right_position_embeddings, + lower_position_embeddings, + h_position_embeddings, + w_position_embeddings, + ], + dim=-1, + ) + return spatial_position_embeddings + + def create_position_ids_from_input_ids(self, input_ids, padding_idx): + """ + 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`. + """ + # 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)) * mask + return incremental_indices.long() + padding_idx + + 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. + """ + 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) + + def forward( + self, + input_ids=None, + bbox=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + ): + if position_ids is None: + if input_ids is not None: + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx).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 = 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 + + spatial_position_embeddings = self.calculate_spatial_position_embeddings(bbox) + + embeddings = embeddings + spatial_position_embeddings + + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +class LayoutLMv3SelfAttention(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.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.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + def cogview_attention(self, attention_scores, alpha=32): + """ + https://huggingface.co/papers/2105.13290 Section 2.4 Stabilization of training: Precision Bottleneck Relaxation + (PB-Relax). A replacement of the original nn.Softmax(dim=-1)(attention_scores). Seems the new attention_probs + will result in a slower speed and a little bias. Can use torch.allclose(standard_attention_probs, + cogview_attention_probs, atol=1e-08) for comparison. The smaller atol (e.g., 1e-08), the better. + """ + scaled_attention_scores = attention_scores / alpha + max_value = scaled_attention_scores.amax(dim=(-1)).unsqueeze(-1) + new_attention_scores = (scaled_attention_scores - max_value) * alpha + return nn.Softmax(dim=-1)(new_attention_scores) + + def forward( + self, + hidden_states, + attention_mask=None, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size = hidden_states.shape[0] + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + + # Take the dot product between "query" and "key" to get the raw attention scores. + # The attention scores QT K/√d could be significantly larger than input elements, and result in overflow. + # Changing the computational order into QT(K/√d) alleviates the problem. (https://huggingface.co/papers/2105.13290) + attention_scores = torch.matmul(query_layer / math.sqrt(self.attention_head_size), key_layer.transpose(-1, -2)) + + if self.has_relative_attention_bias and self.has_spatial_attention_bias: + attention_scores += (rel_pos + rel_2d_pos) / math.sqrt(self.attention_head_size) + elif self.has_relative_attention_bias: + attention_scores += rel_pos / math.sqrt(self.attention_head_size) + + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + # Use the trick of the CogView paper to stabilize training + attention_probs = self.cogview_attention(attention_scores) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer, attention_probs + + +# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfOutput +class LayoutLMv3SelfOutput(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.layoutlmv2.modeling_layoutlmv2.LayoutLMv2Attention with LayoutLMv2->LayoutLMv3 +class LayoutLMv3Attention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = LayoutLMv3SelfAttention(config) + self.output = LayoutLMv3SelfOutput(config) + + def forward( + self, + hidden_states, + attention_mask=None, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + attention_output, _ = self.self( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + attention_output = self.output(attention_output, residual) + return attention_output + + +# Copied from transformers.models.layoutlmv2.modeling_layoutlmv2.LayoutLMv2Layer with LayoutLMv2->LayoutLMv3 +class LayoutLMv3Layer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = LayoutLMv3Attention(config) + self.intermediate = LayoutLMv3Intermediate(config) + self.output = LayoutLMv3Output(config) + + def forward( + self, + hidden_states, + attention_mask=None, + output_attentions=False, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + attention_output = self.attention( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + ) + + layer_output = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output + ) + + return layer_output + + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class LayoutLMv3Encoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([LayoutLMv3Layer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + self.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + if self.has_relative_attention_bias: + self.rel_pos_bins = config.rel_pos_bins + self.max_rel_pos = config.max_rel_pos + self.rel_pos_bias = nn.Linear(self.rel_pos_bins, config.num_attention_heads, bias=False) + + if self.has_spatial_attention_bias: + self.max_rel_2d_pos = config.max_rel_2d_pos + self.rel_2d_pos_bins = config.rel_2d_pos_bins + self.rel_pos_x_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + self.rel_pos_y_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + + def relative_position_bucket(self, relative_position, bidirectional=True, num_buckets=32, max_distance=128): + ret = 0 + if bidirectional: + num_buckets //= 2 + ret += (relative_position > 0).long() * num_buckets + n = torch.abs(relative_position) + else: + n = torch.max(-relative_position, torch.zeros_like(relative_position)) + # now n is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = n < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + val_if_large = max_exact + ( + torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) + ).to(torch.long) + val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1)) + + ret += torch.where(is_small, n, val_if_large) + return ret + + def _cal_1d_pos_emb(self, position_ids): + rel_pos_mat = position_ids.unsqueeze(-2) - position_ids.unsqueeze(-1) + + rel_pos = self.relative_position_bucket( + rel_pos_mat, + num_buckets=self.rel_pos_bins, + max_distance=self.max_rel_pos, + ) + # Since this is a simple indexing operation that is independent of the input, + # no need to track gradients for this operation + # + # Without this no_grad context, training speed slows down significantly + with torch.no_grad(): + rel_pos = self.rel_pos_bias.weight.t()[rel_pos].permute(0, 3, 1, 2) + rel_pos = rel_pos.contiguous() + return rel_pos + + def _cal_2d_pos_emb(self, bbox): + position_coord_x = bbox[:, :, 0] + position_coord_y = bbox[:, :, 3] + rel_pos_x_2d_mat = position_coord_x.unsqueeze(-2) - position_coord_x.unsqueeze(-1) + rel_pos_y_2d_mat = position_coord_y.unsqueeze(-2) - position_coord_y.unsqueeze(-1) + rel_pos_x = self.relative_position_bucket( + rel_pos_x_2d_mat, + num_buckets=self.rel_2d_pos_bins, + max_distance=self.max_rel_2d_pos, + ) + rel_pos_y = self.relative_position_bucket( + rel_pos_y_2d_mat, + num_buckets=self.rel_2d_pos_bins, + max_distance=self.max_rel_2d_pos, + ) + # Since this is a simple indexing operation that is independent of the input, + # no need to track gradients for this operation + # + # Without this no_grad context, training speed slows down significantly + with torch.no_grad(): + rel_pos_x = self.rel_pos_x_bias.weight.t()[rel_pos_x].permute(0, 3, 1, 2) + rel_pos_y = self.rel_pos_y_bias.weight.t()[rel_pos_y].permute(0, 3, 1, 2) + rel_pos_x = rel_pos_x.contiguous() + rel_pos_y = rel_pos_y.contiguous() + rel_2d_pos = rel_pos_x + rel_pos_y + return rel_2d_pos + + def forward( + self, + hidden_states, + bbox=None, + attention_mask=None, + position_ids=None, + patch_height=None, + patch_width=None, + **kwargs: Unpack[TransformersKwargs], + ): + rel_pos = self._cal_1d_pos_emb(position_ids) if self.has_relative_attention_bias else None + rel_2d_pos = self._cal_2d_pos_emb(bbox) if self.has_spatial_attention_bias else None + + for layer_module in self.layer: + hidden_states = layer_module( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +# Copied from transformers.models.roberta.modeling_roberta.RobertaIntermediate +class LayoutLMv3Intermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +# Copied from transformers.models.roberta.modeling_roberta.RobertaOutput +class LayoutLMv3Output(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +@auto_docstring +class LayoutLMv3PreTrainedModel(PreTrainedModel): + config: LayoutLMv3Config + base_model_prefix = "layoutlmv3" + input_modalities = ("image", "text") + _can_record_outputs = {"hidden_states": LayoutLMv3Layer, "attentions": LayoutLMv3SelfAttention} + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, LayoutLMv3Model): + if self.config.visual_embed: + init.zeros_(module.cls_token) + init.zeros_(module.pos_embed) + if hasattr(module, "visual_bbox"): + init.copy_(module.visual_bbox, module.create_visual_bbox(image_size=(module.size, module.size))) + elif isinstance(module, LayoutLMv3TextEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@auto_docstring +class LayoutLMv3Model(LayoutLMv3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.config = config + + if config.text_embed: + self.embeddings = LayoutLMv3TextEmbeddings(config) + + if config.visual_embed: + # use the default pre-training parameters for fine-tuning (e.g., input_size) + # when the input_size is larger in fine-tuning, we will interpolate the position embeddings in forward + self.patch_embed = LayoutLMv3PatchEmbeddings(config) + + self.size = int(config.input_size / config.patch_size) + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.pos_embed = nn.Parameter(torch.zeros(1, self.size * self.size + 1, config.hidden_size)) + self.pos_drop = nn.Dropout(p=0.0) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + if self.config.has_relative_attention_bias or self.config.has_spatial_attention_bias: + self.register_buffer( + "visual_bbox", self.create_visual_bbox(image_size=(self.size, self.size)), persistent=False + ) + + self.norm = nn.LayerNorm(config.hidden_size, eps=1e-6) + + self.encoder = LayoutLMv3Encoder(config) + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def create_visual_bbox(self, image_size=(14, 14), max_len=1000): + """ + Create the bounding boxes for the visual (patch) tokens. + """ + visual_bbox_x = torch.div( + torch.arange(0, max_len * (image_size[1] + 1), max_len), image_size[1], rounding_mode="trunc" + ) + visual_bbox_y = torch.div( + torch.arange(0, max_len * (image_size[0] + 1), max_len), image_size[0], rounding_mode="trunc" + ) + visual_bbox = torch.stack( + [ + visual_bbox_x[:-1].repeat(image_size[0], 1), + visual_bbox_y[:-1].repeat(image_size[1], 1).transpose(0, 1), + visual_bbox_x[1:].repeat(image_size[0], 1), + visual_bbox_y[1:].repeat(image_size[1], 1).transpose(0, 1), + ], + dim=-1, + ).view(-1, 4) + + cls_token_box = torch.tensor([[0 + 1, 0 + 1, max_len - 1, max_len - 1]]) + return torch.cat([cls_token_box, visual_bbox], dim=0) + + def calculate_visual_bbox(self, device, dtype, batch_size): + visual_bbox = self.visual_bbox.repeat(batch_size, 1, 1) + visual_bbox = visual_bbox.to(device).type(dtype) + return visual_bbox + + def forward_image(self, pixel_values): + embeddings = self.patch_embed(pixel_values) + + # add [CLS] token + batch_size, seq_len, _ = embeddings.size() + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add position embeddings + if self.pos_embed is not None: + embeddings = embeddings + self.pos_embed + + embeddings = self.pos_drop(embeddings) + embeddings = self.norm(embeddings) + + return embeddings + + @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, + pixel_values: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, token_sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + bbox (`torch.LongTensor` of shape `(batch_size, token_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. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + token_type_ids (`torch.LongTensor` of shape `(batch_size, token_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. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `(batch_size, token_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]`. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + + [What are position IDs?](../glossary#position-ids) + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, token_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. + + Examples: + + ```python + >>> from transformers import AutoProcessor, AutoModel + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False) + >>> model = AutoModel.from_pretrained("microsoft/layoutlmv3-base") + + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + + >>> encoding = processor(image, words, boxes=boxes, return_tensors="pt") + + >>> outputs = model(**encoding) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + if input_ids is not None: + input_shape = input_ids.size() + batch_size, seq_length = input_shape + device = input_ids.device + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + batch_size, seq_length = input_shape + device = inputs_embeds.device + elif pixel_values is not None: + batch_size = len(pixel_values) + device = pixel_values.device + else: + raise ValueError("You have to specify either input_ids or inputs_embeds or pixel_values") + + if input_ids is not None or inputs_embeds is not None: + 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) + if bbox is None: + bbox = torch.zeros(tuple(list(input_shape) + [4]), dtype=torch.long, device=device) + + embedding_output = self.embeddings( + input_ids=input_ids, + bbox=bbox, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + + final_bbox = final_position_ids = None + patch_height = patch_width = None + if pixel_values is not None: + patch_height, patch_width = ( + torch_int(pixel_values.shape[2] / self.config.patch_size), + torch_int(pixel_values.shape[3] / self.config.patch_size), + ) + visual_embeddings = self.forward_image(pixel_values) + visual_attention_mask = torch.ones( + (batch_size, visual_embeddings.shape[1]), dtype=torch.long, device=device + ) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask, visual_attention_mask], dim=1) + else: + attention_mask = visual_attention_mask + + if self.config.has_relative_attention_bias or self.config.has_spatial_attention_bias: + if self.config.has_spatial_attention_bias: + visual_bbox = self.calculate_visual_bbox(device, dtype=torch.long, batch_size=batch_size) + if bbox is not None: + final_bbox = torch.cat([bbox, visual_bbox], dim=1) + else: + final_bbox = visual_bbox + + visual_position_ids = torch.arange( + 0, visual_embeddings.shape[1], dtype=torch.long, device=device + ).repeat(batch_size, 1) + if input_ids is not None or inputs_embeds is not None: + position_ids = torch.arange(0, input_shape[1], device=device).unsqueeze(0) + position_ids = position_ids.expand(input_shape) + final_position_ids = torch.cat([position_ids, visual_position_ids], dim=1) + else: + final_position_ids = visual_position_ids + + if input_ids is not None or inputs_embeds is not None: + embedding_output = torch.cat([embedding_output, visual_embeddings], dim=1) + else: + embedding_output = visual_embeddings + + embedding_output = self.LayerNorm(embedding_output) + embedding_output = self.dropout(embedding_output) + elif self.config.has_relative_attention_bias or self.config.has_spatial_attention_bias: + if self.config.has_spatial_attention_bias: + final_bbox = bbox + if self.config.has_relative_attention_bias: + position_ids = self.embeddings.position_ids[:, : input_shape[1]] + position_ids = position_ids.expand_as(input_ids) + final_position_ids = position_ids + + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( + attention_mask, None, dtype=embedding_output.dtype + ) + + encoder_outputs = self.encoder( + embedding_output, + bbox=final_bbox, + position_ids=final_position_ids, + attention_mask=extended_attention_mask, + patch_height=patch_height, + patch_width=patch_width, + **kwargs, + ) + + sequence_output = encoder_outputs.last_hidden_state + + return BaseModelOutput( + last_hidden_state=sequence_output, + ) + + +class LayoutLMv3ClassificationHead(nn.Module): + """ + Head for sentence-level classification tasks. Reference: RobertaClassificationHead + """ + + def __init__(self, config, pool_feature=False): + super().__init__() + self.pool_feature = pool_feature + if pool_feature: + self.dense = nn.Linear(config.hidden_size * 3, config.hidden_size) + else: + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + classifier_dropout = ( + config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.out_proj = nn.Linear(config.hidden_size, config.num_labels) + + def forward(self, x): + x = self.dropout(x) + x = self.dense(x) + x = torch.tanh(x) + x = self.dropout(x) + x = self.out_proj(x) + return x + + +@auto_docstring( + custom_intro=""" + LayoutLMv3 Model with a token classification head on top (a linear layer on top of the final hidden states) e.g. + for sequence labeling (information extraction) tasks such as [FUNSD](https://guillaumejaume.github.io/FUNSD/), + [SROIE](https://rrc.cvc.uab.es/?ch=13), [CORD](https://github.com/clovaai/cord) and + [Kleister-NDA](https://github.com/applicaai/kleister-nda). + """ +) +class LayoutLMv3ForTokenClassification(LayoutLMv3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.layoutlmv3 = LayoutLMv3Model(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + if config.num_labels < 10: + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + else: + self.classifier = LayoutLMv3ClassificationHead(config, pool_feature=False) + + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv3.get_input_embeddings() + + def set_input_embeddings(self, value): + self.layoutlmv3.set_input_embeddings(value) + + @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, + pixel_values: 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. + 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 AutoProcessor, AutoModelForTokenClassification + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False) + >>> model = AutoModelForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=7) + + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + >>> word_labels = example["ner_tags"] + + >>> encoding = processor(image, words, boxes=boxes, word_labels=word_labels, return_tensors="pt") + + >>> outputs = model(**encoding) + >>> loss = outputs.loss + >>> logits = outputs.logits + ```""" + outputs = self.layoutlmv3( + input_ids, + bbox=bbox, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + **kwargs, + ) + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + # only take the text part of the output representations + sequence_output = outputs[0][:, :seq_length] + 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 LayoutLMv3ForQuestionAnswering(LayoutLMv3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.layoutlmv3 = LayoutLMv3Model(config) + self.qa_outputs = LayoutLMv3ClassificationHead(config, pool_feature=False) + + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv3.get_input_embeddings() + + def set_input_embeddings(self, value): + self.layoutlmv3.set_input_embeddings(value) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + pixel_values: 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. + + Examples: + + ```python + >>> from transformers import AutoProcessor, AutoModelForQuestionAnswering + >>> from datasets import load_dataset + >>> import torch + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False) + >>> model = AutoModelForQuestionAnswering.from_pretrained("microsoft/layoutlmv3-base") + + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> question = "what's his name?" + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + + >>> encoding = processor(image, question, words, boxes=boxes, return_tensors="pt") + >>> start_positions = torch.tensor([1]) + >>> end_positions = torch.tensor([3]) + + >>> outputs = model(**encoding, start_positions=start_positions, end_positions=end_positions) + >>> loss = outputs.loss + >>> start_scores = outputs.start_logits + >>> end_scores = outputs.end_logits + ```""" + outputs: BaseModelOutput = self.layoutlmv3( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + bbox=bbox, + pixel_values=pixel_values, + **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, + ) + + +@auto_docstring( + custom_intro=""" + LayoutLMv3 Model with a sequence classification head on top (a linear layer on top of the final hidden state of the + [CLS] token) e.g. for document image classification tasks such as the + [RVL-CDIP](https://www.cs.cmu.edu/~aharley/rvl-cdip/) dataset. + """ +) +class LayoutLMv3ForSequenceClassification(LayoutLMv3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + self.layoutlmv3 = LayoutLMv3Model(config) + self.classifier = LayoutLMv3ClassificationHead(config, pool_feature=False) + + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv3.get_input_embeddings() + + def set_input_embeddings(self, value): + self.layoutlmv3.set_input_embeddings(value) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + pixel_values: 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. + + Examples: + + ```python + >>> from transformers import AutoProcessor, AutoModelForSequenceClassification + >>> from datasets import load_dataset + >>> import torch + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False) + >>> model = AutoModelForSequenceClassification.from_pretrained("microsoft/layoutlmv3-base") + + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + + >>> encoding = processor(image, words, boxes=boxes, return_tensors="pt") + >>> sequence_label = torch.tensor([1]) + + >>> outputs = model(**encoding, labels=sequence_label) + >>> loss = outputs.loss + >>> logits = outputs.logits + ```""" + outputs: BaseModelOutput = self.layoutlmv3( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + bbox=bbox, + pixel_values=pixel_values, + **kwargs, + ) + + sequence_output = outputs[0][:, 0, :] + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "LayoutLMv3ForQuestionAnswering", + "LayoutLMv3ForSequenceClassification", + "LayoutLMv3ForTokenClassification", + "LayoutLMv3Model", + "LayoutLMv3PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/processing_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/processing_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..f9433787175155c51756d4d13f14ed732c998eae --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/processing_layoutlmv3.py @@ -0,0 +1,120 @@ +# 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 LayoutLMv3. +""" + +from ...processing_utils import ProcessorMixin +from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy +from ...utils import TensorType, auto_docstring + + +@auto_docstring +class LayoutLMv3Processor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, **kwargs): + super().__init__(image_processor, tokenizer) + + @auto_docstring + def __call__( + self, + images, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchEncoding: + if self.image_processor.apply_ocr and (boxes is not None): + raise ValueError( + "You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True." + ) + + if self.image_processor.apply_ocr and (word_labels is not None): + raise ValueError( + "You cannot provide word labels if you initialized the image processor with apply_ocr set to True." + ) + + # first, apply the image processor + features = self.image_processor(images=images, return_tensors=return_tensors) + + # second, apply the tokenizer + if text is not None and self.image_processor.apply_ocr and text_pair is None: + if isinstance(text, str): + text = [text] # add batch dimension (as the image processor always adds a batch dimension) + text_pair = features["words"] + + encoded_inputs = self.tokenizer( + text=text if text is not None else features["words"], + text_pair=text_pair if text_pair is not None else None, + boxes=boxes if boxes is not None else features["boxes"], + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + return_tensors=return_tensors, + **kwargs, + ) + + # add pixel values + images = features.pop("pixel_values") + if return_overflowing_tokens is True: + images = self.get_overflowing_images(images, encoded_inputs["overflow_to_sample_mapping"]) + encoded_inputs["pixel_values"] = images + + return encoded_inputs + + def get_overflowing_images(self, images, overflow_to_sample_mapping): + # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image + images_with_overflow = [] + for sample_idx in overflow_to_sample_mapping: + images_with_overflow.append(images[sample_idx]) + + if len(images_with_overflow) != len(overflow_to_sample_mapping): + raise ValueError( + "Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got" + f" {len(images_with_overflow)} and {len(overflow_to_sample_mapping)}" + ) + + return images_with_overflow + + @property + def model_input_names(self): + return ["input_ids", "bbox", "attention_mask", "pixel_values"] + + +__all__ = ["LayoutLMv3Processor"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py b/third_party/transformers/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py new file mode 100644 index 0000000000000000000000000000000000000000..cda7c0b8f32439ab8a082303e0229672f9611770 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py @@ -0,0 +1,862 @@ +# Copyright The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization class for LayoutLMv3. Same as LayoutLMv2, but RoBERTa-like BPE tokenization instead of WordPiece.""" + +from tokenizers import Tokenizer, decoders, models, pre_tokenizers, processors + +from ...tokenization_utils_base import ( + BatchEncoding, + EncodedInput, + PaddingStrategy, + PreTokenizedInput, + TensorType, + TextInput, + TextInputPair, + TruncationStrategy, +) +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import add_end_docstrings, logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", + "tokenizer_file": "tokenizer.json", +} + +# Docstring constants for encode methods +LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to encode the sequences with the special tokens relative to their model. + padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + stride (`int`, *optional*, defaults to 0): + If set to a number along with `max_length`, the overflowing tokens returned when + `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence + returned to provide some overlap between truncated and overflowing sequences. The value of this + argument defines the number of overlapping tokens. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). +""" + +LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method +""" + + +class LayoutLMv3Tokenizer(TokenizersBackend): + r""" + Construct a LayoutLMv3 tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level BPE. + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + errors (`str`, *optional*, defaults to `"replace"`): + Paradigm to follow when decoding bytes to UTF-8. See + [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information. + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + add_prefix_space (`bool`, *optional*, defaults to `True`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. + cls_token_box (`list[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [CLS] token. + sep_token_box (`list[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [SEP] token. + pad_token_box (`list[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [PAD] token. + pad_token_label (`int`, *optional*, defaults to -100): + The label to use for padding tokens. Defaults to -100, which is the `ignore_index` of PyTorch's + CrossEntropyLoss. + only_label_first_subword (`bool`, *optional*, defaults to `True`): + Whether or not to only label the first subword, in case word labels are provided. + vocab (`str` or `dict[str, int]`, *optional*): + Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file` when using + `from_pretrained`. + merges (`str` or `list[str]`, *optional*): + Custom merges list. If not provided, merges are loaded from `merges_file` when using `from_pretrained`. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask", "bbox"] + model = models.BPE + + def __init__( + self, + errors="replace", + bos_token="", + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + add_prefix_space=True, + cls_token_box=[0, 0, 0, 0], + sep_token_box=[0, 0, 0, 0], + pad_token_box=[0, 0, 0, 0], + pad_token_label=-100, + only_label_first_subword=True, + vocab: str | dict[str, int] | None = None, + merges: str | list[str] | None = None, + **kwargs, + ): + self.add_prefix_space = add_prefix_space + self._vocab = vocab or {} + self._merges = merges or [] + self._tokenizer = Tokenizer( + models.BPE( + vocab=self._vocab, + merges=self._merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space) + self._tokenizer.decoder = decoders.ByteLevel() + super().__init__( + errors=errors, + bos_token=bos_token, + eos_token=eos_token, + sep_token=sep_token, + cls_token=cls_token, + unk_token=unk_token, + pad_token=pad_token, + mask_token=mask_token, + add_prefix_space=add_prefix_space, + cls_token_box=cls_token_box, + sep_token_box=sep_token_box, + pad_token_box=pad_token_box, + pad_token_label=pad_token_label, + only_label_first_subword=only_label_first_subword, + **kwargs, + ) + + # Now set post_processor with actual token IDs (RoBERTa-style) + cls = str(self.cls_token) + sep = str(self.sep_token) + cls_token_id = self.cls_token_id + sep_token_id = self.sep_token_id + + self._tokenizer.post_processor = processors.RobertaProcessing( + sep=(sep, sep_token_id), + cls=(cls, cls_token_id), + add_prefix_space=add_prefix_space, + trim_offsets=True, + ) + self.cls_token_box = cls_token_box + self.sep_token_box = sep_token_box + self.pad_token_box = pad_token_box + self.pad_token_label = pad_token_label + self.only_label_first_subword = only_label_first_subword + + @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences with word-level normalized bounding boxes and optional labels. + + Args: + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string, a list of strings + (words of a single example or questions of a batch of examples) or a list of list of strings (batch of + words). + text_pair (`List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence should be a list of strings + (pretokenized string). + boxes (`List[List[int]]`, `List[List[List[int]]]`): + Word-level bounding boxes. Each bounding box should be normalized to be on a 0-1000 scale. + word_labels (`List[int]`, `List[List[int]]`, *optional*): + Word-level integer labels (for token classification tasks such as FUNSD, CORD). + """ + + # Input type checking for clearer error + def _is_valid_text_input(t): + if isinstance(t, str): + # Strings are fine + return True + elif isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... empty + return True + elif isinstance(t[0], str): + # ... list of strings + return True + elif isinstance(t[0], (list, tuple)): + # ... list with an empty list or with a list of strings + return len(t[0]) == 0 or isinstance(t[0][0], str) + else: + return False + else: + return False + + if text_pair is not None: + # in case text + text_pair are provided, text = questions, text_pair = words + if not _is_valid_text_input(text): + raise ValueError("text input must of type `str` (single example) or `List[str]` (batch of examples). ") + if not isinstance(text_pair, (list, tuple)): + raise ValueError( + "Words must be of type `List[str]` (single pretokenized example), " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + else: + # in case only text is provided => must be words + if not isinstance(text, (list, tuple)): + raise ValueError( + "Words must be of type `List[str]` (single pretokenized example), " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if text_pair is not None: + is_batched = isinstance(text, (list, tuple)) + else: + is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)) + + words = text if text_pair is None else text_pair + if boxes is None: + raise ValueError("You must provide corresponding bounding boxes") + if is_batched: + if len(words) != len(boxes): + raise ValueError("You must provide words and boxes for an equal amount of examples") + for words_example, boxes_example in zip(words, boxes): + if len(words_example) != len(boxes_example): + raise ValueError("You must provide as many words as there are bounding boxes") + else: + if len(words) != len(boxes): + raise ValueError("You must provide as many words as there are bounding boxes") + + if is_batched: + if text_pair is not None and len(text) != len(text_pair): + raise ValueError( + f"batch length of `text`: {len(text)} does not match batch length of `text_pair`:" + f" {len(text_pair)}." + ) + batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text + is_pair = bool(text_pair is not None) + return self.batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + else: + return self.encode_plus( + text=text, + text_pair=text_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def batch_encode_plus( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def tokenize(self, text: str, pair: str | None = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + batched_input = [(text, pair)] if pair else [text] + encodings = self._tokenizer.encode_batch( + batched_input, add_special_tokens=add_special_tokens, is_pretokenized=False, **kwargs + ) + + return encodings[0].tokens if encodings else [] + + @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def encode_plus( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[int] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a sequence or a pair of sequences. .. warning:: This method is deprecated, + `__call__` should be used instead. + + Args: + text (`str`, `List[str]`, `List[List[str]]`): + The first sequence to be encoded. This can be a string, a list of strings or a list of list of strings. + text_pair (`List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a list of strings (words of a single example) or a + list of list of strings (words of a batch of examples). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._encode_plus( + text=text, + boxes=boxes, + text_pair=text_pair, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + ) -> BatchEncoding: + if not isinstance(batch_text_or_text_pairs, list): + raise TypeError(f"batch_text_or_text_pairs has to be a list (got {type(batch_text_or_text_pairs)})") + + # Set the truncation and padding strategy and restore the initial configuration + self.set_truncation_and_padding( + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + ) + + if is_pair: + batch_text_or_text_pairs = [ + (text.split() if isinstance(text, str) else text, text_pair) + for text, text_pair in batch_text_or_text_pairs + ] + + encodings = self._tokenizer.encode_batch( + batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + is_pretokenized=True, # we set this to True as LayoutLMv3 always expects pretokenized inputs + ) + + # Convert encoding to dict + tokens_and_encodings = [ + self._convert_encoding( + encoding=encoding, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=True + if word_labels is not None + else return_offsets_mapping, # we use offsets to create the labels + return_length=return_length, + verbose=verbose, + ) + for encoding in encodings + ] + + # Convert the output to have dict[list] from list[dict] and remove the additional overflows dimension + sanitized_tokens = {} + for key in tokens_and_encodings[0][0]: + stack = [e for item, _ in tokens_and_encodings for e in item[key]] + sanitized_tokens[key] = stack + sanitized_encodings = [e for _, item in tokens_and_encodings for e in item] + + # If returning overflowing tokens, we need to return a mapping + # from the batch idx to the original sample + if return_overflowing_tokens: + overflow_to_sample_mapping = [] + for i, (toks, _) in enumerate(tokens_and_encodings): + overflow_to_sample_mapping += [i] * len(toks["input_ids"]) + sanitized_tokens["overflow_to_sample_mapping"] = overflow_to_sample_mapping + + for input_ids in sanitized_tokens["input_ids"]: + self._eventual_warn_about_too_long_sequence(input_ids, max_length, verbose) + + # create the token boxes + token_boxes = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + token_boxes_example = [] + for id, sequence_id, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_encodings[batch_index].sequence_ids, + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if is_pair and sequence_id == 0: + token_boxes_example.append(self.pad_token_box) + else: + token_boxes_example.append(boxes[original_index][word_id]) + else: + if id == self.cls_token_id: + token_boxes_example.append(self.cls_token_box) + elif id == self.sep_token_id: + token_boxes_example.append(self.sep_token_box) + elif id == self.pad_token_id: + token_boxes_example.append(self.pad_token_box) + else: + # For RoBERTa, there might be additional sep tokens + token_boxes_example.append(self.sep_token_box) + token_boxes.append(token_boxes_example) + + sanitized_tokens["bbox"] = token_boxes + + # optionally, create the labels + if word_labels is not None: + labels = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + labels_example = [] + previous_word_id = None + for id, offset, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_tokens["offset_mapping"][batch_index], + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if self.only_label_first_subword: + # Check if this is the first token of the word (word_id changed or is first occurrence) + if word_id != previous_word_id: + # Use the real label id for the first token of the word, and padding ids for the remaining tokens + labels_example.append(word_labels[original_index][word_id]) + else: + labels_example.append(self.pad_token_label) + else: + labels_example.append(word_labels[original_index][word_id]) + previous_word_id = word_id + else: + labels_example.append(self.pad_token_label) + previous_word_id = None + labels.append(labels_example) + + sanitized_tokens["labels"] = labels + # finally, remove offsets if the user didn't want them + if not return_offsets_mapping: + del sanitized_tokens["offset_mapping"] + + return BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors) + + def _encode_plus( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[int] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: bool | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # make it a batched input + # 2 options: + # 1) only text, in case text must be a list of str + # 2) text + text_pair, in which case text = str and text_pair a list of str + batched_input = [(text, text_pair)] if text_pair else [text] + batched_boxes = [boxes] + batched_word_labels = [word_labels] if word_labels is not None else None + batched_output = self._batch_encode_plus( + batched_input, + is_pair=bool(text_pair is not None), + boxes=batched_boxes, + word_labels=batched_word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + # Return tensor is None, then we can remove the leading batch axis + # Overflowing tokens are returned as a batch of output so we keep them in this case + if return_tensors is None and not return_overflowing_tokens: + batched_output = BatchEncoding( + { + key: value[0] if len(value) > 0 and isinstance(value[0], list) else value + for key, value in batched_output.items() + }, + batched_output.encodings, + ) + + self._eventual_warn_about_too_long_sequence(batched_output["input_ids"], max_length, verbose) + + return batched_output + + def _pad( + self, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + ) -> dict: + """ + Pad encoded inputs (on left/right and up to predefined length or max length in the batch) + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + padding_strategy: PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in self.padding_side: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + required_input = encoded_inputs[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(required_input) + + if needs_to_be_padded: + difference = max_length - len(required_input) + padding_side = padding_side if padding_side is not None else self.padding_side + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = encoded_inputs["attention_mask"] + [0] * difference + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = ( + encoded_inputs["token_type_ids"] + [self.pad_token_type_id] * difference + ) + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = encoded_inputs["bbox"] + [self.pad_token_box] * difference + if "labels" in encoded_inputs: + encoded_inputs["labels"] = encoded_inputs["labels"] + [self.pad_token_label] * difference + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference + encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[ + "token_type_ids" + ] + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = [self.pad_token_box] * difference + encoded_inputs["bbox"] + if "labels" in encoded_inputs: + encoded_inputs["labels"] = [self.pad_token_label] * difference + encoded_inputs["labels"] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input + else: + raise ValueError("Invalid padding strategy:" + str(padding_side)) + + return encoded_inputs + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. A RoBERTa 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 + + +__all__ = ["LayoutLMv3Tokenizer", "LayoutLMv3TokenizerFast"] + +# Backward alias +LayoutLMv3TokenizerFast = LayoutLMv3Tokenizer diff --git a/third_party/transformers/src/transformers/models/lfm2/__init__.py b/third_party/transformers/src/transformers/models/lfm2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..239ab879832c57465eb85f910cadd0f373abe46a --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_lfm2 import * + from .modeling_lfm2 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/configuration_lfm2.py b/third_party/transformers/src/transformers/models/lfm2/configuration_lfm2.py new file mode 100644 index 0000000000000000000000000000000000000000..f58ce6270d314766d78e671cf8a715dedf38a6c6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2/configuration_lfm2.py @@ -0,0 +1,95 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="LiquidAI/LFM2-1.2B") +@strict +class Lfm2Config(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. + block_multiple_of (`int`, *optional*, defaults to 256): + Multiple for the `intermediate_size`. + block_ffn_dim_multiplier (`float`, *optional*, defaults to 1.0): + Multiplier for the `intermediate_size`. + block_auto_adjust_ff_dim (`bool`, *optional*, defaults to `True`): + Whether to adjust the dim of the `intermediate_size`. + full_attn_idxs (`Optional`, *optional*): + Index of the layers which use attention. + + ```python + >>> from transformers import Lfm2Model, Lfm2Config + + >>> # Initializing a LFM2 model + >>> configuration = Lfm2Config() + + >>> # Initializing a model from the LFM2-1.2B style configuration + >>> model = Lfm2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lfm2" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 1000000.0 + + vocab_size: int = 65536 + hidden_size: int = 2560 + intermediate_size: int = 12288 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + max_position_embeddings: int = 128_000 + initializer_range: float = 0.02 + norm_eps: float = 0.00001 + use_cache: bool = True + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + conv_bias: bool = False + conv_L_cache: int = 3 + block_multiple_of: int = 256 + block_ffn_dim_multiplier: float | int = 1.0 + block_auto_adjust_ff_dim: bool = True + full_attn_idxs: list[int] | None = None + layer_types: list[str] | None = None + + def __post_init__(self, **kwargs): + if self.layer_types is None: + self.full_attn_idxs = ( + self.full_attn_idxs if self.full_attn_idxs is not None else list(range(self.num_hidden_layers)) + ) + self.layer_types = [ + "full_attention" if i in self.full_attn_idxs else "conv" for i in range(self.num_hidden_layers) + ] + + self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings) + self.intermediate_size = kwargs.pop("block_ff_dim", self.intermediate_size) + super().__post_init__(**kwargs) + + +__all__ = ["Lfm2Config"] diff --git a/third_party/transformers/src/transformers/models/lfm2/modeling_lfm2.py b/third_party/transformers/src/transformers/models/lfm2/modeling_lfm2.py new file mode 100644 index 0000000000000000000000000000000000000000..ef753e3b2893a384347745bdd6fe5be92bcde78b --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2/modeling_lfm2.py @@ -0,0 +1,614 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lfm2/modular_lfm2.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.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 ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, 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.import_utils import is_causal_conv1d_available, is_torchdynamo_compiling +from ...utils.output_capturing import capture_outputs +from .configuration_lfm2 import Lfm2Config + + +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 Lfm2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Lfm2RMSNorm 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 Lfm2RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Lfm2Config, 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: Lfm2Config | 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 Lfm2MLP(nn.Module): + def __init__(self, config: Lfm2Config): + super().__init__() + intermediate_size = config.intermediate_size + if config.block_auto_adjust_ff_dim: + intermediate_size = int(2 * intermediate_size / 3) + # custom dim factor multiplier + if config.block_ffn_dim_multiplier is not None: + intermediate_size = int(config.block_ffn_dim_multiplier * intermediate_size) + intermediate_size = config.block_multiple_of * ( + (intermediate_size + config.block_multiple_of - 1) // config.block_multiple_of + ) + self.w1 = nn.Linear(config.hidden_size, intermediate_size, bias=False) + self.w3 = nn.Linear(config.hidden_size, intermediate_size, bias=False) + self.w2 = nn.Linear(intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +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 Lfm2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Lfm2Config, 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 = Lfm2RMSNorm(self.head_dim, eps=config.norm_eps) + self.k_layernorm = Lfm2RMSNorm(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 Lfm2ShortConv(nn.Module): + def __init__( + self, + config: Lfm2Config, + 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 Lfm2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Lfm2Config, layer_idx: int): + super().__init__() + self.is_attention_layer = config.layer_types[layer_idx] == "full_attention" + + if self.is_attention_layer: + self.self_attn = Lfm2Attention(config, layer_idx) + else: + self.conv = Lfm2ShortConv(config, layer_idx) + self.feed_forward = Lfm2MLP(config) + self.operator_norm = Lfm2RMSNorm(config.hidden_size, eps=config.norm_eps) + self.ffn_norm = Lfm2RMSNorm(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 Lfm2PreTrainedModel(PreTrainedModel): + config: Lfm2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Lfm2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Lfm2DecoderLayer, + "attentions": Lfm2Attention, + } + + +@auto_docstring +class Lfm2Model(Lfm2PreTrainedModel): + def __init__(self, config: Lfm2Config): + 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( + [Lfm2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = Lfm2RotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.embedding_norm = Lfm2RMSNorm(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], + ) -> 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) + + 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.rotary_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_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.embedding_norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Lfm2ForCausalLM(Lfm2PreTrainedModel, 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 = Lfm2Model(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, Lfm2ForCausalLM + + >>> model = Lfm2ForCausalLM.from_pretrained("meta-lfm2/Lfm2-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-lfm2/Lfm2-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__ = ["Lfm2ForCausalLM", "Lfm2Model", "Lfm2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/lfm2/modular_lfm2.py b/third_party/transformers/src/transformers/models/lfm2/modular_lfm2.py new file mode 100644 index 0000000000000000000000000000000000000000..c08b61956081d1be794b35996a763caa46d21b94 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2/modular_lfm2.py @@ -0,0 +1,350 @@ +# 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 + +import torch +import torch.nn.functional as F +from torch import nn + +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, logging +from ...utils.import_utils import is_causal_conv1d_available, is_torchdynamo_compiling +from ..bamba.modeling_bamba import apply_mask_to_padding_states +from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaForCausalLM, + LlamaModel, + LlamaPreTrainedModel, + LlamaRMSNorm, + apply_rotary_pos_emb, + eager_attention_forward, +) +from .configuration_lfm2 import Lfm2Config + + +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 Lfm2RMSNorm(LlamaRMSNorm): + pass + + +class Lfm2RotaryEmbedding(Gemma2RotaryEmbedding): + pass + + +class Lfm2MLP(nn.Module): + def __init__(self, config: Lfm2Config): + super().__init__() + intermediate_size = config.intermediate_size + if config.block_auto_adjust_ff_dim: + intermediate_size = int(2 * intermediate_size / 3) + # custom dim factor multiplier + if config.block_ffn_dim_multiplier is not None: + intermediate_size = int(config.block_ffn_dim_multiplier * intermediate_size) + intermediate_size = config.block_multiple_of * ( + (intermediate_size + config.block_multiple_of - 1) // config.block_multiple_of + ) + self.w1 = nn.Linear(config.hidden_size, intermediate_size, bias=False) + self.w3 = nn.Linear(config.hidden_size, intermediate_size, bias=False) + self.w2 = nn.Linear(intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class Lfm2Attention(LlamaAttention): + def __init__(self, config: Lfm2Config, 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.out_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + self.q_layernorm = Lfm2RMSNorm(self.head_dim, eps=config.norm_eps) + self.k_layernorm = Lfm2RMSNorm(self.head_dim, eps=config.norm_eps) + del self.o_proj + del self.attention_dropout + + 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 + + +class Lfm2ShortConv(nn.Module): + def __init__( + self, + config: Lfm2Config, + 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 Lfm2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Lfm2Config, layer_idx: int): + super().__init__() + self.is_attention_layer = config.layer_types[layer_idx] == "full_attention" + + if self.is_attention_layer: + self.self_attn = Lfm2Attention(config, layer_idx) + else: + self.conv = Lfm2ShortConv(config, layer_idx) + self.feed_forward = Lfm2MLP(config) + self.operator_norm = Lfm2RMSNorm(config.hidden_size, eps=config.norm_eps) + self.ffn_norm = Lfm2RMSNorm(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 + + +class Lfm2PreTrainedModel(LlamaPreTrainedModel): + _can_compile_fullgraph = False + + +class Lfm2Model(LlamaModel): + def __init__(self, config: Lfm2Config): + super().__init__(config) + self.embedding_norm = Lfm2RMSNorm(config.hidden_size, eps=config.norm_eps) + del self.norm + + 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) + + 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.rotary_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_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.embedding_norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class Lfm2ForCausalLM(LlamaForCausalLM): + pass + + +__all__ = ["Lfm2ForCausalLM", "Lfm2Model", "Lfm2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/lighton_ocr/__init__.py b/third_party/transformers/src/transformers/models/lighton_ocr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a2091da8c18e450a273122951a99897033d0b5eb --- /dev/null +++ b/third_party/transformers/src/transformers/models/lighton_ocr/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 The LightOn 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_lighton_ocr import * + from .modeling_lighton_ocr import * + from .processing_lighton_ocr import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/lighton_ocr/configuration_lighton_ocr.py b/third_party/transformers/src/transformers/models/lighton_ocr/configuration_lighton_ocr.py new file mode 100644 index 0000000000000000000000000000000000000000..9794e01b8c112336f456e11716e79e0c10860abb --- /dev/null +++ b/third_party/transformers/src/transformers/models/lighton_ocr/configuration_lighton_ocr.py @@ -0,0 +1,103 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lighton_ocr/modular_lighton_ocr.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_lighton_ocr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The LightOn 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, PretrainedConfig +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig + + +@auto_docstring(checkpoint="lightonai/LightOnOCR-1B-1025") +@strict +class LightOnOcrConfig(PretrainedConfig): + r""" + Example: + + ```python + >>> from transformers import LightOnOcrConfig, LightOnOcrForConditionalGeneration + + >>> # Initializing a LightOnOcr configuration + >>> configuration = LightOnOcrConfig() + + >>> # Initializing a model from the configuration + >>> model = LightOnOcrForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lighton_ocr" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + spatial_merge_size: int = 2 + image_token_id: int = 151655 + tie_word_embeddings: bool = True + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = CONFIG_MAPPING["pixtral"]( + attention_dropout=0.0, + head_dim=64, + hidden_act="silu", + hidden_size=1024, + image_size=1540, + initializer_range=0.02, + intermediate_size=4096, + model_type="pixtral", + num_attention_heads=16, + num_channels=3, + num_hidden_layers=24, + patch_size=14, + rope_theta=10000, + ) + elif isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "pixtral") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config) + + if self.text_config is None: + self.text_config = CONFIG_MAPPING["qwen3"]( + attention_dropout=0.0, + head_dim=128, + hidden_act="silu", + hidden_size=1024, + initializer_range=0.02, + intermediate_size=3072, + max_position_embeddings=40960, + num_attention_heads=16, + num_hidden_layers=28, + num_key_value_heads=8, + rms_norm_eps=1e-6, + rope_theta=1000000, + sliding_window=None, + use_cache=True, + vocab_size=151936, + ) + elif isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "qwen3") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + + super().__post_init__(**kwargs) + + +__all__ = ["LightOnOcrConfig"] diff --git a/third_party/transformers/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py b/third_party/transformers/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py new file mode 100644 index 0000000000000000000000000000000000000000..946971c02119d33d2a1fedb8f6486c287ab509bf --- /dev/null +++ b/third_party/transformers/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py @@ -0,0 +1,427 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lighton_ocr/modular_lighton_ocr.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_lighton_ocr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The LightOn 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 dataclasses import dataclass + +import torch +from torch import nn + +from ...cache_utils import Cache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub +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 ...utils.generic import merge_with_config_defaults +from ..auto import AutoModel +from .configuration_lighton_ocr import LightOnOcrConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class LightOnOcrRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + LightOnOcrRMSNorm 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 LightOnOcrPatchMerger(nn.Module): + """ + Learned merging of spatial_merge_size ** 2 patches + """ + + def __init__(self, config: LightOnOcrConfig): + super().__init__() + self.config = config + + hidden_size = config.vision_config.hidden_size + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = self.config.vision_config.patch_size + self.merging_layer = nn.Linear(hidden_size * self.spatial_merge_size**2, hidden_size, bias=False) + + def forward(self, image_features: torch.Tensor, image_sizes: torch.Tensor) -> torch.Tensor: + image_sizes = [ + (image_size[0] // self.patch_size, image_size[1] // self.patch_size) for image_size in image_sizes + ] + + tokens_per_image = [h * w for h, w in image_sizes] + d = image_features.shape[-1] + + permuted_tensor = [] + for image_index, image_tokens in enumerate(image_features.split(tokens_per_image)): + # Reshape image_tokens into a 2D grid + h, w = image_sizes[image_index] + image_grid = image_tokens.view(h, w, d).permute(2, 0, 1).unsqueeze(0) + grid = torch.nn.functional.unfold( + image_grid, kernel_size=self.spatial_merge_size, stride=self.spatial_merge_size + ) + grid = grid.view(d * self.spatial_merge_size**2, -1).t() + permuted_tensor.append(grid) + + image_features = torch.cat(permuted_tensor, dim=0) + image_features = self.merging_layer(image_features) + return image_features + + +class LightOnOcrMultiModalProjector(nn.Module): + def __init__(self, config: LightOnOcrConfig): + super().__init__() + self.config = config + self.norm = LightOnOcrRMSNorm(config.vision_config.hidden_size, eps=config.text_config.rms_norm_eps) + self.patch_merger = LightOnOcrPatchMerger(config) + self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=False) + self.act = nn.GELU() + self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=False) + + def forward(self, image_features: torch.Tensor, image_sizes: torch.Tensor): + image_features = self.norm(image_features) + image_features = self.patch_merger(image_features, image_sizes) + hidden_states = self.linear_1(image_features) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for LightOnOcr outputs, with hidden states and attentions. + """ +) +class LightOnOcrModelOutputWithPast(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 +class LightOnOcrPreTrainedModel(PreTrainedModel): + config: LightOnOcrConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _skip_keys_device_placement = "past_key_values" + + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _supports_flex_attn = True + _supports_attention_backend = True + + +@auto_docstring( + custom_intro=""" + The LightOnOcr model which consists of a vision backbone and a language model, without a language modeling head. + """ +) +class LightOnOcrModel(LightOnOcrPreTrainedModel): + base_model_prefix = "" + + def __init__(self, config: LightOnOcrConfig): + super().__init__(config) + self.vision_encoder = AutoModel.from_config(config.vision_config) + self.vision_projection = LightOnOcrMultiModalProjector(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 + def get_image_features( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_encoder(pixel_values, image_sizes=image_sizes, return_dict=True, **kwargs) + image_features = image_outputs.last_hidden_state + image_features = self.vision_projection(image_features.squeeze(0), image_sizes) + + # Split features per image based on the effective patch size + downsample_ratio = self.config.vision_config.patch_size * self.config.spatial_merge_size + split_sizes = [(height // downsample_ratio) * (width // downsample_ratio) for height, width in image_sizes] + image_features = torch.split(image_features, split_sizes) + 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() + n_image_features = image_features.shape[0] * image_features.shape[1] + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", + ) + return special_image_mask + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + image_sizes: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | LightOnOcrModelOutputWithPast: + 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, image_sizes=image_sizes, return_dict=True + ).pooler_output + image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return LightOnOcrModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for LightOnOcr causal language model (or autoregressive) outputs. + """ +) +class LightOnOcrCausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + The LIGHTON_OCR model which consists of a vision backbone and a language model. + """ +) +class LightOnOcrForConditionalGeneration(LightOnOcrPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + + def __init__(self, config: LightOnOcrConfig): + super().__init__(config) + self.model = LightOnOcrModel(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, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes, **kwargs) + + @merge_with_config_defaults + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + image_sizes: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | LightOnOcrCausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, LightOnOcrForConditionalGeneration + + >>> model = LightOnOcrForConditionalGeneration.from_pretrained("mistralai/Mistral-Small-3.1-24B-Instruct-2503") + >>> processor = AutoProcessor.from_pretrained("mistralai/Mistral-Small-3.1-24B-Instruct-2503") + + >>> prompt = "[INST][IMG]What is the image?[/INST]" + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_new_tokens=15) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is the image?The image depicts two cats lying on a pink blanket." + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + image_sizes=image_sizes, + **kwargs, + ) + + hidden_states = outputs[0] + 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 LightOnOcrCausalLMOutputWithPast( + 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__ = ["LightOnOcrPreTrainedModel", "LightOnOcrForConditionalGeneration", "LightOnOcrModel"] diff --git a/third_party/transformers/src/transformers/models/lighton_ocr/modular_lighton_ocr.py b/third_party/transformers/src/transformers/models/lighton_ocr/modular_lighton_ocr.py new file mode 100644 index 0000000000000000000000000000000000000000..b4be7900c9e7e8bdaca73c8571d04f1214f11684 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lighton_ocr/modular_lighton_ocr.py @@ -0,0 +1,359 @@ +# Copyright 2026 The LightOn 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 numpy as np +import torch +from huggingface_hub.dataclasses import strict +from torch import nn + +from ...cache_utils import Cache +from ...configuration_utils import PreTrainedConfig, PretrainedConfig +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...modeling_outputs import BaseModelOutputWithPooling +from ...modeling_utils import PreTrainedModel +from ...processing_utils import ( + MultiModalData, + ProcessingKwargs, + ProcessorMixin, + Unpack, +) +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel +from ..mistral3.modeling_mistral3 import ( + Mistral3ForConditionalGeneration, + Mistral3Model, + Mistral3ModelOutputWithPast, + Mistral3MultiModalProjector, +) +from ..pixtral.image_processing_pixtral import get_resize_output_image_size + + +@auto_docstring(checkpoint="lightonai/LightOnOCR-1B-1025") +@strict +class LightOnOcrConfig(PretrainedConfig): + r""" + Example: + + ```python + >>> from transformers import LightOnOcrConfig, LightOnOcrForConditionalGeneration + + >>> # Initializing a LightOnOcr configuration + >>> configuration = LightOnOcrConfig() + + >>> # Initializing a model from the configuration + >>> model = LightOnOcrForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lighton_ocr" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + spatial_merge_size: int = 2 + image_token_id: int = 151655 + tie_word_embeddings: bool = True + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = CONFIG_MAPPING["pixtral"]( + attention_dropout=0.0, + head_dim=64, + hidden_act="silu", + hidden_size=1024, + image_size=1540, + initializer_range=0.02, + intermediate_size=4096, + model_type="pixtral", + num_attention_heads=16, + num_channels=3, + num_hidden_layers=24, + patch_size=14, + rope_theta=10000, + ) + elif isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "pixtral") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config) + + if self.text_config is None: + self.text_config = CONFIG_MAPPING["qwen3"]( + attention_dropout=0.0, + head_dim=128, + hidden_act="silu", + hidden_size=1024, + initializer_range=0.02, + intermediate_size=3072, + max_position_embeddings=40960, + num_attention_heads=16, + num_hidden_layers=28, + num_key_value_heads=8, + rms_norm_eps=1e-6, + rope_theta=1000000, + sliding_window=None, + use_cache=True, + vocab_size=151936, + ) + elif isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "qwen3") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + + super().__post_init__(**kwargs) + + +class LightOnOcrProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": False, + }, + "common_kwargs": { + "return_tensors": "pt", + }, + } + + +class LightOnOcrProcessor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + patch_size: int = 14, + spatial_merge_size: int = 2, + chat_template=None, + **kwargs, + ): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + # Calculate effective patch size for image processing + self.effective_patch_size = patch_size * spatial_merge_size + + # Get special tokens and IDs directly from tokenizer attributes + self.image_token = tokenizer.image_token + self.image_break_token = tokenizer.image_break_token + self.image_end_token = tokenizer.image_end_token + self.image_token_id = tokenizer.image_token_id + self.image_break_token_id = tokenizer.image_break_token_id + self.image_end_token_id = tokenizer.image_end_token_id + + self.image_ids = [self.image_token_id, self.image_break_token_id, self.image_end_token_id] + + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[LightOnOcrProcessorKwargs], + ) -> BatchFeature: + if images is None and text is None: + raise ValueError("You must provide either text or images") + output_kwargs = self._merge_kwargs( + LightOnOcrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if images is not None: + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + else: + image_inputs = {} + + 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") + + if image_inputs.get("pixel_values") is not None: + image_sizes_iter = iter(image_inputs["image_sizes"]) + prompt_strings = [] + + for sample in text: + replace_strings = [] + + while self.image_token in sample: + image_height, image_width = next(image_sizes_iter) + num_height_tokens = image_height // self.effective_patch_size + num_width_tokens = image_width // self.effective_patch_size + num_patches = num_height_tokens * num_width_tokens + + replace_str = self.image_token * num_patches + replace_strings.append(replace_str) + + sample = sample.replace(self.image_token, "", 1) + + while "" in sample: + replace_str = replace_strings.pop(0) + sample = sample.replace("", replace_str, 1) + + prompt_strings.append(sample) + else: + prompt_strings = text + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"], return_tensors=None) + self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + vision_data = {} + if image_sizes is not None: + images_kwargs = LightOnOcrProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + + size = images_kwargs.get("size", None) or self.image_processor.size + patch_size = images_kwargs.get("patch_size", None) or self.image_processor.patch_size + if isinstance(patch_size, dict) and "height" in patch_size and "width" in patch_size: + patch_size = (patch_size["height"], patch_size["width"]) + + num_image_tokens = [] + for height, width in image_sizes: + resized_height, resized_width = get_resize_output_image_size( + np.zeros((height, width, 3)), + size=(size["longest_edge"], size["longest_edge"]), + patch_size=patch_size, + ) + num_height_tokens = resized_height // self.effective_patch_size + num_width_tokens = resized_width // self.effective_patch_size + num_image_tokens.append(num_width_tokens * num_height_tokens) + + num_image_patches = [1] * len(image_sizes) + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + +class LightOnOcrMultiModalProjector(Mistral3MultiModalProjector): + def __init__(self, config: LightOnOcrConfig): + self.config = config + super().__init__() + self.act = nn.GELU() + self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=False) + self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=False) + del self.num_feature_layers + + +class LightOnOcrModelOutputWithPast(Mistral3ModelOutputWithPast): + pass + + +class LightOnOcrModel(Mistral3Model): + base_model_prefix = "" + + def __init__(self, config: LightOnOcrConfig): + PreTrainedModel.__init__(self, config) + self.vision_encoder = AutoModel.from_config(config.vision_config) + self.vision_projection = LightOnOcrMultiModalProjector(config) + self.language_model = AutoModel.from_config(config.text_config) + self.post_init() + + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_encoder(pixel_values, image_sizes=image_sizes, return_dict=True, **kwargs) + image_features = image_outputs.last_hidden_state + image_features = self.vision_projection(image_features.squeeze(0), image_sizes) + + # Split features per image based on the effective patch size + downsample_ratio = self.config.vision_config.patch_size * self.config.spatial_merge_size + split_sizes = [(height // downsample_ratio) * (width // downsample_ratio) for height, width in image_sizes] + image_features = torch.split(image_features, split_sizes) + image_outputs.pooler_output = image_features + + return image_outputs + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + image_sizes: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | LightOnOcrModelOutputWithPast: + 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, image_sizes=image_sizes, return_dict=True + ).pooler_output + image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return LightOnOcrModelOutputWithPast( + 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 LightOnOcrForConditionalGeneration(Mistral3ForConditionalGeneration): + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes, **kwargs) + + +__all__ = [ + "LightOnOcrPreTrainedModel", # noqa + "LightOnOcrForConditionalGeneration", + "LightOnOcrModel", + "LightOnOcrConfig", + "LightOnOcrProcessor", +] diff --git a/third_party/transformers/src/transformers/models/lighton_ocr/processing_lighton_ocr.py b/third_party/transformers/src/transformers/models/lighton_ocr/processing_lighton_ocr.py new file mode 100644 index 0000000000000000000000000000000000000000..f7c189c3d849f55852c4fca13570b7f19af161cf --- /dev/null +++ b/third_party/transformers/src/transformers/models/lighton_ocr/processing_lighton_ocr.py @@ -0,0 +1,228 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lighton_ocr/modular_lighton_ocr.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_lighton_ocr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The LightOn 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 numpy as np + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ChannelDimension, ImageInput, get_image_size +from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput + + +class LightOnOcrProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": False, + }, + "common_kwargs": { + "return_tensors": "pt", + }, + } + + +def _num_image_tokens(image_size: tuple[int, int], patch_size: tuple[int, int]) -> int: + """ + Calculate the number of image tokens given the image size and patch size. + + Args: + image_size (`tuple[int, int]`): + The size of the image as `(height, width)`. + patch_size (`tuple[int, int]`): + The patch size as `(height, width)`. + + Returns: + `int`: The number of image tokens. + """ + height, width = image_size + patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size) + num_width_tokens = (width - 1) // patch_width + 1 + num_height_tokens = (height - 1) // patch_height + 1 + return num_height_tokens, num_width_tokens + + +def get_resize_output_image_size( + input_image: ImageInput, + size: int | tuple[int, int] | list[int] | tuple[int], + patch_size: int | tuple[int, int] | list[int] | tuple[int], + input_data_format: str | ChannelDimension | None = None, +) -> tuple: + """ + Find the target (height, width) dimension of the output image after resizing given the input image and the desired + size. + + Args: + input_image (`ImageInput`): + The image to resize. + size (`int` or `tuple[int, int]`): + Max image size an input image can be. Must be a dictionary with the key "longest_edge". + patch_size (`int` or `tuple[int, int]`): + The patch_size as `(height, width)` to use for resizing the image. If patch_size is an integer, `(patch_size, patch_size)` + will be used + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If unset, will use the inferred format from the input. + + Returns: + `tuple`: The target (height, width) dimension of the output image after resizing. + """ + max_height, max_width = size if isinstance(size, (tuple, list)) else (size, size) + patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size) + height, width = get_image_size(input_image, input_data_format) + + ratio = max(height / max_height, width / max_width) + + if ratio > 1: + # Original implementation uses `round` which utilises bankers rounding, which can lead to surprising results + # Here we use floor to ensure the image is always smaller than the given "longest_edge" + height = int(math.floor(height / ratio)) + width = int(math.floor(width / ratio)) + + num_height_tokens, num_width_tokens = _num_image_tokens((height, width), (patch_height, patch_width)) + return num_height_tokens * patch_height, num_width_tokens * patch_width + + +class LightOnOcrProcessor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + patch_size: int = 14, + spatial_merge_size: int = 2, + chat_template=None, + **kwargs, + ): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + # Calculate effective patch size for image processing + self.effective_patch_size = patch_size * spatial_merge_size + + # Get special tokens and IDs directly from tokenizer attributes + self.image_token = tokenizer.image_token + self.image_break_token = tokenizer.image_break_token + self.image_end_token = tokenizer.image_end_token + self.image_token_id = tokenizer.image_token_id + self.image_break_token_id = tokenizer.image_break_token_id + self.image_end_token_id = tokenizer.image_end_token_id + + self.image_ids = [self.image_token_id, self.image_break_token_id, self.image_end_token_id] + + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[LightOnOcrProcessorKwargs], + ) -> BatchFeature: + if images is None and text is None: + raise ValueError("You must provide either text or images") + output_kwargs = self._merge_kwargs( + LightOnOcrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if images is not None: + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + else: + image_inputs = {} + + 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") + + if image_inputs.get("pixel_values") is not None: + image_sizes_iter = iter(image_inputs["image_sizes"]) + prompt_strings = [] + + for sample in text: + replace_strings = [] + + while self.image_token in sample: + image_height, image_width = next(image_sizes_iter) + num_height_tokens = image_height // self.effective_patch_size + num_width_tokens = image_width // self.effective_patch_size + num_patches = num_height_tokens * num_width_tokens + + replace_str = self.image_token * num_patches + replace_strings.append(replace_str) + + sample = sample.replace(self.image_token, "", 1) + + while "" in sample: + replace_str = replace_strings.pop(0) + sample = sample.replace("", replace_str, 1) + + prompt_strings.append(sample) + else: + prompt_strings = text + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"], return_tensors=None) + self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + vision_data = {} + if image_sizes is not None: + images_kwargs = LightOnOcrProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + + size = images_kwargs.get("size", None) or self.image_processor.size + patch_size = images_kwargs.get("patch_size", None) or self.image_processor.patch_size + if isinstance(patch_size, dict) and "height" in patch_size and "width" in patch_size: + patch_size = (patch_size["height"], patch_size["width"]) + + num_image_tokens = [] + for height, width in image_sizes: + resized_height, resized_width = get_resize_output_image_size( + np.zeros((height, width, 3)), + size=(size["longest_edge"], size["longest_edge"]), + patch_size=patch_size, + ) + num_height_tokens = resized_height // self.effective_patch_size + num_width_tokens = resized_width // self.effective_patch_size + num_image_tokens.append(num_width_tokens * num_height_tokens) + + num_image_patches = [1] * len(image_sizes) + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + +__all__ = ["LightOnOcrProcessor"] diff --git a/third_party/transformers/src/transformers/models/llama/__init__.py b/third_party/transformers/src/transformers/models/llama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c15b58a0ae4ede69951aa1a304c93cdcd8feb28 --- /dev/null +++ b/third_party/transformers/src/transformers/models/llama/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_llama import * + from .modeling_llama import * + from .tokenization_llama 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/llama/configuration_llama.py b/third_party/transformers/src/transformers/models/llama/configuration_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..6960a6970592d36182260bf4246b897cdbbb5f3d --- /dev/null +++ b/third_party/transformers/src/transformers/models/llama/configuration_llama.py @@ -0,0 +1,103 @@ +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""LLaMA model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring +from ...utils.type_validators import interval + + +@auto_docstring(checkpoint="meta-llama/Llama-2-7b-hf") +@strict +class LlamaConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import LlamaModel, LlamaConfig + + >>> # Initializing a LLaMA llama-7b style configuration + >>> configuration = LlamaConfig() + + >>> # Initializing a model from the llama-7b style configuration + >>> model = LlamaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "llama" + keys_to_ignore_at_inference = ["past_key_values"] + # Default tensor parallel plan for base model `LlamaModel` + 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 = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = interval(min=0.0, max=1.0)(default=0.02) + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + pretraining_tp: int | None = 1 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: int | float | None = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + + def __post_init__(self, **kwargs): + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +__all__ = ["LlamaConfig"] diff --git a/third_party/transformers/src/transformers/models/llama/convert_llama_weights_to_hf.py b/third_party/transformers/src/transformers/models/llama/convert_llama_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..44906fb633d380e5ef93b3b07b29957c562db3b2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/llama/convert_llama_weights_to_hf.py @@ -0,0 +1,602 @@ +# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import gc +import json +import os +import tempfile +import warnings + +import torch +from tokenizers import AddedToken, processors + +from transformers import GenerationConfig, LlamaConfig, LlamaForCausalLM, LlamaTokenizer, PreTrainedTokenizerFast +from transformers.convert_slow_tokenizer import TikTokenConverter + + +try: + from transformers import LlamaTokenizerFast +except ImportError as e: + warnings.warn(e) + warnings.warn( + "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion" + ) + LlamaTokenizerFast = None + +""" +Sample usage: + +``` +python src/transformers/models/llama/convert_llama_weights_to_hf.py \ + --input_dir /path/to/downloaded/llama/weights --model_size 1B --llama_version 3.2 --output_dir /output/path +``` + +Thereafter, models can be loaded via: + +```py +from transformers import LlamaForCausalLM, LlamaTokenizer + +model = LlamaForCausalLM.from_pretrained("/output/path") +tokenizer = LlamaTokenizer.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). + +If you want your tokenizer to add a bos automatically you should update the tokenizer._tokenizers.post_processor: + +```py +from tokenizers import processors +bos = "<|begin_of_text|>" +tokenizer._tokenizers.post_processor = processors.Sequence( + [ + processors.ByteLevel(trim_offsets=False), + processors.TemplateProcessing( + single=f"{bos}:0 $A:0", + pair=f"{bos}:0 $A:0 {bos}:1 $B:1", + special_tokens=[ + (bos, tokenizer.encode(bos)), + ], + ), + ] +) +``` +""" + +NUM_SHARDS = { + "1B": 1, + "3B": 1, + "7B": 1, + "8B": 1, + "8Bf": 1, + "7Bf": 1, + "13B": 2, + "13Bf": 2, + "34B": 4, + "30B": 4, + "65B": 8, + "70B": 8, + "70Bf": 8, + "405B": 8, + "405B-MP16": 16, +} + +CONTEXT_LENGTH_FOR_VERSION = {"Guard-3": 131072, "3.2": 131072, "3.1": 131072, "3": 8192, "2": 4096, "1": 2048} + +BOS_ADDED_TOKEN = AddedToken( + "<|begin_of_text|>", single_word=False, lstrip=False, rstrip=False, normalized=False, special=True +) +EOS_ADDED_TOKEN = AddedToken( + "<|end_of_text|>", single_word=False, lstrip=False, rstrip=False, normalized=False, special=True +) +EOT_ADDED_TOKEN = AddedToken( + "<|eot_id|>", single_word=False, lstrip=False, rstrip=False, normalized=False, special=True +) + +DEFAULT_LLAMA_SPECIAL_TOKENS = { + "3": [ + "<|begin_of_text|>", + "<|end_of_text|>", + "<|reserved_special_token_0|>", + "<|reserved_special_token_1|>", + "<|reserved_special_token_2|>", + "<|reserved_special_token_3|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|reserved_special_token_4|>", + "<|eot_id|>", # end of turn + ] + + [f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5)], + "3.1": [ + "<|begin_of_text|>", + "<|end_of_text|>", + "<|reserved_special_token_0|>", + "<|reserved_special_token_1|>", + "<|finetune_right_pad_id|>", + "<|reserved_special_token_2|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", # end of message + "<|eot_id|>", # end of turn + "<|python_tag|>", + ] + + [f"<|reserved_special_token_{i}|>" for i in range(3, 256 - 8)], + "3.2": [ + "<|begin_of_text|>", + "<|end_of_text|>", + "<|reserved_special_token_0|>", + "<|reserved_special_token_1|>", + "<|finetune_right_pad_id|>", + "<|reserved_special_token_2|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", # end of message + "<|eot_id|>", # end of turn + "<|python_tag|>", + ] + + [f"<|reserved_special_token_{i}|>" for i in range(3, 256 - 8)], + "Guard-3": [ + "<|begin_of_text|>", + "<|end_of_text|>", + "<|reserved_special_token_0|>", + "<|reserved_special_token_1|>", + "<|finetune_right_pad_id|>", + "<|reserved_special_token_2|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", # end of message + "<|eot_id|>", # end of turn + "<|python_tag|>", + ] + + [f"<|reserved_special_token_{i}|>" for i in range(3, 256 - 8)], +} + + +def is_llama_3(version): + return version in ["3", "3.1", "3.2", "Guard-3"] + + +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=None, + llama_version="1", + vocab_size=None, + num_shards=None, + instruct=False, + push_to_hub=False, +): + print("Converting the model.") + params = read_json(os.path.join(input_base_path, "params.json")) + num_shards = NUM_SHARDS[model_size] if num_shards is None else num_shards + params = params.get("model", params) + n_layers = params["n_layers"] + n_heads = params["n_heads"] + n_heads_per_shard = n_heads // num_shards + dim = params["dim"] + dims_per_head = dim // n_heads + base = params.get("rope_theta", 10000.0) + inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)) + if base > 10000.0 and not is_llama_3(llama_version): + max_position_embeddings = 16384 + else: + max_position_embeddings = CONTEXT_LENGTH_FOR_VERSION[llama_version] + + if params.get("n_kv_heads", None) is not None: + num_key_value_heads = params["n_kv_heads"] # for GQA / MQA + num_key_value_heads_per_shard = num_key_value_heads // num_shards + key_value_dim = dims_per_head * num_key_value_heads + else: # compatibility with other checkpoints + num_key_value_heads = n_heads + num_key_value_heads_per_shard = n_heads_per_shard + key_value_dim = dim + + # permute for sliced rotary + def permute(w, n_heads, dim1=dim, dim2=dim): + return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + with tempfile.TemporaryDirectory() as tmp_model_path: + print(f"Fetching all parameters from the checkpoint at {input_base_path}.") + # Load weights + if num_shards == 1: + # Not sharded + # (The sharded implementation would also work, but this is simpler.) + loaded = torch.load( + os.path.join(input_base_path, "consolidated.00.pth"), map_location="cpu", weights_only=True + ) + else: + # Sharded + checkpoint_list = sorted([file for file in os.listdir(input_base_path) if file.endswith(".pth")]) + print("Loading in order:", checkpoint_list) + loaded = [ + torch.load(os.path.join(input_base_path, file), map_location="cpu", weights_only=True) + for file in checkpoint_list + ] + param_count = 0 + index_dict = {"weight_map": {}} + for layer_i in range(n_layers): + filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin" + if num_shards == 1: + # Unsharded + state_dict = { + f"model.layers.{layer_i}.self_attn.q_proj.weight": permute( + loaded[f"layers.{layer_i}.attention.wq.weight"], n_heads=n_heads + ), + f"model.layers.{layer_i}.self_attn.k_proj.weight": permute( + loaded[f"layers.{layer_i}.attention.wk.weight"], + n_heads=num_key_value_heads, + dim1=key_value_dim, + ), + f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"], + f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"], + f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"], + f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"], + f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"], + f"model.layers.{layer_i}.input_layernorm.weight": loaded[ + f"layers.{layer_i}.attention_norm.weight" + ], + f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[ + f"layers.{layer_i}.ffn_norm.weight" + ], + } + else: + # 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 = { + f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][ + f"layers.{layer_i}.attention_norm.weight" + ].clone(), + f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][ + f"layers.{layer_i}.ffn_norm.weight" + ].clone(), + } + state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"] = permute( + torch.cat( + [ + loaded[i][f"layers.{layer_i}.attention.wq.weight"].view( + n_heads_per_shard, dims_per_head, dim + ) + for i in range(len(loaded)) + ], + dim=0, + ).reshape(dim, dim), + n_heads=n_heads, + ) + state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"] = permute( + torch.cat( + [ + loaded[i][f"layers.{layer_i}.attention.wk.weight"].view( + num_key_value_heads_per_shard, dims_per_head, dim + ) + for i in range(len(loaded)) + ], + dim=0, + ).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"] = torch.cat( + [ + loaded[i][f"layers.{layer_i}.attention.wv.weight"].view( + num_key_value_heads_per_shard, dims_per_head, dim + ) + for i in range(len(loaded)) + ], + dim=0, + ).reshape(key_value_dim, dim) + + state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"] = torch.cat( + [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(len(loaded))], dim=1 + ) + state_dict[f"model.layers.{layer_i}.mlp.gate_proj.weight"] = torch.cat( + [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(len(loaded))], dim=0 + ) + state_dict[f"model.layers.{layer_i}.mlp.down_proj.weight"] = torch.cat( + [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(len(loaded))], dim=1 + ) + state_dict[f"model.layers.{layer_i}.mlp.up_proj.weight"] = torch.cat( + [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(len(loaded))], dim=0 + ) + + state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin" + if num_shards == 1: + # Unsharded + state_dict = { + "model.embed_tokens.weight": loaded["tok_embeddings.weight"], + "model.norm.weight": loaded["norm.weight"], + "lm_head.weight": loaded["output.weight"], + } + else: + concat_dim = 0 if is_llama_3(llama_version) else 1 + state_dict = { + "model.norm.weight": loaded[0]["norm.weight"], + "model.embed_tokens.weight": torch.cat( + [loaded[i]["tok_embeddings.weight"] for i in range(len(loaded))], dim=concat_dim + ), + "lm_head.weight": torch.cat([loaded[i]["output.weight"] for i in range(len(loaded))], dim=0), + } + + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + # Write configs + index_dict["metadata"] = {"total_size": param_count * 2} + write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json")) + ffn_dim_multiplier = params.get("ffn_dim_multiplier", 1) + multiple_of = params.get("multiple_of", 256) + + if is_llama_3(llama_version): + bos_token_id = 128000 + + if instruct: + eos_token_id = [128001, 128008, 128009] + else: + eos_token_id = 128001 + else: + bos_token_id = 1 + eos_token_id = 2 + + if llama_version in ["3.1", "3.2", "Guard-3"]: + rope_parameters = { + "factor": 32.0 if llama_version == "3.2" else 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + } + else: + rope_parameters = None + + config = LlamaConfig( + hidden_size=dim, + intermediate_size=compute_intermediate_size(dim, ffn_dim_multiplier, multiple_of), + num_attention_heads=params["n_heads"], + num_hidden_layers=params["n_layers"], + rms_norm_eps=params["norm_eps"], + num_key_value_heads=num_key_value_heads, + vocab_size=vocab_size, + rope_theta=base, + rope_parameters=rope_parameters, + max_position_embeddings=max_position_embeddings, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=llama_version == "3.2", + ) + + config.save_pretrained(tmp_model_path) + + generation_config = GenerationConfig( + do_sample=True, + temperature=0.6, + top_p=0.9, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + ) + generation_config.save_pretrained(tmp_model_path) + + # Make space so we can load the model properly now. + del state_dict + del loaded + gc.collect() + + print("Loading the checkpoint in a Llama model.") + model = LlamaForCausalLM.from_pretrained(tmp_model_path, dtype=torch.bfloat16) + + # 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.") + if push_to_hub: + print("Pushing to the hub.") + model_name = model_path.split(os.path.sep)[-1] + model.push_to_hub(model_name, private=True) + else: + print("Saving to disk.") + model.save_pretrained(model_path) + + +class Llama3Converter(TikTokenConverter): + def __init__(self, vocab_file, special_tokens=None, instruct=False, llama_version="3.2", **kwargs): + super().__init__(vocab_file, additional_special_tokens=special_tokens, **kwargs) + tokenizer = self.converted() + + # References for chat templates in instruct models + templates_for_version = { + "2": ("meta-llama/Llama-2-7b-chat-hf", "f5db02db724555f92da89c216ac04704f23d4590"), + "3": ("meta-llama/Meta-Llama-3-8B-Instruct", "5f0b02c75b57c5855da9ae460ce51323ea669d8a"), + "3.1": ("meta-llama/Llama-3.1-8B-Instruct", "0e9e39f249a16976918f6564b8830bc894c89659"), + "3.2": ("meta-llama/Llama-3.2-1B-Instruct", "e9f8effbab1cbdc515c11ee6e098e3d5a9f51e14"), + "Guard-3": ("meta-llama/Llama-Guard-3-1B", "acf7aafa60f0410f8f42b1fa35e077d705892029"), + } + + # Add chat_template only if instruct is True. + # Prevents a null chat_template, which triggers + # a parsing warning in the Hub. + additional_kwargs = {} + if instruct or llama_version == "Guard-3": + model_id, revision = templates_for_version.get(llama_version, (None, None)) + if model_id is not None: + from transformers import AutoTokenizer + + t = AutoTokenizer.from_pretrained(model_id, revision=revision) + additional_kwargs["chat_template"] = t.chat_template + + self.converted_tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + bos_token="<|begin_of_text|>", + eos_token="<|end_of_text|>" if not instruct else "<|eot_id|>", + model_input_names=["input_ids", "attention_mask"], + model_max_length=CONTEXT_LENGTH_FOR_VERSION[llama_version], + clean_up_tokenization_spaces=False, + **additional_kwargs, + ) + self.update_post_processor(self.converted_tokenizer) + # finer special_tokens_map.json + self.converted_tokenizer._bos_token = BOS_ADDED_TOKEN + self.converted_tokenizer._eos_token = EOT_ADDED_TOKEN if instruct else EOS_ADDED_TOKEN + + # We can't do this while building the tokenizer because we have no easy access to the bos token id + def update_post_processor(self, tokenizer): + tokenizer._tokenizer.post_processor = processors.Sequence( + [ + processors.ByteLevel(trim_offsets=False), + processors.TemplateProcessing( + single="<|begin_of_text|> $A", + pair="<|begin_of_text|>:0 $A:0 <|begin_of_text|>:1 $B:1", + special_tokens=[ + ("<|begin_of_text|>", tokenizer.convert_tokens_to_ids("<|begin_of_text|>")), + ], + ), + ] + ) + + +def write_tokenizer( + tokenizer_path, input_tokenizer_path, llama_version="2", special_tokens=None, instruct=False, push_to_hub=False +): + print("Converting the tokenizer.") + tokenizer_class = LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast + if is_llama_3(llama_version): + tokenizer = Llama3Converter( + input_tokenizer_path, + special_tokens, + instruct, + llama_version, + ).converted_tokenizer + else: + try: + tokenizer = tokenizer_class(input_tokenizer_path) + except Exception: + raise ValueError( + "Failed to instantiate tokenizer. Please, make sure you have sentencepiece and protobuf installed." + ) + + if push_to_hub: + print(f"Pushing a {tokenizer_class.__name__} to the Hub repo - {tokenizer_path}.") + model_name = tokenizer_path.split(os.path.sep)[-1] + tokenizer.push_to_hub(model_name, private=True) + else: + print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}.") + tokenizer.save_pretrained(tokenizer_path) + return tokenizer + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + help="Location of Llama weights, which contains tokenizer.model and model folders", + ) + parser.add_argument( + "--model_size", + default=None, + help="'f' Deprecated in favor of `num_shards`: models correspond to the finetuned versions, and are specific to the Llama2 official release. For more details on Llama2, check out the original repo: https://huggingface.co/meta-llama", + ) + parser.add_argument( + "--output_dir", + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--push_to_hub", + help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally.", + action="store_true", + default=False, + ) + # Different Llama versions used different default values for max_position_embeddings, hence the need to be able to specify which version is being used. + parser.add_argument( + "--llama_version", + choices=["1", "2", "3", "3.1", "3.2", "Guard-3"], + default="1", + type=str, + help="Version of the Llama model to convert. Currently supports Llama1 and Llama2. Controls the context size", + ) + parser.add_argument( + "--num_shards", + default=None, + type=int, + help="The number of individual shards used for the model. Does not have to be the same as the number of consolidated_xx.pth", + ) + parser.add_argument( + "--special_tokens", + default=None, + type=list[str], + help="The list of special tokens that should be added to the model.", + ) + parser.add_argument( + "--instruct", + action="store_true", + default=False, + help="Whether the model is an instruct model or not. Will affect special tokens and chat template.", + ) + args = parser.parse_args() + if args.model_size is None and args.num_shards is None: + raise ValueError("You have to set at least `num_shards` if you are not giving the `model_size`") + if args.special_tokens is None: + # no special tokens by default + args.special_tokens = DEFAULT_LLAMA_SPECIAL_TOKENS.get(str(args.llama_version), []) + + spm_path = os.path.join(args.input_dir, "tokenizer.model") + vocab_size = len( + write_tokenizer( + args.output_dir, + spm_path, + llama_version=args.llama_version, + special_tokens=args.special_tokens, + instruct=args.instruct, + push_to_hub=args.push_to_hub, + ) + ) + + if args.model_size != "tokenizer_only": + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + model_size=args.model_size, + llama_version=args.llama_version, + vocab_size=vocab_size, + num_shards=args.num_shards, + instruct=args.instruct, + push_to_hub=args.push_to_hub, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/llama/modeling_llama.py b/third_party/transformers/src/transformers/models/llama/modeling_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..9d659c7c6f08cb73cc87e6ee4ecbc8b76bc82f6c --- /dev/null +++ b/third_party/transformers/src/transformers/models/llama/modeling_llama.py @@ -0,0 +1,519 @@ +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import ( + GenericForQuestionAnswering, + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_llama import LlamaConfig + + +logger = logging.get_logger(__name__) + + +@use_kernel_forward_from_hub("RMSNorm") +class LlamaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + LlamaRMSNorm 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 LlamaRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: LlamaConfig, 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: LlamaConfig | 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 + + +class LlamaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +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 LlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class LlamaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx) + + self.mlp = LlamaMLP(config) + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class LlamaPreTrainedModel(PreTrainedModel): + config: LlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _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": LlamaDecoderLayer, + "attentions": LlamaAttention, + } + + +@auto_docstring +class LlamaModel(LlamaPreTrainedModel): + def __init__(self, config: LlamaConfig): + 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( + [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LlamaRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class LlamaForCausalLM(LlamaPreTrainedModel, 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 = LlamaModel(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, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class LlamaForSequenceClassification(GenericForSequenceClassification, LlamaPreTrainedModel): ... + + +class LlamaForQuestionAnswering(GenericForQuestionAnswering, LlamaPreTrainedModel): + base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model` + + +class LlamaForTokenClassification(GenericForTokenClassification, LlamaPreTrainedModel): ... + + +__all__ = [ + "LlamaForCausalLM", + "LlamaModel", + "LlamaPreTrainedModel", + "LlamaForSequenceClassification", + "LlamaForQuestionAnswering", + "LlamaForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/llama/tokenization_llama.py b/third_party/transformers/src/transformers/models/llama/tokenization_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..366e50d74ec27ac19bd0b73c6098f0cfa62e3b83 --- /dev/null +++ b/third_party/transformers/src/transformers/models/llama/tokenization_llama.py @@ -0,0 +1,149 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tokenizers import Tokenizer, decoders, pre_tokenizers +from tokenizers.models import BPE + +from ...tokenization_utils_base import _get_prepend_scheme +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) +VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model", "tokenizer_file": "tokenizer.json"} + +B_INST, E_INST = "[INST]", "[/INST]" +B_SYS, E_SYS = "<>\n", "\n<>\n\n" + +# fmt: off +DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your \ +answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure\ + that your responses are socially unbiased and positive in nature. + +If a question does not make any sense, or is not factually coherent, explain why instead of answering something not \ +correct. If you don't know the answer to a question, please don't share false information.""" +# fmt: on + + +class LlamaTokenizer(TokenizersBackend): + """ + Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding. + + This uses notably ByteFallback and no normalization. + + ```python + >>> from transformers import LlamaTokenizer + + >>> tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") + >>> tokenizer.encode("Hello this is a test") + [1, 15043, 445, 338, 263, 1243] + ``` + + If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or + call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the + values of the first token and final token of an encoded sequence will not be correct). For more details, checkout + [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation. + + + This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab (`str`, `dict` or `list`, *optional*): + Path to the vocabulary file, a dictionary or a list of tokens. + merges (`str` or `list`, *optional*): + Path to the merges file or a list of merges. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like + extra spaces. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The end of sequence token. + add_bos_token (`bool`, *optional*, defaults to `True`): + Whether or not to add an `bos_token` at the start of sequences. + add_eos_token (`bool`, *optional*, defaults to `False`): + Whether or not to add an `eos_token` at the end of sequences. + use_default_system_prompt (`bool`, *optional*, defaults to `False`): + Whether or not the default system prompt for Llama should be used + add_prefix_space (`bool`, *optional*): + Whether or not the tokenizer should automatically add a prefix space + """ + + vocab_files_names = VOCAB_FILES_NAMES + padding_side = "left" + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + vocab: str | dict | list | None = None, + merges: str | list | None = None, + clean_up_tokenization_spaces=False, + unk_token="", + bos_token="", + eos_token="", + use_default_system_prompt=False, + legacy=False, + add_prefix_space=None, + **kwargs, + ): + self.add_prefix_space = add_prefix_space if add_prefix_space is not None else True + self.legacy = legacy + self._vocab = vocab + if vocab is None: + self._vocab = { + str(unk_token): 0, + str(bos_token): 1, + str(eos_token): 2, + } + + self._merges = merges or [] + self._tokenizer = Tokenizer( + BPE(vocab=self._vocab, merges=self._merges, fuse_unk=True, byte_fallback=True, dropout=None) + ) + self._tokenizer.normalizer = None + self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace( + replacement="▁", prepend_scheme=_get_prepend_scheme(self.add_prefix_space, self), split=False + ) + + sequence = [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + + if self.add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + + self._tokenizer.decoder = decoders.Sequence(sequence) + self.use_default_system_prompt = use_default_system_prompt + super().__init__( + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + use_default_system_prompt=use_default_system_prompt, + add_prefix_space=add_prefix_space, + **kwargs, + ) + + +__all__ = ["LlamaTokenizer", "LlamaTokenizerFast"] + +# Backward alias +LlamaTokenizerFast = LlamaTokenizer diff --git a/third_party/transformers/src/transformers/models/luke/__init__.py b/third_party/transformers/src/transformers/models/luke/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db588fb2da1976194806932bf5d39438f9769e19 --- /dev/null +++ b/third_party/transformers/src/transformers/models/luke/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_luke import * + from .modeling_luke import * + from .tokenization_luke 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/luke/configuration_luke.py b/third_party/transformers/src/transformers/models/luke/configuration_luke.py new file mode 100644 index 0000000000000000000000000000000000000000..07689488d6d1a6d43e10b407e9f9007b4033165c --- /dev/null +++ b/third_party/transformers/src/transformers/models/luke/configuration_luke.py @@ -0,0 +1,75 @@ +# Copyright Studio Ousia 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. +"""LUKE configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="studio-ousia/luke-base") +@strict +class LukeConfig(PreTrainedConfig): + r""" + entity_vocab_size (`int`, *optional*, defaults to 500000): + Entity vocabulary size of the LUKE model. Defines the number of different entities that can be represented + by the `entity_ids` passed when calling [`LukeModel`]. + entity_emb_size (`int`, *optional*, defaults to 256): + The number of dimensions of the entity embedding. + use_entity_aware_attention (`bool`, *optional*, defaults to `True`): + Whether or not the model should use the entity-aware self-attention mechanism proposed in [LUKE: Deep + Contextualized Entity Representations with Entity-aware Self-attention (Yamada et + al.)](https://huggingface.co/papers/2010.01057). + + Examples: + + ```python + >>> from transformers import LukeConfig, LukeModel + + >>> # Initializing a LUKE configuration + >>> configuration = LukeConfig() + + >>> # Initializing a model from the configuration + >>> model = LukeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "luke" + + vocab_size: int = 50267 + entity_vocab_size: int = 500000 + hidden_size: int = 768 + entity_emb_size: int = 256 + 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 + use_entity_aware_attention: bool = True + classifier_dropout: float | int | None = None + pad_token_id: int | None = 1 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = True + + +__all__ = ["LukeConfig"] diff --git a/third_party/transformers/src/transformers/models/luke/convert_luke_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/luke/convert_luke_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..af17ae00107a843337fd65030bca29f34c9a6a51 --- /dev/null +++ b/third_party/transformers/src/transformers/models/luke/convert_luke_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,169 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert LUKE checkpoint.""" + +import argparse +import json +import os + +import torch + +from transformers import LukeConfig, LukeModel, LukeTokenizer, RobertaTokenizer +from transformers.tokenization_utils_base import AddedToken + + +@torch.no_grad() +def convert_luke_checkpoint(checkpoint_path, metadata_path, entity_vocab_path, pytorch_dump_folder_path, model_size): + # Load configuration defined in the metadata file + with open(metadata_path) as metadata_file: + metadata = json.load(metadata_file) + config = LukeConfig(use_entity_aware_attention=True, **metadata["model_config"]) + + # Load in the weights from the checkpoint_path + state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + # Load the entity vocab file + entity_vocab = load_entity_vocab(entity_vocab_path) + + tokenizer = RobertaTokenizer.from_pretrained(metadata["model_config"]["bert_model_name"]) + + # Add special tokens to the token vocabulary for downstream tasks + entity_token_1 = AddedToken("", lstrip=False, rstrip=False) + entity_token_2 = AddedToken("", lstrip=False, rstrip=False) + tokenizer.add_special_tokens({"additional_special_tokens": [entity_token_1, entity_token_2]}) + config.vocab_size += 2 + + print(f"Saving tokenizer to {pytorch_dump_folder_path}") + tokenizer.save_pretrained(pytorch_dump_folder_path) + with open(os.path.join(pytorch_dump_folder_path, LukeTokenizer.vocab_files_names["entity_vocab_file"]), "w") as f: + json.dump(entity_vocab, f) + + tokenizer = LukeTokenizer.from_pretrained(pytorch_dump_folder_path) + + # Initialize the embeddings of the special tokens + word_emb = state_dict["embeddings.word_embeddings.weight"] + ent_emb = word_emb[tokenizer.convert_tokens_to_ids(["@"])[0]].unsqueeze(0) + ent2_emb = word_emb[tokenizer.convert_tokens_to_ids(["#"])[0]].unsqueeze(0) + state_dict["embeddings.word_embeddings.weight"] = torch.cat([word_emb, ent_emb, ent2_emb]) + + # Initialize the query layers of the entity-aware self-attention mechanism + for layer_index in range(config.num_hidden_layers): + for matrix_name in ["query.weight", "query.bias"]: + prefix = f"encoder.layer.{layer_index}.attention.self." + state_dict[prefix + "w2e_" + matrix_name] = state_dict[prefix + matrix_name] + state_dict[prefix + "e2w_" + matrix_name] = state_dict[prefix + matrix_name] + state_dict[prefix + "e2e_" + matrix_name] = state_dict[prefix + matrix_name] + + # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks + entity_emb = state_dict["entity_embeddings.entity_embeddings.weight"] + entity_emb[entity_vocab["[MASK2]"]] = entity_emb[entity_vocab["[MASK]"]] + + model = LukeModel(config=config).eval() + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + if not (len(missing_keys) == 1 and missing_keys[0] == "embeddings.position_ids"): + raise ValueError(f"Missing keys {', '.join(missing_keys)}. Expected only missing embeddings.position_ids") + if not (all(key.startswith("entity_predictions") or key.startswith("lm_head") for key in unexpected_keys)): + raise ValueError( + "Unexpected keys" + f" {', '.join([key for key in unexpected_keys if not (key.startswith('entity_predictions') or key.startswith('lm_head'))])}" + ) + + # Check outputs + tokenizer = LukeTokenizer.from_pretrained(pytorch_dump_folder_path, task="entity_classification") + + text = ( + "Top seed Ana Ivanovic said on Thursday she could hardly believe her luck as a fortuitous netcord helped the" + " new world number one avoid a humiliating second- round exit at Wimbledon ." + ) + span = (39, 42) + encoding = tokenizer(text, entity_spans=[span], add_prefix_space=True, return_tensors="pt") + + outputs = model(**encoding) + + # Verify word hidden states + if model_size == "large": + expected_shape = torch.Size((1, 42, 1024)) + expected_slice = torch.tensor( + [[0.0133, 0.0865, 0.0095], [0.3093, -0.2576, -0.7418], [-0.1720, -0.2117, -0.2869]] + ) + else: # base + expected_shape = torch.Size((1, 42, 768)) + expected_slice = torch.tensor([[0.0037, 0.1368, -0.0091], [0.1099, 0.3329, -0.1095], [0.0765, 0.5335, 0.1179]]) + + if not (outputs.last_hidden_state.shape == expected_shape): + raise ValueError( + f"Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}" + ) + if not torch.allclose(outputs.last_hidden_state[0, :3, :3], expected_slice, atol=1e-4): + raise ValueError + + # Verify entity hidden states + if model_size == "large": + expected_shape = torch.Size((1, 1, 1024)) + expected_slice = torch.tensor([[0.0466, -0.0106, -0.0179]]) + else: # base + expected_shape = torch.Size((1, 1, 768)) + expected_slice = torch.tensor([[0.1457, 0.1044, 0.0174]]) + + if not (outputs.entity_last_hidden_state.shape != expected_shape): + raise ValueError( + f"Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is" + f" {expected_shape}" + ) + if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3], expected_slice, atol=1e-4): + raise ValueError + + # Finally, save our PyTorch model and tokenizer + print(f"Saving PyTorch model to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + + +def load_entity_vocab(entity_vocab_path): + entity_vocab = {} + with open(entity_vocab_path, "r", encoding="utf-8") as f: + for index, line in enumerate(f): + title, _ = line.rstrip().split("\t") + entity_vocab[title] = index + + return entity_vocab + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument("--checkpoint_path", type=str, help="Path to a pytorch_model.bin file.") + parser.add_argument( + "--metadata_path", default=None, type=str, help="Path to a metadata.json file, defining the configuration." + ) + parser.add_argument( + "--entity_vocab_path", + default=None, + type=str, + help="Path to an entity_vocab.tsv file, containing the entity vocabulary.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to where to dump the output PyTorch model." + ) + parser.add_argument( + "--model_size", default="base", type=str, choices=["base", "large"], help="Size of the model to be converted." + ) + args = parser.parse_args() + convert_luke_checkpoint( + args.checkpoint_path, + args.metadata_path, + args.entity_vocab_path, + args.pytorch_dump_folder_path, + args.model_size, + ) diff --git a/third_party/transformers/src/transformers/models/luke/modeling_luke.py b/third_party/transformers/src/transformers/models/luke/modeling_luke.py new file mode 100644 index 0000000000000000000000000000000000000000..baa1b9b2c3eafc9145d4a769d9311d10241e30b0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/luke/modeling_luke.py @@ -0,0 +1,2133 @@ +# Copyright Studio Ousia 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 LUKE model.""" + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN, gelu +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_luke import LukeConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the LUKE model. + """ +) +class BaseLukeModelOutputWithPooling(BaseModelOutputWithPooling): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) further processed by a + Linear layer and a Tanh activation function. + entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`): + Sequence of entity hidden-states at the output of the last layer of the model. + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + entity_last_hidden_state: torch.FloatTensor | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model's outputs, with potential hidden states and attentions. + """ +) +class BaseLukeModelOutput(BaseModelOutput): + r""" + entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`): + Sequence of entity hidden-states at the output of the last layer of the model. + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + entity_last_hidden_state: torch.FloatTensor | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model's outputs, with potential hidden states and attentions. + """ +) +class LukeMaskedLMOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + The sum of masked language modeling (MLM) loss and entity prediction loss. + mlm_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Masked language modeling (MLM) loss. + mep_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Masked entity prediction (MEP) 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). + entity_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the entity prediction head (scores for each entity vocabulary token before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + mlm_loss: torch.FloatTensor | None = None + mep_loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + entity_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of entity classification models. + """ +) +class EntityClassificationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification scores (before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of entity pair classification models. + """ +) +class EntityPairClassificationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification scores (before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of entity span classification models. + """ +) +class EntitySpanClassificationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, entity_length, config.num_labels)`): + Classification scores (before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of sentence classification models. + """ +) +class LukeSequenceClassifierOutput(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). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of token classification models. + """ +) +class LukeTokenClassifierOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): + Classification scores (before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of question answering models. + """ +) +class LukeQuestionAnsweringModelOutput(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. + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + start_logits: torch.FloatTensor | None = None + end_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Outputs of multiple choice models. + """ +) +class LukeMultipleChoiceModelOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`): + *num_choices* is the second dimension of the input tensors. (see *input_ids* above). + + Classification scores (before SoftMax). + entity_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 layer) of + shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each + layer plus the initial entity embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + entity_hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +class LukeEmbeddings(nn.Module): + """ + Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. + """ + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # End copy + self.padding_idx = config.pad_token_id + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx + ) + + def forward( + self, + input_ids=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + ): + 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).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 = self.word_embeddings(input_ids) + + position_embeddings = self.position_embeddings(position_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = inputs_embeds + position_embeddings + token_type_embeddings + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + 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 LukeEntityEmbeddings(nn.Module): + def __init__(self, config: LukeConfig): + super().__init__() + self.config = config + + self.entity_embeddings = nn.Embedding(config.entity_vocab_size, config.entity_emb_size, padding_idx=0) + if config.entity_emb_size != config.hidden_size: + self.entity_embedding_dense = nn.Linear(config.entity_emb_size, config.hidden_size, bias=False) + + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward( + self, + entity_ids: torch.LongTensor, + position_ids: torch.LongTensor, + token_type_ids: torch.LongTensor | None = None, + ): + if token_type_ids is None: + token_type_ids = torch.zeros_like(entity_ids) + + entity_embeddings = self.entity_embeddings(entity_ids) + if self.config.entity_emb_size != self.config.hidden_size: + entity_embeddings = self.entity_embedding_dense(entity_embeddings) + + position_embeddings = self.position_embeddings(position_ids.clamp(min=0)) + position_embedding_mask = (position_ids != -1).type_as(position_embeddings).unsqueeze(-1) + position_embeddings = position_embeddings * position_embedding_mask + position_embeddings = torch.sum(position_embeddings, dim=-2) + position_embeddings = position_embeddings / position_embedding_mask.sum(dim=-2).clamp(min=1e-7) + + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = entity_embeddings + position_embeddings + token_type_embeddings + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + + return embeddings + + +class LukeSelfAttention(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.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.use_entity_aware_attention = config.use_entity_aware_attention + + 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) + + if self.use_entity_aware_attention: + self.w2e_query = nn.Linear(config.hidden_size, self.all_head_size) + self.e2w_query = nn.Linear(config.hidden_size, self.all_head_size) + self.e2e_query = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def transpose_for_scores(self, x): + new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) + x = x.view(*new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward( + self, + word_hidden_states, + entity_hidden_states, + attention_mask=None, + output_attentions=False, + ): + word_size = word_hidden_states.size(1) + + if entity_hidden_states is None: + concat_hidden_states = word_hidden_states + else: + concat_hidden_states = torch.cat([word_hidden_states, entity_hidden_states], dim=1) + + key_layer = self.transpose_for_scores(self.key(concat_hidden_states)) + value_layer = self.transpose_for_scores(self.value(concat_hidden_states)) + + if self.use_entity_aware_attention and entity_hidden_states is not None: + # compute query vectors using word-word (w2w), word-entity (w2e), entity-word (e2w), entity-entity (e2e) + # query layers + w2w_query_layer = self.transpose_for_scores(self.query(word_hidden_states)) + w2e_query_layer = self.transpose_for_scores(self.w2e_query(word_hidden_states)) + e2w_query_layer = self.transpose_for_scores(self.e2w_query(entity_hidden_states)) + e2e_query_layer = self.transpose_for_scores(self.e2e_query(entity_hidden_states)) + + # compute w2w, w2e, e2w, and e2e key vectors used with the query vectors computed above + w2w_key_layer = key_layer[:, :, :word_size, :] + e2w_key_layer = key_layer[:, :, :word_size, :] + w2e_key_layer = key_layer[:, :, word_size:, :] + e2e_key_layer = key_layer[:, :, word_size:, :] + + # compute attention scores based on the dot product between the query and key vectors + w2w_attention_scores = torch.matmul(w2w_query_layer, w2w_key_layer.transpose(-1, -2)) + w2e_attention_scores = torch.matmul(w2e_query_layer, w2e_key_layer.transpose(-1, -2)) + e2w_attention_scores = torch.matmul(e2w_query_layer, e2w_key_layer.transpose(-1, -2)) + e2e_attention_scores = torch.matmul(e2e_query_layer, e2e_key_layer.transpose(-1, -2)) + + # combine attention scores to create the final attention score matrix + word_attention_scores = torch.cat([w2w_attention_scores, w2e_attention_scores], dim=3) + entity_attention_scores = torch.cat([e2w_attention_scores, e2e_attention_scores], dim=3) + attention_scores = torch.cat([word_attention_scores, entity_attention_scores], dim=2) + + else: + query_layer = self.transpose_for_scores(self.query(concat_hidden_states)) + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in LukeModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + output_word_hidden_states = context_layer[:, :word_size, :] + if entity_hidden_states is None: + output_entity_hidden_states = None + else: + output_entity_hidden_states = context_layer[:, word_size:, :] + + if output_attentions: + outputs = (output_word_hidden_states, output_entity_hidden_states, attention_probs) + else: + outputs = (output_word_hidden_states, output_entity_hidden_states) + + return outputs + + +# Copied from transformers.models.bert.modeling_bert.BertSelfOutput +class LukeSelfOutput(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 LukeAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = LukeSelfAttention(config) + self.output = LukeSelfOutput(config) + + def forward( + self, + word_hidden_states, + entity_hidden_states, + attention_mask=None, + output_attentions=False, + ): + word_size = word_hidden_states.size(1) + self_outputs = self.self( + word_hidden_states, + entity_hidden_states, + attention_mask, + output_attentions, + ) + if entity_hidden_states is None: + concat_self_outputs = self_outputs[0] + concat_hidden_states = word_hidden_states + else: + concat_self_outputs = torch.cat(self_outputs[:2], dim=1) + concat_hidden_states = torch.cat([word_hidden_states, entity_hidden_states], dim=1) + + attention_output = self.output(concat_self_outputs, concat_hidden_states) + + word_attention_output = attention_output[:, :word_size, :] + if entity_hidden_states is None: + entity_attention_output = None + else: + entity_attention_output = attention_output[:, word_size:, :] + + # add attentions if we output them + outputs = (word_attention_output, entity_attention_output) + self_outputs[2:] + + return outputs + + +# Copied from transformers.models.bert.modeling_bert.BertIntermediate +class LukeIntermediate(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 LukeOutput(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 LukeLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = LukeAttention(config) + self.intermediate = LukeIntermediate(config) + self.output = LukeOutput(config) + + def forward( + self, + word_hidden_states, + entity_hidden_states, + attention_mask=None, + output_attentions=False, + ): + word_size = word_hidden_states.size(1) + + self_attention_outputs = self.attention( + word_hidden_states, + entity_hidden_states, + attention_mask, + output_attentions=output_attentions, + ) + if entity_hidden_states is None: + concat_attention_output = self_attention_outputs[0] + else: + concat_attention_output = torch.cat(self_attention_outputs[:2], dim=1) + + outputs = self_attention_outputs[2:] # 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, concat_attention_output + ) + word_layer_output = layer_output[:, :word_size, :] + if entity_hidden_states is None: + entity_layer_output = None + else: + entity_layer_output = layer_output[:, word_size:, :] + + outputs = (word_layer_output, entity_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 LukeEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([LukeLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + word_hidden_states, + entity_hidden_states, + attention_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ): + all_word_hidden_states = () if output_hidden_states else None + all_entity_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_word_hidden_states = all_word_hidden_states + (word_hidden_states,) + all_entity_hidden_states = all_entity_hidden_states + (entity_hidden_states,) + + layer_outputs = layer_module( + word_hidden_states, + entity_hidden_states, + attention_mask, + output_attentions, + ) + + word_hidden_states = layer_outputs[0] + + if entity_hidden_states is not None: + entity_hidden_states = layer_outputs[1] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[2],) + + if output_hidden_states: + all_word_hidden_states = all_word_hidden_states + (word_hidden_states,) + all_entity_hidden_states = all_entity_hidden_states + (entity_hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + word_hidden_states, + all_word_hidden_states, + all_self_attentions, + entity_hidden_states, + all_entity_hidden_states, + ] + if v is not None + ) + return BaseLukeModelOutput( + last_hidden_state=word_hidden_states, + hidden_states=all_word_hidden_states, + attentions=all_self_attentions, + entity_last_hidden_state=entity_hidden_states, + entity_hidden_states=all_entity_hidden_states, + ) + + +# Copied from transformers.models.bert.modeling_bert.BertPooler +class LukePooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +class EntityPredictionHeadTransform(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.entity_emb_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.entity_emb_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states): + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +class EntityPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.transform = EntityPredictionHeadTransform(config) + self.decoder = nn.Linear(config.entity_emb_size, config.entity_vocab_size, bias=False) + self.bias = nn.Parameter(torch.zeros(config.entity_vocab_size)) + + def forward(self, hidden_states): + hidden_states = self.transform(hidden_states) + hidden_states = self.decoder(hidden_states) + self.bias + + return hidden_states + + +@auto_docstring +class LukePreTrainedModel(PreTrainedModel): + config: LukeConfig + base_model_prefix = "luke" + supports_gradient_checkpointing = True + _no_split_modules = ["LukeAttention", "LukeEntityEmbeddings"] + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + if module.embedding_dim == 1: # embedding for bias parameters + init.zeros_(module.weight) + else: + 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]) + elif isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + + +@auto_docstring( + custom_intro=""" + The bare LUKE model transformer outputting raw hidden-states for both word tokens and entities without any + """ +) +class LukeModel(LukePreTrainedModel): + def __init__(self, config: LukeConfig, add_pooling_layer: bool = True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + self.embeddings = LukeEmbeddings(config) + self.entity_embeddings = LukeEntityEmbeddings(config) + self.encoder = LukeEncoder(config) + + self.pooler = LukePooler(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 + + def get_entity_embeddings(self): + return self.entity_embeddings.entity_embeddings + + def set_entity_embeddings(self, value): + self.entity_embeddings.entity_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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_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, + ) -> tuple | BaseLukeModelOutputWithPooling: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LukeModel + + >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-base") + >>> model = LukeModel.from_pretrained("studio-ousia/luke-base") + # Compute the contextualized entity representation corresponding to the entity mention "Beyoncé" + + >>> text = "Beyoncé lives in Los Angeles." + >>> entity_spans = [(0, 7)] # character-based entity span corresponding to "Beyoncé" + + >>> encoding = tokenizer(text, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt") + >>> outputs = model(**encoding) + >>> word_last_hidden_state = outputs.last_hidden_state + >>> entity_last_hidden_state = outputs.entity_last_hidden_state + # Input Wikipedia entities to obtain enriched contextualized representations of word tokens + + >>> text = "Beyoncé lives in Los Angeles." + >>> entities = [ + ... "Beyoncé", + ... "Los Angeles", + ... ] # Wikipedia entity titles corresponding to the entity mentions "Beyoncé" and "Los Angeles" + >>> entity_spans = [ + ... (0, 7), + ... (17, 28), + ... ] # character-based entity spans corresponding to "Beyoncé" and "Los Angeles" + + >>> encoding = tokenizer( + ... text, entities=entities, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt" + ... ) + >>> outputs = model(**encoding) + >>> word_last_hidden_state = outputs.last_hidden_state + >>> entity_last_hidden_state = outputs.entity_last_hidden_state + ```""" + 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) + if entity_ids is not None: + entity_seq_length = entity_ids.size(1) + if entity_attention_mask is None: + entity_attention_mask = torch.ones((batch_size, entity_seq_length), device=device) + if entity_token_type_ids is None: + entity_token_type_ids = torch.zeros((batch_size, entity_seq_length), dtype=torch.long, device=device) + + # First, compute word embeddings + word_embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + + # Second, compute extended attention mask + extended_attention_mask = self.get_extended_attention_mask(attention_mask, entity_attention_mask) + + # Third, compute entity embeddings and concatenate with word embeddings + if entity_ids is None: + entity_embedding_output = None + else: + entity_embedding_output = self.entity_embeddings(entity_ids, entity_position_ids, entity_token_type_ids) + + # Fourth, send embeddings through the model + encoder_outputs = self.encoder( + word_embedding_output, + entity_embedding_output, + attention_mask=extended_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + # Fifth, get the output. LukeModel outputs the same as BertModel, namely sequence_output of shape (batch_size, seq_len, hidden_size) + sequence_output = encoder_outputs[0] + + # Sixth, we compute the pooled_output, word_sequence_output and entity_sequence_output based on the sequence_output + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseLukeModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + entity_last_hidden_state=encoder_outputs.entity_last_hidden_state, + entity_hidden_states=encoder_outputs.entity_hidden_states, + ) + + def get_extended_attention_mask( + self, word_attention_mask: torch.LongTensor, entity_attention_mask: torch.LongTensor | None + ): + """ + Makes broadcastable attention and causal masks so that future and masked tokens are ignored. + + Arguments: + word_attention_mask (`torch.LongTensor`): + Attention mask for word tokens with ones indicating tokens to attend to, zeros for tokens to ignore. + entity_attention_mask (`torch.LongTensor`, *optional*): + Attention mask for entity tokens with ones indicating tokens to attend to, zeros for tokens to ignore. + + Returns: + `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. + """ + attention_mask = word_attention_mask + if entity_attention_mask is not None: + attention_mask = torch.cat([attention_mask, entity_attention_mask], dim=-1) + + if attention_mask.dim() == 3: + extended_attention_mask = attention_mask[:, None, :, :] + elif attention_mask.dim() == 2: + extended_attention_mask = attention_mask[:, None, None, :] + else: + raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape})") + + extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min + return extended_attention_mask + + +def create_position_ids_from_input_ids(input_ids, padding_idx): + """ + Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols + are ignored. This is modified from fairseq's `utils.make_positions`. + + Args: + x: torch.Tensor x: + + Returns: torch.Tensor + """ + # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA. + mask = input_ids.ne(padding_idx).int() + incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask)) * mask + return incremental_indices.long() + padding_idx + + +# Copied from transformers.models.roberta.modeling_roberta.RobertaLMHead +class LukeLMHead(nn.Module): + """Roberta 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=""" + The LUKE model with a language modeling head and entity prediction head on top for masked language modeling and + masked entity prediction. + """ +) +class LukeForMaskedLM(LukePreTrainedModel): + _tied_weights_keys = { + "entity_predictions.decoder.weight": "luke.entity_embeddings.entity_embeddings.weight", + "lm_head.bias": "lm_head.decoder.bias", + } + + def __init__(self, config): + super().__init__(config) + + self.luke = LukeModel(config) + + self.lm_head = LukeLMHead(config) + self.entity_predictions = EntityPredictionHead(config) + + self.loss_fn = nn.CrossEntropyLoss() + + # 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 + + @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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.LongTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + labels: torch.LongTensor | None = None, + entity_labels: 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, + ) -> tuple | LukeMaskedLMOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + 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]` + entity_labels (`torch.LongTensor` of shape `(batch_size, entity_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.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + loss = None + + mlm_loss = None + logits = self.lm_head(outputs.last_hidden_state) + if labels is not None: + # move labels to correct device + labels = labels.to(logits.device) + mlm_loss = self.loss_fn(logits.view(-1, self.config.vocab_size), labels.view(-1)) + if loss is None: + loss = mlm_loss + + mep_loss = None + entity_logits = None + if outputs.entity_last_hidden_state is not None: + entity_logits = self.entity_predictions(outputs.entity_last_hidden_state) + if entity_labels is not None: + mep_loss = self.loss_fn(entity_logits.view(-1, self.config.entity_vocab_size), entity_labels.view(-1)) + if loss is None: + loss = mep_loss + else: + loss = loss + mep_loss + + if not return_dict: + return tuple( + v + for v in [ + loss, + mlm_loss, + mep_loss, + logits, + entity_logits, + outputs.hidden_states, + outputs.entity_hidden_states, + outputs.attentions, + ] + if v is not None + ) + + return LukeMaskedLMOutput( + loss=loss, + mlm_loss=mlm_loss, + mep_loss=mep_loss, + logits=logits, + entity_logits=entity_logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LUKE model with a classification head on top (a linear layer on top of the hidden state of the first entity + token) for entity classification tasks, such as Open Entity. + """ +) +class LukeForEntityClassification(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.luke = LukeModel(config) + + self.num_labels = config.num_labels + 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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | EntityClassificationOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + labels (`torch.LongTensor` of shape `(batch_size,)` or `(batch_size, num_labels)`, *optional*): + Labels for computing the classification loss. If the shape is `(batch_size,)`, the cross entropy loss is + used for the single-label classification. In this case, labels should contain the indices that should be in + `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, num_labels)`, the binary cross entropy + loss is used for the multi-label classification. In this case, labels should only contain `[0, 1]`, where 0 + and 1 indicate false and true, respectively. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LukeForEntityClassification + + >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-open-entity") + >>> model = LukeForEntityClassification.from_pretrained("studio-ousia/luke-large-finetuned-open-entity") + + >>> text = "Beyoncé lives in Los Angeles." + >>> entity_spans = [(0, 7)] # character-based entity span corresponding to "Beyoncé" + >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") + >>> outputs = model(**inputs) + >>> logits = outputs.logits + >>> predicted_class_idx = logits.argmax(-1).item() + >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) + Predicted class: person + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + feature_vector = outputs.entity_last_hidden_state[:, 0, :] + feature_vector = self.dropout(feature_vector) + logits = self.classifier(feature_vector) + + loss = None + if labels is not None: + # When the number of dimension of `labels` is 1, cross entropy is used as the loss function. The binary + # cross entropy is used otherwise. + # move labels to correct device + labels = labels.to(logits.device) + if labels.ndim == 1: + loss = nn.functional.cross_entropy(logits, labels) + else: + loss = nn.functional.binary_cross_entropy_with_logits(logits.view(-1), labels.view(-1).type_as(logits)) + + if not return_dict: + return tuple( + v + for v in [loss, logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions] + if v is not None + ) + + return EntityClassificationOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LUKE model with a classification head on top (a linear layer on top of the hidden states of the two entity + tokens) for entity pair classification tasks, such as TACRED. + """ +) +class LukeForEntityPairClassification(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.luke = LukeModel(config) + + self.num_labels = config.num_labels + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size * 2, config.num_labels, False) + + # 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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_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, + ) -> tuple | EntityPairClassificationOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + labels (`torch.LongTensor` of shape `(batch_size,)` or `(batch_size, num_labels)`, *optional*): + Labels for computing the classification loss. If the shape is `(batch_size,)`, the cross entropy loss is + used for the single-label classification. In this case, labels should contain the indices that should be in + `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, num_labels)`, the binary cross entropy + loss is used for the multi-label classification. In this case, labels should only contain `[0, 1]`, where 0 + and 1 indicate false and true, respectively. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LukeForEntityPairClassification + + >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-tacred") + >>> model = LukeForEntityPairClassification.from_pretrained("studio-ousia/luke-large-finetuned-tacred") + + >>> text = "Beyoncé lives in Los Angeles." + >>> entity_spans = [ + ... (0, 7), + ... (17, 28), + ... ] # character-based entity spans corresponding to "Beyoncé" and "Los Angeles" + >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") + >>> outputs = model(**inputs) + >>> logits = outputs.logits + >>> predicted_class_idx = logits.argmax(-1).item() + >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) + Predicted class: per:cities_of_residence + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + feature_vector = torch.cat( + [outputs.entity_last_hidden_state[:, 0, :], outputs.entity_last_hidden_state[:, 1, :]], dim=1 + ) + feature_vector = self.dropout(feature_vector) + logits = self.classifier(feature_vector) + + loss = None + if labels is not None: + # When the number of dimension of `labels` is 1, cross entropy is used as the loss function. The binary + # cross entropy is used otherwise. + # move labels to correct device + labels = labels.to(logits.device) + if labels.ndim == 1: + loss = nn.functional.cross_entropy(logits, labels) + else: + loss = nn.functional.binary_cross_entropy_with_logits(logits.view(-1), labels.view(-1).type_as(logits)) + + if not return_dict: + return tuple( + v + for v in [loss, logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions] + if v is not None + ) + + return EntityPairClassificationOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LUKE model with a span classification head on top (a linear layer on top of the hidden states output) for tasks + such as named entity recognition. + """ +) +class LukeForEntitySpanClassification(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.luke = LukeModel(config) + + self.num_labels = config.num_labels + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size * 3, 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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.LongTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + entity_start_positions: torch.LongTensor | None = None, + entity_end_positions: 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, + ) -> tuple | EntitySpanClassificationOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + entity_start_positions (`torch.LongTensor`): + The start positions of entities in the word token sequence. + entity_end_positions (`torch.LongTensor`): + The end positions of entities in the word token sequence. + labels (`torch.LongTensor` of shape `(batch_size, entity_length)` or `(batch_size, entity_length, num_labels)`, *optional*): + Labels for computing the classification loss. If the shape is `(batch_size, entity_length)`, the cross + entropy loss is used for the single-label classification. In this case, labels should contain the indices + that should be in `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, entity_length, + num_labels)`, the binary cross entropy loss is used for the multi-label classification. In this case, + labels should only contain `[0, 1]`, where 0 and 1 indicate false and true, respectively. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LukeForEntitySpanClassification + + >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-conll-2003") + >>> model = LukeForEntitySpanClassification.from_pretrained("studio-ousia/luke-large-finetuned-conll-2003") + + >>> text = "Beyoncé lives in Los Angeles" + # List all possible entity spans in the text + + >>> word_start_positions = [0, 8, 14, 17, 21] # character-based start positions of word tokens + >>> word_end_positions = [7, 13, 16, 20, 28] # character-based end positions of word tokens + >>> entity_spans = [] + >>> for i, start_pos in enumerate(word_start_positions): + ... for end_pos in word_end_positions[i:]: + ... entity_spans.append((start_pos, end_pos)) + + >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") + >>> outputs = model(**inputs) + >>> logits = outputs.logits + >>> predicted_class_indices = logits.argmax(-1).squeeze().tolist() + >>> for span, predicted_class_idx in zip(entity_spans, predicted_class_indices): + ... if predicted_class_idx != 0: + ... print(text[span[0] : span[1]], model.config.id2label[predicted_class_idx]) + Beyoncé PER + Los Angeles LOC + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + hidden_size = outputs.last_hidden_state.size(-1) + + entity_start_positions = entity_start_positions.unsqueeze(-1).expand(-1, -1, hidden_size) + if entity_start_positions.device != outputs.last_hidden_state.device: + entity_start_positions = entity_start_positions.to(outputs.last_hidden_state.device) + start_states = torch.gather(outputs.last_hidden_state, -2, entity_start_positions) + + entity_end_positions = entity_end_positions.unsqueeze(-1).expand(-1, -1, hidden_size) + if entity_end_positions.device != outputs.last_hidden_state.device: + entity_end_positions = entity_end_positions.to(outputs.last_hidden_state.device) + end_states = torch.gather(outputs.last_hidden_state, -2, entity_end_positions) + + feature_vector = torch.cat([start_states, end_states, outputs.entity_last_hidden_state], dim=2) + + feature_vector = self.dropout(feature_vector) + logits = self.classifier(feature_vector) + + loss = None + if labels is not None: + # move labels to correct device + labels = labels.to(logits.device) + # When the number of dimension of `labels` is 2, cross entropy is used as the loss function. The binary + # cross entropy is used otherwise. + if labels.ndim == 2: + loss = nn.functional.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1)) + else: + loss = nn.functional.binary_cross_entropy_with_logits(logits.view(-1), labels.view(-1).type_as(logits)) + + if not return_dict: + return tuple( + v + for v in [loss, logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions] + if v is not None + ) + + return EntitySpanClassificationOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LUKE 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 LukeForSequenceClassification(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.luke = LukeModel(config) + self.dropout = nn.Dropout( + config.classifier_dropout if config.classifier_dropout is not None else 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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | LukeSequenceClassifierOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + 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.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + pooled_output = outputs.pooler_output + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + # move labels to correct device + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + return tuple( + v + for v in [loss, logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions] + if v is not None + ) + + return LukeSequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LUKE Model with a token classification head on top (a linear layer on top of the hidden-states output). To + solve Named-Entity Recognition (NER) task using LUKE, `LukeForEntitySpanClassification` is more suitable than this + class. + """ +) +class LukeForTokenClassification(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.luke = LukeModel(config, add_pooling_layer=False) + self.dropout = nn.Dropout( + config.classifier_dropout if config.classifier_dropout is not None else 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, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | LukeTokenClassifierOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + 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 + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + sequence_output = outputs.last_hidden_state + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + # move labels to correct device + labels = labels.to(logits.device) + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + return tuple( + v + for v in [loss, logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions] + if v is not None + ) + + return LukeTokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class LukeForQuestionAnswering(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + + self.luke = LukeModel(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.FloatTensor | None = None, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_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, + ) -> tuple | LukeQuestionAnsweringModelOutput: + r""" + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + sequence_output = outputs.last_hidden_state + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1) + 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.clamp_(0, ignored_index) + 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: + return tuple( + v + for v in [ + total_loss, + start_logits, + end_logits, + outputs.hidden_states, + outputs.entity_hidden_states, + outputs.attentions, + ] + if v is not None + ) + + return LukeQuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class LukeForMultipleChoice(LukePreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.luke = LukeModel(config) + self.dropout = nn.Dropout( + config.classifier_dropout if config.classifier_dropout is not None else 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, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + entity_ids: torch.LongTensor | None = None, + entity_attention_mask: torch.FloatTensor | None = None, + entity_token_type_ids: torch.LongTensor | None = None, + entity_position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | LukeMultipleChoiceModelOutput: + 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) + entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): + Indices of entity tokens in the entity vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): + Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: + + - 1 for entity tokens that are **not masked**, + - 0 for entity tokens that are **masked**. + entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): + Segment token indices to indicate first and second portions of the entity token inputs. Indices are + selected in `[0, 1]`: + + - 0 corresponds to a *portion A* entity token, + - 1 corresponds to a *portion B* entity token. + entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): + Indices of positions of each input entity in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + 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 + ) + + entity_ids = entity_ids.view(-1, entity_ids.size(-1)) if entity_ids is not None else None + entity_attention_mask = ( + entity_attention_mask.view(-1, entity_attention_mask.size(-1)) + if entity_attention_mask is not None + else None + ) + entity_token_type_ids = ( + entity_token_type_ids.view(-1, entity_token_type_ids.size(-1)) + if entity_token_type_ids is not None + else None + ) + entity_position_ids = ( + entity_position_ids.view(-1, entity_position_ids.size(-2), entity_position_ids.size(-1)) + if entity_position_ids is not None + else None + ) + + outputs = self.luke( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + entity_ids=entity_ids, + entity_attention_mask=entity_attention_mask, + entity_token_type_ids=entity_token_type_ids, + entity_position_ids=entity_position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + pooled_output = outputs.pooler_output + + 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: + # move labels to correct device + labels = labels.to(reshaped_logits.device) + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + return tuple( + v + for v in [ + loss, + reshaped_logits, + outputs.hidden_states, + outputs.entity_hidden_states, + outputs.attentions, + ] + if v is not None + ) + + return LukeMultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=outputs.hidden_states, + entity_hidden_states=outputs.entity_hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "LukeForEntityClassification", + "LukeForEntityPairClassification", + "LukeForEntitySpanClassification", + "LukeForMultipleChoice", + "LukeForQuestionAnswering", + "LukeForSequenceClassification", + "LukeForTokenClassification", + "LukeForMaskedLM", + "LukeModel", + "LukePreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/luke/tokenization_luke.py b/third_party/transformers/src/transformers/models/luke/tokenization_luke.py new file mode 100644 index 0000000000000000000000000000000000000000..bed415c79291d87a0e8eaaa0cc976753a9242756 --- /dev/null +++ b/third_party/transformers/src/transformers/models/luke/tokenization_luke.py @@ -0,0 +1,1664 @@ +# Copyright Studio-Ouisa and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for LUKE.""" + +import itertools +import json +from collections.abc import Mapping + +import numpy as np +from tokenizers import Tokenizer, decoders, pre_tokenizers +from tokenizers.models import BPE + +from ...tokenization_python import PreTrainedTokenizer +from ...tokenization_utils_base import ( + ENCODE_KWARGS_DOCSTRING, + AddedToken, + BatchEncoding, + EncodedInput, + PaddingStrategy, + TensorType, + TextInput, + TextInputPair, + TruncationStrategy, + to_py_obj, +) +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import add_end_docstrings, is_torch_tensor, logging + + +logger = logging.get_logger(__name__) + +EntitySpan = tuple[int, int] +EntitySpanInput = list[EntitySpan] +Entity = str +EntityInput = list[Entity] + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", + "entity_vocab_file": "entity_vocab.json", +} + + +ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) + + - **token_type_ids** -- List of token type ids to be fed to a model (when `return_token_type_ids=True` or + if *"token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) + + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) + + - **entity_ids** -- List of entity ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) + + - **entity_position_ids** -- List of entity positions in the input sequence to be fed to a model. + + - **entity_token_type_ids** -- List of entity token type ids to be fed to a model (when + `return_token_type_ids=True` or if *"entity_token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) + + - **entity_attention_mask** -- List of indices specifying which entities should be attended to by the model + (when `return_attention_mask=True` or if *"entity_attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) + + - **entity_start_positions** -- List of the start positions of entities in the word token sequence (when + `task="entity_span_classification"`). + - **entity_end_positions** -- List of the end positions of entities in the word token sequence (when + `task="entity_span_classification"`). + - **overflowing_tokens** -- List of overflowing tokens sequences (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **num_truncated_tokens** -- Number of tokens truncated (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **special_tokens_mask** -- List of 0s and 1s, with 1 specifying added special tokens and 0 specifying + regular sequence tokens (when `add_special_tokens=True` and `return_special_tokens_mask=True`). + - **length** -- The length of the inputs (when `return_length=True`) + +""" + + +class LukeTokenizer(TokenizersBackend): + """ + Constructs a LUKE tokenizer, derived from the GPT-2 tokenizer, using byte-level Byte-Pair-Encoding. + + This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will + be encoded differently whether it is at the beginning of the sentence (without space) or not: + + ```python + >>> from transformers import LukeTokenizer + + >>> tokenizer = LukeTokenizer.from_pretrained("studio-ousia/luke-base") + >>> tokenizer("Hello world")["input_ids"] + [0, 31414, 232, 2] + + >>> tokenizer(" Hello world")["input_ids"] + [0, 20920, 232, 2] + ``` + + You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you + call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance. + + + + When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one). + + + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. It also creates entity sequences, namely + `entity_ids`, `entity_attention_mask`, `entity_token_type_ids`, and `entity_position_ids` to be used by the LUKE + model. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + merges_file (`str`): + Path to the merges file. + vocab (`str` or `dict[str, int]`, *optional*): + Custom vocabulary dictionary. If not provided, the vocabulary is loaded from `vocab_file`. + merges (`str` or `list[str]`, *optional*): + Custom merges list. If not provided, merges are loaded from `merges_file`. + entity_vocab_file (`str`): + Path to the entity vocabulary file. + task (`str`, *optional*): + Task for which you want to prepare sequences. One of `"entity_classification"`, + `"entity_pair_classification"`, or `"entity_span_classification"`. If you specify this argument, the entity + sequence is automatically created based on the given entity span(s). + max_entity_length (`int`, *optional*, defaults to 32): + The maximum length of `entity_ids`. + max_mention_length (`int`, *optional*, defaults to 30): + The maximum number of tokens inside an entity span. + entity_token_1 (`str`, *optional*, defaults to ``): + The special token used to represent an entity span in a word token sequence. This token is only used when + `task` is set to `"entity_classification"` or `"entity_pair_classification"`. + entity_token_2 (`str`, *optional*, defaults to ``): + The special token used to represent an entity span in a word token sequence. This token is only used when + `task` is set to `"entity_pair_classification"`. + errors (`str`, *optional*, defaults to `"replace"`): + Paradigm to follow when decoding bytes to UTF-8. See + [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information. + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + + + + When building a sequence using special tokens, this is not the token that is used for the beginning of + sequence. The token used is the `cls_token`. + + + + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + + + + When building a sequence using special tokens, this is not the token that is used for the end of sequence. + The token used is the `sep_token`. + + + + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + add_prefix_space (`bool`, *optional*, defaults to `False`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. (LUKE tokenizer detect beginning of words by the preceding space). + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + merges: str | list[str] | None = None, + entity_vocab: str | dict | list | None = None, + errors="replace", + bos_token="", + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + add_prefix_space=False, + task=None, + max_entity_length=32, + max_mention_length=30, + entity_token_1="", + entity_token_2="", + entity_unk_token="[UNK]", + entity_pad_token="[PAD]", + entity_mask_token="[MASK]", + entity_mask2_token="[MASK2]", + **kwargs, + ): + self.add_prefix_space = add_prefix_space + + # Handle entity vocab file for backward compatibility + entity_vocab_file = kwargs.pop("entity_vocab_file", None) + if entity_vocab is None and "entity_vocab" in kwargs: + entity_vocab = kwargs.pop("entity_vocab") + + self._vocab = vocab or {} + self._merges = merges or [] + self._tokenizer = Tokenizer( + BPE( + vocab=self._vocab, + merges=self._merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space) + self._tokenizer.decoder = decoders.ByteLevel() + + # Load entity vocab + if entity_vocab is not None: + self.entity_vocab = entity_vocab + elif entity_vocab_file is not None: + with open(entity_vocab_file, encoding="utf-8") as f: + self.entity_vocab = json.load(f) + else: + # If no entity vocab provided, create a minimal one with required special tokens + self.entity_vocab = { + entity_unk_token: 0, + entity_pad_token: 1, + entity_mask_token: 2, + entity_mask2_token: 3, + } + + # Validate entity special tokens + for entity_special_token in [entity_unk_token, entity_pad_token, entity_mask_token, entity_mask2_token]: + if entity_special_token not in self.entity_vocab: + raise ValueError( + f"Specified entity special token `{entity_special_token}` is not found in entity_vocab." + ) + + self.entity_unk_token_id = self.entity_vocab[entity_unk_token] + self.entity_pad_token_id = self.entity_vocab[entity_pad_token] + self.entity_mask_token_id = self.entity_vocab[entity_mask_token] + self.entity_mask2_token_id = self.entity_vocab[entity_mask2_token] + + # Setup task and max_entity_length + self.task = task + if task is None or task == "entity_span_classification": + self.max_entity_length = max_entity_length + elif task == "entity_classification": + self.max_entity_length = 1 + elif task == "entity_pair_classification": + self.max_entity_length = 2 + else: + raise ValueError( + f"Task {task} not supported. Select task from ['entity_classification', 'entity_pair_classification'," + " 'entity_span_classification'] only." + ) + + self.max_mention_length = max_mention_length + + # Add entity tokens to extra_special_tokens + entity_token_1 = ( + AddedToken(entity_token_1, lstrip=False, rstrip=False) + if isinstance(entity_token_1, str) + else entity_token_1 + ) + entity_token_2 = ( + AddedToken(entity_token_2, lstrip=False, rstrip=False) + if isinstance(entity_token_2, str) + else entity_token_2 + ) + # Handle extra/legacy special tokens (v4 hub files compat) + extra_tokens: list[AddedToken | str] = [] + for key in ("extra_special_tokens", "additional_special_tokens"): + for token in kwargs.pop(key, []) or []: + extra_tokens.append(AddedToken(**token) if isinstance(token, dict) else token) + + # Ensure LUKE entity tokens are present exactly once. + seen = {str(token) for token in extra_tokens} + for token in (entity_token_1, entity_token_2): + token_str = str(token) + if token_str not in seen: + extra_tokens.append(token) + seen.add(token_str) + + kwargs["extra_special_tokens"] = extra_tokens + + # Configure default special token behaviors to match LUKE formatting + token_type_ids_pattern = kwargs.setdefault("token_type_ids_pattern", "all_zeros") + special_tokens_pattern = kwargs.setdefault("special_tokens_pattern", "cls_double_sep") + token_type_ids_include_special_tokens = kwargs.setdefault("token_type_ids_include_special_tokens", True) + self.token_type_ids_pattern = token_type_ids_pattern + self.special_tokens_pattern = special_tokens_pattern + self.token_type_ids_include_special_tokens = token_type_ids_include_special_tokens + + # Set clean_up_tokenization_spaces=True by default to match old Python tokenizer behavior + kwargs.setdefault("clean_up_tokenization_spaces", True) + + super().__init__( + errors=errors, + 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, + add_prefix_space=add_prefix_space, + task=task, + max_entity_length=max_entity_length, + max_mention_length=max_mention_length, + entity_token_1=str(entity_token_1), + entity_token_2=str(entity_token_2), + entity_unk_token=entity_unk_token, + entity_pad_token=entity_pad_token, + entity_mask_token=entity_mask_token, + entity_mask2_token=entity_mask2_token, + entity_vocab=entity_vocab if entity_vocab_file is None else None, # Only store if it was passed as data + **kwargs, + ) + + def build_inputs_with_special_tokens( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + return PreTrainedTokenizer.build_inputs_with_special_tokens(self, token_ids_0, token_ids_1) + + def get_special_tokens_mask( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False + ) -> list[int]: + return PreTrainedTokenizer.get_special_tokens_mask( + self, token_ids_0, token_ids_1, already_has_special_tokens=already_has_special_tokens + ) + + def create_token_type_ids_from_sequences( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + return PreTrainedTokenizer.create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1) + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> str: + text = super()._decode( + token_ids, skip_special_tokens=skip_special_tokens, clean_up_tokenization_spaces=False, **kwargs + ) + + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + text = self.clean_up_tokenization(text) + + return text + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | list[TextInput], + text_pair: TextInput | list[TextInput] | None = None, + entity_spans: EntitySpanInput | list[EntitySpanInput] | None = None, + entity_spans_pair: EntitySpanInput | list[EntitySpanInput] | None = None, + entities: EntityInput | list[EntityInput] | None = None, + entities_pair: EntityInput | list[EntityInput] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + max_entity_length: int | None = None, + stride: int = 0, + is_split_into_words: bool | None = False, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # Check for seq2seq parameters that are not supported with entity-aware encoding + if kwargs.get("text_target") is not None or kwargs.get("text_pair_target") is not None: + if entity_spans is not None or entities is not None or self.task is not None: + raise NotImplementedError( + "text_target and text_pair_target are not supported when using entity-aware encoding. " + "Please use the tokenizer without entities for seq2seq tasks." + ) + # Delegate to parent for seq2seq encoding + return super().__call__( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences, depending on the task you want to prepare them for. + + Args: + text (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence must be a string. Note that this + tokenizer does not support tokenization based on pretokenized strings. + text_pair (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence must be a string. Note that this + tokenizer does not support tokenization based on pretokenized strings. + entity_spans (`list[tuple[int, int]]`, `list[list[tuple[int, int]]]`, *optional*): + The sequence or batch of sequences of entity spans to be encoded. Each sequence consists of tuples each + with two integers denoting character-based start and end positions of entities. If you specify + `"entity_classification"` or `"entity_pair_classification"` as the `task` argument in the constructor, + the length of each sequence must be 1 or 2, respectively. If you specify `entities`, the length of each + sequence must be equal to the length of each sequence of `entities`. + entity_spans_pair (`list[tuple[int, int]]`, `list[list[tuple[int, int]]]`, *optional*): + The sequence or batch of sequences of entity spans to be encoded. Each sequence consists of tuples each + with two integers denoting character-based start and end positions of entities. If you specify the + `task` argument in the constructor, this argument is ignored. If you specify `entities_pair`, the + length of each sequence must be equal to the length of each sequence of `entities_pair`. + entities (`list[str]`, `list[list[str]]`, *optional*): + The sequence or batch of sequences of entities to be encoded. Each sequence consists of strings + representing entities, i.e., special entities (e.g., [MASK]) or entity titles of Wikipedia (e.g., Los + Angeles). This argument is ignored if you specify the `task` argument in the constructor. The length of + each sequence must be equal to the length of each sequence of `entity_spans`. If you specify + `entity_spans` without specifying this argument, the entity sequence or the batch of entity sequences + is automatically constructed by filling it with the [MASK] entity. + entities_pair (`list[str]`, `list[list[str]]`, *optional*): + The sequence or batch of sequences of entities to be encoded. Each sequence consists of strings + representing entities, i.e., special entities (e.g., [MASK]) or entity titles of Wikipedia (e.g., Los + Angeles). This argument is ignored if you specify the `task` argument in the constructor. The length of + each sequence must be equal to the length of each sequence of `entity_spans_pair`. If you specify + `entity_spans_pair` without specifying this argument, the entity sequence or the batch of entity + sequences is automatically constructed by filling it with the [MASK] entity. + max_entity_length (`int`, *optional*): + The maximum length of `entity_ids`. + """ + # Input type checking for clearer error + is_valid_single_text = isinstance(text, str) + is_valid_batch_text = isinstance(text, (list, tuple)) and ( + len(text) == 0 or isinstance(text[0], (str, list, tuple)) + ) + if not (is_valid_single_text or is_valid_batch_text): + raise ValueError( + "text input must be of type `str` (single example), `list[str]` (batch), or `list[tuple]` (batch pairs)." + ) + + is_valid_single_text_pair = isinstance(text_pair, str) + is_valid_batch_text_pair = isinstance(text_pair, (list, tuple)) and ( + len(text_pair) == 0 or isinstance(text_pair[0], str) + ) + if not (text_pair is None or is_valid_single_text_pair or is_valid_batch_text_pair): + raise ValueError("text_pair input must be of type `str` (single example) or `list[str]` (batch).") + + is_batched = bool(isinstance(text, (list, tuple))) + + # Convert padding and truncation to strategies + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + if is_batched: + batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text + if entities is None: + batch_entities_or_entities_pairs = None + else: + batch_entities_or_entities_pairs = ( + list(zip(entities, entities_pair)) if entities_pair is not None else entities + ) + + if entity_spans is None: + batch_entity_spans_or_entity_spans_pairs = None + else: + batch_entity_spans_or_entity_spans_pairs = ( + list(zip(entity_spans, entity_spans_pair)) if entity_spans_pair is not None else entity_spans + ) + + return self._batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + batch_entity_spans_or_entity_spans_pairs=batch_entity_spans_or_entity_spans_pairs, + batch_entities_or_entities_pairs=batch_entities_or_entities_pairs, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + max_entity_length=max_entity_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + else: + return self._encode_plus( + text=text, + text_pair=text_pair, + entity_spans=entity_spans, + entity_spans_pair=entity_spans_pair, + entities=entities, + entities_pair=entities_pair, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + max_entity_length=max_entity_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def _encode_plus( + self, + text: TextInput, + text_pair: TextInput | None = None, + entity_spans: EntitySpanInput | None = None, + entity_spans_pair: EntitySpanInput | None = None, + entities: EntityInput | None = None, + entities_pair: EntityInput | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + max_entity_length: int | None = None, + stride: int = 0, + is_split_into_words: bool | None = False, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # If no entities are provided and task doesn't require them, delegate to parent for proper Encoding support + if ( + entity_spans is None + and entity_spans_pair is None + and entities is None + and entities_pair is None + and self.task is None + ): + # Delegate to parent TokenizersBackend which properly handles Encoding objects + return super()._encode_plus( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + if return_offsets_mapping: + raise NotImplementedError("return_offset_mapping is not available when using entity-aware encoding.") + + if is_split_into_words: + raise NotImplementedError("is_split_into_words is not supported in this tokenizer.") + + ( + first_ids, + second_ids, + first_entity_ids, + second_entity_ids, + first_entity_token_spans, + second_entity_token_spans, + ) = self._create_input_sequence( + text=text, + text_pair=text_pair, + entities=entities, + entities_pair=entities_pair, + entity_spans=entity_spans, + entity_spans_pair=entity_spans_pair, + **kwargs, + ) + + # prepare_for_model will create the attention_mask and token_type_ids + return self.prepare_for_model( + first_ids, + pair_ids=second_ids, + entity_ids=first_entity_ids, + pair_entity_ids=second_entity_ids, + entity_token_spans=first_entity_token_spans, + pair_entity_token_spans=second_entity_token_spans, + add_special_tokens=add_special_tokens, + padding=padding_strategy.value, + truncation=truncation_strategy.value, + max_length=max_length, + max_entity_length=max_entity_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + prepend_batch_axis=True, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair], + batch_entity_spans_or_entity_spans_pairs: list[EntitySpanInput] + | list[tuple[EntitySpanInput, EntitySpanInput]] + | None = None, + batch_entities_or_entities_pairs: list[EntityInput] | list[tuple[EntityInput, EntityInput]] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + max_entity_length: int | None = None, + stride: int = 0, + is_split_into_words: bool | None = False, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # If no entities are provided and task doesn't require them, delegate to parent for proper Encoding support + if ( + batch_entity_spans_or_entity_spans_pairs is None + and batch_entities_or_entities_pairs is None + and self.task is None + ): + # Parent's _encode_plus handles batching internally, so we reconstruct text/text_pair + # from batch_text_or_text_pairs and pass to parent's _encode_plus + # Detect if we have pairs + if batch_text_or_text_pairs and isinstance(batch_text_or_text_pairs[0], (tuple, list)): + # We have pairs + texts, text_pairs = zip(*batch_text_or_text_pairs) + texts = list(texts) + text_pairs = list(text_pairs) + else: + # Just texts + texts = batch_text_or_text_pairs + text_pairs = None + + # Delegate to parent TokenizersBackend which properly handles Encoding objects for batches + return super()._encode_plus( + text=texts, + text_pair=text_pairs, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=kwargs.get("split_special_tokens", self.split_special_tokens), + **kwargs, + ) + + if return_offsets_mapping: + raise NotImplementedError("return_offset_mapping is not available when using entity-aware encoding.") + + if is_split_into_words: + raise NotImplementedError("is_split_into_words is not supported in this tokenizer.") + + # input_ids is a list of tuples (one for each example in the batch) + input_ids = [] + entity_ids = [] + entity_token_spans = [] + for index, text_or_text_pair in enumerate(batch_text_or_text_pairs): + if not isinstance(text_or_text_pair, (list, tuple)): + text, text_pair = text_or_text_pair, None + else: + text, text_pair = text_or_text_pair + + entities, entities_pair = None, None + if batch_entities_or_entities_pairs is not None: + entities_or_entities_pairs = batch_entities_or_entities_pairs[index] + if entities_or_entities_pairs: + if isinstance(entities_or_entities_pairs[0], str): + entities, entities_pair = entities_or_entities_pairs, None + else: + entities, entities_pair = entities_or_entities_pairs + + entity_spans, entity_spans_pair = None, None + if batch_entity_spans_or_entity_spans_pairs is not None: + entity_spans_or_entity_spans_pairs = batch_entity_spans_or_entity_spans_pairs[index] + if len(entity_spans_or_entity_spans_pairs) > 0 and isinstance( + entity_spans_or_entity_spans_pairs[0], list + ): + entity_spans, entity_spans_pair = entity_spans_or_entity_spans_pairs + else: + entity_spans, entity_spans_pair = entity_spans_or_entity_spans_pairs, None + + ( + first_ids, + second_ids, + first_entity_ids, + second_entity_ids, + first_entity_token_spans, + second_entity_token_spans, + ) = self._create_input_sequence( + text=text, + text_pair=text_pair, + entities=entities, + entities_pair=entities_pair, + entity_spans=entity_spans, + entity_spans_pair=entity_spans_pair, + **kwargs, + ) + input_ids.append((first_ids, second_ids)) + entity_ids.append((first_entity_ids, second_entity_ids)) + entity_token_spans.append((first_entity_token_spans, second_entity_token_spans)) + + batch_outputs = self._batch_prepare_for_model( + input_ids, + batch_entity_ids_pairs=entity_ids, + batch_entity_token_spans_pairs=entity_token_spans, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + max_entity_length=max_entity_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=return_tensors, + verbose=verbose, + ) + + return BatchEncoding(batch_outputs) + + def _check_entity_input_format(self, entities: EntityInput | None, entity_spans: EntitySpanInput | None): + if not isinstance(entity_spans, list): + raise TypeError("entity_spans should be given as a list") + elif len(entity_spans) > 0 and not isinstance(entity_spans[0], tuple): + raise ValueError( + "entity_spans should be given as a list of tuples containing the start and end character indices" + ) + + if entities is not None: + if not isinstance(entities, list): + raise ValueError("If you specify entities, they should be given as a list") + + if len(entities) > 0 and not isinstance(entities[0], str): + raise ValueError("If you specify entities, they should be given as a list of entity names") + + if len(entities) != len(entity_spans): + raise ValueError("If you specify entities, entities and entity_spans must be the same length") + + def _create_input_sequence( + self, + text: TextInput, + text_pair: TextInput | None = None, + entities: EntityInput | None = None, + entities_pair: EntityInput | None = None, + entity_spans: EntitySpanInput | None = None, + entity_spans_pair: EntitySpanInput | None = None, + **kwargs, + ) -> tuple[list, list, list, list, list, list]: + def get_input_ids(text): + # Use the underlying tokenizer directly to avoid recursion + encoding = self._tokenizer.encode(text, add_special_tokens=False) + return encoding.ids + + def get_input_ids_and_entity_token_spans(text, entity_spans): + if entity_spans is None: + return get_input_ids(text), None + + cur = 0 + input_ids = [] + entity_token_spans = [None] * len(entity_spans) + + split_char_positions = sorted(frozenset(itertools.chain(*entity_spans))) + char_pos2token_pos = {} + + for split_char_position in split_char_positions: + orig_split_char_position = split_char_position + if ( + split_char_position > 0 and text[split_char_position - 1] == " " + ): # whitespace should be prepended to the following token + split_char_position -= 1 + if cur != split_char_position: + input_ids += get_input_ids(text[cur:split_char_position]) + cur = split_char_position + char_pos2token_pos[orig_split_char_position] = len(input_ids) + + input_ids += get_input_ids(text[cur:]) + + entity_token_spans = [ + (char_pos2token_pos[char_start], char_pos2token_pos[char_end]) for char_start, char_end in entity_spans + ] + + return input_ids, entity_token_spans + + first_ids, second_ids = None, None + first_entity_ids, second_entity_ids = None, None + first_entity_token_spans, second_entity_token_spans = None, None + + if self.task is None: + if entity_spans is None: + first_ids = get_input_ids(text) + else: + self._check_entity_input_format(entities, entity_spans) + + first_ids, first_entity_token_spans = get_input_ids_and_entity_token_spans(text, entity_spans) + if entities is None: + first_entity_ids = [self.entity_mask_token_id] * len(entity_spans) + else: + first_entity_ids = [self.entity_vocab.get(entity, self.entity_unk_token_id) for entity in entities] + + if text_pair is not None: + if entity_spans_pair is None: + second_ids = get_input_ids(text_pair) + else: + self._check_entity_input_format(entities_pair, entity_spans_pair) + + second_ids, second_entity_token_spans = get_input_ids_and_entity_token_spans( + text_pair, entity_spans_pair + ) + if entities_pair is None: + second_entity_ids = [self.entity_mask_token_id] * len(entity_spans_pair) + else: + second_entity_ids = [ + self.entity_vocab.get(entity, self.entity_unk_token_id) for entity in entities_pair + ] + + elif self.task == "entity_classification": + if not (isinstance(entity_spans, list) and len(entity_spans) == 1 and isinstance(entity_spans[0], tuple)): + raise ValueError( + "Entity spans should be a list containing a single tuple " + "containing the start and end character indices of an entity" + ) + first_entity_ids = [self.entity_mask_token_id] + first_ids, first_entity_token_spans = get_input_ids_and_entity_token_spans(text, entity_spans) + + # add special tokens to input ids + entity_token_start, entity_token_end = first_entity_token_spans[0] + first_ids = ( + first_ids[:entity_token_end] + [self.extra_special_tokens_ids[0]] + first_ids[entity_token_end:] + ) + first_ids = ( + first_ids[:entity_token_start] + [self.extra_special_tokens_ids[0]] + first_ids[entity_token_start:] + ) + first_entity_token_spans = [(entity_token_start, entity_token_end + 2)] + + elif self.task == "entity_pair_classification": + if not ( + isinstance(entity_spans, list) + and len(entity_spans) == 2 + and isinstance(entity_spans[0], tuple) + and isinstance(entity_spans[1], tuple) + ): + raise ValueError( + "Entity spans should be provided as a list of two tuples, " + "each tuple containing the start and end character indices of an entity" + ) + + head_span, tail_span = entity_spans + first_entity_ids = [self.entity_mask_token_id, self.entity_mask2_token_id] + first_ids, first_entity_token_spans = get_input_ids_and_entity_token_spans(text, entity_spans) + + head_token_span, tail_token_span = first_entity_token_spans + token_span_with_special_token_ids = [ + (head_token_span, self.extra_special_tokens_ids[0]), + (tail_token_span, self.extra_special_tokens_ids[1]), + ] + if head_token_span[0] < tail_token_span[0]: + first_entity_token_spans[0] = (head_token_span[0], head_token_span[1] + 2) + first_entity_token_spans[1] = (tail_token_span[0] + 2, tail_token_span[1] + 4) + token_span_with_special_token_ids.reverse() + else: + first_entity_token_spans[0] = (head_token_span[0] + 2, head_token_span[1] + 4) + first_entity_token_spans[1] = (tail_token_span[0], tail_token_span[1] + 2) + + for (entity_token_start, entity_token_end), special_token_id in token_span_with_special_token_ids: + first_ids = first_ids[:entity_token_end] + [special_token_id] + first_ids[entity_token_end:] + first_ids = first_ids[:entity_token_start] + [special_token_id] + first_ids[entity_token_start:] + + elif self.task == "entity_span_classification": + if not (isinstance(entity_spans, list) and len(entity_spans) > 0 and isinstance(entity_spans[0], tuple)): + raise ValueError( + "Entity spans should be provided as a list of tuples, " + "each tuple containing the start and end character indices of an entity" + ) + + first_ids, first_entity_token_spans = get_input_ids_and_entity_token_spans(text, entity_spans) + first_entity_ids = [self.entity_mask_token_id] * len(entity_spans) + + else: + raise ValueError(f"Task {self.task} not supported") + + return ( + first_ids, + second_ids, + first_entity_ids, + second_entity_ids, + first_entity_token_spans, + second_entity_token_spans, + ) + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def _batch_prepare_for_model( + self, + batch_ids_pairs: list[tuple[list[int], None]], + batch_entity_ids_pairs: list[tuple[list[int] | None, list[int] | None]], + batch_entity_token_spans_pairs: list[tuple[list[tuple[int, int]] | None, list[tuple[int, int]] | None]], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + max_entity_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_length: bool = False, + verbose: bool = True, + ) -> BatchEncoding: + """ + Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It + adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens + + + Args: + batch_ids_pairs: list of tokenized input ids or input ids pairs + batch_entity_ids_pairs: list of entity ids or entity ids pairs + batch_entity_token_spans_pairs: list of entity spans or entity spans pairs + max_entity_length: The maximum length of the entity sequence. + """ + + batch_outputs = {} + for input_ids, entity_ids, entity_token_span_pairs in zip( + batch_ids_pairs, batch_entity_ids_pairs, batch_entity_token_spans_pairs + ): + first_ids, second_ids = input_ids + first_entity_ids, second_entity_ids = entity_ids + first_entity_token_spans, second_entity_token_spans = entity_token_span_pairs + outputs = self.prepare_for_model( + first_ids, + second_ids, + entity_ids=first_entity_ids, + pair_entity_ids=second_entity_ids, + entity_token_spans=first_entity_token_spans, + pair_entity_token_spans=second_entity_token_spans, + add_special_tokens=add_special_tokens, + padding=PaddingStrategy.DO_NOT_PAD.value, # we pad in batch afterward + truncation=truncation_strategy.value, + max_length=max_length, + max_entity_length=max_entity_length, + stride=stride, + pad_to_multiple_of=None, # we pad in batch afterward + padding_side=None, # we pad in batch afterward + return_attention_mask=False, # we pad in batch afterward + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=None, # We convert the whole batch to tensors at the end + prepend_batch_axis=False, + verbose=verbose, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + batch_outputs = self.pad( + batch_outputs, + padding=padding_strategy.value, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors) + + return batch_outputs + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def prepare_for_model( + self, + ids: list[int], + pair_ids: list[int] | None = None, + entity_ids: list[int] | None = None, + pair_entity_ids: list[int] | None = None, + entity_token_spans: list[tuple[int, int]] | None = None, + pair_entity_token_spans: list[tuple[int, int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + max_entity_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + prepend_batch_axis: bool = False, + **kwargs, + ) -> BatchEncoding: + """ + Prepares a sequence of input id, entity id and entity span, or a pair of sequences of inputs ids, entity ids, + entity spans so that it can be used by the model. It adds special tokens, truncates sequences if overflowing + while taking into account the special tokens and manages a moving window (with user defined stride) for + overflowing tokens. Please Note, for *pair_ids* different than `None` and *truncation_strategy = longest_first* + or `True`, it is not possible to return overflowing tokens. Such a combination of arguments will raise an + error. + + Args: + ids (`list[int]`): + Tokenized input ids of the first sequence. + pair_ids (`list[int]`, *optional*): + Tokenized input ids of the second sequence. + entity_ids (`list[int]`, *optional*): + Entity ids of the first sequence. + pair_entity_ids (`list[int]`, *optional*): + Entity ids of the second sequence. + entity_token_spans (`list[tuple[int, int]]`, *optional*): + Entity spans of the first sequence. + pair_entity_token_spans (`list[tuple[int, int]]`, *optional*): + Entity spans of the second sequence. + max_entity_length (`int`, *optional*): + The maximum length of the entity sequence. + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + # Compute lengths + pair = bool(pair_ids is not None) + len_ids = len(ids) + len_pair_ids = len(pair_ids) if pair else 0 + + if return_token_type_ids and not add_special_tokens: + raise ValueError( + "Asking to return token_type_ids while setting add_special_tokens to False " + "results in an undefined behavior. Please set add_special_tokens to True or " + "set return_token_type_ids to None." + ) + if ( + return_overflowing_tokens + and truncation_strategy == TruncationStrategy.LONGEST_FIRST + and pair_ids is not None + ): + raise ValueError( + "Not possible to return overflowing tokens for pair of sequences with the " + "`longest_first`. Please select another truncation strategy than `longest_first`, " + "for instance `only_second` or `only_first`." + ) + + # Load from model defaults + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + encoded_inputs = {} + + # Compute the total size of the returned word encodings + total_len = len_ids + len_pair_ids + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0) + + # Truncation: Handle max sequence length and max_entity_length + overflowing_tokens = [] + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length: + # truncate words up to max_length + ids, pair_ids, overflowing_tokens = self.truncate_sequences( + ids, + pair_ids=pair_ids, + num_tokens_to_remove=total_len - max_length, + truncation_strategy=truncation_strategy, + stride=stride, + ) + + if return_overflowing_tokens: + encoded_inputs["overflowing_tokens"] = overflowing_tokens + encoded_inputs["num_truncated_tokens"] = total_len - max_length + + # Add special tokens + if add_special_tokens: + sequence = self.build_inputs_with_special_tokens(ids, pair_ids) + token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids) + entity_token_offset = 1 # 1 * token + pair_entity_token_offset = len(ids) + 3 # 1 * token & 2 * tokens + else: + sequence = ids + pair_ids if pair else ids + token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else []) + entity_token_offset = 0 + pair_entity_token_offset = len(ids) + + # Build output dictionary + encoded_inputs["input_ids"] = sequence + if return_token_type_ids: + encoded_inputs["token_type_ids"] = token_type_ids + if return_special_tokens_mask: + if add_special_tokens: + encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask(ids, pair_ids) + else: + encoded_inputs["special_tokens_mask"] = [0] * len(sequence) + + # Set max entity length + if not max_entity_length: + max_entity_length = self.max_entity_length + + if entity_ids is not None: + total_entity_len = 0 + num_invalid_entities = 0 + valid_entity_ids = [ent_id for ent_id, span in zip(entity_ids, entity_token_spans) if span[1] <= len(ids)] + valid_entity_token_spans = [span for span in entity_token_spans if span[1] <= len(ids)] + + total_entity_len += len(valid_entity_ids) + num_invalid_entities += len(entity_ids) - len(valid_entity_ids) + + valid_pair_entity_ids, valid_pair_entity_token_spans = None, None + if pair_entity_ids is not None: + valid_pair_entity_ids = [ + ent_id + for ent_id, span in zip(pair_entity_ids, pair_entity_token_spans) + if span[1] <= len(pair_ids) + ] + valid_pair_entity_token_spans = [span for span in pair_entity_token_spans if span[1] <= len(pair_ids)] + total_entity_len += len(valid_pair_entity_ids) + num_invalid_entities += len(pair_entity_ids) - len(valid_pair_entity_ids) + + if num_invalid_entities != 0: + logger.warning( + f"{num_invalid_entities} entities are ignored because their entity spans are invalid due to the" + " truncation of input tokens" + ) + + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and total_entity_len > max_entity_length: + # truncate entities up to max_entity_length + valid_entity_ids, valid_pair_entity_ids, overflowing_entities = self.truncate_sequences( + valid_entity_ids, + pair_ids=valid_pair_entity_ids, + num_tokens_to_remove=total_entity_len - max_entity_length, + truncation_strategy=truncation_strategy, + stride=stride, + ) + valid_entity_token_spans = valid_entity_token_spans[: len(valid_entity_ids)] + if valid_pair_entity_token_spans is not None: + valid_pair_entity_token_spans = valid_pair_entity_token_spans[: len(valid_pair_entity_ids)] + + if return_overflowing_tokens: + encoded_inputs["overflowing_entities"] = overflowing_entities + encoded_inputs["num_truncated_entities"] = total_entity_len - max_entity_length + + final_entity_ids = valid_entity_ids + valid_pair_entity_ids if valid_pair_entity_ids else valid_entity_ids + encoded_inputs["entity_ids"] = list(final_entity_ids) + entity_position_ids = [] + entity_start_positions = [] + entity_end_positions = [] + for token_spans, offset in ( + (valid_entity_token_spans, entity_token_offset), + (valid_pair_entity_token_spans, pair_entity_token_offset), + ): + if token_spans is not None: + for start, end in token_spans: + start += offset + end += offset + position_ids = list(range(start, end))[: self.max_mention_length] + position_ids += [-1] * (self.max_mention_length - end + start) + entity_position_ids.append(position_ids) + entity_start_positions.append(start) + entity_end_positions.append(end - 1) + + encoded_inputs["entity_position_ids"] = entity_position_ids + if self.task == "entity_span_classification": + encoded_inputs["entity_start_positions"] = entity_start_positions + encoded_inputs["entity_end_positions"] = entity_end_positions + + if return_token_type_ids: + encoded_inputs["entity_token_type_ids"] = [0] * len(encoded_inputs["entity_ids"]) + + # Check lengths + self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose) + + # Padding + if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask: + encoded_inputs = self.pad( + encoded_inputs, + max_length=max_length, + max_entity_length=max_entity_length, + padding=padding_strategy.value, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + if return_length: + encoded_inputs["length"] = len(encoded_inputs["input_ids"]) + + batch_outputs = BatchEncoding( + encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis + ) + + return batch_outputs + + def pad( + self, + encoded_inputs: BatchEncoding + | list[BatchEncoding] + | dict[str, EncodedInput] + | dict[str, list[EncodedInput]] + | list[dict[str, EncodedInput]], + padding: bool | str | PaddingStrategy = True, + max_length: int | None = None, + max_entity_length: int | None = None, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + return_tensors: str | TensorType | None = None, + verbose: bool = True, + ) -> BatchEncoding: + """ + Pad a single encoded input or a batch of encoded inputs up to predefined length or to the max sequence length + in the batch. Padding side (left/right) padding token ids are defined at the tokenizer level (with + `self.padding_side`, `self.pad_token_id` and `self.pad_token_type_id`) .. note:: If the `encoded_inputs` passed + are dictionary of numpy arrays or PyTorch tensors the result will use the same type unless + you provide a different tensor type with `return_tensors`. In the case of PyTorch tensors, you will lose the + specific device of your tensors however. + + Args: + encoded_inputs ([`BatchEncoding`], list of [`BatchEncoding`], `dict[str, list[int]]`, `dict[str, list[list[int]]` or `list[dict[str, list[int]]]`): + Tokenized inputs. Can represent one input ([`BatchEncoding`] or `dict[str, list[int]]`) or a batch of + tokenized inputs (list of [`BatchEncoding`], *dict[str, list[list[int]]]* or *list[dict[str, + list[int]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader + collate function. Instead of `list[int]` you can have tensors (numpy arrays, or PyTorch tensors), + see the note above for the return type. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + max_entity_length (`int`, *optional*): + The maximum length of the entity sequence. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. [What are attention + masks?](../glossary#attention-mask) + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + """ + # If we have a list of dicts, let's convert it in a dict of lists + # We do this to allow using this method as a collate_fn function in PyTorch Dataloader + if isinstance(encoded_inputs, (list, tuple)) and isinstance(encoded_inputs[0], Mapping): + # Call .keys() explicitly for compatibility with TensorDict and other Mapping subclasses + encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()} + + # The model's main input name, usually `input_ids`, has be passed for padding + if self.model_input_names[0] not in encoded_inputs: + raise ValueError( + "You should supply an encoding or a list of encodings to this method " + f"that includes {self.model_input_names[0]}, but you provided {list(encoded_inputs.keys())}" + ) + + required_input = encoded_inputs[self.model_input_names[0]] + + if not required_input: + if return_attention_mask: + encoded_inputs["attention_mask"] = [] + return encoded_inputs + + # If we have PyTorch/NumPy tensors/arrays as inputs, we cast them as python objects + # and rebuild them afterwards if no return_tensors is specified + # Note that we lose the specific device the tensor may be on for PyTorch + + first_element = required_input[0] + if isinstance(first_element, (list, tuple)): + # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element. + index = 0 + while len(required_input[index]) == 0: + index += 1 + if index < len(required_input): + first_element = required_input[index][0] + # At this state, if `first_element` is still a list/tuple, it's an empty one so there is nothing to do. + if not isinstance(first_element, (int, list, tuple)): + if is_torch_tensor(first_element): + return_tensors = "pt" if return_tensors is None else return_tensors + elif isinstance(first_element, np.ndarray): + return_tensors = "np" if return_tensors is None else return_tensors + else: + raise ValueError( + f"type of {first_element} unknown: {type(first_element)}. " + "Should be one of a python, numpy, or pytorch object." + ) + + for key, value in encoded_inputs.items(): + encoded_inputs[key] = to_py_obj(value) + + # Convert padding_strategy in PaddingStrategy + padding_strategy, _, max_length, _ = self._get_padding_truncation_strategies( + padding=padding, max_length=max_length, verbose=verbose + ) + + if max_entity_length is None: + max_entity_length = self.max_entity_length + + required_input = encoded_inputs[self.model_input_names[0]] + if required_input and not isinstance(required_input[0], (list, tuple)): + encoded_inputs = self._pad( + encoded_inputs, + max_length=max_length, + max_entity_length=max_entity_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + return BatchEncoding(encoded_inputs, tensor_type=return_tensors) + + batch_size = len(required_input) + if any(len(v) != batch_size for v in encoded_inputs.values()): + raise ValueError("Some items in the output dictionary have a different batch size than others.") + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = max(len(inputs) for inputs in required_input) + max_entity_length = ( + max(len(inputs) for inputs in encoded_inputs["entity_ids"]) if "entity_ids" in encoded_inputs else 0 + ) + padding_strategy = PaddingStrategy.MAX_LENGTH + + batch_outputs = {} + for i in range(batch_size): + inputs = {k: v[i] for k, v in encoded_inputs.items()} + outputs = self._pad( + inputs, + max_length=max_length, + max_entity_length=max_entity_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def _pad( + self, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, + max_entity_length: int | None = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + ) -> dict: + """ + Pad encoded inputs (on left/right and up to predefined length or max length in the batch) + + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`list[int]`) or batch of tokenized inputs (`list[list[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + max_entity_length: The maximum length of the entity sequence. + padding_strategy: PaddingStrategy to use for padding. + + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in self.padding_side: + + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + entities_provided = bool("entity_ids" in encoded_inputs) + + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(encoded_inputs["input_ids"]) + if entities_provided: + max_entity_length = len(encoded_inputs["entity_ids"]) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + if ( + entities_provided + and max_entity_length is not None + and pad_to_multiple_of is not None + and (max_entity_length % pad_to_multiple_of != 0) + ): + max_entity_length = ((max_entity_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and ( + len(encoded_inputs["input_ids"]) != max_length + or (entities_provided and len(encoded_inputs["entity_ids"]) != max_entity_length) + ) + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(encoded_inputs["input_ids"]) + if entities_provided and return_attention_mask and "entity_attention_mask" not in encoded_inputs: + encoded_inputs["entity_attention_mask"] = [1] * len(encoded_inputs["entity_ids"]) + + if needs_to_be_padded: + difference = max_length - len(encoded_inputs["input_ids"]) + padding_side = padding_side if padding_side is not None else self.padding_side + if entities_provided: + entity_difference = max_entity_length - len(encoded_inputs["entity_ids"]) + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = encoded_inputs["attention_mask"] + [0] * difference + if entities_provided: + encoded_inputs["entity_attention_mask"] = ( + encoded_inputs["entity_attention_mask"] + [0] * entity_difference + ) + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = encoded_inputs["token_type_ids"] + [0] * difference + if entities_provided: + encoded_inputs["entity_token_type_ids"] = ( + encoded_inputs["entity_token_type_ids"] + [0] * entity_difference + ) + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference + encoded_inputs["input_ids"] = encoded_inputs["input_ids"] + [self.pad_token_id] * difference + if entities_provided: + encoded_inputs["entity_ids"] = ( + encoded_inputs["entity_ids"] + [self.entity_pad_token_id] * entity_difference + ) + encoded_inputs["entity_position_ids"] = ( + encoded_inputs["entity_position_ids"] + [[-1] * self.max_mention_length] * entity_difference + ) + if self.task == "entity_span_classification": + encoded_inputs["entity_start_positions"] = ( + encoded_inputs["entity_start_positions"] + [0] * entity_difference + ) + encoded_inputs["entity_end_positions"] = ( + encoded_inputs["entity_end_positions"] + [0] * entity_difference + ) + + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] + if entities_provided: + encoded_inputs["entity_attention_mask"] = [0] * entity_difference + encoded_inputs[ + "entity_attention_mask" + ] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [0] * difference + encoded_inputs["token_type_ids"] + if entities_provided: + encoded_inputs["entity_token_type_ids"] = [0] * entity_difference + encoded_inputs[ + "entity_token_type_ids" + ] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs["input_ids"] = [self.pad_token_id] * difference + encoded_inputs["input_ids"] + if entities_provided: + encoded_inputs["entity_ids"] = [self.entity_pad_token_id] * entity_difference + encoded_inputs[ + "entity_ids" + ] + encoded_inputs["entity_position_ids"] = [ + [-1] * self.max_mention_length + ] * entity_difference + encoded_inputs["entity_position_ids"] + if self.task == "entity_span_classification": + encoded_inputs["entity_start_positions"] = [0] * entity_difference + encoded_inputs[ + "entity_start_positions" + ] + encoded_inputs["entity_end_positions"] = [0] * entity_difference + encoded_inputs[ + "entity_end_positions" + ] + else: + raise ValueError("Invalid padding strategy:" + str(padding_side)) + + return encoded_inputs + + +__all__ = ["LukeTokenizer"] diff --git a/third_party/transformers/src/transformers/models/mamba/__init__.py b/third_party/transformers/src/transformers/models/mamba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a8cd57ed3f2bd9ddd52437b42186635279917eeb --- /dev/null +++ b/third_party/transformers/src/transformers/models/mamba/__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_mamba import * + from .modeling_mamba 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/mamba/configuration_mamba.py b/third_party/transformers/src/transformers/models/mamba/configuration_mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..bc91025ead2e7a5b8f2105b52fe0642182843cdb --- /dev/null +++ b/third_party/transformers/src/transformers/models/mamba/configuration_mamba.py @@ -0,0 +1,106 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MAMBA configuration""" + +import math + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="state-spaces/mamba-2.8b") +@strict +class MambaConfig(PreTrainedConfig): + r""" + layer_norm_epsilon (`float`, *optional*, defaults to 1e-05): + The epsilon to use in the layer normalization layers. + expand (`int`, *optional*, defaults to 2): + Expanding factor used to determine the intermediate size. + use_bias (`bool`, *optional*, defaults to `False`): + Whether or not to use bias in ["in_proj", "out_proj"] of the mixer block + use_conv_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use bias in the convolution layer of the mixer block. + residual_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model + rescale_prenorm_residual (`bool`, *optional*, defaults to `False`): + Whether or not to rescale `out_proj` weights when initializing. + use_mambapy (`bool`, *optional*, defaults to `False`): + Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, + the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive + version if memory is limited. + use_associative_scan (`bool`, *optional*, defaults to `True`): + Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive + sequential implementation. The associative scan is only active during `torch.compile` tracing and + requires torch >= 2.9.0. Both paths are tested to produce numerically identical results (see + `test_associative_scan_matches_sequential`). Set to `False` to fall back to the sequential loop. + + Example: + + ```python + >>> from transformers import MambaConfig, MambaModel + + >>> # Initializing a Mamba configuration + >>> configuration = MambaConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = MambaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mamba" + + vocab_size: int = 50280 + hidden_size: int = 768 + state_size: int = 16 + num_hidden_layers: int = 32 + layer_norm_epsilon: float = 1e-5 + pad_token_id: int | None = 0 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 0 + expand: int = 2 + conv_kernel: int = 4 + use_bias: bool = False + use_conv_bias: bool = True + hidden_act: str = "silu" + initializer_range: float = 0.1 + residual_in_fp32: bool = True + time_step_rank: str | int = "auto" + time_step_scale: float = 1.0 + time_step_min: float = 0.001 + time_step_max: float = 0.1 + time_step_init_scheme: str = "random" + time_step_floor: float = 1e-4 + rescale_prenorm_residual: bool = False + use_cache: bool = True + use_mambapy: bool = False + use_associative_scan: bool = True + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + self.intermediate_size = int(self.expand * self.hidden_size) + self.time_step_rank = ( + math.ceil(self.hidden_size / 16) if self.time_step_rank == "auto" else self.time_step_rank + ) + super().__post_init__(**kwargs) + + @property + def layer_types(self): + return ["mamba"] * self.num_hidden_layers + + +__all__ = ["MambaConfig"] diff --git a/third_party/transformers/src/transformers/models/mamba/convert_mamba_ssm_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/mamba/convert_mamba_ssm_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..96dfbbc1d4a7e515939312e71dd29832e2682451 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mamba/convert_mamba_ssm_checkpoint_to_pytorch.py @@ -0,0 +1,151 @@ +# Copyright 2024 state-spaces/mamba org and 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. +"""This script can be used to convert checkpoints provided in the `mamba_ssm` library into the format provided in HuggingFace `transformers`. It depends on the `mamba_ssm` package to be installed.""" + +import argparse +import json +import math + +import torch + +from transformers import AutoTokenizer, MambaConfig, MambaForCausalLM +from transformers.utils import logging +from transformers.utils.import_utils import is_mamba_ssm_available + + +if is_mamba_ssm_available(): + from mamba_ssm.models.config_mamba import MambaConfig as MambaConfigSSM + from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel + + def convert_ssm_config_to_hf_config(config_ssm: MambaConfigSSM) -> MambaConfig: + """Convert a MambaConfig from mamba_ssm to a MambaConfig from transformers.""" + hf_config = MambaConfig() + # Set config hidden size, num hidden layers, and vocab size directly from the original config + hf_config.hidden_size = config_ssm.d_model + hf_config.intermediate_size = config_ssm.d_model * 2 + hf_config.time_step_rank = math.ceil(config_ssm.d_model / 16) + + hf_config.num_hidden_layers = config_ssm.n_layer + vocab_size = config_ssm.vocab_size + pad_vocab_size_multiple = config_ssm.pad_vocab_size_multiple + if (vocab_size % pad_vocab_size_multiple) != 0: + vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple) + hf_config.vocab_size = vocab_size + return hf_config + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def convert_mamba_ssm_checkpoint_to_huggingface_model( + original_state_dict: dict, original_ssm_config_dict: dict +) -> tuple[MambaForCausalLM, AutoTokenizer]: + if not is_mamba_ssm_available(): + raise ImportError( + "Calling convert_mamba_ssm_checkpoint_to_huggingface_model requires the mamba_ssm library to be installed. Please install it with `pip install mamba_ssm`." + ) + original_ssm_config = MambaConfigSSM(**original_ssm_config_dict) + + # Convert mamba_ssm config to huggingface MambaConfig + hf_config = convert_ssm_config_to_hf_config(original_ssm_config) + + # No weights need to be renamed between the two models. + converted_state_dict = original_state_dict + + # Load reshaped state dict into a huggingface model. + hf_model = MambaForCausalLM(hf_config) + tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") + hf_model.load_state_dict(converted_state_dict) + return (hf_model, tokenizer) + + +def validate_converted_model( + original_state_dict: dict, original_ssm_config_dict: dict, hf_model: MambaForCausalLM, tokenizer: AutoTokenizer +) -> None: + """Validate the converted model returns the same output as the original model.""" + torch_device = "cuda" + + original_config = MambaConfigSSM(**original_ssm_config_dict) + original_model = MambaLMHeadModel(original_config).to(torch_device) + original_model.load_state_dict(original_state_dict) + + hf_model = hf_model.to(torch_device) + input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"].to(torch_device) + # Assert model logits are close + with torch.no_grad(): + original_model_logits = original_model(input_ids).logits + hf_model_logits = hf_model(input_ids).logits + if not torch.allclose(original_model_logits, hf_model_logits, atol=1e-3): + raise ValueError("The converted model did not return the same logits as the original model.") + + logger.info("Model conversion validated successfully.") + + +def convert_mamba_checkpoint_file_to_huggingface_model_file( + mamba_checkpoint_path: str, config_json_file: str, output_dir: str +) -> None: + if not is_mamba_ssm_available(): + raise ImportError( + "Calling convert_mamba_checkpoint_file_to_huggingface_model_file requires the mamba_ssm library to be installed. Please install it with `pip install mamba_ssm`." + ) + if not torch.cuda.is_available(): + raise ValueError( + "This script is to be run with a CUDA device, as the original mamba_ssm model does not support cpu." + ) + logger.info(f"Loading model from {mamba_checkpoint_path} based on config from {config_json_file}") + # Load weights and config from paths + original_state_dict = torch.load(mamba_checkpoint_path, map_location="cpu", weights_only=True) + with open(config_json_file, "r", encoding="utf-8") as json_file: + original_ssm_config_dict = json.load(json_file) + + # Convert the model + hf_model, tokenizer = convert_mamba_ssm_checkpoint_to_huggingface_model( + original_state_dict, original_ssm_config_dict + ) + + # Validate the conversion + validate_converted_model(original_state_dict, original_ssm_config_dict, hf_model, tokenizer) + + logger.info(f"Model converted successfully. Saving model to {output_dir}") + + # Save new model to pytorch_dump_path + hf_model.save_pretrained(output_dir) + tokenizer.save_pretrained(output_dir) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-i", + "--mamba_checkpoint_file", + type=str, + required=True, + help="Path to a `pytorch_model.bin` mamba_ssm checkpoint file to be converted.", + ) + parser.add_argument( + "-c", + "--config_json_file", + type=str, + required=True, + help="Path to a `config.json` file corresponding to a MambaConfig of the original mamba_ssm model.", + ) + parser.add_argument( + "-o", "--output_dir", type=str, required=True, help="Path to directory to save the converted output model to." + ) + args = parser.parse_args() + + convert_mamba_checkpoint_file_to_huggingface_model_file( + args.mamba_checkpoint_file, args.config_json_file, args.output_dir + ) diff --git a/third_party/transformers/src/transformers/models/mamba/modeling_mamba.py b/third_party/transformers/src/transformers/models/mamba/modeling_mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..2e5695c1d4fafd885d4c92206f5192bbe9c29b49 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mamba/modeling_mamba.py @@ -0,0 +1,720 @@ +# Copyright 2024 state-spaces/mamba org and 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 MAMBA 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 +from ...generation import GenerationMixin +from ...integrations import lazy_load_kernel +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_utils import PreTrainedModel +from ...utils import ( + ModelOutput, + auto_docstring, + logging, +) +from ...utils.import_utils import ( + is_mambapy_available, + is_torch_greater_or_equal, + is_tracing, + resolve_internal_import, +) +from .configuration_mamba import MambaConfig + + +logger = logging.get_logger(__name__) + +if is_torch_greater_or_equal("2.9.0"): + from torch._higher_order_ops.associative_scan import associative_scan +else: + associative_scan = None + + +if is_mambapy_available(): + from mambapy.pscan import pscan +else: + pscan = None + + +class MambaMixer(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: MambaConfig, layer_idx: int, initialize_mixer_weights: bool = True): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.ssm_state_size = config.state_size + self.conv_kernel_size = config.conv_kernel + self.intermediate_size = config.intermediate_size + self.time_step_rank = int(config.time_step_rank) + self.layer_idx = layer_idx + self.use_conv_bias = config.use_conv_bias + self.conv1d = nn.Conv1d( + in_channels=self.intermediate_size, + out_channels=self.intermediate_size, + bias=config.use_conv_bias, + kernel_size=config.conv_kernel, + groups=self.intermediate_size, + padding=config.conv_kernel - 1, + ) + + self.activation = config.hidden_act + self.act = ACT2FN[config.hidden_act] + + self.use_mambapy = config.use_mambapy + self.use_associative_scan = config.use_associative_scan + + # projection of the input hidden states + self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias) + # selective projection used to make dt, B and C input dependent + self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False) + # time step projection (discretization) + self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True) + + # 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 + self.A_log = nn.Parameter(torch.empty(self.intermediate_size, self.ssm_state_size)) + self.D = nn.Parameter(torch.empty(self.intermediate_size)) + if initialize_mixer_weights and self.dt_proj.weight.device.type != "meta": + self.init_mamba_weights() + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) + self.use_bias = config.use_bias + + global causal_conv1d, 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 mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn + 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" + ) + selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) + mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None) + + self.warn_slow_implementation() + + @torch.no_grad() + def init_mamba_weights(self): + A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32, device=self.A_log.device)[None, :] + A = A.expand(self.intermediate_size, -1).contiguous() + init.copy_(self.A_log, torch.log(A)) + init.ones_(self.D) + + dt_init_std = self.config.time_step_rank**-0.5 * self.config.time_step_scale + if self.config.time_step_init_scheme == "constant": + init.constant_(self.dt_proj.weight, dt_init_std) + elif self.config.time_step_init_scheme == "random": + init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std) + + dt = torch.exp( + torch.rand(self.intermediate_size, device=self.dt_proj.bias.device, dtype=torch.float32) + * (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)) + init.copy_(self.dt_proj.bias, inv_dt) + + def warn_slow_implementation(self): + is_fast_path_available = all( + (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn) + ) + if not is_fast_path_available: + if self.use_mambapy: + if is_mambapy_available(): + logger.warning_once( + "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" + " is None. Falling back to the mamba.py backend. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and" + " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d" + ) + else: + raise ImportError( + "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py." + ) + else: + logger.warning_once( + "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" + " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and" + " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d. For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py." + ) + + def cuda_kernels_forward( + self, + hidden_states: torch.Tensor, + cache_params: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + ): + # 1. Gated MLP's linear projection + projected_states = self.in_proj(hidden_states).transpose(1, 2) + + if self.training and cache_params is None: # Doesn't support outputting the states -> used for training + contextualized_states = mamba_inner_fn( + projected_states, + self.conv1d.weight, + self.conv1d.bias if self.use_conv_bias else None, + self.x_proj.weight, + self.dt_proj.weight, + self.out_proj.weight, + self.out_proj.bias.float() if self.use_bias else None, + -torch.exp(self.A_log.float()), + None, # input-dependent B + None, # input-dependent C + self.D.float(), + delta_bias=self.dt_proj.bias.float(), + delta_softplus=True, + ) + + else: + hidden_states, gate = projected_states.chunk(2, dim=1) + + if attention_mask is not None: + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + is_decoding = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + + # 2. Convolution sequence transformation + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) + if is_decoding: + hidden_states = causal_conv1d_update( + hidden_states.squeeze(-1), + cache_params.layers[self.layer_idx].conv_states, + conv_weights, + self.conv1d.bias, + self.activation, + ) + hidden_states = hidden_states.unsqueeze(-1) + else: + if cache_params is not None: + conv_states = nn.functional.pad( + hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0) + ) + cache_params.update_conv_state(conv_states, self.layer_idx) + hidden_states = causal_conv1d_fn( + hidden_states, conv_weights, self.conv1d.bias, activation=self.activation + ) + + if attention_mask is not None: + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 3. State Space Model sequence transformation + # 3.a. input varying initialization of time_step, B and C + ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) + time_step, B, C = torch.split( + ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + ) + discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + + A = -torch.exp(self.A_log.float()) + # 3.c perform the recurrence y ← SSM(A, B, C)(x) + time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None + if is_decoding: + scan_outputs = selective_state_update( + cache_params.layers[self.layer_idx].recurrent_states, + hidden_states[..., 0], + discrete_time_step[..., 0], + A, + B[:, 0], + C[:, 0], + self.D, + gate[..., 0], + time_proj_bias, + dt_softplus=True, + ).unsqueeze(-1) + else: + scan_outputs, ssm_state = selective_scan_fn( + hidden_states, + discrete_time_step, + A, + B.transpose(1, 2), + C.transpose(1, 2), + self.D.float(), + gate, + time_proj_bias, + delta_softplus=True, + return_last_state=True, + ) + if ssm_state is not None and cache_params is not None: + cache_params.update_recurrent_state(ssm_state, self.layer_idx) + + # 4. Final linear projection + contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) + return contextualized_states + + # fmt: off + def slow_forward(self, input_states, cache_params: Cache | None=None, attention_mask: torch.LongTensor | None = None): + batch_size, seq_len, _ = input_states.shape + dtype = input_states.dtype + # 1. Gated MLP's linear projection + projected_states = self.in_proj(input_states).transpose(1, 2) # [batch, 2 * intermediate_size, seq_len] + hidden_states, gate = projected_states.chunk(2, dim=1) + + if attention_mask is not None: + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + if cache_params is not None and cache_params.has_previous_state(self.layer_idx): + ssm_state = cache_params.layers[self.layer_idx].recurrent_states.clone() + else: + ssm_state = torch.zeros( + (batch_size, self.intermediate_size, self.ssm_state_size), + device=hidden_states.device, dtype=dtype + ) + + # 2. Convolution sequence transformation + if cache_params is not None: + if not cache_params.has_previous_state(self.layer_idx): + conv_state = nn.functional.pad( + hidden_states, + (self.conv_kernel_size - hidden_states.shape[-1], 0) + ) + + cache_params.update_conv_state(conv_state, self.layer_idx) + hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len] + else: + conv_state = cache_params.update_conv_state(hidden_states, self.layer_idx) + conv_state = conv_state.to(self.conv1d.weight.device) + 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).unsqueeze(-1) # [batch, intermediate_size, 1] : decoding + else: + hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len] + + if attention_mask is not None: + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 3. State Space Model sequence transformation + # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] + ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) + time_step, B, C = torch.split( + ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + ) + discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size] + discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) # [batch, intermediate_size, seq_len] + + # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) + A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size] + discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) # [batch, intermediate_size, seq_len, ssm_state_size] + discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() # [batch, intermediate_size, seq_len, ssm_state_size] + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + # 3.c perform the recurrence y ← SSM(A, B, C)(x) + if self.use_mambapy and self.training and cache_params is None: + hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) # [batch, seq_len, intermediate_size, ssm_state_size] + + scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len] + scan_output = scan_output + hidden_states * self.D[None, :, None] + scan_output = scan_output * self.act(gate) + else: + # Use associative_scan for parallel computation when available + if self.use_associative_scan and associative_scan is not None and is_tracing(hidden_states) and cache_params is None: + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return (a_left * a_right, a_right * b_left + b_right) + + combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" + _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode) + # all_h: [B, D, S, N] -> output: [B, D, S] + scan_output = torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1) + ssm_state = all_h[:, :, -1, :] + else: + # Sequential loop for decoding or when associative_scan unavailable + scan_outputs = [] + for i in range(seq_len): + ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] # [batch, intermediate_size, ssm_state] + scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) # [batch, intermediate_size, 1] + scan_outputs.append(scan_output[:, :, 0]) + scan_output = torch.stack(scan_outputs, dim=-1) # [batch, intermediate_size, seq_len] + + scan_output = scan_output + (hidden_states * self.D[None, :, None]) + scan_output = (scan_output * self.act(gate)) + + if cache_params is not None: + cache_params.update_recurrent_state(ssm_state, self.layer_idx) + + # 4. Final linear projection + contextualized_states = self.out_proj(scan_output.transpose(1, 2)) # [batch, seq_len, hidden_size] + return contextualized_states + # fmt: on + + def forward( + self, + hidden_states, + cache_params: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + **kwargs, + ): + is_fast_path_available = all( + (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn) + ) + if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states): + return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) + return self.slow_forward(hidden_states, cache_params, attention_mask) + + +class MambaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MambaRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{self.weight.shape[0]}, eps={self.variance_epsilon}" + + +class MambaBlock(GradientCheckpointingLayer): + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.residual_in_fp32 = config.residual_in_fp32 + self.norm = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.mixer = MambaMixer(config, layer_idx=layer_idx, initialize_mixer_weights=False) + + def forward( + self, + hidden_states, + cache_params: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + **kwargs, + ): + residual = hidden_states + hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) + if self.residual_in_fp32: + residual = residual.to(torch.float32) + + hidden_states = self.mixer(hidden_states, cache_params=cache_params, attention_mask=attention_mask) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class MambaPreTrainedModel(PreTrainedModel): + config: MambaConfig + base_model_prefix = "backbone" + _no_split_modules = ["MambaBlock", "MambaMixer"] + supports_gradient_checkpointing = True + _is_stateful = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + std = self.config.initializer_range + if isinstance(module, MambaMixer): + # 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 + module.init_mamba_weights() + + init.kaiming_uniform_(module.conv1d.weight, a=math.sqrt(5)) + if module.conv1d.bias is not None: + init.zeros_(module.conv1d.bias) + init.kaiming_uniform_(module.out_proj.weight, a=math.sqrt(5)) + + 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 + # 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 + p = module.out_proj.weight + p /= math.sqrt(self.config.num_hidden_layers) + + if isinstance(module, nn.Linear): + init.normal_(module.weight, std=std) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, MambaRMSNorm): + init.ones_(module.weight) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, std=std) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for the MAMBA model outputs. + """ +) +class MambaOutput(ModelOutput): + r""" + cache_params (`Cache`): + The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to + avoid providing the old `input_ids`. + + Includes both the State space model state matrices after the selective scan, and the Convolutional states + """ + + last_hidden_state: torch.FloatTensor | None = None + cache_params: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for causal language model (or autoregressive) outputs. + """ +) +class MambaCausalLMOutput(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). + cache_params (`Cache`): + The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to + avoid providing the old `input_ids`. + + Includes both the State space model state matrices after the selective scan, and the Convolutional states + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + cache_params: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + + +@auto_docstring +class MambaModel(MambaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([MambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.gradient_checkpointing = False + self.norm_f = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + # Initialize weights and apply final processing + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + for k in state_dict: + if "embedding." in k: + state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) + break + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.LongTensor | None = None, + cache_params: Cache | None = None, + use_cache: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + attention_mask: torch.LongTensor | None = None, + **kwargs, + ) -> tuple | MambaOutput: + r""" + cache_params (`Cache`, *optional*): + If passed along, the model uses the previous state in all the blocks (which will give the output for the + `input_ids` provided as if the model add `state_input_ids + input_ids` as context). + use_cache (`bool`, *optional*): + If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits. + """ + 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 if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): # ^ 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 self.gradient_checkpointing and self.training and use_cache: + use_cache = False + + if use_cache and cache_params is None: + cache_params = DynamicCache(config=self.config) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + for mixer_block in self.layers: + hidden_states = mixer_block( + hidden_states, + cache_params=cache_params, + attention_mask=attention_mask, + ) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None) + + return MambaOutput( + last_hidden_state=hidden_states, + cache_params=cache_params if use_cache else None, + hidden_states=all_hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + The MAMBA Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """ +) +class MambaForCausalLM(MambaPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "backbone.embeddings.weight"} + + def __init__(self, config): + super().__init__(config) + self.backbone = MambaModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.backbone.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.backbone.set_input_embeddings(new_embeddings) + + def prepare_inputs_for_generation( + self, + input_ids, + inputs_embeds=None, + use_cache=None, + cache_params: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + is_first_iteration: bool | None = False, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + cache_params=cache_params, + attention_mask=attention_mask, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if use_cache and not is_first_iteration: + model_inputs["attention_mask"] = None + + return model_inputs + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + cache_params: Cache | None = None, + labels: torch.LongTensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, # for now we need this for generation + ) -> tuple | MambaCausalLMOutput: + r""" + cache_params (`Cache`, *optional*): + If passed along, the model uses the previous state in all the blocks (which will give the output for the + `input_ids` provided as if the model add `state_input_ids + input_ids` as context). + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + use_cache (`bool`, *optional*): + If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + mamba_outputs = self.backbone( + input_ids, + cache_params=cache_params, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + use_cache=use_cache, + attention_mask=attention_mask, + ) + + hidden_states = mamba_outputs[0] + # Only compute necessary logits + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :].to(self.lm_head.weight.dtype)).float() + + loss = None + if labels is not None: + # move labels to correct device + labels = labels.to(logits.device) + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + if not return_dict: + output = (logits,) + mamba_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return MambaCausalLMOutput( + loss=loss, + logits=logits, + cache_params=mamba_outputs.cache_params, + hidden_states=mamba_outputs.hidden_states, + ) + + +__all__ = ["MambaForCausalLM", "MambaModel", "MambaPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/mask2former/__init__.py b/third_party/transformers/src/transformers/models/mask2former/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c37a74a766325896ff4577b3df32cbbd4a335382 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/__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_mask2former import * + from .image_processing_mask2former import * + from .image_processing_pil_mask2former import * + from .modeling_mask2former 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/mask2former/configuration_mask2former.py b/third_party/transformers/src/transformers/models/mask2former/configuration_mask2former.py new file mode 100644 index 0000000000000000000000000000000000000000..3be9b5d681bf64764987aee07b7a508cd9edf5c5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/configuration_mask2former.py @@ -0,0 +1,134 @@ +# Copyright 2022 Meta Platforms, Inc.and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Mask2Former 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, logging +from ..auto import AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="facebook/mask2former-swin-small-coco-instance") +@strict +class Mask2FormerConfig(PreTrainedConfig): + r""" + feature_size (`int`, *optional*, defaults to 256): + The features (channels) of the resulting feature maps. + mask_feature_size (`int`, *optional*, defaults to 256): + The masks' features size, this value will also be used to specify the Feature Pyramid Network features' + size. + encoder_feedforward_dim (`int`, *optional*, defaults to 1024): + Dimension of feedforward network for deformable detr encoder used as part of pixel decoder. + dim_feedforward (`int`, *optional*, defaults to 2048): + Feature dimension in feedforward network for transformer decoder. + pre_norm (`bool`, *optional*, defaults to `False`): + Whether to use pre-LayerNorm or not for transformer decoder. + enforce_input_projection (`bool`, *optional*, defaults to `False`): + Whether to add an input projection 1x1 convolution even if the input channels and hidden dim are identical + in the Transformer decoder. + common_stride (`int`, *optional*, defaults to 4): + Parameter used for determining number of FPN levels used as part of pixel decoder. + ignore_value (`int`, *optional*, defaults to 255): + Category id to be ignored during training. + num_queries (`int`, *optional*, defaults to 100): + Number of queries for the decoder. + train_num_points (`str` or `function`, *optional*, defaults to 12544): + Number of points used for sampling during loss calculation. + oversample_ratio (`float`, *optional*, defaults to 3.0): + Oversampling parameter used for calculating no. of sampled points + importance_sample_ratio (`float`, *optional*, defaults to 0.75): + Ratio of points that are sampled via importance sampling. + feature_strides (`list[int]`, *optional*, defaults to `[4, 8, 16, 32]`): + Feature strides corresponding to features generated from backbone network. + output_auxiliary_logits (`bool`, *optional*): + Should the model output its `auxiliary_logits` or not. + + Examples: + + ```python + >>> from transformers import Mask2FormerConfig, Mask2FormerModel + + >>> # Initializing a Mask2Former facebook/mask2former-swin-small-coco-instance configuration + >>> configuration = Mask2FormerConfig() + + >>> # Initializing a model (with random weights) from the facebook/mask2former-swin-small-coco-instance style configuration + >>> model = Mask2FormerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + + """ + + model_type = "mask2former" + sub_configs = {"backbone_config": AutoConfig} + backbones_supported = ["swin"] + attribute_map = {"hidden_size": "hidden_dim", "num_hidden_layers": "decoder_layers"} + + backbone_config: dict | PreTrainedConfig | None = None + feature_size: int = 256 + mask_feature_size: int = 256 + hidden_dim: int = 256 + encoder_feedforward_dim: int = 1024 + activation_function: str = "relu" + encoder_layers: int = 6 + decoder_layers: int = 10 + num_attention_heads: int = 8 + dropout: float | int = 0.0 + dim_feedforward: int = 2048 + pre_norm: bool = False + enforce_input_projection: bool = False + common_stride: int = 4 + ignore_value: int = 255 + num_queries: int = 100 + no_object_weight: float = 0.1 + class_weight: float = 2.0 + mask_weight: float = 5.0 + dice_weight: float = 5.0 + train_num_points: int = 12544 + oversample_ratio: float = 3.0 + importance_sample_ratio: float = 0.75 + init_std: float = 0.02 + init_xavier_std: float = 1.0 + use_auxiliary_loss: bool = True + feature_strides: list[int] | tuple[int, ...] = (4, 8, 16, 32) + output_auxiliary_logits: bool | None = None + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="swin", + default_config_kwargs={ + "depths": [2, 2, 18, 2], + "drop_path_rate": 0.3, + "out_features": ["stage1", "stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + if self.backbone_config.model_type not in self.backbones_supported: + logger.warning_once( + f"Backbone {self.backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. " + f"Supported model types: {','.join(self.backbones_supported)}" + ) + + super().__post_init__(**kwargs) + + +__all__ = ["Mask2FormerConfig"] diff --git a/third_party/transformers/src/transformers/models/mask2former/convert_mask2former_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/mask2former/convert_mask2former_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b51ae601263e3de6156354e482e00bb8482709 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/convert_mask2former_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,1020 @@ +# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +import sys +from argparse import ArgumentParser +from collections.abc import Iterator +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from pprint import pformat +from typing import Any + +import httpx +import torch +import torchvision.transforms as T +from detectron2.checkpoint import DetectionCheckpointer +from detectron2.config import get_cfg +from detectron2.projects.deeplab import add_deeplab_config +from huggingface_hub import hf_hub_download +from PIL import Image +from torch import Tensor, nn + +from transformers import ( + Mask2FormerConfig, + Mask2FormerForUniversalSegmentation, + Mask2FormerImageProcessor, + Mask2FormerModel, + SwinConfig, +) +from transformers.models.mask2former.modeling_mask2former import ( + Mask2FormerForUniversalSegmentationOutput, + Mask2FormerModelOutput, +) +from transformers.utils import logging + + +StateDict = dict[str, Tensor] + +logging.set_verbosity_info() +logger = logging.get_logger() + +torch.manual_seed(0) + + +class TrackedStateDict: + def __init__(self, to_track: dict): + """This class "tracks" a python dictionary by keeping track of which item is accessed. + + Args: + to_track (Dict): The dictionary we wish to track + """ + self.to_track = to_track + self._seen: set[str] = set() + + def __getitem__(self, key: str) -> Any: + return self.to_track[key] + + def __setitem__(self, key: str, item: Any): + self._seen.add(key) + self.to_track[key] = item + + def diff(self) -> list[str]: + """This method returns a set difference between the keys in the tracked state dict and the one we have access so far. + This is an effective method to check if we have update all the keys + + Returns: + list[str]: List of keys not yet updated + """ + return set(self.to_track.keys()) - self._seen + + def copy(self) -> dict: + # proxy the call to the internal dictionary + return self.to_track.copy() + + +# 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 + + +@dataclass +class Args: + """Fake command line arguments needed by mask2former/detectron implementation""" + + config_file: str + + +def setup_cfg(args: Args): + # load config from file and command-line arguments + cfg = get_cfg() + add_deeplab_config(cfg) + add_maskformer2_config(cfg) + cfg.merge_from_file(args.config_file) + cfg.freeze() + return cfg + + +class OriginalMask2FormerConfigToOursConverter: + def __call__(self, original_config: object) -> Mask2FormerConfig: + model = original_config.MODEL + + repo_id = "huggingface/label-files" + if model.SEM_SEG_HEAD.NUM_CLASSES == 847: + filename = "mask2former-ade20k-full-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 150: + filename = "ade20k-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 80: + filename = "coco-detection-mmdet-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 171: + filename = "mask2former-coco-stuff-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 133: + filename = "coco-panoptic-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 19: + filename = "cityscapes-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 8: + filename = "cityscapes-instance-id2label.json" + elif model.SEM_SEG_HEAD.NUM_CLASSES == 65: + filename = "mapillary-vistas-id2label.json" + + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + label2id = {label: idx for idx, label in id2label.items()} + + if model.SWIN.EMBED_DIM == 96: + backbone_config = SwinConfig.from_pretrained( + "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"] + ) + elif model.SWIN.EMBED_DIM == 128: + backbone_config = SwinConfig( + embed_dim=128, + window_size=12, + depths=(2, 2, 18, 2), + num_heads=(4, 8, 16, 32), + out_features=["stage1", "stage2", "stage3", "stage4"], + ) + + elif model.SWIN.EMBED_DIM == 192: + backbone_config = SwinConfig.from_pretrained( + "microsoft/swin-large-patch4-window12-384", out_features=["stage1", "stage2", "stage3", "stage4"] + ) + else: + raise ValueError(f"embed dim {model.SWIN.EMBED_DIM} not supported for Swin!") + + backbone_config.drop_path_rate = model.SWIN.DROP_PATH_RATE + backbone_config.attention_probs_dropout_prob = model.SWIN.ATTN_DROP_RATE + backbone_config.depths = model.SWIN.DEPTHS + + config: Mask2FormerConfig = Mask2FormerConfig( + ignore_value=model.SEM_SEG_HEAD.IGNORE_VALUE, + num_labels=model.SEM_SEG_HEAD.NUM_CLASSES, + num_queries=model.MASK_FORMER.NUM_OBJECT_QUERIES, + no_object_weight=model.MASK_FORMER.NO_OBJECT_WEIGHT, + class_weight=model.MASK_FORMER.CLASS_WEIGHT, + mask_weight=model.MASK_FORMER.MASK_WEIGHT, + dice_weight=model.MASK_FORMER.DICE_WEIGHT, + train_num_points=model.MASK_FORMER.TRAIN_NUM_POINTS, + oversample_ratio=model.MASK_FORMER.OVERSAMPLE_RATIO, + importance_sample_ratio=model.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO, + init_std=0.02, + init_xavier_std=1.0, + use_auxiliary_loss=model.MASK_FORMER.DEEP_SUPERVISION, + feature_strides=[4, 8, 16, 32], + backbone_config=backbone_config, + id2label=id2label, + label2id=label2id, + feature_size=model.SEM_SEG_HEAD.CONVS_DIM, + mask_feature_size=model.SEM_SEG_HEAD.MASK_DIM, + hidden_dim=model.MASK_FORMER.HIDDEN_DIM, + encoder_layers=model.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS, + encoder_feedforward_dim=1024, + decoder_layers=model.MASK_FORMER.DEC_LAYERS, + num_attention_heads=model.MASK_FORMER.NHEADS, + dropout=model.MASK_FORMER.DROPOUT, + dim_feedforward=model.MASK_FORMER.DIM_FEEDFORWARD, + pre_norm=model.MASK_FORMER.PRE_NORM, + enforce_input_proj=model.MASK_FORMER.ENFORCE_INPUT_PROJ, + common_stride=model.SEM_SEG_HEAD.COMMON_STRIDE, + ) + return config + + +class OriginalMask2FormerConfigToImageProcessorConverter: + def __call__(self, original_config: object) -> Mask2FormerImageProcessor: + model = original_config.MODEL + model_input = original_config.INPUT + + return Mask2FormerImageProcessor( + image_mean=(torch.tensor(model.PIXEL_MEAN) / 255).tolist(), + image_std=(torch.tensor(model.PIXEL_STD) / 255).tolist(), + size=model_input.MIN_SIZE_TEST, + max_size=model_input.MAX_SIZE_TEST, + num_labels=model.SEM_SEG_HEAD.NUM_CLASSES, + ignore_index=model.SEM_SEG_HEAD.IGNORE_VALUE, + size_divisibility=32, + ) + + +class OriginalMask2FormerCheckpointToOursConverter: + def __init__(self, original_model: nn.Module, config: Mask2FormerConfig): + self.original_model = original_model + self.config = config + + def pop_all(self, renamed_keys: list[tuple[str, str]], dst_state_dict: StateDict, src_state_dict: StateDict): + for src_key, dst_key in renamed_keys: + dst_state_dict[dst_key] = src_state_dict.pop(src_key) + + def replace_maskformer_swin_backbone( + self, dst_state_dict: StateDict, src_state_dict: StateDict, config: Mask2FormerConfig + ): + dst_prefix: str = "pixel_level_module.encoder" + src_prefix: str = "backbone" + + renamed_keys = [ + ( + f"{src_prefix}.patch_embed.proj.weight", + f"{dst_prefix}.model.embeddings.patch_embeddings.projection.weight", + ), + (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.model.embeddings.patch_embeddings.projection.bias"), + (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.model.embeddings.norm.weight"), + (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.model.embeddings.norm.bias"), + ] + num_layers = len(config.backbone_config.depths) + for layer_idx in range(num_layers): + for block_idx in range(config.backbone_config.depths[layer_idx]): + renamed_keys.extend( + [ # src, dst + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table", + ), + ] + ) + # now we need to handle the attentions + # read in weights + bias of input projection layer of cross-attention + + src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"] + src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"] + + size = src_att_weight.shape[0] + offset = size // 3 + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight" + ] = src_att_weight[:offset, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias" + ] = src_att_bias[:offset] + + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight" + ] = src_att_weight[offset : offset * 2, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias" + ] = src_att_bias[offset : offset * 2] + + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight" + ] = src_att_weight[-offset:, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias" + ] = src_att_bias[-offset:] + + # let's pop them + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight") + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias") + # proj + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias", + ), + ] + ) + + # second norm + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias", + ), + ] + ) + + # mlp + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias", + ), + ] + ) + + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index", + ) + ] + ) + + if layer_idx < num_layers - 1: + # patch merging + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.reduction.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.bias", + ), + ] + ) + + # hidden states norms + renamed_keys.extend( + [ + ( + f"{src_prefix}.norm{layer_idx}.weight", + f"{dst_prefix}.hidden_states_norms.{layer_idx}.weight", + ), + ( + f"{src_prefix}.norm{layer_idx}.bias", + f"{dst_prefix}.hidden_states_norms.{layer_idx}.bias", + ), + ] + ) + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_swin_backbone(self, dst_state_dict: StateDict, src_state_dict: StateDict, config: Mask2FormerConfig): + dst_prefix: str = "pixel_level_module.encoder" + src_prefix: str = "backbone" + + renamed_keys = [ + ( + f"{src_prefix}.patch_embed.proj.weight", + f"{dst_prefix}.embeddings.patch_embeddings.projection.weight", + ), + (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.embeddings.patch_embeddings.projection.bias"), + (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.embeddings.norm.weight"), + (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.embeddings.norm.bias"), + ] + + for layer_idx in range(len(config.backbone_config.depths)): + for block_idx in range(config.backbone_config.depths[layer_idx]): + renamed_keys.extend( + [ # src, dst + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table", + ), + ] + ) + # now we need to handle the attentions + # read in weights + bias of input projection layer of cross-attention + + src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"] + src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"] + + size = src_att_weight.shape[0] + offset = size // 3 + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight" + ] = src_att_weight[:offset, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias" + ] = src_att_bias[:offset] + + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight" + ] = src_att_weight[offset : offset * 2, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias" + ] = src_att_bias[offset : offset * 2] + + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight" + ] = src_att_weight[-offset:, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias" + ] = src_att_bias[-offset:] + + # let's pop them + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight") + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias") + # proj + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias", + ), + ] + ) + + # second norm + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias", + ), + ] + ) + + # mlp + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias", + ), + ] + ) + + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index", + ) + ] + ) + + if layer_idx < 3: + # patch merging + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.reduction.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.bias", + ), + ] + ) + + # hidden states norms + renamed_keys.extend( + [ + ( + f"{src_prefix}.norm{layer_idx}.weight", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.weight", + ), + ( + f"{src_prefix}.norm{layer_idx}.bias", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.bias", + ), + ] + ) + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + # Backbone + Pixel Decoder + def replace_pixel_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "pixel_level_module.decoder" + src_prefix: str = "sem_seg_head.pixel_decoder" + + self.replace_swin_backbone(dst_state_dict, src_state_dict, self.config) + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + def rename_keys_for_self_attn(src_prefix: str, dst_prefix: str): + self_attn_keys = [] + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.attention_weights", f"{dst_prefix}.attention_weights") + ) + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.output_proj", f"{dst_prefix}.output_proj") + ) + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.sampling_offsets", f"{dst_prefix}.sampling_offsets") + ) + self_attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.value_proj", f"{dst_prefix}.value_proj")) + + return self_attn_keys + + def rename_keys_for_encoder_layer(src_prefix: str, dst_prefix: str): + encoder_keys = [] + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear1", f"{dst_prefix}.fc1")) + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear2", f"{dst_prefix}.fc2")) + encoder_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.norm1", f"{dst_prefix}.self_attn_layer_norm") + ) + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm2", f"{dst_prefix}.final_layer_norm")) + encoder_keys.extend(rename_keys_for_self_attn(f"{src_prefix}.self_attn", f"{dst_prefix}.self_attn")) + + return encoder_keys + + # convolution layer for final features + renamed_keys = [ + (f"{src_prefix}.adapter_1.weight", f"{dst_prefix}.adapter_1.0.weight"), + (f"{src_prefix}.adapter_1.norm.weight", f"{dst_prefix}.adapter_1.1.weight"), + (f"{src_prefix}.adapter_1.norm.bias", f"{dst_prefix}.adapter_1.1.bias"), + ] + + renamed_keys.extend( + [ + (f"{src_prefix}.layer_1.weight", f"{dst_prefix}.layer_1.0.weight"), + (f"{src_prefix}.layer_1.norm.weight", f"{dst_prefix}.layer_1.1.weight"), + (f"{src_prefix}.layer_1.norm.bias", f"{dst_prefix}.layer_1.1.bias"), + ] + ) + + # proj layers + for i in range(3): + for j in range(2): + renamed_keys.extend( + [ + (f"{src_prefix}.input_proj.{i}.{j}.weight", f"{dst_prefix}.input_projections.{i}.{j}.weight"), + (f"{src_prefix}.input_proj.{i}.{j}.bias", f"{dst_prefix}.input_projections.{i}.{j}.bias"), + ] + ) + + renamed_keys.extend([(f"{src_prefix}.transformer.level_embed", f"{dst_prefix}.level_embed")]) + + # layers + for layer_idx in range(self.config.encoder_layers): + renamed_keys.extend( + rename_keys_for_encoder_layer( + f"{src_prefix}.transformer.encoder.layers.{layer_idx}", f"{dst_prefix}.encoder.layers.{layer_idx}" + ) + ) + + # proj + renamed_keys.extend( + [ + (f"{src_prefix}.mask_features.weight", f"{dst_prefix}.mask_projection.weight"), + (f"{src_prefix}.mask_features.bias", f"{dst_prefix}.mask_projection.bias"), + ] + ) + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + # Transformer Decoder + def rename_keys_in_masked_attention_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder" + src_prefix: str = "sem_seg_head.predictor" + + rename_keys = [] + for i in range(self.config.decoder_layers - 1): + rename_keys.append( + ( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.out_proj.weight", + f"{dst_prefix}.layers.{i}.self_attn.out_proj.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.out_proj.bias", + f"{dst_prefix}.layers.{i}.self_attn.out_proj.bias", + ) + ) + + rename_keys.append( + ( + f"{src_prefix}.transformer_self_attention_layers.{i}.norm.weight", + f"{dst_prefix}.layers.{i}.self_attn_layer_norm.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_self_attention_layers.{i}.norm.bias", + f"{dst_prefix}.layers.{i}.self_attn_layer_norm.bias", + ) + ) + + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.in_proj_weight", + f"{dst_prefix}.layers.{i}.cross_attn.in_proj_weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.in_proj_bias", + f"{dst_prefix}.layers.{i}.cross_attn.in_proj_bias", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.out_proj.weight", + f"{dst_prefix}.layers.{i}.cross_attn.out_proj.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.out_proj.bias", + f"{dst_prefix}.layers.{i}.cross_attn.out_proj.bias", + ) + ) + + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.norm.weight", + f"{dst_prefix}.layers.{i}.cross_attn_layer_norm.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_cross_attention_layers.{i}.norm.bias", + f"{dst_prefix}.layers.{i}.cross_attn_layer_norm.bias", + ) + ) + + rename_keys.append( + (f"{src_prefix}.transformer_ffn_layers.{i}.linear1.weight", f"{dst_prefix}.layers.{i}.fc1.weight") + ) + rename_keys.append( + (f"{src_prefix}.transformer_ffn_layers.{i}.linear1.bias", f"{dst_prefix}.layers.{i}.fc1.bias") + ) + rename_keys.append( + (f"{src_prefix}.transformer_ffn_layers.{i}.linear2.weight", f"{dst_prefix}.layers.{i}.fc2.weight") + ) + rename_keys.append( + (f"{src_prefix}.transformer_ffn_layers.{i}.linear2.bias", f"{dst_prefix}.layers.{i}.fc2.bias") + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_ffn_layers.{i}.norm.weight", + f"{dst_prefix}.layers.{i}.final_layer_norm.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.transformer_ffn_layers.{i}.norm.bias", + f"{dst_prefix}.layers.{i}.final_layer_norm.bias", + ) + ) + + return rename_keys + + def replace_masked_attention_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder" + src_prefix: str = "sem_seg_head.predictor" + + renamed_keys = self.rename_keys_in_masked_attention_decoder(dst_state_dict, src_state_dict) + + # add more + renamed_keys.extend( + [ + (f"{src_prefix}.decoder_norm.weight", f"{dst_prefix}.layernorm.weight"), + (f"{src_prefix}.decoder_norm.bias", f"{dst_prefix}.layernorm.bias"), + ] + ) + + mlp_len = 3 + for i in range(mlp_len): + renamed_keys.extend( + [ + ( + f"{src_prefix}.mask_embed.layers.{i}.weight", + f"{dst_prefix}.mask_predictor.mask_embedder.{i}.0.weight", + ), + ( + f"{src_prefix}.mask_embed.layers.{i}.bias", + f"{dst_prefix}.mask_predictor.mask_embedder.{i}.0.bias", + ), + ] + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_keys_qkv_transformer_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder.layers" + src_prefix: str = "sem_seg_head.predictor" + for i in range(self.config.decoder_layers - 1): + # read in weights + bias of input projection layer of self-attention + in_proj_weight = src_state_dict.pop( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_weight" + ) + in_proj_bias = src_state_dict.pop( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_bias" + ) + # next, add query, keys and values (in that order) to the state dict + dst_state_dict[f"{dst_prefix}.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:] + + def replace_transformer_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module" + src_prefix: str = "sem_seg_head.predictor" + + self.replace_masked_attention_decoder(dst_state_dict, src_state_dict) + + renamed_keys = [ + (f"{src_prefix}.query_embed.weight", f"{dst_prefix}.queries_embedder.weight"), + (f"{src_prefix}.query_feat.weight", f"{dst_prefix}.queries_features.weight"), + (f"{src_prefix}.level_embed.weight", f"{dst_prefix}.level_embed.weight"), + ] + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + self.replace_keys_qkv_transformer_decoder(dst_state_dict, src_state_dict) + + def replace_universal_segmentation_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "" + src_prefix: str = "sem_seg_head.predictor" + + renamed_keys = [ + (f"{src_prefix}.class_embed.weight", f"{dst_prefix}class_predictor.weight"), + (f"{src_prefix}.class_embed.bias", f"{dst_prefix}class_predictor.bias"), + ] + + logger.info(f"Replacing keys {pformat(renamed_keys)}") + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def convert(self, mask2former: Mask2FormerModel) -> Mask2FormerModel: + dst_state_dict = TrackedStateDict(mask2former.state_dict()) + src_state_dict = self.original_model.state_dict() + + self.replace_pixel_module(dst_state_dict, src_state_dict) + self.replace_transformer_module(dst_state_dict, src_state_dict) + + logger.info(f"Missed keys are {pformat(dst_state_dict.diff())}") + logger.info(f"Not copied keys are {pformat(src_state_dict.keys())}") + logger.info("🙌 Done") + + state_dict = {key: dst_state_dict[key] for key in dst_state_dict.to_track} + mask2former.load_state_dict(state_dict) + return mask2former + + def convert_universal_segmentation( + self, mask2former: Mask2FormerForUniversalSegmentation + ) -> Mask2FormerForUniversalSegmentation: + dst_state_dict = TrackedStateDict(mask2former.state_dict()) + src_state_dict = self.original_model.state_dict() + + self.replace_universal_segmentation_module(dst_state_dict, src_state_dict) + + state_dict = {key: dst_state_dict[key] for key in dst_state_dict.to_track} + mask2former.load_state_dict(state_dict) + + return mask2former + + @staticmethod + def using_dirs(checkpoints_dir: Path, config_dir: Path) -> Iterator[tuple[object, Path, Path]]: + checkpoints: list[Path] = checkpoints_dir.glob("**/*.pkl") + + for checkpoint in checkpoints: + logger.info(f"Converting {checkpoint.stem}") + # find associated config file + + # dataset_name e.g 'coco' + dataset_name = checkpoint.parents[2].stem + if dataset_name == "ade": + dataset_name = dataset_name.replace("ade", "ade20k") + + # task type e.g 'instance-segmentation' + segmentation_task = checkpoint.parents[1].stem + + # config file corresponding to checkpoint + config_file_name = f"{checkpoint.parents[0].stem}.yaml" + + config: Path = config_dir / dataset_name / segmentation_task / "swin" / config_file_name + yield config, checkpoint + + +def test( + original_model, + our_model: Mask2FormerForUniversalSegmentation, + image_processor: Mask2FormerImageProcessor, + tolerance: float, +): + with torch.no_grad(): + original_model = original_model.eval() + our_model = our_model.eval() + + im = prepare_img() + x = image_processor(images=im, return_tensors="pt")["pixel_values"] + + original_model_backbone_features = original_model.backbone(x.clone()) + our_model_output: Mask2FormerModelOutput = our_model.model(x.clone(), output_hidden_states=True) + + # Test backbone + for original_model_feature, our_model_feature in zip( + original_model_backbone_features.values(), our_model_output.encoder_hidden_states + ): + assert torch.allclose(original_model_feature, our_model_feature, atol=tolerance), ( + "The backbone features are not the same." + ) + + # Test pixel decoder + mask_features, _, multi_scale_features = original_model.sem_seg_head.pixel_decoder.forward_features( + original_model_backbone_features + ) + + for original_model_feature, our_model_feature in zip( + multi_scale_features, our_model_output.pixel_decoder_hidden_states + ): + assert torch.allclose(original_model_feature, our_model_feature, atol=tolerance), ( + "The pixel decoder feature are not the same" + ) + + # Let's test the full model + tr_complete = T.Compose( + [T.Resize((384, 384)), T.ToTensor()], + ) + y = (tr_complete(im) * 255.0).to(torch.int).float() + + # modify original Mask2Former code to return mask and class logits + original_class_logits, original_mask_logits = original_model([{"image": y.clone().squeeze(0)}]) + + our_model_out: Mask2FormerForUniversalSegmentationOutput = our_model(x.clone()) + our_mask_logits = our_model_out.masks_queries_logits + our_class_logits = our_model_out.class_queries_logits + + assert original_mask_logits.shape == our_mask_logits.shape, "Output masks shapes are not matching." + assert original_class_logits.shape == our_class_logits.shape, "Output class logits shapes are not matching." + assert torch.allclose(original_class_logits, our_class_logits, atol=tolerance), ( + "The class logits are not the same." + ) + assert torch.allclose(original_mask_logits, our_mask_logits, atol=tolerance), ( + "The predicted masks are not the same." + ) + + logger.info("Test passed!") + + +def get_model_name(checkpoint_file: Path): + # model_name_raw is something like maskformer2_swin_small_bs16_50ep + model_name_raw: str = checkpoint_file.parents[0].stem + + # `segmentation_task_type` must be one of the following: `instance-segmentation`, `panoptic-segmentation`, `semantic-segmentation` + segmentation_task_name: str = checkpoint_file.parents[1].stem + if segmentation_task_name not in ["instance-segmentation", "panoptic-segmentation", "semantic-segmentation"]: + raise ValueError( + f"{segmentation_task_name} must be wrong since acceptable values are: instance-segmentation," + " panoptic-segmentation, semantic-segmentation." + ) + + # dataset name must be one of the following: `coco`, `ade`, `cityscapes`, `mapillary-vistas` + dataset_name: str = checkpoint_file.parents[2].stem + if dataset_name not in ["coco", "ade", "cityscapes", "mapillary-vistas"]: + raise ValueError( + f"{dataset_name} must be wrong since we didn't find 'coco' or 'ade' or 'cityscapes' or 'mapillary-vistas'" + " in it " + ) + + backbone = "swin" + backbone_types = ["tiny", "small", "base_IN21k", "base", "large"] + backbone_type = list(filter(lambda x: x in model_name_raw, backbone_types))[0].replace("_", "-") + + model_name = f"mask2former-{backbone}-{backbone_type}-{dataset_name}-{segmentation_task_name.split('-')[0]}" + + return model_name + + +if __name__ == "__main__": + parser = ArgumentParser( + description="Command line to convert the original mask2formers (with swin backbone) to our implementations." + ) + + parser.add_argument( + "--checkpoints_dir", + type=Path, + help=( + "A directory containing the model's checkpoints. The directory has to have the following structure:" + " ///.pkl" + ), + ) + parser.add_argument( + "--configs_dir", + type=Path, + help=( + "A directory containing the model's configs, see detectron2 doc. The directory has to have the following" + " structure: ///.yaml" + ), + ) + parser.add_argument( + "--mask2former_dir", + required=True, + type=Path, + help=( + "A path to Mask2Former's original implementation directory. You can download from here:" + " https://github.com/facebookresearch/Mask2Former" + ), + ) + + args = parser.parse_args() + + checkpoints_dir: Path = args.checkpoints_dir + config_dir: Path = args.configs_dir + mask2former_dir: Path = args.mask2former_dir + # append the path to the parents to mask2former dir + sys.path.append(str(mask2former_dir.parent)) + # import original Mask2Former config and model from original source code repo + from Mask2Former.mask2former.config import add_maskformer2_config + from Mask2Former.mask2former.maskformer_model import MaskFormer as OriginalMask2Former + + for config_file, checkpoint_file in OriginalMask2FormerCheckpointToOursConverter.using_dirs( + checkpoints_dir, config_dir + ): + model_name = get_model_name(checkpoint_file) + image_processor = OriginalMask2FormerConfigToImageProcessorConverter()( + setup_cfg(Args(config_file=config_file)) + ) + image_processor.size = {"height": 384, "width": 384} + + original_config = setup_cfg(Args(config_file=config_file)) + mask2former_kwargs = OriginalMask2Former.from_config(original_config) + original_model = OriginalMask2Former(**mask2former_kwargs).eval() + + DetectionCheckpointer(original_model).load(str(checkpoint_file)) + + config: Mask2FormerConfig = OriginalMask2FormerConfigToOursConverter()(original_config) + mask2former = Mask2FormerModel(config=config).eval() + + converter = OriginalMask2FormerCheckpointToOursConverter(original_model, config) + mask2former = converter.convert(mask2former) + + mask2former_for_segmentation = Mask2FormerForUniversalSegmentation(config=config).eval() + mask2former_for_segmentation.model = mask2former + + mask2former_for_segmentation = converter.convert_universal_segmentation(mask2former_for_segmentation) + + tolerance = 3e-1 + high_tolerance_models = [ + "mask2former-swin-base-IN21k-coco-instance", + "mask2former-swin-base-coco-instance", + "mask2former-swin-small-cityscapes-semantic", + ] + + if model_name in high_tolerance_models: + tolerance = 3e-1 + + logger.info(f"Testing {model_name}...") + test(original_model, mask2former_for_segmentation, image_processor, tolerance) + logger.info(f"Pushing {model_name} to hub...") + + image_processor.push_to_hub(model_name) + mask2former_for_segmentation.push_to_hub(model_name) diff --git a/third_party/transformers/src/transformers/models/mask2former/image_processing_mask2former.py b/third_party/transformers/src/transformers/models/mask2former/image_processing_mask2former.py new file mode 100644 index 0000000000000000000000000000000000000000..a72540fe7a03c51154434762a00de3aa8cc6908f --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/image_processing_mask2former.py @@ -0,0 +1,822 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/mask2former/modular_mask2former.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_mask2former.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Any, Optional, Union + +import numpy as np +import torch +from torch import nn +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 get_size_with_aspect_ratio, group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size_for_max_height_width, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +class Mask2FormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ignore_index (`int`, *optional*): + Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels + denoted with 0 (background) will be replaced with `ignore_index`. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). + The background label will be replaced by `ignore_index`. + num_labels (`int`, *optional*): + The number of labels in the segmentation map. + size_divisor (`int`, *optional*, defaults to `32`): + Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in + Swin Transformer. + pad_size (`SizeDict`, *optional*): + The size to pad the images to. Must be larger than any image size provided for preprocessing. If `pad_size` + is not provided, images will be padded to the largest height and width in the batch. + """ + + ignore_index: int | None + do_reduce_labels: bool + num_labels: int | None + size_divisor: int + pad_size: SizeDict | None + + +# Helper functions for post-processing (PyTorch-based) +def binary_mask_to_rle(mask: "torch.Tensor | np.ndarray") -> list[int]: + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `np.ndarray`): + A binary mask 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. + """ + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask) + + pixels = mask.flatten() + zero = torch.zeros(1, device=pixels.device, dtype=pixels.dtype) + pixels = torch.cat([zero, pixels, zero]) + runs = torch.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return runs.tolist() + + +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`): + 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 + + +def convert_segmentation_map_to_binary_masks_fast( + segmentation_map: "torch.Tensor", + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, +): + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = torch.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = torch.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] # drop background label if applicable + + binary_masks = [(segmentation_map == i) for i in all_labels] + if binary_masks: + binary_masks = torch.stack(binary_masks, dim=0) + else: + binary_masks = torch.zeros((0, *segmentation_map.shape), device=segmentation_map.device) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = torch.zeros(all_labels.shape[0], device=segmentation_map.device) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(label.item() + 1 if do_reduce_labels else label.item())] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels + return binary_masks.float(), labels.long() + + +@auto_docstring +class Mask2FormerImageProcessor(TorchvisionBackend): + valid_kwargs = Mask2FormerImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_pad = True + model_input_names = ["pixel_values", "pixel_mask"] + size_divisor = 32 + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[Mask2FormerImageProcessorKwargs]) -> None: + size = kwargs.pop("size", None) + max_size = kwargs.pop("max_size", None) + + if size is None and max_size is not None: + size = self.size.copy() + size["longest_edge"] = max_size + elif size is None: + size = self.size + + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + super().__init__(**kwargs) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the + `_max_size` attribute from the dictionary. + """ + image_processor_dict = super().to_dict() + image_processor_dict.pop("_max_size", None) + return image_processor_dict + + def reduce_label(self, labels: list["torch.Tensor"]): + for idx in range(len(labels)): + label = labels[idx] + label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype), label) + label = label - 1 + label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype), label) + labels[idx] = label + + def resize( + self, + image: torch.Tensor, + size: SizeDict, + size_divisor: int = 0, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = 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. + size_divisor (`int`, *optional*, defaults to 0): + If `size_divisor` is given, the output image size will be divisible by the number. + resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*): + 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.size()[-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.size()[-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}." + ) + if size_divisor > 0: + height, width = new_size + height = int(math.ceil(height / size_divisor) * size_divisor) + width = int(math.ceil(width / size_divisor) * size_divisor) + new_size = (height, width) + + image = super().resize( + image, size=SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs + ) + return image + + def pad( + self, + images: torch.Tensor, + padded_size: tuple[int, int], + segmentation_maps: torch.Tensor | None = None, + fill: int = 0, + ignore_index: int = 255, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + original_size = images.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] + images = tvF.pad(images, padding, fill=fill) + if segmentation_maps is not None: + segmentation_maps = [tvF.pad(mask, padding, fill=ignore_index) for mask in segmentation_maps] + + # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + pixel_mask = torch.zeros((images.shape[0], *padded_size), dtype=torch.int64, device=images.device) + pixel_mask[:, : original_size[0], : original_size[1]] = 1 + + return images, pixel_mask, segmentation_maps + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[Mask2FormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + device: Union[str, "torch.device"] | None = None, + **kwargs: Unpack[Mask2FormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + segmentation_maps: Optional["torch.Tensor"], + instance_id_to_semantic_id: dict[int, int] | None, + do_resize: bool | None, + size: SizeDict | None, + pad_size: SizeDict | None, + size_divisor: int | None, + resample: Union["PILImageResampling", "tvF.InterpolationMode"] | None, + do_rescale: bool | None, + rescale_factor: float | None, + do_normalize: bool | None, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + if segmentation_maps is not None and len(images) != len(segmentation_maps): + raise ValueError("Images and segmentation maps must have the same length.") + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + if segmentation_maps is not None: + grouped_segmentation_maps, grouped_segmentation_maps_index = group_images_by_shape( + segmentation_maps, disable_grouping=disable_grouping + ) + resized_segmentation_maps_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_images, size=size, size_divisor=size_divisor, resample=resample + ) + if segmentation_maps is not None: + stacked_segmentation_maps = grouped_segmentation_maps[shape] + if do_resize: + stacked_segmentation_maps = self.resize( + image=stacked_segmentation_maps, + size=size, + size_divisor=size_divisor, + resample=tvF.InterpolationMode.NEAREST_EXACT, + ) + resized_images_grouped[shape] = stacked_images + if segmentation_maps is not None: + resized_segmentation_maps_grouped[shape] = stacked_segmentation_maps + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + if segmentation_maps is not None: + resized_segmentation_maps = reorder_images( + resized_segmentation_maps_grouped, grouped_segmentation_maps_index + ) + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(resized_images) + + if segmentation_maps is not None: + mask_labels = [] + class_labels = [] + # Convert to list of binary masks and labels + for idx, segmentation_map in enumerate(resized_segmentation_maps): + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + # Use instance2class_id mapping per image + masks, classes = convert_segmentation_map_to_binary_masks_fast( + segmentation_map.squeeze(0), + instance_id, + ignore_index=ignore_index, + do_reduce_labels=do_reduce_labels, + ) + mask_labels.append(masks) + class_labels.append(classes) + + if segmentation_maps is not None: + # group mask_labels as paired inputs and not images so as not to stack them + grouped_images, grouped_segmentation_maps, grouped_images_index = group_images_by_shape( + resized_images, mask_labels, disable_grouping=disable_grouping + ) + processed_segmentation_maps_grouped = {} + else: + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_pixel_masks_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 + ) + padded_images, pixel_masks, padded_segmentation_maps = self.pad( + images=stacked_images, + segmentation_maps=grouped_segmentation_maps[shape] if segmentation_maps is not None else None, + padded_size=padded_size, + ignore_index=ignore_index, + ) + processed_images_grouped[shape] = padded_images + processed_pixel_masks_grouped[shape] = pixel_masks + if segmentation_maps is not None: + processed_segmentation_maps_grouped[shape] = padded_segmentation_maps + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + processed_pixel_masks = reorder_images(processed_pixel_masks_grouped, grouped_images_index) + encoded_inputs = BatchFeature( + data={"pixel_values": processed_images, "pixel_mask": processed_pixel_masks}, + tensor_type=return_tensors, + ) + if segmentation_maps is not None: + mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_images_index) + # we cannot batch them since they don't share a common class size + encoded_inputs["mask_labels"] = mask_labels + encoded_inputs["class_labels"] = class_labels + + return encoded_inputs + + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`Mask2FormerForUniversalSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *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 left to None, 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.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # 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 = torch.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, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`Mask2FormerForUniversalSegmentationOutput`] into instance segmentation predictions. + Only supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + 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 left to None, 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. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + 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`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps 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 return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros((384, 384)) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + 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 [`Mask2FormerForUniversalSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentationOutput`]): + The outputs from [`Mask2FormerForUniversalSegmentation`]. + 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 left to None, 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`, set + to `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("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + 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__ = ["Mask2FormerImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/mask2former/image_processing_pil_mask2former.py b/third_party/transformers/src/transformers/models/mask2former/image_processing_pil_mask2former.py new file mode 100644 index 0000000000000000000000000000000000000000..8358a3601bed0d4f84731b253b26c52efc439eaa --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/image_processing_pil_mask2former.py @@ -0,0 +1,862 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/mask2former/modular_mask2former.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_mask2former.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math +from typing import Any + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import PaddingMode, get_size_with_aspect_ratio +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, + get_image_size_for_max_height_width, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available, logging, requires_backends +from ...utils.import_utils import requires + + +if is_torch_available(): + import torch + from torch import nn + +logger = logging.get_logger(__name__) + + +class Mask2FormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ignore_index (`int`, *optional*): + Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels + denoted with 0 (background) will be replaced with `ignore_index`. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). + The background label will be replaced by `ignore_index`. + num_labels (`int`, *optional*): + The number of labels in the segmentation map. + size_divisor (`int`, *optional*, defaults to `32`): + Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in + Swin Transformer. + pad_size (`SizeDict`, *optional*): + The size to pad the images to. Must be larger than any image size provided for preprocessing. If `pad_size` + is not provided, images will be padded to the largest height and width in the batch. + """ + + ignore_index: int | None + do_reduce_labels: bool + num_labels: int | None + size_divisor: int + pad_size: SizeDict | None + + +def convert_segmentation_map_to_binary_masks( + segmentation_map: np.ndarray, + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, +): + """Convert segmentation map to binary masks using NumPy operations.""" + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = np.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] + + binary_masks = [(segmentation_map == i) for i in all_labels] + if binary_masks: + binary_masks = np.stack(binary_masks, axis=0) + else: + binary_masks = np.zeros((0, *segmentation_map.shape), dtype=np.float32) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = np.zeros(all_labels.shape[0], dtype=np.int64) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(int(label) + 1 if do_reduce_labels else int(label))] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels.astype(np.int64) + return binary_masks.astype(np.float32), labels + + +# Adapted from transformers.models.mask2former.image_processing_mask2former.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.mask2former.image_processing_mask2former.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.mask2former.image_processing_mask2former.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, +): + 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.mask2former.image_processing_mask2former.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`): + 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 + + +# Adapted from transformers.models.mask2former.image_processing_mask2former.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] + + +@requires(backends=("torch",)) +class Mask2FormerImageProcessorPil(PilBackend): + valid_kwargs = Mask2FormerImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_pad = True + model_input_names = ["pixel_values", "pixel_mask"] + size_divisor = 32 + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[Mask2FormerImageProcessorKwargs]) -> None: + size = kwargs.pop("size", None) + max_size = kwargs.pop("max_size", None) + + # Store max_size as private attribute for backward compatibility + self._max_size = max_size if max_size is not None else 1333 + + if size is None and max_size is not None: + size = self.size.copy() + size["longest_edge"] = max_size + elif size is None: + size = self.size + + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + super().__init__(**kwargs) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the + `_max_size` attribute from the dictionary. + """ + image_processor_dict = super().to_dict() + image_processor_dict.pop("_max_size", None) + return image_processor_dict + + def reduce_label(self, labels: list[np.ndarray]): + """Reduce label values by 1, replacing 0 with 255.""" + for idx in range(len(labels)): + label = labels[idx].copy() + label[label == 0] = 255 + label = label - 1 + label[label == 254] = 255 + labels[idx] = label + + def resize( + self, + image: np.ndarray, + size: SizeDict, + size_divisor: int = 0, + resample: PILImageResampling | None = None, + **kwargs, + ) -> np.ndarray: + """ + Resize the image to the given size with optional size_divisor. + + Args: + image (`np.ndarray`): + Image to resize. + size (`SizeDict`): + Size of the image's `(height, width)` dimensions after resizing. + size_divisor (`int`, *optional*, defaults to 0): + If `size_divisor` is given, the output image size will be divisible by the number. + resample (`PILImageResampling | int | None`, *optional*): + Resampling filter to use if resizing the image. + """ + + if size.shortest_edge and size.longest_edge: + # Get current image size + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + new_size = get_size_with_aspect_ratio((height, width), size.shortest_edge, size.longest_edge) + elif size.max_height and size.max_width: + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + new_size = get_image_size_for_max_height_width((height, width), 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}." + ) + if size_divisor > 0: + height, width = new_size + height = int(math.ceil(height / size_divisor) * size_divisor) + width = int(math.ceil(width / size_divisor) * size_divisor) + new_size = (height, width) + + return super().resize(image, size=SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def pad( + self, + images: list[np.ndarray], + padded_size: tuple[int, int], + segmentation_maps: list[np.ndarray] | None = None, + fill: int = 0, + ignore_index: int = 255, + ) -> tuple[list[np.ndarray], list[np.ndarray], list[np.ndarray] | None]: + """ + Pad images and optionally segmentation maps to the given size. + + Args: + images (`list[np.ndarray]`): + Images to pad. + padded_size (`tuple[int, int]`): + Target size (height, width) to pad to. + segmentation_maps (`list[np.ndarray]`, *optional*): + Segmentation maps to pad. + fill (`int`, *optional*, defaults to 0): + Fill value for images. + ignore_index (`int`, *optional*, defaults to 255): + Fill value for segmentation maps. + + Returns: + `tuple`: (padded_images, pixel_masks, padded_segmentation_maps) + """ + padded_images = [] + pixel_masks = [] + + for image in images: + 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}." + ) + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + image = np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=fill, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + padded_images.append(image) + + # Make a pixel mask for the image + pixel_mask = np.zeros(padded_size, dtype=np.int64) + pixel_mask[: original_size[0], : original_size[1]] = 1 + pixel_masks.append(pixel_mask) + + padded_segmentation_maps = None + if segmentation_maps is not None: + padded_segmentation_maps = [] + for mask in segmentation_maps: + original_size = mask.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + mask = np_pad( + mask, + padding, + mode=PaddingMode.CONSTANT, + constant_values=ignore_index, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + padded_segmentation_maps.append(mask) + + return padded_images, pixel_masks, padded_segmentation_maps + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[Mask2FormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + **kwargs: Unpack[Mask2FormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + segmentation_maps: list[np.ndarray] | None, + instance_id_to_semantic_id: dict[int, int] | None, + do_resize: bool | None, + size: SizeDict | None, + pad_size: SizeDict | None, + size_divisor: int | None, + resample: PILImageResampling | None, + do_rescale: bool | None, + rescale_factor: float | None, + do_normalize: bool | None, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + if segmentation_maps is not None and len(images) != len(segmentation_maps): + raise ValueError("Images and segmentation maps must have the same length.") + + # Process images one by one (no batching in PIL backend) + resized_images = [] + resized_segmentation_maps = None + if segmentation_maps is not None: + resized_segmentation_maps = [] + + for idx, image in enumerate(images): + if do_resize: + image = self.resize(image=image, size=size, size_divisor=size_divisor, resample=resample) + resized_images.append(image) + + if segmentation_maps is not None: + seg_map = segmentation_maps[idx] + if do_resize: + seg_map = self.resize( + image=seg_map, size=size, size_divisor=size_divisor, resample=PILImageResampling.NEAREST + ) + resized_segmentation_maps.append(seg_map) + + # Determine padded size + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(resized_images, input_data_format=ChannelDimension.FIRST) + + # Convert segmentation maps to binary masks if provided + mask_labels = None + class_labels = None + if segmentation_maps is not None: + mask_labels = [] + class_labels = [] + for idx, segmentation_map in enumerate(resized_segmentation_maps): + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + # Squeeze channel dimension if present + if segmentation_map.ndim == 3 and segmentation_map.shape[0] == 1: + segmentation_map = segmentation_map.squeeze(0) + masks, classes = convert_segmentation_map_to_binary_masks( + segmentation_map, instance_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels + ) + mask_labels.append(masks) + class_labels.append(classes) + + # Process images: rescale, normalize, pad + processed_images = [] + for image in resized_images: + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + + # Pad images and create pixel masks (also pad mask_labels to match padded image size) + padded_images, pixel_masks, padded_mask_labels = self.pad( + images=processed_images, + padded_size=padded_size, + segmentation_maps=mask_labels, + fill=0, + ignore_index=ignore_index, # Match Torchvision backend for cross-backend equivalence + ) + + encoded_inputs = BatchFeature( + data={"pixel_values": padded_images, "pixel_mask": pixel_masks}, tensor_type=return_tensors + ) + # we cannot batch them since they don't share a common class size + if segmentation_maps is not None: + encoded_inputs["mask_labels"] = [ + torch.from_numpy(mask_label) if return_tensors == "pt" else mask_label + for mask_label in padded_mask_labels + ] + encoded_inputs["class_labels"] = [ + torch.from_numpy(class_label) if return_tensors == "pt" else class_label + for class_label in class_labels + ] + + return encoded_inputs + + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`Mask2FormerForUniversalSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *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 left to None, 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.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # 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 = torch.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, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`Mask2FormerForUniversalSegmentationOutput`] into instance segmentation predictions. + Only supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + 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 left to None, 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. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + 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`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps 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`. + """ + requires_backends(self, ["torch"]) + if return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros((384, 384)) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + # Adapted from transformers.models.mask2former.image_processing_mask2former.Mask2FormerImageProcessor.post_process_panoptic_segmentation + 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 [`Mask2FormerForUniversalSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentationOutput`]): + The outputs from [`Mask2FormerForUniversalSegmentation`]. + 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 left to None, 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`, set + to `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`. + """ + requires_backends(self, ["torch"]) + if label_ids_to_fuse is None: + logger.warning("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + 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__ = ["Mask2FormerImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/mask2former/modeling_mask2former.py b/third_party/transformers/src/transformers/models/mask2former/modeling_mask2former.py new file mode 100644 index 0000000000000000000000000000000000000000..b20646bc2cb4c00f499be905155523c8e3079356 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/modeling_mask2former.py @@ -0,0 +1,2507 @@ +# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Mask2Former model.""" + +import math +import warnings +from dataclasses import dataclass + +import numpy as np +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import load_backbone +from ...file_utils import ModelOutput, is_scipy_available, requires_backends +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import auto_docstring, is_accelerate_available, logging, torch_compilable_check +from .configuration_mask2former import Mask2FormerConfig + + +if is_scipy_available(): + from scipy.optimize import linear_sum_assignment + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import reduce + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Mask2Former's pixel decoder module output, practically a Multi-Scale Deformable Attention based decoder. It returns + the mask features and the multiscale features. + """ +) +class Mask2FormerPixelDecoderOutput(ModelOutput): + r""" + multi_scale_features (`tuple(torch.FloatTensor)`): + Tuple of multi-scale features of scales [1/8, 1/16, 1/32] and shape `(batch_size, num_channels, height, + width)`from the Multi-Scale Deformable Attenntion based Pixel Decoder. + mask_features (`torch.FloatTensor`): + Tensor of shape `(batch_size, num_channels, height, width)`, 1/4 scale features from the last Pixel Decoder + Layer. + attentions (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights from pixel decoder. Returned when `output_attentions=True` is passed + or when `config.output_attentions=True` + """ + + multi_scale_features: tuple[torch.FloatTensor] | None = None + mask_features: torch.FloatTensor | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the Transformer decoder. This class adds two attributes to + BaseModelOutputWithCrossAttentions for mask predictions logits and a tuple of intermediate decoder activations, + i.e. the output of each decoder layer, each of them gone through a layernorm. + """ +) +class Mask2FormerMaskedAttentionDecoderOutput(BaseModelOutputWithCrossAttentions): + r""" + hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer + plus the initial embedding outputs. Returned when `output_hidden_states=True`. + attentions (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in + the self-attention heads. Returned when `output_attentions=True`. + masks_queries_logits (`tuple(torch.FloatTensor)` of shape `(batch_size, num_queries, height, width)`): + Tuple of mask predictions from all layers of the transformer decoder. + intermediate_hidden_states (`tuple(torch.FloatTensor)` of shape `(num_queries, 1, hidden_size)`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + """ + + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: torch.FloatTensor | None = None + masks_queries_logits: tuple[torch.FloatTensor] | None = None + intermediate_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Mask2Former's pixel level module output. It returns the output of the encoder (optional) and all hidden states + (multi-scale features) from the `decoder`. By default, the `encoder` is a Swin Backbone and the `decoder` is a + Multi-Scale Deformable Attention based decoder. + + The `decoder_last_hidden_state` are the **per-pixel embeddings** while `decoder_hidden_states` refer to multi-scale + feature maps produced using **multi-scaling strategy** defined in the paper. + """ +) +class Mask2FormerPixelLevelModuleOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor`): + Last hidden states (final feature map of shape `(batch_size, num_channels, height, width)`) of the last + stage of the encoder. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden states (also + called feature maps) of the model at the output of each stage. Returned if output_hidden_states is set to + True. + decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)): + 1/4 scale features from the last Pixel Decoder Layer. + decoder_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden states (also + called feature maps) of the model at the output of each stage. + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_last_hidden_state: torch.FloatTensor | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`Mask2FormerModel`]. This class returns all the needed hidden states to compute the logits. + """ +) +class Mask2FormerModelOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`, *optional*): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). Returned when + `output_hidden_states=True` is passed. + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`, *optional*): + Last hidden states (final feature map) of the last stage of the pixel decoder model. + transformer_decoder_last_hidden_state (`tuple(torch.FloatTensor)`): + Final output of the transformer decoder `(batch_size, sequence_length, hidden_size)`. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder + model at the output of each stage. Returned when `output_hidden_states=True` is passed. + pixel_decoder_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, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel + decoder model at the output of each stage. Returned when `output_hidden_states=True` is passed. + transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the + transformer decoder at the output of each stage. Returned when `output_hidden_states=True` is passed. + transformer_decoder_intermediate_states (`tuple(torch.FloatTensor)` of shape `(num_queries, 1, hidden_size)`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + masks_queries_logits (`tuple(torch.FloatTensor)` of shape `(batch_size, num_queries, height, width)`) + Mask Predictions from each layer in the transformer decoder. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed): + Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Self attentions weights from transformer decoder. + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_intermediate_states: tuple[torch.FloatTensor] | None = None + masks_queries_logits: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`Mask2FormerForUniversalSegmentationOutput`]. + + This output can be directly passed to [`~Mask2FormerImageProcessor.post_process_semantic_segmentation`] or + [`~Mask2FormerImageProcessor.post_process_instance_segmentation`] or + [`~Mask2FormerImageProcessor.post_process_panoptic_segmentation`] to compute final segmentation maps. Please, see + [`~Mask2FormerImageProcessor] for details regarding usage. + """ +) +class Mask2FormerForUniversalSegmentationOutput(ModelOutput): + r""" + loss (`torch.Tensor`, *optional*): + The computed loss, returned when labels are present. + class_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each + query. Note the `+ 1` is needed because we incorporate the null class. + masks_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each + query. + auxiliary_logits (`list[Dict(str, torch.FloatTensor)]`, *optional*): + List of class and mask predictions from each layer of the transformer decoder. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the pixel decoder model. + transformer_decoder_last_hidden_state (`tuple(torch.FloatTensor)`): + Final output of the transformer decoder `(batch_size, sequence_length, hidden_size)`. + encoder_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, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder + model at the output of each stage. + pixel_decoder_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, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel + decoder model at the output of each stage. + transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the + transformer decoder at the output of each stage. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Self and Cross Attentions weights from transformer decoder. + """ + + loss: torch.FloatTensor | None = None + class_queries_logits: torch.FloatTensor | None = None + masks_queries_logits: torch.FloatTensor | None = None + auxiliary_logits: list[dict[str, torch.FloatTensor]] | None = None + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: torch.FloatTensor | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py +def sample_point( + input_features: torch.Tensor, point_coordinates: torch.Tensor, add_dim=False, **kwargs +) -> torch.Tensor: + """ + A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors. + + Args: + input_features (`torch.Tensor` of shape (batch_size, channels, height, width)): + A tensor that contains features map on a height * width grid + point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,: + 2)): + A tensor that contains [0, 1] * [0, 1] normalized point coordinates + add_dim (`bool`): + boolean value to keep track of added dimension + + Returns: + point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels, + height_grid, width_grid): + A tensor that contains features for points in `point_coordinates`. + """ + if point_coordinates.dim() == 3: + add_dim = True + point_coordinates = point_coordinates.unsqueeze(2) + + # use nn.function.grid_sample to get features for points in `point_coordinates` via bilinear interpolation + point_features = torch.nn.functional.grid_sample(input_features, 2.0 * point_coordinates - 1.0, **kwargs) + if add_dim: + point_features = point_features.squeeze(3) + + return point_features + + +# Copied from transformers.models.maskformer.modeling_maskformer.dice_loss +def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor: + r""" + Compute the DICE loss, similar to generalized IOU for masks as follows: + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$ + + In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$ + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + num_masks (`int`): + The number of masks present in the current batch, used for normalization. + + Returns: + `torch.Tensor`: The computed loss. + """ + probs = inputs.sigmoid().flatten(1) + numerator = 2 * (probs * labels).sum(-1) + denominator = probs.sum(-1) + labels.sum(-1) + loss = 1 - (numerator + 1) / (denominator + 1) + loss = loss.sum() / num_masks + return loss + + +def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor: + r""" + Args: + inputs (`torch.Tensor`): + A float tensor of arbitrary shape. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss. + """ + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss = criterion(inputs, labels) + + loss = cross_entropy_loss.mean(1).sum() / num_masks + return loss + + +# Copied from transformers.models.maskformer.modeling_maskformer.pair_wise_dice_loss +def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor: + """ + A pair wise version of the dice loss, see `dice_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + `torch.Tensor`: The computed loss between each pairs. + """ + inputs = inputs.sigmoid().flatten(1) + numerator = 2 * torch.matmul(inputs, labels.T) + # using broadcasting to get a [num_queries, NUM_CLASSES] matrix + denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :] + loss = 1 - (numerator + 1) / (denominator + 1) + return loss + + +def pair_wise_sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + r""" + A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss between each pairs. + """ + + height_and_width = inputs.shape[1] + + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) + cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) + + loss_pos = torch.matmul(cross_entropy_loss_pos / height_and_width, labels.T) + loss_neg = torch.matmul(cross_entropy_loss_neg / height_and_width, (1 - labels).T) + loss = loss_pos + loss_neg + return loss + + +# Adapted from https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/matcher.py +class Mask2FormerHungarianMatcher(nn.Module): + """This class computes an assignment between the labels and the predictions of the network. + + For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more + predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are + un-matched (and thus treated as non-objects). + """ + + def __init__( + self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0, num_points: int = 12544 + ): + """Creates the matcher + + Params: + cost_class (`float`, *optional*, defaults to 1.0): + Relative weight of the classification error in the matching cost. + cost_mask (`float`, *optional*, defaults to 1.0): + This is the relative weight of the focal loss of the binary mask in the matching cost. + cost_dice (`float`, *optional*, defaults to 1.0): + This is the relative weight of the dice loss of the binary mask in the matching cost. + num_points (`int`, *optional*, defaults to 12544): + No. of points to sample on which the mask loss will be calculated. The same set of K points are + uniformly sampled for all prediction and ground truth masks to construct the cost matrix for bipartite + matching. + """ + super().__init__() + if cost_class == 0 and cost_mask == 0 and cost_dice == 0: + raise ValueError("All costs can't be 0") + + self.num_points = num_points + self.cost_class = cost_class + self.cost_mask = cost_mask + self.cost_dice = cost_dice + + @torch.no_grad() + def forward( + self, + masks_queries_logits: torch.Tensor, + class_queries_logits: torch.Tensor, + mask_labels: torch.Tensor, + class_labels: torch.Tensor, + ) -> list[tuple[Tensor]]: + """ + Params: + masks_queries_logits (`torch.Tensor`): + A tensor of dim `batch_size, num_queries, num_labels` with the classification logits. + class_queries_logits (`torch.Tensor`): + A tensor of dim `batch_size, num_queries, height, width` with the predicted masks. + class_labels (`torch.Tensor`): + A tensor of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the + target) containing the class labels. + mask_labels (`torch.Tensor`): + A tensor of dim `num_target_boxes, height, width` containing the target masks. + + Returns: + matched_indices (`list[tuple[Tensor]]`): A list of size batch_size, containing tuples of (index_i, index_j) + where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected labels (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes). + """ + indices: list[tuple[np.array]] = [] + + # iterate through batch size + batch_size = masks_queries_logits.shape[0] + for i in range(batch_size): + pred_probs = class_queries_logits[i].softmax(-1) + pred_mask = masks_queries_logits[i] + + # Compute the classification cost. Contrary to the loss, we don't use the NLL, but approximate it in 1 - proba[target class]. The 1 is a constant that doesn't change the matching, it can be omitted. + cost_class = -pred_probs[:, class_labels[i]] + target_mask = mask_labels[i].to(pred_mask) + target_mask = target_mask[:, None] + pred_mask = pred_mask[:, None] + + # Sample ground truth and predicted masks + point_coordinates = torch.rand(1, self.num_points, 2, device=pred_mask.device) + + target_coordinates = point_coordinates.repeat(target_mask.shape[0], 1, 1) + target_mask = sample_point(target_mask, target_coordinates, align_corners=False).squeeze(1) + + pred_coordinates = point_coordinates.repeat(pred_mask.shape[0], 1, 1) + pred_mask = sample_point(pred_mask, pred_coordinates, align_corners=False).squeeze(1) + + # compute the cross entropy loss between each mask pairs -> shape (num_queries, num_labels) + cost_mask = pair_wise_sigmoid_cross_entropy_loss(pred_mask, target_mask) + # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels) + cost_dice = pair_wise_dice_loss(pred_mask, target_mask) + # final cost matrix + cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice + # eliminate infinite values in cost_matrix to avoid the error ``ValueError: cost matrix is infeasible`` + cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10)) + cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10)) + cost_matrix = torch.nan_to_num(cost_matrix, 0) + # do the assignment using the hungarian algorithm in scipy + assigned_indices: tuple[np.array] = linear_sum_assignment(cost_matrix.cpu()) + indices.append(assigned_indices) + + # It could be stacked in one tensor + matched_indices = [ + (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices + ] + return matched_indices + + +# Adapted from https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/criterion.py +class Mask2FormerLoss(nn.Module): + def __init__(self, config: Mask2FormerConfig, weight_dict: dict[str, float]): + """ + The Mask2Former Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we + compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair + of matched ground-truth / prediction (supervise class and mask) + + Args: + config (`Mask2FormerConfig`): + The configuration for Mask2Former model also containing loss calculation specific parameters. + weight_dict (`dict[str, float]`): + A dictionary of weights to be applied to the different losses. + """ + super().__init__() + requires_backends(self, ["scipy"]) + self.num_labels = config.num_labels + self.weight_dict = weight_dict + + # Weight to apply to the null class + self.eos_coef = config.no_object_weight + empty_weight = torch.ones(self.num_labels + 1) + empty_weight[-1] = self.eos_coef + self.register_buffer("empty_weight", empty_weight) + + # pointwise mask loss parameters + self.num_points = config.train_num_points + self.oversample_ratio = config.oversample_ratio + self.importance_sample_ratio = config.importance_sample_ratio + + self.matcher = Mask2FormerHungarianMatcher( + cost_class=config.class_weight, + cost_dice=config.dice_weight, + cost_mask=config.mask_weight, + num_points=self.num_points, + ) + + def _max_by_axis(self, sizes: list[list[int]]) -> list[int]: + maxes = sizes[0] + for sublist in sizes[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + # Adapted from nested_tensor_from_tensor_list() in original implementation + def _pad_images_to_max_in_batch(self, tensors: list[Tensor]) -> tuple[Tensor, Tensor]: + # get the maximum size in the batch + max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors]) + # compute final size + batch_shape = [len(tensors)] + max_size + batch_size, _, height, width = batch_shape + dtype = tensors[0].dtype + device = tensors[0].device + padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device) + padding_masks = torch.ones((batch_size, height, width), dtype=torch.bool, device=device) + # pad the tensors to the size of the biggest one + for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks): + padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor) + padding_mask[: tensor.shape[1], : tensor.shape[2]] = False + + return padded_tensors, padding_masks + + def loss_labels( + self, class_queries_logits: Tensor, class_labels: list[Tensor], indices: tuple[np.array] + ) -> dict[str, Tensor]: + """Compute the losses related to the labels using cross entropy. + + Args: + class_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, num_labels` + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + """ + pred_logits = class_queries_logits + batch_size, num_queries, _ = pred_logits.shape + criterion = nn.CrossEntropyLoss(weight=self.empty_weight) + idx = self._get_predictions_permutation_indices(indices) # shape of (batch_size, num_queries) + target_classes_o = torch.cat( + [target[j] for target, (_, j) in zip(class_labels, indices)] + ) # shape of (batch_size, num_queries) + target_classes = torch.full( + (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device + ) + target_classes[idx] = target_classes_o + # Permute target_classes (batch_size, num_queries, num_labels) -> (batch_size, num_labels, num_queries) + pred_logits_transposed = pred_logits.transpose(1, 2) + loss_ce = criterion(pred_logits_transposed, target_classes) + losses = {"loss_cross_entropy": loss_ce} + return losses + + def loss_masks( + self, + masks_queries_logits: torch.Tensor, + mask_labels: list[torch.Tensor], + indices: tuple[np.array], + num_masks: int, + ) -> dict[str, torch.Tensor]: + """Compute the losses related to the masks using sigmoid_cross_entropy_loss and dice loss. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, height, width)`. + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + num_masks (`int)`: + The number of masks, used for normalization. + + Returns: + losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing two keys: + - **loss_mask** -- The loss computed using sigmoid cross entropy loss on the predicted and ground truth. + masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth, + masks. + """ + src_idx = self._get_predictions_permutation_indices(indices) + tgt_idx = self._get_targets_permutation_indices(indices) + # shape (batch_size * num_queries, height, width) + pred_masks = masks_queries_logits[src_idx] + # shape (batch_size, num_queries, height, width) + # pad all and stack the targets to the num_labels dimension + target_masks, _ = self._pad_images_to_max_in_batch(mask_labels) + target_masks = target_masks[tgt_idx] + + # No need to upsample predictions as we are using normalized coordinates + pred_masks = pred_masks[:, None] + target_masks = target_masks[:, None] + + # Sample point coordinates + with torch.no_grad(): + point_coordinates = self.sample_points_using_uncertainty( + pred_masks, + lambda logits: self.calculate_uncertainty(logits), + self.num_points, + self.oversample_ratio, + self.importance_sample_ratio, + ) + + point_labels = sample_point(target_masks, point_coordinates, align_corners=False).squeeze(1) + + point_logits = sample_point(pred_masks, point_coordinates, align_corners=False).squeeze(1) + + losses = { + "loss_mask": sigmoid_cross_entropy_loss(point_logits, point_labels, num_masks), + "loss_dice": dice_loss(point_logits, point_labels, num_masks), + } + + del pred_masks + del target_masks + return losses + + def _get_predictions_permutation_indices(self, indices): + # Permute predictions following indices + batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) + predictions_indices = torch.cat([src for (src, _) in indices]) + return batch_indices, predictions_indices + + def _get_targets_permutation_indices(self, indices): + # Permute labels following indices + batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) + target_indices = torch.cat([tgt for (_, tgt) in indices]) + return batch_indices, target_indices + + def calculate_uncertainty(self, logits: torch.Tensor) -> torch.Tensor: + """ + In Mask2Former paper, uncertainty is estimated as L1 distance between 0.0 and the logit prediction in 'logits' + for the foreground class in `classes`. + + Args: + logits (`torch.Tensor`): + A tensor of shape (R, 1, ...) for class-specific or class-agnostic, where R is the total number of predicted masks in all images and C is: + the number of foreground classes. The values are logits. + + Returns: + scores (`torch.Tensor`): A tensor of shape (R, 1, ...) that contains uncertainty scores with the most + uncertain locations having the highest uncertainty score. + """ + uncertainty_scores = -(torch.abs(logits)) + return uncertainty_scores + + def sample_points_using_uncertainty( + self, + logits: torch.Tensor, + uncertainty_function, + num_points: int, + oversample_ratio: int, + importance_sample_ratio: float, + ) -> torch.Tensor: + """ + This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The + uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit + prediction as input. + + Args: + logits (`float`): + Logit predictions for P points. + uncertainty_function: + A function that takes logit predictions for P points and returns their uncertainties. + num_points (`int`): + The number of points P to sample. + oversample_ratio (`int`): + Oversampling parameter. + importance_sample_ratio (`float`): + Ratio of points that are sampled via importance sampling. + + Returns: + point_coordinates (`torch.Tensor`): + Coordinates for P sampled points. + """ + + num_boxes = logits.shape[0] + num_points_sampled = int(num_points * oversample_ratio) + + # Get random point coordinates + point_coordinates = torch.rand(num_boxes, num_points_sampled, 2, device=logits.device) + # Get sampled prediction value for the point coordinates + point_logits = sample_point(logits, point_coordinates, align_corners=False) + # Calculate the uncertainties based on the sampled prediction values of the points + point_uncertainties = uncertainty_function(point_logits) + + num_uncertain_points = int(importance_sample_ratio * num_points) + num_random_points = num_points - num_uncertain_points + + idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] + shift = num_points_sampled * torch.arange(num_boxes, dtype=torch.long, device=logits.device) + idx += shift[:, None] + point_coordinates = point_coordinates.view(-1, 2)[idx.view(-1), :].view(num_boxes, num_uncertain_points, 2) + + if num_random_points > 0: + point_coordinates = torch.cat( + [point_coordinates, torch.rand(num_boxes, num_random_points, 2, device=logits.device)], + dim=1, + ) + return point_coordinates + + def forward( + self, + masks_queries_logits: torch.Tensor, + class_queries_logits: torch.Tensor, + mask_labels: list[torch.Tensor], + class_labels: list[torch.Tensor], + auxiliary_predictions: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: + """ + This performs the loss computation. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, height, width)`. + class_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, num_labels)`. + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*): + if `use_auxiliary_loss` was set to `true` in [`Mask2FormerConfig`], then it contains the logits from + the inner layers of the Mask2FormerMaskedAttentionDecoder. + + Returns: + losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing three keys: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + - **loss_mask** -- The loss computed using sigmoid cross_entropy loss on the predicted and ground truth + masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth + masks. + if `use_auxiliary_loss` was set to `true` in [`Mask2FormerConfig`], the dictionary contains additional + losses for each auxiliary predictions. + """ + + # retrieve the matching between the outputs of the last layer and the labels + indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + # compute the average number of target masks for normalization purposes + num_masks = self.get_num_masks(class_labels, device=class_labels[0].device) + # get all the losses + losses: dict[str, Tensor] = { + **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), + **self.loss_labels(class_queries_logits, class_labels, indices), + } + # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if auxiliary_predictions is not None: + for idx, aux_outputs in enumerate(auxiliary_predictions): + masks_queries_logits = aux_outputs["masks_queries_logits"] + class_queries_logits = aux_outputs["class_queries_logits"] + loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} + losses.update(loss_dict) + + return losses + + def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Computes the average number of target masks across the batch, for normalization purposes. + """ + num_masks = sum(len(classes) for classes in class_labels) + num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device) + world_size = 1 + if is_accelerate_available(): + if PartialState._shared_state != {}: + num_masks = reduce(num_masks) + world_size = PartialState().num_processes + + num_masks = torch.clamp(num_masks / world_size, min=1) + return num_masks + + +# Copied from transformers.models.oneformer.modeling_oneformer.multi_scale_deformable_attention +def multi_scale_deformable_attention( + value: Tensor, + value_spatial_shapes: Tensor | list[tuple], + sampling_locations: Tensor, + attention_weights: Tensor, +) -> Tensor: + batch_size, _, num_heads, hidden_dim = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([height * width for height, width in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level_id, (height, width) in enumerate(value_spatial_shapes): + # batch_size, height*width, num_heads, hidden_dim + # -> batch_size, height*width, num_heads*hidden_dim + # -> batch_size, num_heads*hidden_dim, height*width + # -> batch_size*num_heads, hidden_dim, height, width + value_l_ = ( + value_list[level_id].flatten(2).transpose(1, 2).reshape(batch_size * num_heads, hidden_dim, height, width) + ) + # batch_size, num_queries, num_heads, num_points, 2 + # -> batch_size, num_heads, num_queries, num_points, 2 + # -> batch_size*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1) + # batch_size*num_heads, hidden_dim, num_queries, num_points + sampling_value_l_ = nn.functional.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (batch_size, num_queries, num_heads, num_levels, num_points) + # -> (batch_size, num_heads, num_queries, num_levels, num_points) + # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + batch_size * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(batch_size, num_heads * hidden_dim, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +# Copied from transformers.models.maskformer.modeling_maskformer.MaskFormerSinePositionEmbedding with MaskFormer->Mask2Former +class Mask2FormerSinePositionEmbedding(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_pos_feats: 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_pos_feats = num_pos_feats + 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: Tensor | None = None, + ) -> Tensor: + if mask is None: + mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool) + not_mask = (~mask).to(dtype) + y_embed = not_mask.cumsum(1) + x_embed = not_mask.cumsum(2) + 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_pos_feats, dtype=torch.int64, device=device).to(dtype) + dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_pos_feats) + + 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) + return pos + + +# Modified from transformers.models.detr.modeling_deformable_detr.DeformableDetrMultiscaleDeformableAttention +class Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention(nn.Module): + """ + Multiscale deformable attention as proposed in Deformable DETR. + """ + + def __init__(self, embed_dim: int, num_heads: int, n_levels: int, n_points: int): + super().__init__() + if embed_dim % num_heads != 0: + raise ValueError( + f"embed_dim (d_model) must be divisible by num_heads, but got {embed_dim} and {num_heads}" + ) + dim_per_head = embed_dim // num_heads + # check if dim_per_head is power of 2 + if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0): + warnings.warn( + "You'd better set embed_dim (d_model) in DeformableDetrMultiscaleDeformableAttention to make the" + " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA" + " implementation." + ) + + self.im2col_step = 128 + + self.d_model = embed_dim + self.n_levels = n_levels + self.n_heads = num_heads + self.n_points = n_points + + self.sampling_offsets = nn.Linear(embed_dim, num_heads * n_levels * n_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_heads * n_levels * n_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): + return tensor if position_embeddings is None else tensor + position_embeddings + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states=None, + encoder_attention_mask=None, + position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes_list=None, + level_start_index=None, + output_attentions: bool = False, + ): + # add position embeddings to the hidden states before projecting to queries and keys + if position_embeddings is not None: + hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + + batch_size, num_queries, _ = hidden_states.shape + batch_size, sequence_length, _ = encoder_hidden_states.shape + total_elements = sum(height * width for height, width in spatial_shapes_list) + torch_compilable_check( + total_elements == sequence_length, + "Make sure to align the spatial shapes with the sequence length of the encoder hidden states", + ) + + value = self.value_proj(encoder_hidden_states) + if attention_mask is not None: + # we invert the attention_mask + value = value.masked_fill(attention_mask[..., None], float(0)) + value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads) + sampling_offsets = self.sampling_offsets(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2 + ) + attention_weights = self.attention_weights(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels * self.n_points + ) + attention_weights = nn.functional.softmax(attention_weights, -1).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points + ) + # batch_size, num_queries, n_heads, n_levels, n_points, 2 + if reference_points.shape[-1] == 2: + offset_normalizer = torch.tensor( + [[shape[1], shape[0]] for shape in spatial_shapes_list], + dtype=torch.long, + device=reference_points.device, + ) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) + else: + raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}") + + output = multi_scale_deformable_attention(value, spatial_shapes_list, sampling_locations, attention_weights) + output = self.output_proj(output) + + return output, attention_weights + + +class Mask2FormerPixelDecoderEncoderLayer(nn.Module): + def __init__(self, config: Mask2FormerConfig): + super().__init__() + self.embed_dim = config.feature_size + self.self_attn = Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention( + embed_dim=self.embed_dim, + num_heads=config.num_attention_heads, + n_levels=3, + n_points=4, + ) + + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = nn.functional.relu + self.activation_dropout = config.dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_feedforward_dim) + self.fc2 = nn.Linear(config.encoder_feedforward_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes_list=None, + level_start_index=None, + output_attentions: bool = False, + ): + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Input to the layer. + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Attention mask. + position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings, to be added to `hidden_states`. + reference_points (`torch.FloatTensor`, *optional*): + Reference points. + spatial_shapes_list (`list` of `tuple`): + Spatial shapes of the backbone feature maps as a list of tuples. + level_start_index (`torch.LongTensor`, *optional*): + Level start index. + 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 + + # Apply Multi-scale Deformable Attention Module on the multi-scale feature maps. + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + position_embeddings=position_embeddings, + reference_points=reference_points, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + 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) + + 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 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) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights.transpose(1, 0),) + + return outputs + + +# Modified from from transformers.models.detr.modeling_deformable_detr.DeformableDetrEncoder with DeformableDetrEncoder->Mask2FormerPixelDecoderEncoderOnly +class Mask2FormerPixelDecoderEncoderOnly(nn.Module): + """ + Transformer encoder consisting of *config.encoder_layers* deformable attention layers. Each layer is a + [`Mask2FormerPixelDecoderEncoderLayer`]. The encoder updates the flattened multi-scale feature maps through + multiple deformable attention layers. + + Args: + config: Mask2FormerConfig + """ + + def __init__(self, config: Mask2FormerConfig): + super().__init__() + + self.config = config + self.dropout = config.dropout + self.layers = nn.ModuleList( + [Mask2FormerPixelDecoderEncoderLayer(config) for _ in range(config.encoder_layers)] + ) + + @staticmethod + def get_reference_points(spatial_shapes_list, valid_ratios, device): + """ + Get reference points for each feature map. Used in decoder. + + Args: + spatial_shapes_list (`list` of `tuple`): + Spatial shapes of the backbone feature maps as a list of tuples. + valid_ratios (`torch.FloatTensor`): + Valid ratios of each feature map, has shape of `(batch_size, num_feature_levels, 2)`. + device (`torch.device`): + Device on which to create the tensors. + Returns: + `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)` + """ + reference_points_list = [] + for lvl, (height, width) in enumerate(spatial_shapes_list): + ref_y, ref_x = torch.meshgrid( + torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device), + torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device), + indexing="ij", + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * height) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * width) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + + return reference_points + + def forward( + self, + inputs_embeds=None, + attention_mask=None, + position_embeddings=None, + spatial_shapes_list=None, + level_start_index=None, + valid_ratios=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + 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) + position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Position embeddings that are added to the queries and keys in each self-attention layer. + spatial_shapes_list (`list` of `tuple`): + Spatial shapes of each feature map as a list of tuples. + level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`): + Starting index of each feature map. + valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): + Ratio of valid area in each feature level. + 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 [`~file_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 + + hidden_states = inputs_embeds + reference_points = self.get_reference_points(spatial_shapes_list, valid_ratios, device=inputs_embeds.device) + + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + for i, encoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states.transpose(1, 0),) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + position_embeddings=position_embeddings, + reference_points=reference_points, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states += (hidden_states.transpose(1, 0),) + + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions + ) + + +# Modified from from transformers.models.detr.modeling_deformable_detr.DeformableDetrModel with DeformableDetrModel->Mask2FormerPixelDecoder +class Mask2FormerPixelDecoder(nn.Module): + def __init__(self, config: Mask2FormerConfig, feature_channels): + super().__init__() + + self.config = config + + feature_dim = config.feature_size + mask_dim = config.mask_feature_size + num_pos_features = feature_dim // 2 + + self.position_embedding = Mask2FormerSinePositionEmbedding(num_pos_feats=num_pos_features, normalize=True) + self.num_feature_levels = 3 + transformer_in_channels = feature_channels[-self.num_feature_levels :] + + self.transformer_feature_strides = config.feature_strides[-self.num_feature_levels :] + self.feature_channels = feature_channels + self.level_embed = nn.Parameter(torch.Tensor(self.num_feature_levels, feature_dim)) + + # Create input projection layers + if self.num_feature_levels > 1: + input_projections_list = [] + for in_channels in transformer_in_channels[::-1]: + input_projections_list.append( + nn.Sequential( + nn.Conv2d(in_channels, feature_dim, kernel_size=1), + nn.GroupNorm(32, feature_dim), + ) + ) + self.input_projections = nn.ModuleList(input_projections_list) + else: + self.input_projections = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(transformer_in_channels[-1], feature_dim, kernel_size=1), + nn.GroupNorm(32, feature_dim), + ) + ] + ) + + self.encoder = Mask2FormerPixelDecoderEncoderOnly(config) + self.mask_projection = nn.Conv2d(feature_dim, mask_dim, kernel_size=1, stride=1, padding=0) + + # Extra FPN levels + stride = min(self.transformer_feature_strides) + self.common_stride = config.common_stride + self.num_fpn_levels = int(np.log2(stride) - np.log2(self.common_stride)) + + lateral_convs = [] + output_convs = [] + + for idx, in_channels in enumerate(self.feature_channels[: self.num_fpn_levels]): + lateral_conv = nn.Sequential( + nn.Conv2d(in_channels, feature_dim, kernel_size=1, bias=False), + nn.GroupNorm(32, feature_dim), + ) + + output_conv = nn.Sequential( + nn.Conv2d(feature_dim, feature_dim, kernel_size=3, stride=1, padding=1, bias=False), + nn.GroupNorm(32, feature_dim), + nn.ReLU(), + ) + self.add_module(f"adapter_{idx + 1}", lateral_conv) + self.add_module(f"layer_{idx + 1}", output_conv) + + lateral_convs.append(lateral_conv) + output_convs.append(output_conv) + + # Order convolutional layers from low to high resolution + self.lateral_convolutions = lateral_convs[::-1] + self.output_convolutions = output_convs[::-1] + + def get_valid_ratio(self, mask, dtype=torch.float32): + """Get the valid ratio of all feature maps.""" + + _, height, width = mask.shape + valid_height = torch.sum(~mask[:, :, 0], 1) + valid_width = torch.sum(~mask[:, 0, :], 1) + valid_ratio_height = valid_height.to(dtype) / height + valid_ratio_width = valid_width.to(dtype) / width + valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1) + return valid_ratio + + def forward( + self, + features, + encoder_outputs=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + 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 + ) + + # Apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) + inputs_embeds = [] + position_embeddings = [] + for level, x in enumerate(features[::-1][: self.num_feature_levels]): + inputs_embeds.append(self.input_projections[level](x)) + position_embeddings.append(self.position_embedding(x.shape, x.device, x.dtype)) + + masks = [ + torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in inputs_embeds + ] + + # Prepare encoder inputs (by flattening) + spatial_shapes_list = [(embed.shape[2], embed.shape[3]) for embed in inputs_embeds] + input_embeds_flat = torch.cat([embed.flatten(2).transpose(1, 2) for embed in inputs_embeds], 1) + spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=input_embeds_flat.device) + masks_flat = torch.cat([mask.flatten(1) for mask in masks], 1) + + position_embeddings = [embed.flatten(2).transpose(1, 2) for embed in position_embeddings] + level_pos_embed_flat = [x + self.level_embed[i].view(1, 1, -1) for i, x in enumerate(position_embeddings)] + level_pos_embed_flat = torch.cat(level_pos_embed_flat, 1) + + level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) + valid_ratios = torch.stack([self.get_valid_ratio(mask, dtype=input_embeds_flat.dtype) for mask in masks], 1) + + # Send input_embeds_flat + masks_flat + level_pos_embed_flat (backbone + proj layer output) through encoder + if encoder_outputs is None: + encoder_outputs = self.encoder( + inputs_embeds=input_embeds_flat, + attention_mask=masks_flat, + position_embeddings=level_pos_embed_flat, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + valid_ratios=valid_ratios, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + batch_size = last_hidden_state.shape[0] + + # We compute level_start_index_list separately from the tensor version level_start_index + # to avoid iterating over a tensor which breaks torch.compile/export. + level_start_index_list = [0] + for height, width in spatial_shapes_list[:-1]: + level_start_index_list.append(level_start_index_list[-1] + height * width) + split_sizes = [None] * self.num_feature_levels + for i in range(self.num_feature_levels): + if i < self.num_feature_levels - 1: + split_sizes[i] = level_start_index_list[i + 1] - level_start_index_list[i] + else: + split_sizes[i] = last_hidden_state.shape[1] - level_start_index_list[i] + + encoder_output = torch.split(last_hidden_state, split_sizes, dim=1) + + # Compute final features + outputs = [ + x.transpose(1, 2).view(batch_size, -1, spatial_shapes_list[i][0], spatial_shapes_list[i][1]) + for i, x in enumerate(encoder_output) + ] + + # Append extra FPN levels to outputs, ordered from low to high resolution + for idx, feature in enumerate(features[: self.num_fpn_levels][::-1]): + lateral_conv = self.lateral_convolutions[idx] + output_conv = self.output_convolutions[idx] + current_fpn = lateral_conv(feature) + + # Following FPN implementation, we use nearest upsampling here + out = current_fpn + nn.functional.interpolate( + outputs[-1], size=current_fpn.shape[-2:], mode="bilinear", align_corners=False + ) + out = output_conv(out) + outputs.append(out) + + num_cur_levels = 0 + multi_scale_features = [] + + for out in outputs: + if num_cur_levels < self.num_feature_levels: + multi_scale_features.append(out) + num_cur_levels += 1 + + return Mask2FormerPixelDecoderOutput( + mask_features=self.mask_projection(outputs[-1]), + multi_scale_features=tuple(multi_scale_features), + attentions=encoder_outputs.attentions, + ) + + +class Mask2FormerPixelLevelModule(nn.Module): + def __init__(self, config: Mask2FormerConfig): + """ + Pixel Level Module proposed in [Masked-attention Mask Transformer for Universal Image + Segmentation](https://huggingface.co/papers/2112.01527). It runs the input image through a backbone and a pixel + decoder, generating multi-scale feature maps and pixel embeddings. + + Args: + config ([`Mask2FormerConfig`]): + The configuration used to instantiate this model. + """ + super().__init__() + + self.encoder = load_backbone(config) + self.decoder = Mask2FormerPixelDecoder(config, feature_channels=self.encoder.channels) + + def forward(self, pixel_values: Tensor, output_hidden_states: bool = False) -> Mask2FormerPixelLevelModuleOutput: + backbone_features = self.encoder(pixel_values).feature_maps + decoder_output = self.decoder(backbone_features, output_hidden_states=output_hidden_states) + + return Mask2FormerPixelLevelModuleOutput( + encoder_last_hidden_state=backbone_features[-1], + encoder_hidden_states=tuple(backbone_features) if output_hidden_states else None, + decoder_last_hidden_state=decoder_output.mask_features, + decoder_hidden_states=decoder_output.multi_scale_features, + ) + + +# Modified from transformers.models.detr.modeling_detr.DetrAttention with Detr->Mask2Former +class Mask2FormerAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Here, we add position embeddings to the queries and + keys (as explained in the DETR paper). + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + ): + 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.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 _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int): + return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): + return tensor if position_embeddings is None else tensor + position_embeddings + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + key_value_states: torch.Tensor | None = None, + key_value_position_embeddings: torch.Tensor | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + hidden_states = hidden_states.permute(1, 0, 2) if hidden_states is not None else None + position_embeddings = position_embeddings.permute(1, 0, 2) if position_embeddings is not None else None + key_value_states = key_value_states.permute(1, 0, 2) if key_value_states is not None else None + key_value_position_embeddings = ( + key_value_position_embeddings.permute(1, 0, 2) if key_value_position_embeddings is not None else None + ) + + # 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 + batch_size, target_len, embed_dim = hidden_states.size() + + # add position embeddings to the hidden states before projecting to queries and keys + if position_embeddings is not None: + hidden_states_original = hidden_states + hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + + # add key-value position embeddings to the key value states + if key_value_position_embeddings is not None: + key_value_states_original = key_value_states + key_value_states = self.with_pos_embed(key_value_states, key_value_position_embeddings) + + # get query proj + query_states = self.q_proj(hidden_states) * self.scaling + # get key, value proj + if is_cross_attention: + # cross_attentions + key_states = self._shape(self.k_proj(key_value_states), -1, batch_size) + value_states = self._shape(self.v_proj(key_value_states_original), -1, batch_size) + else: + # self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, batch_size) + value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size) + + proj_shape = (batch_size * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_states = value_states.view(*proj_shape) + + source_len = key_states.size(1) + + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) + + if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len): + raise ValueError( + f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (batch_size * self.num_heads, target_len, source_len): + raise ValueError( + f"Attention mask should be of size {(target_len, batch_size * self.num_heads, source_len)}, but is" + f" {attention_mask.size()}" + ) + attn_weights += attention_mask + + 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 reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_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() != (batch_size * self.num_heads, target_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(batch_size, target_len, embed_dim) + + attn_output = self.out_proj(attn_output).permute(1, 0, 2) + + return attn_output, attn_weights_reshaped + + +class Mask2FormerMaskedAttentionDecoderLayer(GradientCheckpointingLayer): + """ + The Mask2FormerMaskedAttentionDecoderLayer is made up of self-attention, cross (masked) attention as well as FFN + blocks. The cross attention block used as part of `Mask2FormerMaskedAttentionDecoderLayer` is actually a `masked + attention` block that restricts the attention to localized features centered around predicted segments which leads + to faster convergence and improved performance. The order of self and cross (i.e. masked) attention blocks have + also been swapped in Mask2FormerMaskedAttentionDecoder compared to a standard DetrDecoder as an optimization + improvement. + + Args: + config (`Mask2FormerConfig`): + The configuration used to initialize the Mask2FormerMaskedAttentionDecoder. + """ + + def __init__(self, config: Mask2FormerConfig): + super().__init__() + self.config = config + self.embed_dim = self.config.hidden_dim + self.pre_norm = self.config.pre_norm + self.self_attn = Mask2FormerAttention( + embed_dim=self.embed_dim, + num_heads=config.num_attention_heads, + dropout=config.dropout, + is_decoder=True, + ) + + self.dropout = self.config.dropout + self.activation_fn = ACT2FN[self.config.activation_function] + self.activation_dropout = self.config.dropout + + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.cross_attn = nn.MultiheadAttention(self.embed_dim, self.config.num_attention_heads, self.config.dropout) + self.cross_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.fc1 = nn.Linear(self.embed_dim, self.config.dim_feedforward) + self.fc2 = nn.Linear(self.config.dim_feedforward, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def with_pos_embed(self, tensor, pos: Tensor | None): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + hidden_states: torch.Tensor, + level_index: int | None = None, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ): + # Masked(Cross)-Attention Block + cross_attn_weights = None + self_attn_weights = None + + residual = hidden_states + + hidden_states, cross_attn_weights = self.cross_attn( + query=self.with_pos_embed(hidden_states, query_position_embeddings), + key=self.with_pos_embed(encoder_hidden_states[level_index], position_embeddings[level_index]), + value=encoder_hidden_states[level_index], + attn_mask=encoder_attention_mask, + key_padding_mask=None, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.cross_attn_layer_norm(hidden_states) + + # Self Attention Block + residual = hidden_states + + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=query_position_embeddings, + attention_mask=None, + output_attentions=True, + ) + + 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) + + # 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) + + return outputs + + def forward_pre( + self, + hidden_states: torch.Tensor, + level_index: int | None = None, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ): + # Masked(Cross)-Attention Block + cross_attn_weights = None + self_attn_weights = None + + residual = hidden_states + + hidden_states = self.cross_attn_layer_norm(hidden_states) + + hidden_states, cross_attn_weights = self.cross_attn( + query=self.with_pos_embed(hidden_states, query_position_embeddings), + key=self.with_pos_embed(encoder_hidden_states[level_index], position_embeddings[level_index]), + value=encoder_hidden_states[level_index], + attn_mask=encoder_attention_mask, + key_padding_mask=None, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Self Attention Block + residual = hidden_states + + hidden_states = self.self_attn_layer_norm(hidden_states) + + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=query_position_embeddings, + attention_mask=None, + output_attentions=True, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights, cross_attn_weights) + + return outputs + + def forward( + self, + hidden_states: torch.Tensor, + level_index: int | None = None, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ): + """ + Args: + hidden_states (`torch.FloatTensor`): + Input to the layer of shape `(seq_len, batch, embed_dim)`. + attention_mask (`torch.FloatTensor`): + Attention mask of shape `(1, seq_len, tgt_len, src_len)`. + position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings that are added to the keys in the masked-attention layer. + query_position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings 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`(1, seq_len, tgt_len, src_len)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + + if self.pre_norm: + outputs = self.forward_pre( + hidden_states=hidden_states, + level_index=level_index, + position_embeddings=position_embeddings, + query_position_embeddings=query_position_embeddings, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + ) + else: + outputs = self.forward_post( + hidden_states=hidden_states, + level_index=level_index, + position_embeddings=position_embeddings, + query_position_embeddings=query_position_embeddings, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + ) + + return outputs + + +class Mask2FormerMaskedAttentionDecoder(nn.Module): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a + [`Mask2FormerMaskedAttentionDecoderLayer`]. The decoder updates the query embeddings through multiple cross + (masked) and self-attention layers. The decoder uses a new **masked attention** mechanism instead of the standard + cross-attention, which extracts localized features by constraining cross-attention to within the foreground region + of the predicted mask for each query, instead of attending to the full feature map. + + Args: + config (`Mask2FormerConfig`): + Configuration used to instantiate Mask2FormerMaskedAttentionDecoder. + """ + + def __init__(self, config: Mask2FormerConfig): + super().__init__() + + self.config = config + self.mask_feature_size = config.mask_feature_size + self.dropout = config.dropout + self.layerdrop = config.dropout + self.num_feature_levels = 3 # level embedding (3 scales) + self.decoder_layers = config.decoder_layers - 1 + + self.layers = nn.ModuleList( + [Mask2FormerMaskedAttentionDecoderLayer(self.config) for _ in range(self.decoder_layers)] + ) + self.layernorm = nn.LayerNorm(config.hidden_dim) + + self.mask_predictor = Mask2FormerMaskPredictor( + hidden_size=config.hidden_dim, + num_heads=config.num_attention_heads, + mask_feature_size=self.mask_feature_size, + ) + + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds: torch.Tensor | None = None, + multi_stage_positional_embeddings: torch.Tensor | None = None, + pixel_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + feature_size_list: list | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ): + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(num_queries, batch_size, hidden_size)`): + The query embeddings that are passed into the decoder. + multi_stage_positional_embeddings (`torch.FloatTensor` of shape `(height*width, batch_size, num_channels)`): + Position embeddings that are added to the keys in each cross(masked)-attention layer. + pixel_embeddings (`torch.FloatTensor`): + Tensor of shape `(batch_size, num_channels, height, width)`, 1/4 scale features from the last Pixel + Decoder. + query_position_embeddings (`torch.FloatTensor` of shape `(num_queries, batch_size, hidden_size)`): + , *optional*): Position embeddings that are added to the queries and keys in each self-attention layer. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the + cross(masked)-attention of the decoder. + feature_size_list (`list[torch.Size]`): + This is a list containing shapes (height & width) of multi-scale features from the Pixel Decoder. + 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 + + if inputs_embeds is not None: + hidden_states = inputs_embeds + + # intermediate hidden states with layernorm applied - required for predicting class logits + intermediate = () + + # decoder layers + all_hidden_states = () if output_hidden_states else None + attentions = () if output_attentions else None + + # intermediate mask predictions from transformer decoder layers + intermediate_mask_predictions = () + + intermediate_hidden_states = self.layernorm(inputs_embeds) + intermediate += (intermediate_hidden_states,) + + predicted_mask, attention_mask = self.mask_predictor( + intermediate_hidden_states, pixel_embeddings, feature_size_list[0] + ) + intermediate_mask_predictions += (predicted_mask,) + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + dropout_probability = torch.rand([]) + + if self.training and (dropout_probability < self.layerdrop): + continue + + level_index = idx % self.num_feature_levels + + where = (attention_mask.sum(-1) != attention_mask.shape[-1]).to(attention_mask.dtype) + # Multiply the attention mask instead of indexing to avoid issue in torch.export. + attention_mask = attention_mask * where.unsqueeze(-1) + + layer_outputs = decoder_layer( + hidden_states, + level_index, + None, # attention_mask + multi_stage_positional_embeddings, + query_position_embeddings, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=attention_mask, + output_attentions=output_attentions, + ) + + intermediate_hidden_states = self.layernorm(layer_outputs[0]) + + predicted_mask, attention_mask = self.mask_predictor( + intermediate_hidden_states, + pixel_embeddings, + feature_size_list[(idx + 1) % self.num_feature_levels], + ) + + intermediate_mask_predictions += (predicted_mask,) + + # add intermediate hidden states with layer norm applied which will be used for predicting class logits + intermediate += (intermediate_hidden_states,) + + hidden_states = layer_outputs[0] + + if output_attentions: + attentions += (layer_outputs[1],) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states = hidden_states.transpose(1, 0) + if not return_dict: + outputs = [hidden_states, all_hidden_states, attentions, intermediate, intermediate_mask_predictions] + return tuple(v for v in outputs if v is not None) + + return Mask2FormerMaskedAttentionDecoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=attentions, + intermediate_hidden_states=intermediate, + masks_queries_logits=intermediate_mask_predictions, + ) + + +# Copied from transformers.models.maskformer.modeling_maskformer.PredictionBlock with MaskFormer->Mask2Former +class Mask2FormerPredictionBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None: + super().__init__() + self.layers = [nn.Linear(in_dim, out_dim), activation] + # Maintain submodule indexing as if part of a Sequential block + for i, layer in enumerate(self.layers): + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class Mask2FormerMLPPredictionHead(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3): + """ + A classic Multi Layer Perceptron (MLP). + + Args: + input_dim (`int`): + The input dimensions. + hidden_dim (`int`): + The hidden dimensions. + output_dim (`int`): + The output dimensions. + num_layers (int, *optional*, defaults to 3): + The number of layers. + """ + super().__init__() + in_dims = [input_dim] + [hidden_dim] * (num_layers - 1) + out_dims = [hidden_dim] * (num_layers - 1) + [output_dim] + + self.layers = [] + for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)): + activation = nn.ReLU() if i < num_layers - 1 else nn.Identity() + layer = Mask2FormerPredictionBlock(in_dim, out_dim, activation=activation) + self.layers.append(layer) + # Provide backwards compatibility from when the class inherited from nn.Sequential + # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. + # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. + # self.my_layer_name = Layer() + # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register + # explicitly + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class Mask2FormerMaskPredictor(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mask_feature_size: torch.Tensor): + """ + This class is used to get the predicted mask for a given Mask2FormerMaskedAttentionDecoder layer. It also + generates the binarized attention mask associated with the given predicted mask. The attention mask obtained + using predicted mask of the (l-1)th decoder layer is fed to the cross(masked)-attention block of the next + decoder layer as input. + + Args: + hidden_size (`int`): + The feature dimension of the Mask2FormerMaskedAttentionDecoder + num_heads (`int`): + The number of heads used in the Mask2FormerMaskedAttentionDecoder + mask_feature_size (`torch.Tensor`): + one of the output dimensions of the predicted masks for each query + """ + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + + self.mask_embedder = Mask2FormerMLPPredictionHead(self.hidden_size, self.hidden_size, mask_feature_size) + + def forward( + self, outputs: torch.Tensor, pixel_embeddings: torch.Tensor, attention_mask_target_size: int | None = None + ): + mask_embeddings = self.mask_embedder(outputs.transpose(0, 1)) + + # Sum up over the channels + outputs_mask = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings) + + attention_mask = nn.functional.interpolate( + outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False + ) + + attention_mask = attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1) + attention_mask = (attention_mask.flatten(0, 1) < 0.5).bool() + attention_mask = attention_mask.detach() + + return outputs_mask, attention_mask + + +class Mask2FormerTransformerModule(nn.Module): + """ + The Mask2Former's transformer module. + """ + + def __init__(self, in_features: int, config: Mask2FormerConfig): + super().__init__() + hidden_dim = config.hidden_dim + self.num_feature_levels = 3 + self.position_embedder = Mask2FormerSinePositionEmbedding(num_pos_feats=hidden_dim // 2, normalize=True) + self.queries_embedder = nn.Embedding(config.num_queries, hidden_dim) + self.queries_features = nn.Embedding(config.num_queries, hidden_dim) + self.input_projections = [] + + for _ in range(self.num_feature_levels): + if in_features != hidden_dim or config.enforce_input_projection: + self.input_projections.append(nn.Conv2d(in_features, hidden_dim, kernel_size=1)) + else: + self.input_projections.append(nn.Sequential()) + + self.decoder = Mask2FormerMaskedAttentionDecoder(config=config) + self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim) + + def forward( + self, + multi_scale_features: list[Tensor], + mask_features: Tensor, + output_hidden_states: bool = False, + output_attentions: bool = False, + ) -> Mask2FormerMaskedAttentionDecoderOutput: + multi_stage_features = [] + multi_stage_positional_embeddings = [] + size_list = [] + + for i in range(self.num_feature_levels): + size_list.append(multi_scale_features[i].shape[-2:]) + multi_stage_positional_embeddings.append( + self.position_embedder( + multi_scale_features[i].shape, multi_scale_features[i].device, multi_scale_features[i].dtype, None + ).flatten(2) + ) + multi_stage_features.append( + self.input_projections[i](multi_scale_features[i]).flatten(2) + + self.level_embed.weight[i][None, :, None] + ) + + # Flatten (batch_size, num_channels, height, width) -> (height*width, batch_size, num_channels) + multi_stage_positional_embeddings[-1] = multi_stage_positional_embeddings[-1].permute(2, 0, 1) + multi_stage_features[-1] = multi_stage_features[-1].permute(2, 0, 1) + + _, batch_size, _ = multi_stage_features[0].shape + + # [num_queries, batch_size, num_channels] + query_embeddings = self.queries_embedder.weight.unsqueeze(1).repeat(1, batch_size, 1) + query_features = self.queries_features.weight.unsqueeze(1).repeat(1, batch_size, 1) + + decoder_output = self.decoder( + inputs_embeds=query_features, + multi_stage_positional_embeddings=multi_stage_positional_embeddings, + pixel_embeddings=mask_features, + encoder_hidden_states=multi_stage_features, + query_position_embeddings=query_embeddings, + feature_size_list=size_list, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + ) + + return decoder_output + + +@auto_docstring +class Mask2FormerPreTrainedModel(PreTrainedModel): + config: Mask2FormerConfig + base_model_prefix = "model" + main_input_name = "pixel_values" + input_modalities = ("image",) + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + xavier_std = self.config.init_xavier_std + std = self.config.init_std + + if isinstance(module, Mask2FormerTransformerModule): + if module.input_projections is not None: + for input_projection in module.input_projections: + if not isinstance(input_projection, nn.Sequential): + init.xavier_uniform_(input_projection.weight, gain=xavier_std) + init.constant_(input_projection.bias, 0) + + elif isinstance(module, Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention): + init.constant_(module.sampling_offsets.weight, 0.0) + thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (2.0 * math.pi / module.n_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(module.n_heads, 1, 1, 2) + .repeat(1, module.n_levels, module.n_points, 1) + ) + for i in range(module.n_points): + grid_init[:, :, i, :] *= i + 1 + + init.copy_(module.sampling_offsets.bias, grid_init.view(-1)) + + init.constant_(module.attention_weights.weight, 0.0) + init.constant_(module.attention_weights.bias, 0.0) + init.xavier_uniform_(module.value_proj.weight) + init.constant_(module.value_proj.bias, 0.0) + init.xavier_uniform_(module.output_proj.weight) + init.constant_(module.output_proj.bias, 0.0) + + elif isinstance(module, Mask2FormerMaskedAttentionDecoderLayer): + for p in module.parameters(): + if p.dim() > 1: + init.xavier_uniform_(p, gain=xavier_std) + init.zeros_(module.cross_attn.in_proj_bias) + + elif isinstance(module, Mask2FormerPixelDecoder): + init.normal_(module.level_embed, std=0) + + elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + if getattr(module, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.ones_(module.weight) + 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, Mask2FormerLoss): + empty_weight = torch.ones(module.num_labels + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + + if hasattr(module, "reference_points"): + init.xavier_uniform_(module.reference_points.weight, gain=1.0) + init.constant_(module.reference_points.bias, 0.0) + + +@auto_docstring +class Mask2FormerModel(Mask2FormerPreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: Mask2FormerConfig): + super().__init__(config) + self.pixel_level_module = Mask2FormerPixelLevelModule(config) + self.transformer_module = Mask2FormerTransformerModule(in_features=config.feature_size, config=config) + + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + pixel_mask: Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> Mask2FormerModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + batch_size, _, height, width = pixel_values.shape + + if pixel_mask is None: + pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device) + + pixel_level_module_output = self.pixel_level_module( + pixel_values=pixel_values, output_hidden_states=output_hidden_states + ) + + transformer_module_output = self.transformer_module( + multi_scale_features=pixel_level_module_output.decoder_hidden_states, + mask_features=pixel_level_module_output.decoder_last_hidden_state, + output_hidden_states=True, + output_attentions=output_attentions, + ) + + encoder_hidden_states = None + pixel_decoder_hidden_states = None + transformer_decoder_hidden_states = None + transformer_decoder_intermediate_states = None + + if output_hidden_states: + encoder_hidden_states = pixel_level_module_output.encoder_hidden_states + pixel_decoder_hidden_states = pixel_level_module_output.decoder_hidden_states + transformer_decoder_hidden_states = transformer_module_output.hidden_states + transformer_decoder_intermediate_states = transformer_module_output.intermediate_hidden_states + + output = Mask2FormerModelOutput( + encoder_last_hidden_state=pixel_level_module_output.encoder_last_hidden_state, + pixel_decoder_last_hidden_state=pixel_level_module_output.decoder_last_hidden_state, + transformer_decoder_last_hidden_state=transformer_module_output.last_hidden_state, + encoder_hidden_states=encoder_hidden_states, + pixel_decoder_hidden_states=pixel_decoder_hidden_states, + transformer_decoder_hidden_states=transformer_decoder_hidden_states, + transformer_decoder_intermediate_states=transformer_decoder_intermediate_states, + attentions=transformer_module_output.attentions, + masks_queries_logits=transformer_module_output.masks_queries_logits, + ) + + if not return_dict: + output = tuple(v for v in output.values() if v is not None) + + return output + + +@auto_docstring( + custom_intro=""" + The Mask2Former Model with heads on top for instance/semantic/panoptic segmentation. + """ +) +class Mask2FormerForUniversalSegmentation(Mask2FormerPreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: Mask2FormerConfig): + super().__init__(config) + self.model = Mask2FormerModel(config) + + self.weight_dict: dict[str, float] = { + "loss_cross_entropy": config.class_weight, + "loss_mask": config.mask_weight, + "loss_dice": config.dice_weight, + } + + self.class_predictor = nn.Linear(config.hidden_dim, config.num_labels + 1) + + self.criterion = Mask2FormerLoss(config=config, weight_dict=self.weight_dict) + self.post_init() + + def get_loss_dict( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: Tensor, + class_labels: Tensor, + auxiliary_predictions: dict[str, Tensor], + ) -> dict[str, Tensor]: + loss_dict: dict[str, Tensor] = self.criterion( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + auxiliary_predictions=auxiliary_predictions, + ) + + # weight each loss by `self.weight_dict[]` including auxiliary losses + for key, weight in self.weight_dict.items(): + for loss_key, loss in loss_dict.items(): + if key in loss_key: + loss *= weight + + return loss_dict + + def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: + return sum(loss_dict.values()) + + def get_auxiliary_logits(self, classes: torch.Tensor, output_masks: torch.Tensor): + auxiliary_logits: list[dict[str, Tensor]] = [] + + for aux_binary_masks, aux_classes in zip(output_masks[:-1], classes[:-1]): + auxiliary_logits.append({"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes}) + + return auxiliary_logits + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + pixel_mask: Tensor | None = None, + output_hidden_states: bool | None = None, + output_auxiliary_logits: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> Mask2FormerForUniversalSegmentationOutput: + r""" + mask_labels (`list[torch.Tensor]`, *optional*): + List of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + output_auxiliary_logits (`bool`, *optional*): + Whether or not to output auxiliary logits. + + Examples: + + Instance segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> import torch + + >>> # Load Mask2Former trained on COCO instance segmentation dataset + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-coco-instance") + >>> model = Mask2FormerForUniversalSegmentation.from_pretrained( + ... "facebook/mask2former-swin-small-coco-instance" + ... ) + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # Perform post-processing to get instance segmentation map + >>> pred_instance_map = image_processor.post_process_instance_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0] + >>> print(pred_instance_map.shape) + torch.Size([480, 640]) + ``` + + Semantic segmentation example: + ```python + >>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> import torch + + >>> # Load Mask2Former trained on ADE20k semantic segmentation dataset + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-ade-semantic") + >>> model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-small-ade-semantic") + + >>> url = ( + ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg" + ... ) + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # Perform post-processing to get semantic segmentation map + >>> pred_semantic_map = image_processor.post_process_semantic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0] + >>> print(pred_semantic_map.shape) + torch.Size([512, 683]) + ``` + + Panoptic segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> import torch + + >>> # Load Mask2Former trained on CityScapes panoptic segmentation dataset + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-cityscapes-panoptic") + >>> model = Mask2FormerForUniversalSegmentation.from_pretrained( + ... "facebook/mask2former-swin-small-cityscapes-panoptic" + ... ) + + >>> url = "https://cdn-media.huggingface.co/Inference-API/Sample-results-on-the-Cityscapes-dataset-The-above-images-show-how-our-method-can-handle.png" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # Model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # Perform post-processing to get panoptic segmentation map + >>> pred_panoptic_map = image_processor.post_process_panoptic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0]["segmentation"] + >>> print(pred_panoptic_map.shape) + torch.Size([338, 676]) + ``` + """ + 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 + + outputs = self.model( + pixel_values=pixel_values, + pixel_mask=pixel_mask, + output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, + output_attentions=output_attentions, + return_dict=True, + ) + + loss, loss_dict, auxiliary_logits = None, None, None + class_queries_logits = () + + for decoder_output in outputs.transformer_decoder_intermediate_states: + class_prediction = self.class_predictor(decoder_output.transpose(0, 1)) + class_queries_logits += (class_prediction,) + + masks_queries_logits = outputs.masks_queries_logits + + auxiliary_logits = self.get_auxiliary_logits(class_queries_logits, masks_queries_logits) + + if mask_labels is not None and class_labels is not None: + loss_dict = self.get_loss_dict( + masks_queries_logits=masks_queries_logits[-1], + class_queries_logits=class_queries_logits[-1], + mask_labels=mask_labels, + class_labels=class_labels, + auxiliary_predictions=auxiliary_logits, + ) + loss = self.get_loss(loss_dict) + + encoder_hidden_states = None + pixel_decoder_hidden_states = None + transformer_decoder_hidden_states = None + + if output_hidden_states: + encoder_hidden_states = outputs.encoder_hidden_states + pixel_decoder_hidden_states = outputs.pixel_decoder_hidden_states + transformer_decoder_hidden_states = outputs.transformer_decoder_hidden_states + + output_auxiliary_logits = ( + self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits + ) + if not output_auxiliary_logits: + auxiliary_logits = None + + output = Mask2FormerForUniversalSegmentationOutput( + loss=loss, + class_queries_logits=class_queries_logits[-1], + masks_queries_logits=masks_queries_logits[-1], + auxiliary_logits=auxiliary_logits, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + pixel_decoder_last_hidden_state=outputs.pixel_decoder_last_hidden_state, + transformer_decoder_last_hidden_state=outputs.transformer_decoder_last_hidden_state, + encoder_hidden_states=encoder_hidden_states, + pixel_decoder_hidden_states=pixel_decoder_hidden_states, + transformer_decoder_hidden_states=transformer_decoder_hidden_states, + attentions=outputs.attentions, + ) + + if not return_dict: + output = tuple(v for v in output.values() if v is not None) + if loss is not None: + output = (loss) + output + return output + + +__all__ = ["Mask2FormerForUniversalSegmentation", "Mask2FormerModel", "Mask2FormerPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/mask2former/modular_mask2former.py b/third_party/transformers/src/transformers/models/mask2former/modular_mask2former.py new file mode 100644 index 0000000000000000000000000000000000000000..87f2b834991f8a1dcfff75b260fd76cb1526e5d4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mask2former/modular_mask2former.py @@ -0,0 +1,613 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from torch import nn + +from ...image_utils import SizeDict +from ...processing_utils import ImagesKwargs +from ...utils import ( + TensorType, + logging, + requires_backends, +) +from ...utils.import_utils import requires +from ..maskformer.image_processing_maskformer import MaskFormerImageProcessor +from ..maskformer.image_processing_pil_maskformer import MaskFormerImageProcessorPil +from .image_processing_mask2former import ( + compute_segments, + convert_segmentation_to_rle, + remove_low_and_no_objects, +) + + +logger = logging.get_logger(__name__) + + +class Mask2FormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ignore_index (`int`, *optional*): + Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels + denoted with 0 (background) will be replaced with `ignore_index`. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). + The background label will be replaced by `ignore_index`. + num_labels (`int`, *optional*): + The number of labels in the segmentation map. + size_divisor (`int`, *optional*, defaults to `32`): + Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in + Swin Transformer. + pad_size (`SizeDict`, *optional*): + The size to pad the images to. Must be larger than any image size provided for preprocessing. If `pad_size` + is not provided, images will be padded to the largest height and width in the batch. + """ + + ignore_index: int | None + do_reduce_labels: bool + num_labels: int | None + size_divisor: int + pad_size: SizeDict | None + + +class Mask2FormerImageProcessor(MaskFormerImageProcessor): + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`Mask2FormerForUniversalSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *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 left to None, 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.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # 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 = torch.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, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`Mask2FormerForUniversalSegmentationOutput`] into instance segmentation predictions. + Only supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + 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 left to None, 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. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + 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`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps 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 return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros((384, 384)) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + 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 [`Mask2FormerForUniversalSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentationOutput`]): + The outputs from [`Mask2FormerForUniversalSegmentation`]. + 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 left to None, 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`, set + to `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("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + 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 + + +@requires(backends=("torch",)) +class Mask2FormerImageProcessorPil(MaskFormerImageProcessorPil): + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`Mask2FormerForUniversalSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *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 left to None, 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.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # 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 = torch.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, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`Mask2FormerForUniversalSegmentationOutput`] into instance segmentation predictions. + Only supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`Mask2FormerForUniversalSegmentation`]): + 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 left to None, 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. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + 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`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps 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`. + """ + requires_backends(self, ["torch"]) + if return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros((384, 384)) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + 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 [`Mask2FormerForUniversalSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`Mask2FormerForUniversalSegmentationOutput`]): + The outputs from [`Mask2FormerForUniversalSegmentation`]. + 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 left to None, 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`, set + to `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`. + """ + requires_backends(self, ["torch"]) + if label_ids_to_fuse is None: + logger.warning("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Scale back to preprocessed image size - (384, 384) for all models + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False + ) + + 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__ = ["Mask2FormerImageProcessor", "Mask2FormerImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/mllama/__init__.py b/third_party/transformers/src/transformers/models/mllama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6b64f2a8779692c561e05c1e6deb6d39ec927458 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/__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_mllama import * + from .image_processing_mllama import * + from .image_processing_pil_mllama import * + from .modeling_mllama import * + from .processing_mllama 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/mllama/configuration_mllama.py b/third_party/transformers/src/transformers/models/mllama/configuration_mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..17658a6e5512a618d7f947b61bea36e4427e9a1d --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/configuration_mllama.py @@ -0,0 +1,200 @@ +# Copyright 2024 HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Mllama 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="meta-llama/Llama-3.2-11B-Vision") +@strict +class MllamaVisionConfig(PreTrainedConfig): + r""" + num_global_layers (`int`, *optional*, defaults to 8): + Number of global layers in the Transformer encoder. Vision model has a second transformer encoder, called global. + vision_output_dim (`int`, *optional*, defaults to 7680): + Dimensionality of the vision model output. Includes output of transformer + encoder with intermediate layers and global transformer encoder. + max_num_tiles (`int`, *optional*, defaults to 4): + Maximum number of tiles for image splitting. + intermediate_layers_indices (`list[int]`, *optional*, defaults to [3, 7, 15, 23, 30]): + Indices of intermediate layers of transformer encoder from which to extract and output features. + These output features are concatenated with final hidden state of transformer encoder. + supported_aspect_ratios (`list[list[int]]`, *optional*): + List of supported aspect ratios for image splitting. If not specified, the default supported aspect ratios + are [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]] for `max_num_tiles=4`. + + Example: + + ```python + >>> from transformers import MllamaVisionConfig, MllamaVisionModel + + >>> # Initializing a Llama config + >>> config = MllamaVisionConfig() + + >>> # Initializing a vision model from the mllama-11b style configuration + >>> model = MllamaVisionModel(config) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mllama_vision_model" + base_config_key = "vision_config" + attribute_map = {"num_attention_heads": "attention_heads"} + + hidden_size: int = 1280 + hidden_act: str = "gelu" + num_hidden_layers: int = 32 + num_global_layers: int = 8 + attention_heads: int = 16 + num_channels: int = 3 + intermediate_size: int = 5120 + vision_output_dim: int = 7680 + image_size: int | list[int] | tuple[int, int] = 448 + patch_size: int | list[int] | tuple[int, int] = 14 + norm_eps: float = 1e-5 + max_num_tiles: int = 4 + intermediate_layers_indices: list[int] | None = None + supported_aspect_ratios: list[list[int]] | None = None + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + if self.supported_aspect_ratios is None: + self.supported_aspect_ratios = [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]] + + if self.intermediate_layers_indices is None: + self.intermediate_layers_indices = [3, 7, 15, 23, 30] + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if ( + self.supported_aspect_ratios == [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]] + and self.max_num_tiles != 4 + ): + raise ValueError("max_num_tiles must be 4 for default supported aspect ratios") + + @property + def max_aspect_ratio_id(self) -> int: + return len(self.supported_aspect_ratios) + + +@auto_docstring(checkpoint="meta-llama/Llama-3.2-11B-Vision") +@strict +class MllamaTextConfig(PreTrainedConfig): + r""" + cross_attention_layers (`list[int]`, *optional*): + Indices of the cross attention layers. If not specified, will default to [3, 8, 13, 18, 23, 28, 33, 38]. + + Example: + + ```python + >>> from transformers import MllamaTextModel, MllamaTextConfig + + >>> # Initializing a Mllama text config + >>> config = MllamaTextConfig() + + >>> # Initializing a model from the Mllama text configuration + >>> model = MllamaTextModel(config) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mllama_text_model" + base_config_key = "text_config" + default_theta = 500000.0 + + vocab_size: int = 128256 + hidden_size: int = 4096 + hidden_act: str = "silu" + num_hidden_layers: int = 40 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + intermediate_size: int = 14_336 + rope_parameters: dict | None = None + rms_norm_eps: float = 1e-5 + max_position_embeddings: int = 131_072 + initializer_range: float = 0.02 + use_cache: bool = True + tie_word_embeddings: bool = False + cross_attention_layers: list[int] | None = None + dropout: float | int = 0.0 + bos_token_id: int = 128000 + eos_token_id: int | list[int] | None = 128001 + pad_token_id: int | None = 128004 + + def __post_init__(self, **kwargs): + if self.cross_attention_layers is None: + self.cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38] + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="meta-llama/Llama-3.2-11B-Vision") +@strict +class MllamaConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import MllamaForConditionalGeneration, MllamaConfig, MllamaVisionConfig, MllamaTextConfig + + >>> # Initializing a CLIP-vision config + >>> vision_config = MllamaVisionConfig() + + >>> # Initializing a Llama config + >>> text_config = MllamaTextConfig() + + >>> # Initializing a mllama-11b style configuration + >>> configuration = MllamaConfig(vision_config, text_config) + + >>> # Initializing a model from the mllama-11b style configuration + >>> model = MllamaForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mllama" + attribute_map = { + "image_token_id": "image_token_index", + } + sub_configs = {"text_config": MllamaTextConfig, "vision_config": MllamaVisionConfig} + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_token_index: int = 128256 + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = MllamaVisionConfig() + logger.info("vision_config is None, using default mllama vision config") + elif isinstance(self.vision_config, dict): + self.vision_config = MllamaVisionConfig(**self.vision_config) + + if self.text_config is None: + self.text_config = MllamaTextConfig() + logger.info("text_config is None, using default mllama text config") + elif isinstance(self.text_config, dict): + self.text_config = MllamaTextConfig(**self.text_config) + + super().__post_init__(**kwargs) + + +__all__ = ["MllamaConfig"] diff --git a/third_party/transformers/src/transformers/models/mllama/convert_mllama_weights_to_hf.py b/third_party/transformers/src/transformers/models/mllama/convert_mllama_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..857987f65ac92d98881d58c3b68673c12d828734 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/convert_mllama_weights_to_hf.py @@ -0,0 +1,638 @@ +# 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. + +import argparse +import gc +import json +import math +import os + +import regex as re +import torch +import torch.nn.functional as F + +from transformers import ( + GenerationConfig, + MllamaConfig, + MllamaForConditionalGeneration, + MllamaImageProcessor, + PreTrainedTokenizerFast, +) +from transformers.convert_slow_tokenizer import TikTokenConverter +from transformers.models.mllama.configuration_mllama import MllamaTextConfig, MllamaVisionConfig +from transformers.models.mllama.image_processing_mllama import get_all_supported_aspect_ratios + + +# fmt: off +# If a weight needs to be split in two or more keys, use `|` to indicate it. ex: +# r"text_model.layers.(\d+).attention.wqkv.weight": r"language_model.model.layers.\1.self_attn.q|k|v|_proj.weight" +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + r"text_model.norm.weight": r"language_model.model.norm.weight", + r"text_model.output.weight": r"language_model.lm_head.weight", + r"text_model.tok_embeddings": r"language_model.model.embed_tokens", + r"text_model.learnable_embedding": r"language_model.model.learnable_embedding", + r"text_model.rope.freqs": None, # meaning we skip it and don't want it + # For every cross attention layer, the layer needs to be updated + r"text_model.cross_attention_layers.(\d+).gate_attn": r"language_model.model.layers.\1.cross_attn_attn_gate", + r"text_model.cross_attention_layers.(\d+).gate_ffwd": r"language_model.model.layers.\1.cross_attn_mlp_gate", + # special key, wqkv needs to be split afterwards + r"text_model.cross_attention_layers.(\d+).attention.w(q|k|v|o)": r"language_model.model.layers.\1.cross_attn.\2_proj", + r"text_model.cross_attention_layers.(\d+).attention.(q|k)_norm": r"language_model.model.layers.\1.cross_attn.\2_norm", + r"text_model.cross_attention_layers.(\d+).attention_norm.weight": r"language_model.model.layers.\1.input_layernorm.weight", + r"text_model.cross_attention_layers.(\d+).attention.wk.layer_norm_weight": r"language_model.model.layers.\1.post_attention_layernorm.weight", + r"text_model.cross_attention_layers.(\d+).feed_forward.w1.weight": r"language_model.model.layers.\1.mlp.gate_proj.weight", + r"text_model.cross_attention_layers.(\d+).feed_forward.w2.weight": r"language_model.model.layers.\1.mlp.down_proj.weight", + r"text_model.cross_attention_layers.(\d+).feed_forward.w3.weight": r"language_model.model.layers.\1.mlp.up_proj.weight", + r"text_model.cross_attention_layers.(\d+).ffn_norm.weight": r"language_model.model.layers.\1.post_attention_layernorm.weight", + # self attention layers + r"text_model.layers.(\d+).attention.w(q|k|v|o).weight": r"language_model.model.layers.\1.self_attn.\2_proj.weight", + r"text_model.layers.(\d+).attention_norm.weight": r"language_model.model.layers.\1.input_layernorm.weight", + r"text_model.layers.(\d+).feed_forward.w1.": r"language_model.model.layers.\1.mlp.gate_proj.", + r"text_model.layers.(\d+).feed_forward.w2.": r"language_model.model.layers.\1.mlp.down_proj.", + r"text_model.layers.(\d+).feed_forward.w3.": r"language_model.model.layers.\1.mlp.up_proj.", + r"text_model.layers.(\d+).ffn_norm.weight": r"language_model.model.layers.\1.post_attention_layernorm.weight", + # Vision encoder mapping + r"vision_model.vision_encoder.conv1._linear": r"vision_model.patch_embedding", + r'vision_model.vision_projection.': r"multi_modal_projector.", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).attn.wq": r"vision_model.\1.layers.\2.self_attn.q_proj", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).attn.wk": r"vision_model.\1.layers.\2.self_attn.k_proj", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).attn.wv": r"vision_model.\1.layers.\2.self_attn.v_proj", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).attn.wo": r"vision_model.\1.layers.\2.self_attn.o_proj", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).mlp.c_fc": r"vision_model.\1.layers.\2.mlp.fc1", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).mlp.c_proj": r"vision_model.\1.layers.\2.mlp.fc2", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).ln_1": r"vision_model.\1.layers.\2.input_layernorm", + r"vision_model.vision_encoder.(global_transformer|transformer).resblocks.(\d+).ln_2": r"vision_model.\1.layers.\2.post_attention_layernorm", + r"vision_model.vision_encoder.global_transformer.resblocks.(\d+).(gate_ffn|gate_attn)": r"vision_model.global_transformer.layers.\1.\2", + r'vision_model.vision_encoder.ln_(pre|post).(weight|bias)': r'vision_model.vision_encoder.layernorm_\1.\2', + r'vision_model.vision_encoder.positional_embedding\b': r'vision_model.gated_positional_embedding.embedding', + r'vision_model.vision_encoder.gated_positional_embedding\b': r'vision_model.gated_positional_embedding.tile_embedding.weight', + r'vision_model.vision_encoder.gated_positional_embedding_gate': r'vision_model.gated_positional_embedding.gate', + r"vision_model.vision_encoder.pre_tile_pos_embed.embedding": r"vision_model.pre_tile_positional_embedding.embedding.weight", + r"vision_model.vision_encoder.post_tile_pos_embed.embedding": r"vision_model.post_tile_positional_embedding.embedding.weight", + r"vision_model.vision_encoder.pre_tile_pos_embed.gate": r"vision_model.pre_tile_positional_embedding.gate", + r"vision_model.vision_encoder.post_tile_pos_embed.gate": r"vision_model.post_tile_positional_embedding.gate", + r"vision_model.vision_encoder.(?=\w)": r"vision_model.", +} +# fmt: on + +CONTEXT_LENGTH = 131072 + + +def convert_old_keys_to_new_keys(state_dict_keys: dict | None = None): + """ + This function should be applied only once, on the concatenated keys to efficiently rename using + the key mappings. + """ + output_dict = {} + if state_dict_keys is not None: + old_text = "\n".join(state_dict_keys) + new_text = old_text + for pattern, replacement in ORIGINAL_TO_CONVERTED_KEY_MAPPING.items(): + if replacement is None: + new_text = re.sub(pattern, "", new_text) # an empty line + continue + new_text = re.sub(pattern, replacement, new_text) + output_dict = dict(zip(old_text.split("\n"), new_text.split("\n"))) + return output_dict + + +def permute_for_rope(input_tensor, n_heads, dim1, dim2): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + input_tensor = input_tensor.reshape(dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) + return input_tensor + + +def pre_compute_positional_embedding(embedding): + """ + Instead of iterating of the batch of images, and the ratios inside, we pre-compute the + positional embeddings depending on the aspect ratio id. This is done to support `torch.compile` + and efficient inference / training with different aspect ratios. + """ + max_num_tiles, *shapes = embedding.shape + hidden_size = shapes[-1] + supported_aspect_ratios = get_all_supported_aspect_ratios(max_num_tiles) + max_aspect_ratio_id = len(supported_aspect_ratios) # we keep 0 index for padding + # tile embedding does not have patches + num_patches = 1 if len(shapes) == 2 else shapes[1] + precomputed_embeddings = torch.zeros( + max_aspect_ratio_id + 1, + max_num_tiles, + num_patches, + hidden_size, + device=embedding.device, + dtype=embedding.dtype, + ) + + for i, (height, width) in enumerate(supported_aspect_ratios): + aspect_ratio_id = i + 1 # we keep 0 index for padding + current_embedding = embedding[:height, :width].reshape(height * width, num_patches, hidden_size) + precomputed_embeddings[aspect_ratio_id, : height * width] = current_embedding + precomputed_embeddings = precomputed_embeddings.flatten(1) + return precomputed_embeddings + + +def is_param_different_across_shards(key): + """ + Return `True` if the parameter is different across checkpoint shards + and needs to be concatenated. + """ + patterns = [r"vision_model.patch_embedding.weight",r"vision_model.(transformer|global_transformer).layers.(\d+).self_attn.(q|k|v|o)_proj.weight",r"vision_model.(transformer|global_transformer).layers.(\d+).mlp.fc1.(weight|bias)",r"vision_model.(transformer|global_transformer).layers.(\d+).mlp.fc2.weight", r"multi_modal_projector.(weight|bias)",r"language_model.model.embed_tokens.weight",r"language_model.lm_head.weight",r"language_model.model.layers.(\d+).self_attn.(q|k|v|o)_proj.weight",r"language_model.model.layers.(\d+).cross_attn.(q|k|v|o)_proj.weight",r"language_model.model.layers.(\d+).mlp.(up|down|gate)_proj.weight",r"language_model.model.learnable_embedding.weight"] # fmt: skip + return any(re.search(pattern, key) for pattern in patterns) + + +def get_concat_dim(key): + """ + Return the dimension to concatenate the weights on. + """ + concat_dim_1 = [r"vision_model.(transformer|global_transformer).layers.(\d+).mlp.fc2.weight",r"vision_model.(transformer|global_transformer).layers.(\d+).self_attn.o_proj.weight",r"language_model.model.layers.(\d+).cross_attn.o_proj.weight",r"language_model.model.layers.(\d+).self_attn.o_proj.weight",r"language_model.model.layers.(\d+).mlp.down_proj.weight"] # fmt: off + if any(re.search(pattern, key) for pattern in concat_dim_1): + return 1 + return 0 + + +def compute_intermediate_size(hidden_dim, multiple_of=1024, ffn_dim_multiplier=1.3): + hidden_dim = 4 * int(2 * hidden_dim / 3) + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + return hidden_dim + + +def interpolate_positional_embedding( + embeddings: torch.Tensor, vision_tile_size: int, vision_patch_size: int +) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position embeddings, to be able to use the model on higher resolution + images. + """ + cls_embedding, positional_embedding = embeddings[:1], embeddings[1:] + total_num_patches, dim = positional_embedding.shape + + # compute current and target number of patches for height and width + num_patches = int(round(total_num_patches**0.5)) + new_num_patches = vision_tile_size // vision_patch_size + + # Check if the number of patches is already the desired size + if num_patches == new_num_patches: + return embeddings + + positional_embedding = positional_embedding.transpose(0, 1) + positional_embedding = positional_embedding.reshape(1, dim, num_patches, num_patches) + positional_embedding = F.interpolate( + positional_embedding, + size=(new_num_patches, new_num_patches), + mode="bicubic", + align_corners=False, + ) + positional_embedding = positional_embedding.reshape(dim, -1).transpose(0, 1) + + embeddings = torch.cat([cls_embedding, positional_embedding], dim=0) + return embeddings + + +def write_model( + model_path, + input_base_path, + num_shards, + instruct=False, +): + os.makedirs(model_path, exist_ok=True) + + with open(os.path.join(input_base_path, "params.json"), "r") as f: + params = json.load(f) + + params = params.get("model", params) + dtype = "bfloat16" + + # ------------------------------------------------------------ + # Text model params and config + # ------------------------------------------------------------ + + # params from config + text_vocab_size = params["vocab_size"] + text_num_layers = params["n_layers"] + text_dim = params["dim"] + text_num_heads = params["n_heads"] + text_rms_norm_eps = params["norm_eps"] + text_rope_theta = params["rope_theta"] + cross_attention_num_layers = params["vision_num_cross_attention_layers"] + + # some constants from original code + rope_parameters = { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + } + max_position_embeddings = CONTEXT_LENGTH + + # compute additional params for weight conversion + text_num_heads_per_shard = text_num_heads // num_shards + text_dim_per_head = text_dim // text_num_heads + text_intermediate_size = compute_intermediate_size(text_dim, multiple_of=params["multiple_of"]) + + if params.get("n_kv_heads", None) is not None: + text_num_key_value_heads = params["n_kv_heads"] # for GQA / MQA + text_num_key_value_heads_per_shard = text_num_key_value_heads // num_shards + text_key_value_dim = text_dim_per_head * text_num_key_value_heads + else: # compatibility with other checkpoints + text_num_key_value_heads = text_num_heads + text_num_key_value_heads_per_shard = text_num_heads_per_shard + text_key_value_dim = text_dim + + # cross-attention layers: 20 for 90B, 8 for 11B + cross_attention_frequency = math.ceil(text_num_layers / cross_attention_num_layers) + text_num_total_layers = text_num_layers + cross_attention_num_layers + cross_attention_layers_shift = list( + range(cross_attention_frequency - 1, text_num_total_layers, cross_attention_frequency + 1) + ) + self_attention_layers_shift = [k for k in range(text_num_total_layers) if k not in cross_attention_layers_shift] + + bos_token_id = 128000 + eos_token_id = [128001, 128008, 128009] if instruct else 128001 + pad_token_id = 128004 + + text_config = MllamaTextConfig( + num_attention_heads=text_num_heads, + vocab_size=text_vocab_size, + hidden_size=text_dim, + rms_norm_eps=text_rms_norm_eps, + rope_theta=text_rope_theta, + num_hidden_layers=text_num_total_layers, + cross_attention_layers=cross_attention_layers_shift, + intermediate_size=text_intermediate_size, + max_position_embeddings=max_position_embeddings, + rope_parameters=rope_parameters, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=False, # Constant set to False + dtype=dtype, + ) + + # ------------------------------------------------------------ + # Vision model params and config + # ------------------------------------------------------------ + + # params from config + vision_tile_size = params["vision_chunk_size"] + vision_max_num_tiles = params["vision_max_num_chunks"] + + # some constants from original code + vision_patch_size = 14 + vision_num_channels = 3 + vision_num_layers = 32 + vision_num_layers_global = 8 + vision_dim = 1280 + vision_num_heads = 16 + vision_intermediate_layers_indices = [3, 7, 15, 23, 30] + + # compute additional params for weight conversion + vision_dim_per_head = vision_dim // vision_num_heads + vision_num_heads_per_shard = vision_num_heads // num_shards + vision_intermediate_size = vision_dim * 4 + vision_supported_aspect_ratios = get_all_supported_aspect_ratios(vision_max_num_tiles) + + vision_config = MllamaVisionConfig( + hidden_size=vision_dim, + patch_size=vision_patch_size, + num_channels=vision_num_channels, + intermediate_size=vision_intermediate_size, + num_hidden_layers=vision_num_layers, + num_attention_heads=vision_num_heads, + num_global_layers=vision_num_layers_global, + intermediate_layers_indices=vision_intermediate_layers_indices, + image_size=vision_tile_size, + max_num_tiles=vision_max_num_tiles, + supported_aspect_ratios=vision_supported_aspect_ratios, + dtype=dtype, + ) + + # save config + config = MllamaConfig(vision_config=vision_config, text_config=text_config, dtype=dtype) + config.architectures = ["MllamaForConditionalGeneration"] + config.save_pretrained(model_path) + print("Model config saved successfully...") + + # ------------------------------------------------------------ + # Convert weights + # ------------------------------------------------------------ + + print(f"Fetching all parameters from the checkpoint at {input_base_path}...") + if num_shards == 1: + if os.path.exists(os.path.join(input_base_path, "consolidated.00.pth")): + path = os.path.join(input_base_path, "consolidated.00.pth") + else: + path = os.path.join(input_base_path, "consolidated.pth") + loaded = [torch.load(path, map_location="cpu", mmap=True, weights_only=True)] + else: + loaded = [ + torch.load( + os.path.join(input_base_path, f"consolidated.{i:02d}.pth"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + for i in range(num_shards) + ] + + print("Converting model...") + all_keys = list(loaded[0].keys()) + new_keys = convert_old_keys_to_new_keys(all_keys) + + state_dict = {} + for key in all_keys: + new_key = new_keys[key] + + # In the original model, self-attention layers and cross-attention layers are different lists of layers. + # In the converted model, they are merged into one list with corresponding index shift to preserve the order. + if ("cross_attention" in key or "text_model.layers" in key) and "language_model" in new_key: + shift = cross_attention_layers_shift if "cross_attention" in key else self_attention_layers_shift + new_key = re.sub(r"layers.(\d+).", lambda _match: f"layers.{shift[int(_match.groups()[0])]}.", new_key) + + current_parameter = [chunk.pop(key).contiguous().clone() for chunk in loaded] + if not is_param_different_across_shards(new_key): + current_parameter = current_parameter[0] + + concat_dim = get_concat_dim(new_key) + + # Post-process the current_parameter. + if re.search("(k|v|q)_proj.weight", new_key) and "language_model" in new_key: + if "q_proj" in new_key: + param_num_heads = text_num_heads + param_num_head_per_shard = text_num_heads_per_shard + param_dim = text_dim + else: + param_num_heads = text_num_key_value_heads + param_num_head_per_shard = text_num_key_value_heads_per_shard + param_dim = text_key_value_dim + shards = [param.view(param_num_head_per_shard, text_dim_per_head, text_dim) for param in current_parameter] + current_parameter = torch.cat(shards, dim=concat_dim) + if "cross_attn" not in new_key and "v_proj.weight" not in new_key: + current_parameter = permute_for_rope(current_parameter, param_num_heads, param_dim, text_dim) + state_dict[new_key] = current_parameter.reshape(param_num_heads * text_dim_per_head, text_dim) + + elif "vision_model" in new_key and re.search("(k|v|q)_proj", new_key): + shards = [ + param.view(vision_num_heads_per_shard, vision_dim_per_head, vision_dim) for param in current_parameter + ] + param = torch.cat(shards, dim=concat_dim) + state_dict[new_key] = param.reshape(vision_num_heads * vision_dim_per_head, vision_dim) + + elif new_key == "vision_model.patch_embedding.weight": + current_parameter = torch.cat(current_parameter, dim=concat_dim) + state_dict[new_key] = current_parameter.reshape( + -1, vision_num_channels, vision_patch_size, vision_patch_size + ) + + elif new_key.endswith("gate"): + state_dict[new_key] = current_parameter[0].view(1) + + elif "vision_model.gated_positional_embedding.embedding" in new_key: + current_parameter = interpolate_positional_embedding( + current_parameter, vision_tile_size, vision_patch_size + ) + state_dict[new_key] = current_parameter + + elif "vision_model.gated_positional_embedding.tile_embedding.weight" in new_key: + current_parameter = current_parameter.permute(2, 0, 1, 3).flatten(1) + current_parameter = interpolate_positional_embedding( + current_parameter, vision_tile_size, vision_patch_size + ) + current_parameter = current_parameter.reshape( + -1, vision_max_num_tiles, vision_max_num_tiles, vision_dim + ).permute(1, 2, 0, 3) + state_dict[new_key] = pre_compute_positional_embedding(current_parameter) + + elif "tile_positional_embedding.embedding" in new_key: + state_dict[new_key] = pre_compute_positional_embedding(current_parameter) + + elif new_key != "": + if isinstance(current_parameter, list): + current_parameter = torch.cat(current_parameter, dim=concat_dim) + state_dict[new_key] = current_parameter + + state_dict["language_model.model.embed_tokens.weight"] = torch.cat( + [ + state_dict["language_model.model.embed_tokens.weight"], + state_dict.pop("language_model.model.learnable_embedding.weight"), + ], + dim=0, + ) + del loaded + gc.collect() + + print("Loading the checkpoint in a Mllama model.") + with torch.device("meta"): + model = MllamaForConditionalGeneration(config) + model.load_state_dict(state_dict, strict=True, assign=True) + print("Checkpoint loaded successfully.") + del model.config._name_or_path + + print("Saving the model.") + model.save_pretrained(model_path) + del state_dict, model + + # Safety check: reload the converted model + gc.collect() + print("Reloading the model to check if it's saved correctly.") + MllamaForConditionalGeneration.from_pretrained(model_path, dtype=torch.bfloat16, device_map="auto") + print("Model reloaded successfully.") + + # generation config + if instruct: + print("Saving generation config...") + generation_config = GenerationConfig( + do_sample=True, + temperature=0.6, + top_p=0.9, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + ) + generation_config.save_pretrained(model_path) + + +class MllamaConverter(TikTokenConverter): + def __init__( + self, + vocab_file, + special_tokens: list[str], + pattern: str, + model_max_length: int, + chat_template: str | None = None, + **kwargs, + ): + super().__init__(vocab_file, pattern=pattern) + self.additional_special_tokens = special_tokens + tokenizer = self.converted() + if chat_template is not None: + kwargs["chat_template"] = chat_template + self.tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + model_input_names=["input_ids", "attention_mask"], + model_max_length=model_max_length, + **kwargs, + ) + + +def write_tokenizer(tokenizer_path: str, save_dir: str, instruct: bool = False): + model_max_length = CONTEXT_LENGTH + pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" + + # Special tokens + num_reserved_special_tokens = 256 + special_tokens = [ + "<|begin_of_text|>", + "<|end_of_text|>", + "<|reserved_special_token_0|>", + "<|reserved_special_token_1|>", + "<|finetune_right_pad_id|>", + "<|step_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", # end of message + "<|eot_id|>", # end of turn + "<|python_tag|>", + ] + special_tokens += [ + f"<|reserved_special_token_{i + 2}|>" for i in range(num_reserved_special_tokens - len(special_tokens)) + ] + # original tokenizer has <|image|> with 128011 token_id, + # however, later in the code it is replaced with 128256 token_id + special_tokens.append("<|image|>") + + # Chat template + chat_template = ( + "{% for message in messages %}" + "{% if loop.index0 == 0 %}" + "{{ bos_token }}" + "{% endif %}" + "{{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }}" + "{% if message['content'] is string %}" + "{{ message['content'] }}" + "{% else %}" + "{% for content in message['content'] %}" + "{% if content['type'] == 'image' %}" + "{{ '<|image|>' }}" + "{% elif content['type'] == 'text' %}" + "{{ content['text'] }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + "{{ '<|eot_id|>' }}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}" + "{% endif %}" + ) + + converter = MllamaConverter( + vocab_file=tokenizer_path, + pattern=pattern, + special_tokens=special_tokens, + model_max_length=model_max_length, + chat_template=chat_template if instruct else None, + bos_token="<|begin_of_text|>", + eos_token="<|end_of_text|>" if not instruct else "<|eot_id|>", + pad_token="<|finetune_right_pad_id|>", + ) + tokenizer = converter.tokenizer + tokenizer.save_pretrained(save_dir) + + if instruct: + print("Saving chat template...") + chat_template_path = os.path.join(save_dir, "chat_template.json") + with open(chat_template_path, "w") as f: + json.dump({"chat_template": chat_template}, f, indent=2) + + +def write_image_processor(config_path: str, save_dir: str): + with open(config_path, "r") as f: + params = json.load(f) + + tile_size = params["vision_chunk_size"] + max_image_tiles = params["vision_max_num_chunks"] + + image_processor = MllamaImageProcessor( + do_resize=True, + size={"height": tile_size, "width": tile_size}, + do_rescale=True, + rescale_factor=1 / 255, + do_normalize=True, + image_mean=[0.48145466, 0.4578275, 0.40821073], + image_std=[0.26862954, 0.26130258, 0.27577711], + do_pad=True, + max_image_tiles=max_image_tiles, + ) + + image_processor.save_pretrained(save_dir) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + default="Llama-3.2-11B-Vision/original", + help="Location of LLaMA weights, which contains tokenizer.model and model folders", + ) + parser.add_argument( + "--output_dir", + default="Llama-3.2-11B-Vision", + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--special_tokens", + default=None, + type=list[str], + help="The list of special tokens that should be added to the model.", + ) + parser.add_argument( + "--num_shards", + default=1, + type=int, + help="The number of individual shards used for the model. Does not have to be the same as the number of consolidated_xx.pth", + ) + parser.add_argument( + "--instruct", + action="store_true", + help="Whether the model is an instruct model", + ) + args = parser.parse_args() + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + num_shards=args.num_shards, + instruct=args.instruct, + ) + + write_tokenizer( + tokenizer_path=os.path.join(args.input_dir, "tokenizer.model"), + save_dir=args.output_dir, + instruct=args.instruct, + ) + + write_image_processor( + config_path=os.path.join(args.input_dir, "params.json"), + save_dir=args.output_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/mllama/image_processing_mllama.py b/third_party/transformers/src/transformers/models/mllama/image_processing_mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..42ad79c8ad93624e5999a0457a863bbafa7b735b --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/image_processing_mllama.py @@ -0,0 +1,545 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from functools import lru_cache +from typing import Optional, Union + +import numpy as np +import torch + +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, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_vision_available + + +if is_vision_available(): + from PIL import Image + +from torchvision.transforms.v2 import functional as tvF + + +class MllamaImageProcessorKwargs(ImagesKwargs, total=False): + """ + max_image_tiles (`int`, *optional*): + The maximum number of tiles allowed. + """ + + max_image_tiles: int + + +@lru_cache(maxsize=10) +def get_all_supported_aspect_ratios(max_image_tiles: int) -> list[tuple[int, int]]: + """ + Computes all allowed aspect ratios for a given maximum number of input tiles. + + This function calculates all possible arrangements of tiles that can be formed + within the constraint of the maximum number of tiles. Each arrangement is + represented by its aspect ratio (width/height) and the corresponding tile configuration. + + Args: + max_image_tiles (`int`): + The maximum number of tiles allowed. + + Returns: + `list[tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height) + configuration in terms of number of tiles. + + Example: + >>> get_all_supported_aspect_ratios(4) + [(1, 1), (1, 2), (1, 3), (1, 4), (2, 1), (2, 2), (3, 1), (4, 1)] + + """ + aspect_ratios = [] + for width in range(1, max_image_tiles + 1): + for height in range(1, max_image_tiles + 1): + if width * height <= max_image_tiles: + aspect_ratios.append((width, height)) + return aspect_ratios + + +def get_image_size_fit_to_canvas( + image_height: int, + image_width: int, + canvas_height: int, + canvas_width: int, + tile_size: int, +) -> tuple[int, int]: + """ + Calculates the new size of an image to fit within a canvas while maintaining aspect ratio. + + This function calculates the optimal size for an image to fit within a canvas defined by + canvas_height and canvas_width, while ensuring that the image dimensions are not smaller than + tile_size. If the image is larger than the canvas, the returned size will fit within the canvas. + If the image already fits within the canvas, the size remains unchanged. + The aspect ratio of the original image is preserved as much as possible. + + Args: + image_height (`int`): + The height of the original image. + image_width (`int`): + The width of the original image. + canvas_height (`int`): + The height of the canvas. + canvas_width (`int`): + The width of the canvas. + tile_size (`int`): + The tile size. + + Returns: + `tuple[int, int]`: A tuple containing the new height and width of the image. + + """ + # Set target image size in between `tile_size` and canvas_size + target_width = np.clip(image_width, tile_size, canvas_width) + target_height = np.clip(image_height, tile_size, canvas_height) + + scale_h = target_height / image_height + scale_w = target_width / image_width + + if scale_w < scale_h: + new_width = target_width + # minimum height is 1 to avoid invalid height of 0 + new_height = min(math.floor(image_height * scale_w) or 1, target_height) + else: + new_height = target_height + # minimum width is 1 to avoid invalid width of 0 + new_width = min(math.floor(image_width * scale_h) or 1, target_width) + + return new_height, new_width + + +@lru_cache(maxsize=100) +def get_optimal_tiled_canvas( + image_height: int, + image_width: int, + max_image_tiles: int, + tile_size: int, +) -> tuple[int, int]: + r""" + Determines the best canvas based on image and tile size and maximum number of tiles. + + First, calculates possible resolutions based on the maximum number of tiles and tile size. + For each possible resolution, calculates the scaling factors for width and height, and selects + the smallest one. If upscaling is possible, picks the smallest upscaling factor > 1. + If upscaling is not possible, picks the largest scaling factor <= 1. + If there are multiple resolutions with the same max scale, picks the one with the lowest area. + + Args: + image_height (`int`): + The height of the image. + image_width (`int`): + The width of the image. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + tile_size (`int`): + The tile size. + + Returns: + `tuple[int, int]`: The best canvas resolution [height, width] for the given image. + """ + possible_tile_arrangements = get_all_supported_aspect_ratios(max_image_tiles) + possible_canvas_sizes = np.array(possible_tile_arrangements) * tile_size + + target_heights, target_widths = np.array(possible_canvas_sizes).T + + scale_h = target_heights / image_height + scale_w = target_widths / image_width + + scales = np.where(scale_w > scale_h, scale_h, scale_w) + + upscaling_options = scales[scales >= 1] + if len(upscaling_options) > 0: + selected_scale = np.min(upscaling_options) + else: + downscaling_options = scales[scales < 1] + selected_scale = np.max(downscaling_options) + + chosen_canvas = possible_canvas_sizes[scales == selected_scale] + + if len(chosen_canvas) > 1: + areas = chosen_canvas[:, 0] * chosen_canvas[:, 1] + optimal_idx = np.argmin(areas) + optimal_canvas = chosen_canvas[optimal_idx] + else: + optimal_canvas = chosen_canvas[0] + + return tuple(optimal_canvas) + + +def _validate_size(size: SizeDict) -> None: + if not (size.height and size.width): + raise ValueError(f"Argument `size` must be a dictionary with keys 'height' and 'width'. Got: {size}") + if size.height != size.width: + raise ValueError(f"Argument `size` must have the same height and width, got {size}") + + +def _validate_mllama_preprocess_arguments(do_resize, size, do_pad, max_image_tiles): + if not do_pad: + raise ValueError("MllamaImageProcessor doesn't support `do_pad=False` mode.") + if not do_resize: + raise ValueError("MllamaImageProcessor doesn't support `do_resize=False` mode.") + if max_image_tiles is None or max_image_tiles <= 0: + raise ValueError(f"MllamaImageProcessor `max_image_tiles` must be a positive integer, got {max_image_tiles}.") + _validate_size(size) + + +def build_aspect_ratio_mask( + aspect_ratios: list[list[tuple[int, int]]], + max_image_tiles: int, + device: Optional["torch.device"] = None, +) -> "torch.Tensor": + """ + Builds a mask for the aspect ratios of the images. + + Args: + aspect_ratios (`List[List[Tuple[int, int]]]`): + A list of lists containing aspect ratios for each image in the batch. + Each aspect ratio is represented as a tuple of (width, height) in terms of number of tiles. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + device (`torch.device`, *optional*): + The device to create the tensor on. Defaults to CPU. + + Returns: + `torch.Tensor`: A 3D torch.Tensor of shape (batch_size, max_num_images, max_image_tiles). + The mask contains 1s for valid tiles and 0s for padding. + """ + batch_size = len(aspect_ratios) + max_num_images = max(len(row) for row in aspect_ratios) + + aspect_ratio_mask = torch.zeros((batch_size, max_num_images, max_image_tiles), dtype=torch.long, device=device) + + # Set the first tile to 1 for all aspect ratios + # because in original implementation aspect ratios are padded with (1, 1), + # but original code examples are not built to handle batches, so we might remove it later + aspect_ratio_mask[:, :, 0] = 1 + + # Set the aspect ratio mask for the rest of the tiles + for i, sample_aspect_ratios in enumerate(aspect_ratios): + for j, (num_tiles_w, num_tiles_h) in enumerate(sample_aspect_ratios): + aspect_ratio_mask[i, j, : num_tiles_w * num_tiles_h] = 1 + + return aspect_ratio_mask + + +def pad_batches_and_tiles( + batch_images: list[list["torch.Tensor"]], + max_image_tiles: int, +) -> tuple["torch.Tensor", list[list[int]]]: + """ + Stack a list of lists of images with variable lengths into a torch.Tensor, applying zero padding as needed. + Each list in the input represents a batch sample, and each image within a list is expected to be + pre-split into tiles. The resulting array will have a shape of + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width). + + Args: + batch_images (`List[List[torch.Tensor]]`): + A list of lists of image tiles. Each inner list represents + a batch sample containing multiple images, where each image is pre-split into tiles. + The shape of each tile array is (num_tiles, channels, tile_height, tile_width). + max_image_tiles (int): + The maximum number of tiles any image was potantially split. + + Returns: + `Tuple[torch.Tensor, List[List[int]]]`: A tuple containing: + - stacked_images (`torch.Tensor`): + A numpy array of stacked images with shape + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width). + - all_num_tiles (`List[List[int]]`): + A list of lists containing the number of tiles + for each image in each batch sample. + """ + # Determine output shape + batch_size = len(batch_images) + max_num_images = max(len(images) for images in batch_images) + shapes = [image.shape for images in batch_images for image in images] + _, channels, tile_height, tile_width = shapes[0] + + # Initialize the stacked images array with zeros + stacked_images = torch.zeros( + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width), + dtype=torch.float32, + device=batch_images[0][0].device, + ) + + # Fill the stacked images array with the tiled images from the batch + all_num_tiles = [] + for i, images in enumerate(batch_images): + num_sample_tiles = [] + for j, image in enumerate(images): + num_tiles = image.shape[0] + stacked_images[i, j, :num_tiles] = image + num_sample_tiles.append(num_tiles) + all_num_tiles.append(num_sample_tiles) + + return stacked_images, all_num_tiles + + +def convert_aspect_ratios_to_ids( + aspect_ratios: list[list[tuple[int, int]]], + max_image_tiles: int, + device: Optional["torch.device"] = None, +) -> "torch.Tensor": + """ + Convert aspect ratio tuples to unique ids. + + For batch padding we use 0, because there might be different number of images in each batch. + The aspect ratio ids start from 1, with 1 corresponding to the first supported aspect ratio. + + Args: + aspect_ratios (`List[List[Tuple[int, int]]]`): + A list of aspect ratios for each image in the batch. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + device (`torch.device`, *optional*): + The device to create the tensor on. Defaults to CPU. + + Returns: + `torch.Tensor`: + The aspect ratios ids as a numpy array with shape (batch_size, max_num_images). + Each id corresponds to the index of the aspect ratio in the list of supported aspect ratios, + offset by 1 (so 0 can be used for padding). + """ + batch_size = len(aspect_ratios) + max_num_images = max(len(row) for row in aspect_ratios) + supported_aspect_ratios = get_all_supported_aspect_ratios(max_image_tiles) + + aspect_ratios_ids = torch.zeros((batch_size, max_num_images), dtype=torch.long, device=device) + for i, sample_aspect_ratios in enumerate(aspect_ratios): + for j, (num_tiles_h, num_tiles_w) in enumerate(sample_aspect_ratios): + aspect_ratios_ids[i, j] = supported_aspect_ratios.index((num_tiles_h, num_tiles_w)) + 1 + return aspect_ratios_ids + + +# Copied from transformers.models.idefics2.image_processing_idefics2.convert_to_rgb +def convert_to_rgb(image: ImageInput) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + """ + if not is_vision_available() or not isinstance(image, Image.Image): + return image + + if image.mode == "RGB": + return image + + image_rgba = image.convert("RGBA") + background = Image.new("RGBA", image_rgba.size, (255, 255, 255)) + alpha_composite = Image.alpha_composite(background, image_rgba) + alpha_composite = alpha_composite.convert("RGB") + return alpha_composite + + +@auto_docstring +class MllamaImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + do_pad = True + max_image_tiles = 4 + valid_kwargs = MllamaImageProcessorKwargs + model_input_names = ["pixel_values", "num_tiles", "aspect_ratio_ids", "aspect_ratio_mask"] + + def __init__(self, **kwargs: Unpack[MllamaImageProcessorKwargs]): + super().__init__(**kwargs) + _validate_mllama_preprocess_arguments(self.do_resize, self.size, self.do_pad, self.max_image_tiles) + + def _validate_preprocess_kwargs(self, **kwargs): + super()._validate_preprocess_kwargs(**kwargs) + _validate_mllama_preprocess_arguments(self.do_resize, self.size, self.do_pad, self.max_image_tiles) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[MllamaImageProcessorKwargs]) -> 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.""" + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def convert_to_rgb(self, image: ImageInput) -> ImageInput: + """Converts an image to RGB format.""" + return convert_to_rgb(image) + + def pad( + self, + image: "torch.Tensor", + size: SizeDict, + aspect_ratio: tuple[int, int], + ) -> "torch.Tensor": + """ + Pad an image to the `size` x `aspect_ratio`. For example, if size is {height: 224, width: 224} and aspect ratio is + (1, 2), the image will be padded to 224x448. + + Args: + image (`torch.Tensor`): + Image to pad. + size (`Dict[str, int]`): + Size of the output image. + aspect_ratio (`Tuple[int, int]`): + The aspect ratio of the image. + + Returns: + `torch.Tensor`: The padded image. + """ + image_height, image_width = image.shape[-2:] + num_tiles_height, num_tiles_width = aspect_ratio + padded_height = num_tiles_height * size.height + padded_width = num_tiles_width * size.width + pad_size = (0, 0, padded_width - image_width, padded_height - image_height) + + image = tvF.pad(image, pad_size, fill=0) + return image + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + max_image_tiles: int, + resample: Optional["tvF.InterpolationMode"] = None, + antialias: bool = True, + ) -> Union["torch.Tensor", tuple[int, int]]: + """ + Resizes an image to fit within a tiled canvas while maintaining its aspect ratio. + The optimal canvas size is calculated based on the maximum number of tiles and the tile size. + + The function first determines the best tile arrangement for the image, then resizes the image + to fit within this canvas. The resized image and the number of tiles along the height and width + dimensions are returned. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`Dict[str, int]`): + Size of the output image. + max_image_tiles (`int`): + The maximum number of tiles to split the image into. + resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`): + Resampling filter to use when resizing the image. + + Returns: + `Union[torch.Tensor, Tuple[int, int]]`: The resized image and a tuple containing the number of tiles + along the height and width dimensions. + """ + image_height, image_width = image.shape[-2:] + tile_size = size.height + + canvas_height, canvas_width = get_optimal_tiled_canvas( + image_height=image_height, + image_width=image_width, + max_image_tiles=max_image_tiles, + tile_size=tile_size, + ) + num_tiles_height = canvas_height // tile_size + num_tiles_width = canvas_width // tile_size + + new_height, new_width = get_image_size_fit_to_canvas( + image_height=image_height, + image_width=image_width, + canvas_height=canvas_height, + canvas_width=canvas_width, + tile_size=tile_size, + ) + + image = super().resize( + image, SizeDict(height=new_height, width=new_width), resample=resample, antialias=antialias + ) + + return image, (num_tiles_height, num_tiles_width) + + def _preprocess( + self, + images: list[list["torch.Tensor"]], + 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, + max_image_tiles: int | None, + return_tensors: str | TensorType | None, + disable_grouping: bool | None, + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing (nested structure) + grouped_images, grouped_images_index = group_images_by_shape( + images, is_nested=True, disable_grouping=disable_grouping + ) + split_images_grouped = {} + aspect_ratio_grouped = {} + for shape, stacked_images in grouped_images.items(): + stacked_images, aspect_ratio = self.resize( + image=stacked_images, size=size, resample=resample, max_image_tiles=max_image_tiles + ) + stacked_images = self.pad( + image=stacked_images, + size=size, + aspect_ratio=aspect_ratio, + ) + num_tiles_height, num_tiles_width = aspect_ratio + aspect_ratio_grouped[shape] = [aspect_ratio] * len(stacked_images) + # same aspect ratio for all images in the batch + split_images = split_to_tiles(stacked_images, num_tiles_height, num_tiles_width) + + # Rescale and normalize + split_images = self.rescale_and_normalize( + split_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + split_images_grouped[shape] = split_images + + split_images = reorder_images(split_images_grouped, grouped_images_index, is_nested=True) + aspect_ratios = reorder_images(aspect_ratio_grouped, grouped_images_index, is_nested=True) + + split_images, num_tiles = pad_batches_and_tiles(split_images, max_image_tiles) + + # Use the same device as the processed image tiles so that all output tensors are consistent. + pixel_values_device = split_images.device + aspect_ratio_ids = convert_aspect_ratios_to_ids( + aspect_ratios, max_image_tiles=max_image_tiles, device=pixel_values_device + ) + aspect_ratio_mask = build_aspect_ratio_mask( + aspect_ratios, max_image_tiles=max_image_tiles, device=pixel_values_device + ) + + encoded_inputs = BatchFeature( + data={ + "pixel_values": split_images, + "aspect_ratio_ids": aspect_ratio_ids, + "aspect_ratio_mask": aspect_ratio_mask, + }, + tensor_type=return_tensors, + ) + encoded_inputs["num_tiles"] = num_tiles + + return encoded_inputs + + +__all__ = ["MllamaImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/mllama/image_processing_pil_mllama.py b/third_party/transformers/src/transformers/models/mllama/image_processing_pil_mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..397590c2b00c11854935cdc7ef8c0b3d044870a2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/image_processing_pil_mllama.py @@ -0,0 +1,541 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PIL Image processor class for Mllama.""" + +import math +from functools import lru_cache + +import numpy as np +from PIL import Image + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import PaddingMode, get_image_size +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_vision_available + + +def split_to_tiles_np(image: np.ndarray, num_tiles_height: int, num_tiles_width: int) -> np.ndarray: + """Split an image into tiles (numpy version).""" + num_channels, height, width = image.shape + tile_height = height // num_tiles_height + tile_width = width // num_tiles_width + + image = image.reshape(num_channels, num_tiles_height, tile_height, num_tiles_width, tile_width) + image = image.transpose(1, 3, 0, 2, 4) + image = image.reshape(num_tiles_width * num_tiles_height, num_channels, tile_height, tile_width) + + return np.ascontiguousarray(image) + + +def build_aspect_ratio_mask_np(aspect_ratios: list[list[tuple[int, int]]], max_image_tiles: int) -> np.ndarray: + """ + Build aspect ratio mask (numpy version). + + Args: + aspect_ratios (`List[List[Tuple[int, int]]]`): + A list of lists containing aspect ratios for each image in the batch. + Each aspect ratio is represented as a tuple of (width, height) in terms of number of tiles. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + + Returns: + `np.ndarray`: A 3D array of shape (batch_size, max_num_images, max_image_tiles). + The mask contains 1s for valid tiles and 0s for padding. + """ + batch_size = len(aspect_ratios) + max_num_images = max(len(row) for row in aspect_ratios) + + aspect_ratio_mask = np.zeros((batch_size, max_num_images, max_image_tiles), dtype=np.int64) + # Set the first tile to 1 for all aspect ratios + # because in original implementation aspect ratios are padded with (1, 1), + # but original code examples are not built to handle batches, so we might remove it later + aspect_ratio_mask[:, :, 0] = 1 + + # Set the aspect ratio mask for the rest of the tiles + for i, sample_aspect_ratios in enumerate(aspect_ratios): + for j, (num_tiles_w, num_tiles_h) in enumerate(sample_aspect_ratios): + aspect_ratio_mask[i, j, : num_tiles_w * num_tiles_h] = 1 + + return aspect_ratio_mask + + +def pack_images( + batch_images: list[list[np.ndarray]], + max_image_tiles: int, +) -> tuple[np.ndarray, list[list[int]]]: + """ + Stack nested image tiles into a padded array. + Each list in the input represents a batch sample, and each image within a list is expected to be + pre-split into tiles. The resulting array will have a shape of + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width). + + Args: + batch_images (`List[List[np.ndarray]]`): + A list of lists of image tiles. Each inner list represents + a batch sample containing multiple images, where each image is pre-split into tiles. + The shape of each tile array is (num_tiles, channels, tile_height, tile_width). + max_image_tiles (int): + The maximum number of tiles any image was potentially split. + + Returns: + `Tuple[np.ndarray, List[List[int]]]`: A tuple containing: + - stacked_images (`np.ndarray`): + An array of stacked images with shape + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width). + - all_num_tiles (`List[List[int]]`): + A list of lists containing the number of tiles + for each image in each batch sample. + """ + # Determine output shape + batch_size = len(batch_images) + max_num_images = max(len(images) for images in batch_images) + shapes = [image.shape for images in batch_images for image in images] + _, channels, tile_height, tile_width = shapes[0] + + # Initialize the stacked images array with zeros + stacked_images = np.zeros( + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width), + dtype=np.float32, + ) + + # Fill the stacked images array with the tiled images from the batch + all_num_tiles = [] + for i, images in enumerate(batch_images): + num_sample_tiles = [] + for j, image in enumerate(images): + num_tiles = image.shape[0] + stacked_images[i, j, :num_tiles] = image + num_sample_tiles.append(num_tiles) + all_num_tiles.append(num_sample_tiles) + + return stacked_images, all_num_tiles + + +def convert_aspect_ratios_to_ids_np(aspect_ratios: list[list[tuple[int, int]]], max_image_tiles: int) -> np.ndarray: + """ + Convert aspect ratio tuples to ids (numpy version). + + For batch padding we use 0, because there might be different number of images in each batch. + The aspect ratio ids start from 1, with 1 corresponding to the first supported aspect ratio. + + Args: + aspect_ratios (`List[List[Tuple[int, int]]]`): + A list of aspect ratios for each image in the batch. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + + Returns: + `np.ndarray`: + The aspect ratios ids as an array with shape (batch_size, max_num_images). + Each id corresponds to the index of the aspect ratio in the list of supported aspect ratios, + offset by 1 (so 0 can be used for padding). + """ + batch_size = len(aspect_ratios) + max_num_images = max(len(row) for row in aspect_ratios) + supported_aspect_ratios = get_all_supported_aspect_ratios(max_image_tiles) + + aspect_ratios_ids = np.zeros((batch_size, max_num_images), dtype=np.int64) + for i, sample_aspect_ratios in enumerate(aspect_ratios): + for j, (num_tiles_h, num_tiles_w) in enumerate(sample_aspect_ratios): + aspect_ratios_ids[i, j] = supported_aspect_ratios.index((num_tiles_h, num_tiles_w)) + 1 + return aspect_ratios_ids + + +# Adapted from transformers.models.mllama.image_processing_mllama.MllamaImageProcessorKwargs +class MllamaImageProcessorKwargs(ImagesKwargs, total=False): + """ + max_image_tiles (`int`, *optional*): + The maximum number of tiles allowed. + """ + + max_image_tiles: int + + +# Adapted from transformers.models.mllama.image_processing_mllama._validate_size +def _validate_size(size: SizeDict) -> None: + if not (size.height and size.width): + raise ValueError(f"Argument `size` must be a dictionary with keys 'height' and 'width'. Got: {size}") + if size.height != size.width: + raise ValueError(f"Argument `size` must have the same height and width, got {size}") + + +# Adapted from transformers.models.mllama.image_processing_mllama._validate_mllama_preprocess_arguments +def _validate_mllama_preprocess_arguments(do_resize, size, do_pad, max_image_tiles): + if not do_pad: + raise ValueError("MllamaImageProcessor doesn't support `do_pad=False` mode.") + if not do_resize: + raise ValueError("MllamaImageProcessor doesn't support `do_resize=False` mode.") + if max_image_tiles is None or max_image_tiles <= 0: + raise ValueError(f"MllamaImageProcessor `max_image_tiles` must be a positive integer, got {max_image_tiles}.") + _validate_size(size) + + +# Adapted from transformers.models.idefics2.image_processing_idefics2.convert_to_rgb +def convert_to_rgb(image: ImageInput) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + """ + if not is_vision_available() or not isinstance(image, Image.Image): + return image + + if image.mode == "RGB": + return image + + image_rgba = image.convert("RGBA") + background = Image.new("RGBA", image_rgba.size, (255, 255, 255)) + alpha_composite = Image.alpha_composite(background, image_rgba) + alpha_composite = alpha_composite.convert("RGB") + return alpha_composite + + +# Adapted from transformers.models.mllama.image_processing_mllama.get_all_supported_aspect_ratios +@lru_cache(maxsize=10) +def get_all_supported_aspect_ratios(max_image_tiles: int) -> list[tuple[int, int]]: + """ + Computes all allowed aspect ratios for a given maximum number of input tiles. + + This function calculates all possible arrangements of tiles that can be formed + within the constraint of the maximum number of tiles. Each arrangement is + represented by its aspect ratio (width/height) and the corresponding tile configuration. + + Args: + max_image_tiles (`int`): + The maximum number of tiles allowed. + + Returns: + `list[tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height) + configuration in terms of number of tiles. + + Example: + >>> get_all_supported_aspect_ratios(4) + [(1, 1), (1, 2), (1, 3), (1, 4), (2, 1), (2, 2), (3, 1), (4, 1)] + + """ + aspect_ratios = [] + for width in range(1, max_image_tiles + 1): + for height in range(1, max_image_tiles + 1): + if width * height <= max_image_tiles: + aspect_ratios.append((width, height)) + return aspect_ratios + + +# Adapted from transformers.models.mllama.image_processing_mllama.get_image_size_fit_to_canvas +def get_image_size_fit_to_canvas( + image_height: int, + image_width: int, + canvas_height: int, + canvas_width: int, + tile_size: int, +) -> tuple[int, int]: + """ + Calculates the new size of an image to fit within a canvas while maintaining aspect ratio. + + This function calculates the optimal size for an image to fit within a canvas defined by + canvas_height and canvas_width, while ensuring that the image dimensions are not smaller than + tile_size. If the image is larger than the canvas, the returned size will fit within the canvas. + If the image already fits within the canvas, the size remains unchanged. + The aspect ratio of the original image is preserved as much as possible. + + Args: + image_height (`int`): + The height of the original image. + image_width (`int`): + The width of the original image. + canvas_height (`int`): + The height of the canvas. + canvas_width (`int`): + The width of the canvas. + tile_size (`int`): + The tile size. + + Returns: + `tuple[int, int]`: A tuple containing the new height and width of the image. + + """ + # Set target image size in between `tile_size` and canvas_size + target_width = np.clip(image_width, tile_size, canvas_width) + target_height = np.clip(image_height, tile_size, canvas_height) + + scale_h = target_height / image_height + scale_w = target_width / image_width + + if scale_w < scale_h: + new_width = target_width + # minimum height is 1 to avoid invalid height of 0 + new_height = min(math.floor(image_height * scale_w) or 1, target_height) + else: + new_height = target_height + # minimum width is 1 to avoid invalid width of 0 + new_width = min(math.floor(image_width * scale_h) or 1, target_width) + + return new_height, new_width + + +# Adapted from transformers.models.mllama.image_processing_mllama.get_optimal_tiled_canvas +@lru_cache(maxsize=100) +def get_optimal_tiled_canvas( + image_height: int, + image_width: int, + max_image_tiles: int, + tile_size: int, +) -> tuple[int, int]: + r""" + Determines the best canvas based on image and tile size and maximum number of tiles. + + First, calculates possible resolutions based on the maximum number of tiles and tile size. + For each possible resolution, calculates the scaling factors for width and height, and selects + the smallest one. If upscaling is possible, picks the smallest upscaling factor > 1. + If upscaling is not possible, picks the largest scaling factor <= 1. + If there are multiple resolutions with the same max scale, picks the one with the lowest area. + + Args: + image_height (`int`): + The height of the image. + image_width (`int`): + The width of the image. + max_image_tiles (`int`): + The maximum number of tiles any image can be split into. + tile_size (`int`): + The tile size. + + Returns: + `tuple[int, int]`: The best canvas resolution [height, width] for the given image. + """ + possible_tile_arrangements = get_all_supported_aspect_ratios(max_image_tiles) + possible_canvas_sizes = np.array(possible_tile_arrangements) * tile_size + + target_heights, target_widths = np.array(possible_canvas_sizes).T + + scale_h = target_heights / image_height + scale_w = target_widths / image_width + + scales = np.where(scale_w > scale_h, scale_h, scale_w) + + upscaling_options = scales[scales >= 1] + if len(upscaling_options) > 0: + selected_scale = np.min(upscaling_options) + else: + downscaling_options = scales[scales < 1] + selected_scale = np.max(downscaling_options) + + chosen_canvas = possible_canvas_sizes[scales == selected_scale] + + if len(chosen_canvas) > 1: + areas = chosen_canvas[:, 0] * chosen_canvas[:, 1] + optimal_idx = np.argmin(areas) + optimal_canvas = chosen_canvas[optimal_idx] + else: + optimal_canvas = chosen_canvas[0] + + return tuple(optimal_canvas) + + +@auto_docstring +class MllamaImageProcessorPil(PilBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + do_pad = True + max_image_tiles = 4 + valid_kwargs = MllamaImageProcessorKwargs + model_input_names = ["pixel_values", "num_tiles", "aspect_ratio_ids", "aspect_ratio_mask"] + + def __init__(self, **kwargs: Unpack[MllamaImageProcessorKwargs]): + super().__init__(**kwargs) + _validate_mllama_preprocess_arguments(self.do_resize, self.size, self.do_pad, self.max_image_tiles) + + def _validate_preprocess_kwargs(self, **kwargs): + super()._validate_preprocess_kwargs(**kwargs) + _validate_mllama_preprocess_arguments(self.do_resize, self.size, self.do_pad, self.max_image_tiles) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[MllamaImageProcessorKwargs]) -> 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.""" + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def convert_to_rgb(self, image: ImageInput) -> ImageInput: + """Converts an image to RGB format.""" + return convert_to_rgb(image) + + def pad( + self, + image: np.ndarray, + size: dict[str, int], + aspect_ratio: tuple[int, int], + ) -> np.ndarray: + """ + Pad an image to the `size` x `aspect_ratio`. For example, if size is {height: 224, width: 224} and aspect ratio is + (1, 2), the image will be padded to 224x448. + + Args: + image (`np.ndarray`): + Image to pad. + size (`Dict[str, int]`): + Size of the output image. + aspect_ratio (`Tuple[int, int]`): + The aspect ratio of the image. + + Returns: + `np.ndarray`: The padded image. + """ + image_height, image_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + num_tiles_height, num_tiles_width = aspect_ratio + padded_height = num_tiles_height * size["height"] + padded_width = num_tiles_width * size["width"] + # Spatial padding: ((height_before, height_after), (width_before, width_after)) + # np_pad will add channel dimension padding for channels_first + padding = ((0, padded_height - image_height), (0, padded_width - image_width)) + + image = np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=0, + data_format="channels_first", + input_data_format="channels_first", + ) + return image + + def resize( + self, + image: np.ndarray, + size: SizeDict, + max_image_tiles: int, + resample: PILImageResampling | None = None, + ) -> tuple[np.ndarray, tuple[int, int]]: + """ + Resizes an image to fit within a tiled canvas while maintaining its aspect ratio. + The optimal canvas size is calculated based on the maximum number of tiles and the tile size. + + The function first determines the best tile arrangement for the image, then resizes the image + to fit within this canvas. The resized image and the number of tiles along the height and width + dimensions are returned. + + Args: + image (`np.ndarray`): + Image to resize. + size (`Dict[str, int]`): + Size of the output image. + max_image_tiles (`int`): + The maximum number of tiles to split the image into. + resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`): + Resampling filter to use when resizing the image. + + Returns: + `Union[np.ndarray, Tuple[int, int]]`: The resized image and a tuple containing the number of tiles + along the height and width dimensions. + """ + image_height, image_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + tile_size = size.height if hasattr(size, "height") else size["height"] + + canvas_height, canvas_width = get_optimal_tiled_canvas( + image_height=image_height, + image_width=image_width, + max_image_tiles=max_image_tiles, + tile_size=tile_size, + ) + num_tiles_height = canvas_height // tile_size + num_tiles_width = canvas_width // tile_size + + new_height, new_width = get_image_size_fit_to_canvas( + image_height=image_height, + image_width=image_width, + canvas_height=canvas_height, + canvas_width=canvas_width, + tile_size=tile_size, + ) + + image = super().resize(image, SizeDict(height=new_height, width=new_width), resample=resample) + return image, (num_tiles_height, num_tiles_width) + + def _preprocess( + self, + images: list[list[np.ndarray]], + 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, + max_image_tiles: int | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + batch_images = [] + batch_aspect_ratios = [] + + size_dict = size if isinstance(size, dict) else {"height": size.height, "width": size.width} + + for sample_images in images: + sample_tiles = [] + sample_aspect_ratios = [] + + for image in sample_images: + image, aspect_ratio = self.resize(image, size=size, max_image_tiles=max_image_tiles, resample=resample) + image = self.pad(image, size=size_dict, aspect_ratio=aspect_ratio) + + # Rescale and normalize + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + + num_tiles_height, num_tiles_width = aspect_ratio + tiles = split_to_tiles_np(image, num_tiles_height, num_tiles_width) + sample_tiles.append(tiles) + sample_aspect_ratios.append(aspect_ratio) + + batch_images.append(sample_tiles) + batch_aspect_ratios.append(sample_aspect_ratios) + + stacked_images, num_tiles = pack_images(batch_images, max_image_tiles) + aspect_ratio_ids = convert_aspect_ratios_to_ids_np(batch_aspect_ratios, max_image_tiles) + aspect_ratio_mask = build_aspect_ratio_mask_np(batch_aspect_ratios, max_image_tiles) + + encoded_inputs = BatchFeature( + data={ + "pixel_values": stacked_images, + "aspect_ratio_ids": aspect_ratio_ids, + "aspect_ratio_mask": aspect_ratio_mask, + }, + tensor_type=return_tensors, + ) + encoded_inputs["num_tiles"] = num_tiles + + return encoded_inputs + + +__all__ = ["MllamaImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/mllama/modeling_mllama.py b/third_party/transformers/src/transformers/models/mllama/modeling_mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..3b9d12b9a225473ca3931e5522b89412631e26bb --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/modeling_mllama.py @@ -0,0 +1,1634 @@ +# 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. +"""PyTorch Mllama model.""" + +import math +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 ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, 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, logging +from ...utils.generic import ( + maybe_autocast, + merge_with_config_defaults, +) +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_mllama import MllamaConfig, MllamaTextConfig, MllamaVisionConfig + + +logger = logging.get_logger(__name__) + + +def _prepare_cross_attention_mask( + cross_attention_mask: torch.Tensor, + num_vision_tokens: int, + dtype: str, +) -> tuple[torch.Tensor, torch.Tensor]: + # reshape so it can be used by attn module + batch_size, text_total_length, *_ = cross_attention_mask.shape + cross_attention_mask = cross_attention_mask.repeat_interleave(num_vision_tokens, dim=3) + cross_attention_mask = cross_attention_mask.view(batch_size, text_total_length, -1) + cross_attention_mask = cross_attention_mask.unsqueeze(1) + + # invert the mask + inverted_cross_attn_mask = (1.0 - cross_attention_mask).to(dtype) + cross_attention_mask = inverted_cross_attn_mask.masked_fill( + inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min + ) + + # apply full-row bias, which return 4D tensor of shape [B, H, S1, 1] where value is 0 if the a full row in cross attn mask's + # last dimension contains negative infinity values, otherwise it's 1 + negative_inf_value = torch.finfo(dtype).min + full_text_row_masked_out_mask = ( + (cross_attention_mask != negative_inf_value).any(dim=-1).type_as(cross_attention_mask)[..., None] + ) + cross_attention_mask *= full_text_row_masked_out_mask + + return cross_attention_mask, full_text_row_masked_out_mask + + +def _prepare_aspect_ratio_attention_mask( + aspect_ratio_mask: torch.Tensor, + num_patches: int, + target_length: int, + dtype: torch.dtype, +) -> torch.Tensor: + # Expand aspect ratio mask to target_length + batch_size, max_num_tiles = aspect_ratio_mask.shape + attention_mask = aspect_ratio_mask.view(batch_size, max_num_tiles, 1, 1).to(dtype) + attention_mask = attention_mask.repeat(1, 1, target_length, 1) + + # Mask padding patches + pad_patches = target_length - num_patches + attention_mask[:, :, -pad_patches:] = 0 + + # Invert the mask (0 -> 1, 1 -> 0) + attention_mask = 1 - attention_mask + + # Reshape to 2D and create 4D attention mask + # (batch_size, 1, max_num_tiles * target_length, max_num_tiles * target_length) + attention_mask = attention_mask.reshape(batch_size, max_num_tiles * target_length, 1) + attention_mask = attention_mask @ attention_mask.transpose(-1, -2) * torch.finfo(dtype).min + attention_mask = attention_mask.unsqueeze(1) + + return attention_mask + + +class MllamaPrecomputedAspectRatioEmbedding(nn.Module): + def __init__(self, config: MllamaVisionConfig, is_gated: bool = True): + super().__init__() + self.max_num_tiles = config.max_num_tiles + self.hidden_size = config.hidden_size + self.max_aspect_ratio_id = config.max_aspect_ratio_id + self.is_gated = is_gated + + self.embedding = nn.Embedding(self.max_aspect_ratio_id + 1, self.max_num_tiles * self.hidden_size) + if is_gated: + self.gate = nn.Parameter(torch.zeros(1)) + + def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor: + embeddings = self.embedding(aspect_ratio_ids) + embeddings = embeddings.reshape(-1, self.max_num_tiles, 1, self.hidden_size) + + if self.is_gated: + embeddings = embeddings * self.gate.tanh() + + hidden_state = hidden_state + embeddings + return hidden_state + + +class MllamaPrecomputedPositionEmbedding(nn.Module): + def __init__(self, config: MllamaVisionConfig): + super().__init__() + self.max_num_tiles = config.max_num_tiles + self.max_aspect_ratio_id = config.max_aspect_ratio_id + self.num_patches = (config.image_size // config.patch_size) ** 2 + 1 + self.hidden_size = config.hidden_size + self.scale = config.hidden_size**-0.5 + + self.gate = nn.Parameter(torch.zeros(1)) + + # position embedding + position_embedding = torch.randn(self.num_patches, self.hidden_size) + self.embedding = nn.Parameter(self.scale * position_embedding) + + # tile position embedding + self.tile_embedding = nn.Embedding( + self.max_aspect_ratio_id + 1, self.max_num_tiles * self.num_patches * self.hidden_size + ) + + def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor: + # position embeddings + gated_position_embedding = (1 - self.gate.tanh()) * self.embedding + hidden_state = hidden_state + gated_position_embedding.view(1, 1, self.num_patches, self.hidden_size) + + # precomputed tile position embeddings + tile_position_embedding = self.tile_embedding(aspect_ratio_ids) + batch_size = hidden_state.shape[0] + tile_position_embedding = tile_position_embedding.reshape( + batch_size, self.max_num_tiles, self.num_patches, self.hidden_size + ) + gated_tile_position_embedding = self.gate.tanh() * tile_position_embedding + hidden_state = hidden_state + gated_tile_position_embedding + + return hidden_state + + +# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->MllamaVision +class MllamaVisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.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: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class MllamaVisionAttention(nn.Module): + def __init__(self, config: MllamaVisionConfig): + super().__init__() + + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.attention_heads + self.head_dim = config.hidden_size // config.attention_heads + self.scaling = self.head_dim**-0.5 + self.num_key_value_groups = 1 + + self.q_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=False) + + def forward( + self, + hidden_state: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + query = self.q_proj(hidden_state) + key = self.k_proj(hidden_state) + value = self.v_proj(hidden_state) + + batch_size, q_seq_len, _ = query.shape + _, kv_seq_len, _ = key.shape + + query = query.view(batch_size, q_seq_len, self.num_heads, self.head_dim).transpose(1, 2) + key = key.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2) + value = value.view(batch_size, kv_seq_len, 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, + query, + key, + value, + attention_mask, + dropout=0.0, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(batch_size, q_seq_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class MllamaVisionEncoderLayer(nn.Module): + def __init__(self, config: MllamaVisionConfig, is_gated: bool = False): + super().__init__() + + self.hidden_size = config.hidden_size + self.num_attention_heads = config.attention_heads + self.is_gated = is_gated + self.intermediate_size = config.intermediate_size + + self.self_attn = MllamaVisionAttention(config) + self.mlp = MllamaVisionMLP(config) + + self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps) + self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps) + + if is_gated: + self.gate_attn = nn.Parameter(torch.ones(1) * math.pi / 4) + self.gate_ffn = nn.Parameter(torch.ones(1) * math.pi / 4) + + def forward( + self, + hidden_state: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ): + # Self Attention + residual = hidden_state + hidden_state = self.input_layernorm(hidden_state) + hidden_state, attn_weights = self.self_attn(hidden_state, attention_mask=attention_mask) + if self.is_gated: + hidden_state = self.gate_attn.tanh() * hidden_state + hidden_state = residual + hidden_state + + # Feed forward + residual = hidden_state + hidden_state = self.post_attention_layernorm(hidden_state) + hidden_state = self.mlp(hidden_state) + if self.is_gated: + hidden_state = self.gate_ffn.tanh() * hidden_state + hidden_state = residual + hidden_state + + return hidden_state + + +class MllamaVisionEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`MllamaEncoderLayer`]. + + Args: + config: MllamaConfig + """ + + def __init__(self, config: MllamaVisionConfig, num_layers=32, is_gated=False): + super().__init__() + self.config = config + self.layers = nn.ModuleList([MllamaVisionEncoderLayer(config, is_gated) for _ in range(num_layers)]) + self.gradient_checkpointing = False + self.config = config + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + """ + encoder_states = () + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_state=hidden_states, + attention_mask=attention_mask, + ) + encoder_states = encoder_states + (hidden_states,) + + return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MllamaText +class MllamaTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + MllamaTextRMSNorm 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 MllamaTextCrossAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + config: MllamaTextConfig | None = None, + layer_idx: int | None = None, + ): + super().__init__() + self.config = config + self.num_heads = self.config.num_attention_heads + self.num_key_value_heads = self.config.num_key_value_heads + self.dropout = config.dropout + self.hidden_size = config.hidden_size + self.head_dim = config.hidden_size // self.num_heads + self.layer_idx = layer_idx + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.q_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + bsz, q_len, _ = hidden_states.size() + query_states = self.q_proj(hidden_states) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + query_states = self.q_norm(query_states) + + if cross_attention_states is not None: + key_states = self.k_proj(cross_attention_states) + value_states = self.v_proj(cross_attention_states) + key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + key_states = self.k_norm(key_states) + if past_key_values is not None: + # if we have a new image + new tokens, we only computed key_states on that new image + # we still update the cross key states, past_image, new_image. And use it! + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + elif past_key_values is not None and past_key_values.get_seq_length() > 0: + key_states, value_states = ( + past_key_values.layers[self.layer_idx].keys, + past_key_values.layers[self.layer_idx].values, + ) + else: + raise ValueError( + "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!" + ) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class MllamaTextSelfAttention(nn.Module): + def __init__(self, config: MllamaTextConfig, layer_idx: int): + super().__init__() + self.config = config + self.num_heads = config.num_attention_heads + self.dropout = config.dropout + self.hidden_size = config.hidden_size + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.hidden_size // self.num_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + + self.layer_idx = layer_idx + self.is_causal = True + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_embeddings: torch.Tensor, + past_key_values=None, + **kwargs, + ): + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + 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.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +# Copied from transformers.models.gemma2.modeling_gemma2.Gemma2MLP with Gemma2->MllamaText +class MllamaTextMLP(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) + # Ignore copy + 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 + + +# Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer +class MllamaSelfAttentionDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MllamaTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = MllamaTextSelfAttention(config=config, layer_idx=layer_idx) + + self.mlp = MllamaTextMLP(config) + self.input_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.layer_idx = layer_idx + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + """ + 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. + + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class MllamaCrossAttentionDecoderLayer(GradientCheckpointingLayer): + """Cross-attention transformer block with tanh-gated attention and feedforward.""" + + def __init__(self, config: MllamaTextConfig, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + self.cross_attn = MllamaTextCrossAttention(config, layer_idx=layer_idx) + + self.input_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn_attn_gate = torch.nn.Parameter(torch.zeros(1)) + + self.mlp = MllamaTextMLP(config) + self.post_attention_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn_mlp_gate = torch.nn.Parameter(torch.zeros(1)) + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor, + cross_attention_mask: torch.Tensor, + attention_mask: torch.Tensor, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor], + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, attn_weights = self.cross_attn( + hidden_states=hidden_states, + attention_mask=cross_attention_mask, + cross_attention_states=cross_attention_states, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + if full_text_row_masked_out_mask is not None: + hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states # type: ignore + hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states + + return hidden_states + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with LlamaConfig->MllamaTextConfig,Llama->Mllama +class MllamaRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: MllamaTextConfig, 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: MllamaTextConfig | 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 + + # Ignore copy + @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.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@auto_docstring +class MllamaPreTrainedModel(PreTrainedModel): + config: MllamaConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = [ + "MllamaVisionEncoderLayer", + "MllamaCrossAttentionDecoderLayer", + "MllamaSelfAttentionDecoderLayer", + ] + _can_compile_fullgraph = False # static cache cannot have different shapes for each layer + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": [MllamaSelfAttentionDecoderLayer, MllamaCrossAttentionDecoderLayer], + "attentions": [ + OutputRecorder(MllamaTextSelfAttention, index=1, layer_name="self_attn"), + OutputRecorder(MllamaTextSelfAttention, index=1, layer_name="cross_attn"), + OutputRecorder(MllamaTextCrossAttention, index=1, layer_name="cross_attn"), + ], + } + + @torch.no_grad() + def _init_weights(self, module): + std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + elif isinstance(module, nn.LayerNorm): + init.ones_(module.weight) + init.zeros_(module.bias) + elif isinstance(module, MllamaTextRMSNorm): + init.ones_(module.weight) + elif isinstance(module, MllamaVisionModel): + init.normal_(module.class_embedding, std=std) + elif isinstance(module, MllamaPrecomputedPositionEmbedding): + init.normal_(module.embedding, std=std) + init.zeros_(module.gate) + elif isinstance(module, MllamaVisionEncoderLayer) and module.is_gated: + init.normal_(module.gate_attn, std=std) + init.normal_(module.gate_ffn, std=std) + elif isinstance(module, MllamaCrossAttentionDecoderLayer): + init.zeros_(module.cross_attn_attn_gate) + init.zeros_(module.cross_attn_mlp_gate) + elif isinstance(module, MllamaPrecomputedAspectRatioEmbedding): + if module.is_gated: + init.zeros_(module.gate) + elif isinstance(module, MllamaRotaryEmbedding): + rope_fn = ( + ROPE_INIT_FUNCTIONS[module.rope_type] + if module.rope_type != "default" + else module.compute_default_rope_parameters + ) + buffer_value, _ = rope_fn(module.config) + init.copy_(module.inv_freq, buffer_value) + init.copy_(module.original_inv_freq, buffer_value) + + +@auto_docstring( + custom_intro=""" + The Mllama Vision Model which consists of two vision encoders. + """ +) +class MllamaVisionModel(MllamaPreTrainedModel): + config: MllamaVisionConfig + base_model_prefix = "vision_model" + input_modalities = ("image",) + + def __init__(self, config: MllamaVisionConfig): + super().__init__(config) + self.image_size = config.image_size + self.patch_size = config.patch_size + self.max_num_tiles = config.max_num_tiles + self.hidden_size = config.hidden_size + self.num_channels = config.num_channels + self.intermediate_layers_indices = config.intermediate_layers_indices + + self.num_patches = (self.image_size // self.patch_size) ** 2 + 1 + self.scale = config.hidden_size**-0.5 + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.hidden_size, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + bias=False, + ) + + self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size)) + self.gated_positional_embedding = MllamaPrecomputedPositionEmbedding(config) + + self.pre_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding(config, is_gated=True) + self.post_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding(config, is_gated=True) + + # layer norms + self.layernorm_pre = nn.LayerNorm(self.hidden_size) + self.layernorm_post = nn.LayerNorm(self.hidden_size) + + # encoders + self.transformer = MllamaVisionEncoder(config, config.num_hidden_layers, is_gated=False) + self.global_transformer = MllamaVisionEncoder(config, config.num_global_layers, is_gated=True) + + self.post_init() + + def get_input_embeddings(self): + """ + This function is used to fetch the first embedding layer to activate grads on inputs. + """ + return self.patch_embedding + + def apply_class_embedding(self, hidden_state: torch.Tensor) -> torch.Tensor: + batch_size, _, hidden_size = hidden_state.shape + class_embedding = self.class_embedding.expand(batch_size, 1, hidden_size) + hidden_state = torch.cat([class_embedding, hidden_state], dim=1) + return hidden_state + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, pixel_values: torch.Tensor, aspect_ratio_ids: torch.Tensor, aspect_ratio_mask: torch.Tensor, **kwargs + ) -> BaseModelOutput: + r""" + aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*): + Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image. + These ids correspond to indices in the model's list of supported aspect ratios, offset by 1. + + For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]: + - An image with aspect ratio [1, 1] would have ID 1 + - An image with aspect ratio [1, 2] would have ID 2 + - An image with aspect ratio [2, 1] would have ID 3 + + The id 0 is reserved for padding (i.e., no image). + + If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2. + aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*): + Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`: + + - 1 for tiles that are **not masked**, + - 0 for tiles that are **masked**. + + Example: + + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, MllamaVisionModel + + >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" + >>> model = MllamaVisionModel.from_pretrained(checkpoint) + >>> processor = AutoProcessor.from_pretrained(checkpoint) + + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = processor(images=image, return_tensors="pt") + + >>> output = model(**inputs) + + >>> print(output.last_hidden_state.shape) + torch.Size([1, 1, 4, 1025, 7680]) + ``` + """ + batch_size, num_concurrent_media, num_tiles, num_channels, height, width = pixel_values.shape + + pixel_values = pixel_values.reshape(batch_size * num_concurrent_media * num_tiles, num_channels, height, width) + aspect_ratio_ids = aspect_ratio_ids.reshape(batch_size * num_concurrent_media, -1) + + # Patch embedding + target_dtype = self.patch_embedding.weight.dtype + target_device = self.patch_embedding.weight.device + patch_embeds = self.patch_embedding(pixel_values.to(target_device, target_dtype)) + hidden_state = patch_embeds.flatten(2).transpose(1, 2) + + # Tile embeddings + _, num_patches, dim = hidden_state.shape + hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, -1, dim) + hidden_state = self.pre_tile_positional_embedding(hidden_state, aspect_ratio_ids) + + # Add cls token + hidden_state = hidden_state.reshape(batch_size * num_concurrent_media * num_tiles, num_patches, dim) + hidden_state = self.apply_class_embedding(hidden_state) + num_patches += 1 + + # Position embeddings + hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, num_patches, dim) + hidden_state = self.gated_positional_embedding(hidden_state, aspect_ratio_ids) + + hidden_state = self.layernorm_pre(hidden_state) + + # Compute the number of tokens to pad + num_padding_patches = (8 - (hidden_state.shape[-2] % 8)) % 8 + # Compute padding tuple for pad function + padding = (0, 0, 0, num_padding_patches) # (pad_left, pad_right, pad_left for dim -2, pad_right for dim -2) + # Pad the tensor + hidden_state = F.pad(hidden_state, padding, mode="constant", value=0) + slice_index = -num_padding_patches if num_padding_patches > 0 else None + + # Prepare attention mask + attention_mask = aspect_ratio_mask.reshape(batch_size * num_concurrent_media, -1) + attention_mask = _prepare_aspect_ratio_attention_mask( + aspect_ratio_mask=attention_mask, + num_patches=self.num_patches, + target_length=hidden_state.shape[2], + dtype=self.dtype, + ) + + # Apply encoder + hidden_state = hidden_state.view(batch_size * num_concurrent_media, -1, dim) + output = self.transformer( + hidden_state, + attention_mask=attention_mask, + ) + hidden_state = output.last_hidden_state + + hidden_state = self.layernorm_post(hidden_state) + + # Apply global encoder + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim + ) + hidden_state = self.post_tile_positional_embedding(hidden_state, aspect_ratio_ids) + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, num_tiles * (num_patches + num_padding_patches), dim + ) + global_output = self.global_transformer( + hidden_state, + attention_mask=attention_mask, + ) + hidden_state = global_output.last_hidden_state + + # Remove padding form hidden state + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim + ) + hidden_state = hidden_state[:, :, :slice_index] + hidden_state = hidden_state.reshape(batch_size, num_concurrent_media, num_tiles, num_patches, dim) + + # Collect intermediate layer outputs from encoder output + all_intermediate_hidden_states = [output.hidden_states[i] for i in self.intermediate_layers_indices] + intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) + + # Remove padding from intermediate hidden states + intermediate_hidden_states = intermediate_hidden_states.reshape( + batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, -1 + ) + intermediate_hidden_states = intermediate_hidden_states[:, :, :slice_index] + intermediate_hidden_states = intermediate_hidden_states.reshape( + batch_size, num_concurrent_media, num_tiles, num_patches, -1 + ) + + # Concatenate final hidden state and intermediate hidden states + hidden_state = torch.cat([hidden_state, intermediate_hidden_states], dim=-1) + + return BaseModelOutput(last_hidden_state=hidden_state) + + +@auto_docstring( + custom_intro=""" + The Mllama Text Model which consists of transformer with self and cross attention layers. + """ +) +class MllamaTextModel(MllamaPreTrainedModel): + config: MllamaTextConfig + base_model_prefix = "language_model.model" + input_modalities = ("text",) + + def __init__(self, config: MllamaTextConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size + 8, config.hidden_size, self.padding_idx) + self.cross_attention_layers = config.cross_attention_layers + + layers = [] + for layer_idx in range(config.num_hidden_layers): + if layer_idx in self.cross_attention_layers: + layers.append(MllamaCrossAttentionDecoderLayer(config, layer_idx)) + else: + layers.append(MllamaSelfAttentionDecoderLayer(config, layer_idx)) + + self.layers = nn.ModuleList(layers) + self.norm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = MllamaRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @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, + cross_attention_states: torch.FloatTensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + r""" + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*): + A tuple containing two tensors that mask out rows in the cross-attention mechanism: + - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1. + A value of 0 indicates that the corresponding text token's entire row in the cross-attention + matrix should be masked out (all image tokens ignored). + - The second tensor has the same shape and is used internally to apply the masking during + the forward pass of cross-attention layers. + This mask is derived from the cross_attention_mask and is used to handle cases where a text token + should not attend to any image token. + + Example: + + ```python + >>> from transformers import AutoProcessor, MllamaTextModel + + >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" + >>> model = MllamaTextModel.from_pretrained(checkpoint) + >>> processor = AutoProcessor.from_pretrained(checkpoint) + + >>> text = "<|image|>If I had to write a haiku for this one" + >>> inputs = processor(text=text, return_tensors="pt") + + >>> output = model(**inputs) + + >>> print(output.last_hidden_state.shape) + torch.Size([1, 13, 4096]) + ``` + """ + use_cache = use_cache if use_cache is not None else self.config.use_cache + + 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) + + hidden_states = inputs_embeds + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(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, + ) + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + # decoder layers + for idx, decoder_layer in enumerate(self.layers): + # For text-only path we should skip cross attention layers. + # Let's check if the layer is cross attention layer and if we have cross attention states + # or cached cross attention states. + is_cross_attention_layer = idx in self.cross_attention_layers + is_cross_attention_cache_empty = past_key_values is None or ( + past_key_values is not None and past_key_values.get_seq_length(idx) == 0 + ) + + if is_cross_attention_layer and cross_attention_states is None and is_cross_attention_cache_empty: + continue + + hidden_states = decoder_layer( + hidden_states, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + attention_mask=causal_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_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 BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring( + custom_intro=""" + The Mllama Text Model with a language modeling head on top. + """ +) +class MllamaForCausalLM(MllamaPreTrainedModel, GenerationMixin): + config: MllamaTextConfig + _can_compile_fullgraph = True # only the LLM without cross attn can do compile + base_model_prefix = "language_model" + + def __init__(self, config): + super().__init__(config.get_text_config()) + self.text_config = config.get_text_config() + self.vocab_size = self.text_config.vocab_size + self.model = MllamaTextModel._from_config(self.text_config) + self.lm_head = nn.Linear(self.text_config.hidden_size, self.vocab_size, bias=False) + + 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, + cross_attention_states: torch.LongTensor | None = None, + cross_attention_mask: torch.LongTensor | None = None, + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*): + A tuple containing two tensors that mask out rows in the cross-attention mechanism: + - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1. + A value of 0 indicates that the corresponding text token's entire row in the cross-attention + matrix should be masked out (all image tokens ignored). + - The second tensor has the same shape and is used internally to apply the masking during + the forward pass of cross-attention layers. + This mask is derived from the cross_attention_mask and is used to handle cases where a text token + should not attend to any image token. + 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, MllamaForCausalLM + + >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision") + >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision") + + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + cross_attention_states=cross_attention_states, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + 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, :]).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, + ) + + +@auto_docstring( + custom_intro=""" + The Mllama model which consists of a vision encoder and a language model without language modeling head. + """ +) +class MllamaModel(MllamaPreTrainedModel): + def __init__(self, config: MllamaConfig): + super().__init__(config) + self.vocab_size = config.text_config.vocab_size + self.hidden_size = config.text_config.hidden_size + self.max_num_tiles = config.vision_config.max_num_tiles + self.vision_output_dim = config.vision_config.vision_output_dim + + self.vision_model = MllamaVisionModel._from_config(config.vision_config) + self.language_model = MllamaTextModel._from_config(config.text_config) + self.multi_modal_projector = nn.Linear( + config.vision_config.vision_output_dim, + config.text_config.hidden_size, + bias=True, + ) + 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 + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + aspect_ratio_mask: torch.Tensor | None = None, + aspect_ratio_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + cross_attention_states: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + r""" + aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*): + Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`: + + - 1 for tiles that are **not masked**, + - 0 for tiles that are **masked**. + aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*): + Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image. + These ids correspond to indices in the model's list of supported aspect ratios, offset by 1. + + For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]: + - An image with aspect ratio [1, 1] would have ID 1 + - An image with aspect ratio [1, 2] would have ID 2 + - An image with aspect ratio [2, 1] would have ID 3 + + The id 0 is reserved for padding (i.e., no image). + + If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and cross_attention_states is not None: + raise ValueError("`pixel_values` and `cross_attention_states` cannot be provided simultaneously") + + if pixel_values is not None: + if aspect_ratio_ids is None: + raise ValueError("`aspect_ratio_ids` must be provided if `pixel_values` is provided") + # get vision tokens from vision model + vision_outputs = self.vision_model( + pixel_values=pixel_values, + aspect_ratio_ids=aspect_ratio_ids, + aspect_ratio_mask=aspect_ratio_mask, + ) + cross_attention_states = vision_outputs.last_hidden_state + cross_attention_states = self.multi_modal_projector(cross_attention_states).reshape( + -1, cross_attention_states.shape[-2], self.hidden_size + ) + + if cross_attention_mask is not None: + cross_attention_mask, full_text_row_masked_out_mask = _prepare_cross_attention_mask( + cross_attention_mask, + num_vision_tokens=self.vision_model.num_patches, + dtype=self.dtype, + ) + else: + full_text_row_masked_out_mask = None + + if cross_attention_mask is not None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + seq_len = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + device = input_ids.device if input_ids is not None else inputs_embeds.device + current_pos = torch.arange(seq_len, device=device) + past_seen_tokens + + cross_attention_mask = cross_attention_mask[:, :, current_pos] + full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, current_pos] + + outputs = self.language_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + past_key_values=past_key_values, + use_cache=use_cache, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + return BaseModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The Mllama model which consists of a vision encoder and a language model. + """, +) +class MllamaForConditionalGeneration(MllamaPreTrainedModel, GenerationMixin): + # _tied_weights_keys = {"lm_head.weight": "model.language_moddel.embed_tokens.weight"} + + def __init__(self, config: MllamaConfig): + super().__init__(config) + self.model = MllamaModel(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) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + aspect_ratio_mask: torch.Tensor | None = None, + aspect_ratio_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + cross_attention_states: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*): + Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`: + + - 1 for tiles that are **not masked**, + - 0 for tiles that are **masked**. + aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*): + Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image. + These ids correspond to indices in the model's list of supported aspect ratios, offset by 1. + + For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]: + - An image with aspect ratio [1, 1] would have ID 1 + - An image with aspect ratio [1, 2] would have ID 2 + - An image with aspect ratio [2, 1] would have ID 3 + + The id 0 is reserved for padding (i.e., no image). + + If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + 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, MllamaForConditionalGeneration + + >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" + >>> model = MllamaForConditionalGeneration.from_pretrained(checkpoint) + >>> processor = AutoProcessor.from_pretrained(checkpoint) + + >>> prompt = "<|image|>If I had to write a haiku for this one" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(text=prompt, images=image, return_tensors="pt") + + >>> # Generate + >>> output = model.generate(**inputs, max_new_tokens=15) + + >>> prompt_len = inputs.input_ids.shape[-1] + >>> generated_ids = output[:, prompt_len:] + >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) + >>> print(generated_text) + [', it would be:.\\nA stop sign in Chinatown.\\n'] + ``` + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + aspect_ratio_mask=aspect_ratio_mask, + aspect_ratio_ids=aspect_ratio_ids, + cross_attention_mask=cross_attention_mask, + cross_attention_states=cross_attention_states, + 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, labels, self.config.text_config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + position_ids=None, + pixel_values=None, + aspect_ratio_ids=None, + aspect_ratio_mask=None, + cross_attention_mask=None, + past_key_values=None, + use_cache=False, + 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, + use_cache=use_cache, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + aspect_ratio_ids=aspect_ratio_ids, + aspect_ratio_mask=aspect_ratio_mask, + cross_attention_mask=cross_attention_mask, + logits_to_keep=logits_to_keep, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + # If we're in pre-fill or cacheless decoding step, then we need pixel_values and aspect ratios + # to compute image hidden states, otherwise they are cached within each cross attn layer + if not is_first_iteration and use_cache: + model_inputs["pixel_values"] = None + model_inputs["aspect_ratio_ids"] = None + model_inputs["aspect_ratio_mask"] = None + + return model_inputs + + def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder, **kwargs): + cross_attention_mask_prev = model_kwargs.get("cross_attention_mask", None) + model_kwargs = super()._update_model_kwargs_for_generation( + outputs=outputs, + model_kwargs=model_kwargs, + is_encoder_decoder=is_encoder_decoder, + **kwargs, + ) + + # add cross-attn mask for new token + if cross_attention_mask_prev is not None: + model_kwargs["cross_attention_mask"] = torch.cat( + [cross_attention_mask_prev, cross_attention_mask_prev[:, -1:, ...]], dim=1 + ) + return model_kwargs + + +__all__ = [ + "MllamaForConditionalGeneration", + "MllamaForCausalLM", + "MllamaTextModel", + "MllamaVisionModel", + "MllamaPreTrainedModel", + "MllamaModel", +] diff --git a/third_party/transformers/src/transformers/models/mllama/processing_mllama.py b/third_party/transformers/src/transformers/models/mllama/processing_mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..2a604b4cf0b0eaa1449482e6af48611dec5875ab --- /dev/null +++ b/third_party/transformers/src/transformers/models/mllama/processing_mllama.py @@ -0,0 +1,311 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Processor class for Mllama.""" + +import numpy as np + +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 PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +class MllamaProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "image_kwargs": { + "max_image_tiles": 4, + }, + } + + +def get_cross_attention_token_mask(input_ids: list[int], image_token_id: int) -> list[list[int]]: + """ + Generate a cross-attention token mask for image tokens in the input sequence. + + This function identifies the positions of image tokens in the input sequence and creates + a mask that defines which subsequent tokens each image token should attend to. + + Args: + input_ids (list[int]): A list of token ids representing the input sequence. + image_token_id (int): The id of the token used to represent images in the sequence. + + Returns: + list[list[int]]: A list of [start, end] pairs, where each pair represents the range + of tokens an image token should attend to. + + Notes: + - If no image tokens are present, an empty list is returned. + - For a single image token, it attends to all subsequent tokens until the end of the sequence. + - For multiple image tokens, each attends to tokens up to the next image token or the end of the sequence. + - Consecutive image tokens are treated as a group and attend to all subsequent tokens together. + """ + + image_token_locations = [i for i, token in enumerate(input_ids) if token == image_token_id] + + if len(image_token_locations) == 0: + return [] + + # only one image present, unmask until end of sequence + if len(image_token_locations) == 1: + return [[image_token_locations[0], -1]] + + vision_masks = [[loc1, loc2] for loc1, loc2 in zip(image_token_locations[:-1], image_token_locations[1:])] + + # last image will attend to all subsequent text + vision_masks.append([image_token_locations[-1], len(input_ids)]) + + # if there are two or more consecutive vision tokens, + # they should all attend to all subsequent + # text present + last_mask_end = vision_masks[-1][1] + for vision_mask in vision_masks[::-1]: + if vision_mask[0] == vision_mask[1] - 1: + vision_mask[1] = last_mask_end + last_mask_end = vision_mask[1] + + return vision_masks + + +def convert_sparse_cross_attention_mask_to_dense( + cross_attention_token_mask: list[list[list[int]]], + num_tiles: list[list[int]], + max_num_tiles: int, + length: int, +) -> np.ndarray: + """ + Convert the cross attention mask indices to a cross attention mask 4D array. + + This function takes a sparse representation of cross attention masks and converts it to a dense 4D numpy array. + The sparse representation is a nested list structure that defines attention ranges for each image in each batch item. + + Args: + cross_attention_token_mask (list[list[list[int]]]): A nested list structure where: + - The outer list represents the batch dimension. + - The middle list represents different images within each batch item. + - The inner list contains pairs of integers [start, end] representing token ranges for each image. + num_tiles (list[list[int]]): A nested list structure specifying the number of tiles for each image in each batch item. + max_num_tiles (int): The maximum possible number of tiles. + length (int): The total sequence length of the input. + + Returns: + np.ndarray: A 4D numpy array of shape (batch_size, length, max_num_images, max_num_tiles) + The array contains `1` where attention is allowed and `0` where it is not. + + Note: + - Special handling is done for cases where the end token is -1, which is interpreted as attending to the end of the sequence. + """ + + batch_size = len(cross_attention_token_mask) + max_num_images = max(len(masks) for masks in cross_attention_token_mask) + + cross_attention_mask = np.zeros( + shape=(batch_size, length, max_num_images, max_num_tiles), + dtype=np.int64, + ) + + for sample_idx, (sample_masks, sample_num_tiles) in enumerate(zip(cross_attention_token_mask, num_tiles)): + for mask_idx, (locations, mask_num_tiles) in enumerate(zip(sample_masks, sample_num_tiles)): + if len(locations) == 2: + start, end = locations + end = min(end, length) + if end == -1: + end = length + cross_attention_mask[sample_idx, start:end, mask_idx, :mask_num_tiles] = 1 + return cross_attention_mask + + +def build_string_from_input(prompt: str, bos_token: str, image_token: str) -> str: + """ + Builds a string from the input prompt by adding `bos_token` if not already present. + + Args: + prompt (`str`): + The input prompt string. + bos_token (`str`): + The beginning of sentence token to be added. + image_token (`str`): + The image token used to identify the start of an image sequence. + + Returns: + str: The modified prompt string with the `bos_token` added if necessary. + + Examples: + >>> build_string_from_input("Hello world", "", "<|image|>") + 'Hello world' + + >>> build_string_from_input("<|image|>Hello world", "", "<|image|>") + '<|image|>Hello world' + + >>> build_string_from_input("Hello world", "", "<|image|>") + 'Hello world' + """ + + if bos_token in prompt: + return prompt + + num_image_tokens_on_start = 0 + while prompt.startswith(image_token): + prompt = prompt[len(image_token) :] + num_image_tokens_on_start += 1 + + return f"{image_token * num_image_tokens_on_start}{bos_token}{prompt}" + + +@auto_docstring +class MllamaProcessor(ProcessorMixin): + def __init__(self, image_processor, tokenizer, chat_template=None): + if not hasattr(tokenizer, "image_token"): + self.image_token = "<|image|>" + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + else: + self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id + + self.python_token = "<|python_tag|>" + self.python_token_id = tokenizer.convert_tokens_to_ids(self.python_token) + self.bos_token = tokenizer.bos_token + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + **kwargs: Unpack[MllamaProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + TODO: add aspect_ratio_ids and aspect_ratio_mask and cross_attention_mask + """ + if text is None and images is None: + raise ValueError("You must specify either text or images.") + + output_kwargs = self._merge_kwargs( + MllamaProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + + data = {} + if text is not None: + if isinstance(text, str): + text = [text] + elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): + raise ValueError("Invalid input text. Please provide a string, or a list of strings") + n_images_in_text = [t.count(self.image_token) for t in text] + text = [build_string_from_input(text_item, self.bos_token, self.image_token) for text_item in text] + encoding = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, encoding, modalities=["image"]) + n_images_in_ids = [token_ids.count(self.image_token_id) for token_ids in encoding["input_ids"]] + data.update(encoding) + + n_images_in_images = [0] + if images is not None: + images = self.image_processor.fetch_images(images) + images = make_nested_list_of_images(images) + n_images_in_images = [len(sample) for sample in images] + + if text is not None: + if any(batch_img == 0 for batch_img in n_images_in_text) and not all( + batch_img == 0 for batch_img in n_images_in_text + ): + raise ValueError( + "If a batch of text is provided, there should be either no images or at least one image per sample" + ) + if sum(n_images_in_text) > 0 and ( + n_images_in_images != n_images_in_text or n_images_in_ids != n_images_in_images + ): + if images is None: + raise ValueError("No image were provided, but there are image tokens in the prompt") + else: + add_message = "" + if sum(n_images_in_images) == sum(n_images_in_text) and n_images_in_images != n_images_in_text: + add_message = "Make sure to pass your images as a nested list, where each sub-list holds images per batch" + elif n_images_in_ids != n_images_in_images: + add_message = "If you activated truncation with `max_length`, increase the `max_length` so image tokens aren't cropped." + + raise ValueError( + f"The number of image tokens in each text ({n_images_in_text}) should be the same as the " + f"number of provided images per batch ({n_images_in_images}). {add_message}" + ) + + if images is not None: + image_features = self.image_processor(images, **output_kwargs["images_kwargs"]) + num_tiles = image_features.pop("num_tiles") + data.update(image_features) + + # Create cross attention mask + if images is not None and text is not None: + cross_attention_token_mask = [ + get_cross_attention_token_mask(token_ids, self.image_token_id) for token_ids in encoding["input_ids"] + ] + cross_attention_mask = convert_sparse_cross_attention_mask_to_dense( + cross_attention_token_mask, + num_tiles=num_tiles, + max_num_tiles=self.image_processor.max_image_tiles, + length=max(len(input_ids) for input_ids in encoding["input_ids"]), + ) + data["cross_attention_mask"] = cross_attention_mask + + return BatchFeature(data=data, tensor_type=return_tensors) + + def post_process_image_text_to_text( + self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs + ): + """ + Post-process the output of the model to decode the text. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + or `(sequence_length,)`. + skip_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `list[str]`: The decoded text. + """ + return self.tokenizer.batch_decode( + generated_outputs, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + + # Remove `num_tiles`, it is popped and used only when processing. Make a copy of list when removing + # otherwise `self.image_processor.model_input_names` is also modified + image_processor_input_names = [name for name in image_processor_input_names if name != "num_tiles"] + return list(tokenizer_input_names + image_processor_input_names + ["cross_attention_mask"]) + + +__all__ = ["MllamaProcessor"] diff --git a/third_party/transformers/src/transformers/models/mt5/__init__.py b/third_party/transformers/src/transformers/models/mt5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f04d056d6e08afd696c8f2eb93cbe75dd84f713d --- /dev/null +++ b/third_party/transformers/src/transformers/models/mt5/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_mt5 import * + from .modeling_mt5 import * + from .tokenization_mt5 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/mt5/configuration_mt5.py b/third_party/transformers/src/transformers/models/mt5/configuration_mt5.py new file mode 100644 index 0000000000000000000000000000000000000000..e71270d835cec8f0fbbbd097e2a922d6f890ffb1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mt5/configuration_mt5.py @@ -0,0 +1,95 @@ +# Copyright 2020, The T5 Authors and HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""mT5 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/mt5-small") +@strict +class MT5Config(PreTrainedConfig): + r""" + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. + relative_attention_max_distance (`int`, *optional*, defaults to 128): + The maximum distance of the longer sequences for the bucket separation. + feed_forward_proj (`str`, *optional*, defaults to `"gated-gelu"`): + Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. + """ + + model_type = "mt5" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + "head_dim": "d_kv", + } + + vocab_size: int = 250112 + d_model: int = 512 + d_kv: int = 64 + d_ff: int = 1024 + num_layers: int = 8 + num_decoder_layers: int | None = None + num_heads: int = 6 + relative_attention_num_buckets: int = 32 + relative_attention_max_distance: int = 128 + dropout_rate: float | int = 0.1 + layer_norm_epsilon: float = 1e-6 + initializer_factor: float = 1.0 + feed_forward_proj: str = "gated-gelu" + is_encoder_decoder: bool = True + use_cache: bool = True + tie_word_embeddings: bool = True + bos_token_id: int | None = None + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + decoder_start_token_id: int | None = 0 + classifier_dropout: float | int = 0.0 + is_decoder: bool = False + + def __post_init__(self, **kwargs): + self.num_decoder_layers = ( + self.num_decoder_layers if self.num_decoder_layers is not None else self.num_layers + ) # default = symmetry + + act_info = self.feed_forward_proj.split("-") + self.dense_act_fn = act_info[-1] + self.is_gated_act = act_info[0] == "gated" + + if self.feed_forward_proj == "gated-gelu": + self.dense_act_fn = "gelu_new" + + # Force because official weights have False serialized, but we have to tie always + kwargs.pop("tie_word_embeddings", None) + self.tie_word_embeddings = True + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + act_info = self.feed_forward_proj.split("-") + + if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2: + raise ValueError( + f"`feed_forward_proj`: {self.feed_forward_proj} is not a valid activation function of the dense layer. " + "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. " + "'gated-gelu' or 'relu'" + ) + + +__all__ = ["MT5Config"] diff --git a/third_party/transformers/src/transformers/models/mt5/modeling_mt5.py b/third_party/transformers/src/transformers/models/mt5/modeling_mt5.py new file mode 100644 index 0000000000000000000000000000000000000000..90a9515397bb3deb556bb75fe1a29ea7dd438122 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mt5/modeling_mt5.py @@ -0,0 +1,1682 @@ +# Copyright 2020 Mesh TensorFlow authors, T5 Authors and 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 mT5 model.""" + +import copy +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 ...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 ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, + Seq2SeqQuestionAnsweringModelOutput, + Seq2SeqSequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import DUMMY_INPUTS, DUMMY_MASK, auto_docstring, logging, torch_compilable_check +from .configuration_mt5 import MT5Config + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->MT5 +class MT5LayerNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Construct a layernorm module in the MT5 style. No bias and no subtraction of mean. + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # MT5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean + # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated + # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for + # half-precision inputs is done in fp32 + + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->MT5 +class MT5DenseActDense(nn.Module): + def __init__(self, config: MT5Config): + super().__init__() + self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_states = self.wi(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.dropout(hidden_states) + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + hidden_states = self.wo(hidden_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->MT5 +class MT5DenseGatedActDense(nn.Module): + def __init__(self, config: MT5Config): + super().__init__() + self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->MT5 +class MT5LayerFF(nn.Module): + def __init__(self, config: MT5Config): + super().__init__() + if config.is_gated_act: + self.DenseReluDense = MT5DenseGatedActDense(config) + else: + self.DenseReluDense = MT5DenseActDense(config) + + self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, hidden_states): + forwarded_states = self.layer_norm(hidden_states) + forwarded_states = self.DenseReluDense(forwarded_states) + hidden_states = hidden_states + self.dropout(forwarded_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->MT5 +class MT5Attention(nn.Module): + def __init__( + self, + config: MT5Config, + has_relative_attention_bias=False, + layer_idx: int | None = None, + ): + super().__init__() + self.is_decoder = config.is_decoder + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + self.d_model = config.d_model + self.key_value_proj_dim = config.d_kv + self.n_heads = config.num_heads + self.dropout = config.dropout_rate + self.inner_dim = self.n_heads * self.key_value_proj_dim + self.layer_idx = layer_idx + if layer_idx is None and self.is_decoder: + logger.warning_once( + f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and " + "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) + + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) + + self.gradient_checkpointing = False + + @staticmethod + def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + + Translate relative position to a bucket number for relative attention. The relative position is defined as + memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to + position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for + small absolute relative_position and larger buckets for larger absolute relative_positions. All relative + positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. + This should allow for more graceful generalization to longer sequences than the model has been trained on + + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + + Returns: + a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) + """ + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) + # now relative_position is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) + ) + + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) + return relative_buckets + + def compute_bias(self, query_length, key_length, device=None, past_seen_tokens=0): + """Compute binned relative position bias""" + if device is None: + device = self.relative_attention_bias.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + past_seen_tokens + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] + relative_position = memory_position - context_position # shape (query_length, key_length) + relative_position_bucket = self._relative_position_bucket( + relative_position, # shape (query_length, key_length) + bidirectional=(not self.is_decoder), + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) + values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) + return values + + def forward( + self, + hidden_states, + mask=None, + key_value_states=None, + position_bias=None, + past_key_values=None, + output_attentions=False, + **kwargs, + ): + """ + Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). + """ + # Input is (batch_size, seq_length, dim) + # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder) + batch_size, seq_length = hidden_states.shape[:2] + past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0 + # We clone here for StaticCache, as we get the value before updating it, but use it after and it's the same ref + past_seen_tokens = past_seen_tokens.clone() if isinstance(past_seen_tokens, torch.Tensor) else past_seen_tokens + + # 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 + + q_input_shape = (batch_size, seq_length, -1, self.key_value_proj_dim) + query_states = self.q(hidden_states).view(*q_input_shape).transpose(1, 2) + + # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` + is_updated = False + 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: + kv_input_shape = (batch_size, current_states.shape[1], -1, self.key_value_proj_dim) + key_states = self.k(current_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v(current_states).view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 + scores = torch.matmul(query_states, key_states.transpose(3, 2)) + + if position_bias is None: + key_length = key_states.shape[-2] + if not self.has_relative_attention_bias: + position_bias = torch.zeros( + (1, query_states.shape[1], seq_length, key_length), device=scores.device, dtype=scores.dtype + ) + if self.gradient_checkpointing and self.training: + position_bias.requires_grad = True + else: + position_bias = self.compute_bias( + seq_length, key_length, device=scores.device, past_seen_tokens=past_seen_tokens + ) + + if mask is not None: + causal_mask = mask[:, :, :, : key_states.shape[-2]] + position_bias = position_bias + causal_mask + + position_bias_masked = position_bias + scores += position_bias_masked + + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(batch_size, seq_length, -1) + attn_output = self.o(attn_output) + + outputs = (attn_output, position_bias) + + if output_attentions: + outputs = outputs + (attn_weights,) + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->MT5 +class MT5LayerSelfAttention(nn.Module): + def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): + super().__init__() + self.SelfAttention = MT5Attention( + config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx + ) + self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + past_key_values=None, + use_cache=False, + output_attentions=False, + **kwargs, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.SelfAttention( + normed_hidden_states, + mask=attention_mask, + position_bias=position_bias, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = hidden_states + self.dropout(attention_output[0]) + outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->MT5 +class MT5LayerCrossAttention(nn.Module): + def __init__(self, config, layer_idx: int | None = None): + super().__init__() + self.EncDecAttention = MT5Attention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + key_value_states, + attention_mask=None, + position_bias=None, + past_key_values=None, + output_attentions=False, + **kwargs, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.EncDecAttention( + normed_hidden_states, + mask=attention_mask, + key_value_states=key_value_states, + position_bias=position_bias, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + layer_output = hidden_states + self.dropout(attention_output[0]) + outputs = (layer_output,) + attention_output[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 +class MT5Block(GradientCheckpointingLayer): + def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): + super().__init__() + self.is_decoder = config.is_decoder + self.layer = nn.ModuleList() + self.layer.append( + MT5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx) + ) + if self.is_decoder: + self.layer.append(MT5LayerCrossAttention(config, layer_idx=layer_idx)) + + self.layer.append(MT5LayerFF(config)) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + encoder_decoder_position_bias=None, + past_key_values=None, + use_cache=False, + output_attentions=False, + return_dict=True, + **kwargs, + ): + self_attention_outputs = self.layer[0]( + hidden_states, + attention_mask=attention_mask, + position_bias=position_bias, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = self_attention_outputs[0] + attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + do_cross_attention = self.is_decoder and encoder_hidden_states is not None + if do_cross_attention: + cross_attention_outputs = self.layer[1]( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_bias=encoder_decoder_position_bias, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + hidden_states = cross_attention_outputs[0] + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + # Keep cross-attention outputs and relative position weights + attention_outputs = attention_outputs + cross_attention_outputs[1:] + + # Apply Feed Forward layer + hidden_states = self.layer[-1](hidden_states) + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + return ( + outputs + attention_outputs + ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + + +# Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 +class MT5ClassificationHead(nn.Module): + """Head for sentence-level classification tasks.""" + + def __init__(self, config: MT5Config): + super().__init__() + self.dense = nn.Linear(config.d_model, config.d_model) + self.dropout = nn.Dropout(p=config.classifier_dropout) + self.out_proj = nn.Linear(config.d_model, config.num_labels) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dropout(hidden_states) + hidden_states = self.dense(hidden_states) + hidden_states = torch.tanh(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.out_proj(hidden_states) + return hidden_states + + +@auto_docstring +# Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel with T5->MT5, t5->mt5 +class MT5PreTrainedModel(PreTrainedModel): + config: MT5Config + base_model_prefix = "transformer" + supports_gradient_checkpointing = True + _can_compile_fullgraph = True + + _no_split_modules = ["MT5Block"] + _keep_in_fp32_modules = ["wo"] + + @property + def dummy_inputs(self): + input_ids = torch.tensor(DUMMY_INPUTS) + input_mask = torch.tensor(DUMMY_MASK) + dummy_inputs = { + "decoder_input_ids": input_ids, + "input_ids": input_ids, + "decoder_attention_mask": input_mask, + } + return dummy_inputs + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor # Used for testing weights initialization + if isinstance(module, MT5LayerNorm): + init.constant_(module.weight, factor * 1.0) + elif isinstance( + module, + (MT5Model, MT5ForConditionalGeneration, MT5EncoderModel, MT5ForQuestionAnswering), + ): + init.normal_(module.shared.weight, mean=0.0, std=factor * 1.0) + if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: + init.normal_(module.lm_head.weight, mean=0.0, std=factor * 1.0) + if hasattr(module, "qa_outputs"): + init.normal_(module.qa_outputs.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + init.zeros_(module.qa_outputs.bias) + elif isinstance(module, MT5ForTokenClassification): + if hasattr(module, "classifier"): + init.normal_(module.classifier.weight, mean=0.0, std=factor * 1.0) + init.zeros_(module.classifier.bias) + elif isinstance(module, MT5ClassificationHead): + init.normal_(module.dense.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.dense, "bias") and module.dense.bias is not None: + init.zeros_(module.dense.bias) + init.normal_(module.out_proj.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: + init.zeros_(module.out_proj.bias) + elif isinstance(module, MT5DenseActDense): + init.normal_(module.wi.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi, "bias") and module.wi.bias is not None: + init.zeros_(module.wi.bias) + init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + init.zeros_(module.wo.bias) + elif isinstance(module, MT5DenseGatedActDense): + init.normal_(module.wi_0.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: + init.zeros_(module.wi_0.bias) + init.normal_(module.wi_1.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: + init.zeros_(module.wi_1.bias) + init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + init.zeros_(module.wo.bias) + elif isinstance(module, MT5Attention): + d_model = self.config.d_model + key_value_proj_dim = self.config.d_kv + n_heads = self.config.num_heads + init.normal_(module.q.weight, mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) + init.normal_(module.k.weight, mean=0.0, std=factor * (d_model**-0.5)) + init.normal_(module.v.weight, mean=0.0, std=factor * (d_model**-0.5)) + init.normal_(module.o.weight, mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) + if module.has_relative_attention_bias: + init.normal_(module.relative_attention_bias.weight, mean=0.0, std=factor * ((d_model) ** -0.5)) + + def _shift_right(self, input_ids): + decoder_start_token_id = self.config.decoder_start_token_id + pad_token_id = self.config.pad_token_id + + if decoder_start_token_id is None: + raise ValueError( + "self.model.config.decoder_start_token_id has to be defined. In MT5 it is usually set to the pad_token_id. " + "See MT5 docs for more information." + ) + + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +# Copied from transformers.models.t5.modeling_t5.T5Stack with T5->MT5 +class MT5Stack(MT5PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model) + self.is_decoder = config.is_decoder + + self.block = nn.ModuleList( + [MT5Block(config, has_relative_attention_bias=bool(i == 0), layer_idx=i) for i in range(config.num_layers)] + ) + self.final_layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + # Initialize weights and apply final processing + self.post_init() + self.gradient_checkpointing = False + + def set_input_embeddings(self, new_embeddings): + self.embed_tokens = new_embeddings + + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + inputs_embeds=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + use_cache = use_cache if use_cache is not None else self.config.use_cache + 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: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError( + f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}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: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if inputs_embeds is None: + if self.embed_tokens is None: + raise ValueError("You have to initialize the model with valid token embeddings") + inputs_embeds = self.embed_tokens(input_ids) + + batch_size, seq_length = input_shape + + if use_cache is True: + if not self.is_decoder: + raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") + + if self.is_decoder: + if use_cache and past_key_values is None: + if self.config.is_encoder_decoder: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + else: + past_key_values = DynamicCache(config=self.config) + elif not self.is_decoder: + # do not pass cache object down the line for encoder stack + # it messes indexing later in decoder-stack because cache object is modified in-place + past_key_values = None + + if self.config.is_decoder: + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values.self_attention_cache + if isinstance(past_key_values, EncoderDecoderCache) + else past_key_values, + ) + else: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + encoder_extended_attention_mask = None + if self.is_decoder and encoder_hidden_states is not None: + encoder_extended_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + all_cross_attentions = () if (output_attentions and self.is_decoder) else None + position_bias = None + encoder_decoder_position_bias = None + + hidden_states = self.dropout(inputs_embeds) + + for layer_module in self.block: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_outputs = layer_module( + hidden_states, + attention_mask, + position_bias, + encoder_hidden_states, + encoder_extended_attention_mask, + encoder_decoder_position_bias, # as a positional argument for gradient checkpointing + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + return_dict=return_dict, + ) + + hidden_states = layer_outputs[0] + + # We share the position biases between the layers - the first layer store them + # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), + # (cross-attention position bias), (cross-attention weights) + position_bias = layer_outputs[1] + if self.is_decoder and encoder_hidden_states is not None: + encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[2],) + if self.is_decoder: + all_cross_attentions = all_cross_attentions + (layer_outputs[4],) + + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + # Add last layer + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + past_key_values, + all_hidden_states, + all_attentions, + 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_attentions, + cross_attentions=all_cross_attentions, + ) + + +@auto_docstring +class MT5Model(MT5PreTrainedModel): + r""" + Examples: + + ```python + >>> from transformers import MT5Model, AutoTokenizer + + >>> model = MT5Model.from_pretrained("google/mt5-small") + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." + >>> summary = "Weiter Verhandlung in Syrien." + >>> inputs = tokenizer(article, return_tensors="pt") + >>> labels = tokenizer(text_target=summary, return_tensors="pt") + + >>> outputs = model(input_ids=inputs["input_ids"], decoder_input_ids=labels["input_ids"]) + >>> hidden_states = outputs.last_hidden_state + ```""" + + model_type = "mt5" + config: MT5Config + _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + } + + # Copied from transformers.models.t5.modeling_t5.T5Model.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + self.encoder = MT5Stack(encoder_config) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.num_layers = config.num_decoder_layers + self.decoder = MT5Stack(decoder_config) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.t5.modeling_t5.T5Model.get_input_embeddings + def get_input_embeddings(self): + return self.shared + + # Copied from transformers.models.t5.modeling_t5.T5Model.set_input_embeddings + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + @auto_docstring + # Copied from transformers.models.t5.modeling_t5.T5Model.forward with google-t5/->google/, T5->MT5, t5->mt5 + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.BoolTensor | None = None, + encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor] | Seq2SeqModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + MT5 uses the `pad_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`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 + Training](./mt5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + + Example: + + ```python + >>> from transformers import AutoTokenizer, MT5Model + + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> model = MT5Model.from_pretrained("google/mt5-small") + + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 + + >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for MT5Model. + >>> # This is not needed for torch's MT5ForConditionalGeneration as it does this internally using labels arg. + >>> decoder_input_ids = model._shift_right(decoder_input_ids) + + >>> # forward pass + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + 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 + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and 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, + ) + + hidden_states = encoder_outputs[0] + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + 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 Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + MT5 Model with a `language modeling` head on top. + """ +) +class MT5ForConditionalGeneration(MT5PreTrainedModel, GenerationMixin): + r""" + Examples: + + ```python + >>> from transformers import MT5ForConditionalGeneration, AutoTokenizer + + >>> model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." + >>> summary = "Weiter Verhandlung in Syrien." + >>> inputs = tokenizer(article, text_target=summary, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> loss = outputs.loss + ```""" + + model_type = "mt5" + config: MT5Config + _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + "lm_head.weight": "shared.weight", + } + + # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.model_dim = config.d_model + + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + self.encoder = MT5Stack(encoder_config) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.num_layers = config.num_decoder_layers + self.decoder = MT5Stack(decoder_config) + + self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.BoolTensor | None = None, + encoder_outputs: tuple[tuple[torch.Tensor]] | 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.FloatTensor] | Seq2SeqLMOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + MT5 uses the `pad_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`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 + Training](./mt5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., + config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for + labels in `[0, ..., config.vocab_size]` + + Examples: + + ```python + >>> from transformers import AutoTokenizer, MT5ForConditionalGeneration + + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") + + >>> # training + >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids + >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids + >>> outputs = model(input_ids=input_ids, labels=labels) + >>> loss = outputs.loss + >>> logits = outputs.logits + + >>> # inference + >>> input_ids = tokenizer( + ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> outputs = model.generate(input_ids) + >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) + >>> # studies have shown that owning a dog is good for you. + ```""" + 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 + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + # Convert encoder inputs in embeddings if needed + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and 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, + ) + + hidden_states = encoder_outputs[0] + + if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: + # get decoder inputs from shifting lm labels to the right + decoder_input_ids = self._shift_right(labels) + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = decoder_outputs[0] + + lm_logits = self.lm_head(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss(ignore_index=-100) + # move labels to correct device to enable PP + labels = labels.to(lm_logits.device) + loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) + + if not return_dict: + output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + 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, + ) + + # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_decoder_input_ids_from_labels + def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): + return self._shift_right(labels) + + +@auto_docstring +class MT5EncoderModel(MT5PreTrainedModel): + r""" + Examples: + + ```python + >>> from transformers import MT5EncoderModel, AutoTokenizer + + >>> model = MT5EncoderModel.from_pretrained("google/mt5-small") + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." + >>> input_ids = tokenizer(article, return_tensors="pt").input_ids + >>> outputs = model(input_ids) + >>> hidden_state = outputs.last_hidden_state + ```""" + + model_type = "mt5" + config: MT5Config + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + } + + # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = config + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = MT5Stack(encoder_config) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_input_embeddings + def get_input_embeddings(self): + return self.shared + + # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.set_input_embeddings + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + + @auto_docstring + # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.forward with google-t5/->google/, T5->MT5, t5->mt5 + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | 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, + ) -> tuple[torch.FloatTensor] | BaseModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). + + Example: + + ```python + >>> from transformers import AutoTokenizer, MT5EncoderModel + + >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") + >>> model = MT5EncoderModel.from_pretrained("google/mt5-small") + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + return encoder_outputs + + +@auto_docstring( + custom_intro=""" + MT5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE + tasks. + """ +) +class MT5ForSequenceClassification(MT5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] + + # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.transformer = MT5Model(config) + self.classification_head = MT5ClassificationHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.forward with T5->MT5, t5->mt5 + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | Seq2SeqSequenceClassifierOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + MT5 uses the `pad_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`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 + Training](./mt5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + if labels is not None: + use_cache = False + + if input_ids is None and inputs_embeds is not None: + raise NotImplementedError( + f"Passing input embeddings is currently not supported for {self.__class__.__name__}" + ) + + # Copied from models.bart.modeling_bart.BartModel.forward different to other models, MT5 automatically creates + # decoder_input_ids from input_ids if no decoder_input_ids are provided + if decoder_input_ids is None and decoder_inputs_embeds is None: + if input_ids is None: + raise ValueError( + "If no `decoder_input_ids` or `decoder_inputs_embeds` are " + "passed, `input_ids` cannot be `None`. Please pass either " + "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." + ) + decoder_input_ids = self._shift_right(input_ids) + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + + eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) + + torch_compilable_check( + torch.unique_consecutive(eos_mask.sum(1)).numel() == 1, + "All examples must have the same number of tokens.", + ) + batch_size, _, hidden_size = sequence_output.shape + sentence_representation = sequence_output[eos_mask, :].view(batch_size, -1, hidden_size)[:, -1, :] + logits = self.classification_head(sentence_representation) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.config.num_labels == 1: + self.config.problem_type = "regression" + elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.config.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return Seq2SeqSequenceClassifierOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + +@auto_docstring +class MT5ForTokenClassification(MT5PreTrainedModel): + # Copied from transformers.models.t5.modeling_t5.T5ForTokenClassification.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = MT5EncoderModel(config) + self.dropout = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + # Copied from transformers.models.t5.modeling_t5.T5ForTokenClassification.forward with T5->MT5 + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor] | TokenClassifierOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./t5#training). + 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.transformer( + input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.dropout(hidden_states) + logits = self.classifier(hidden_states) + + 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:-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 MT5ForQuestionAnswering(MT5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + } + + # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.__init__ with T5->MT5 + def __init__(self, config: MT5Config): + super().__init__(config) + self.model_dim = config.d_model + + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + self.encoder = MT5Stack(encoder_config) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.num_layers = config.num_decoder_layers + self.decoder = MT5Stack(decoder_config) + + self.num_labels = config.num_labels + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_input_embeddings + def get_input_embeddings(self): + return self.shared + + # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.set_input_embeddings + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + @auto_docstring + # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.forward + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.BoolTensor | None = None, + encoder_outputs: tuple[tuple[torch.Tensor]] | 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.FloatTensor] | Seq2SeqQuestionAnsweringModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + T5 uses the `pad_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`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 + Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + use_cache = use_cache if use_cache is not None else self.config.use_cache + if start_positions is not None and end_positions is not None: + use_cache = False + + # Copied from models.bart.modeling_bart.BartModel.forward + # different to other models, T5 automatically creates decoder_input_ids from + # input_ids if no decoder_input_ids are provided + if decoder_input_ids is None and decoder_inputs_embeds is None: + if input_ids is None: + raise ValueError( + "If no `decoder_input_ids` or `decoder_inputs_embeds` are " + "passed, `input_ids` cannot be `None`. Please pass either " + "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." + ) + decoder_input_ids = self._shift_right(input_ids) + + 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 + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and 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, + ) + + hidden_states = encoder_outputs[0] + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=None, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = decoder_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).to(start_logits.device) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1).to(end_logits.device) + # 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) + decoder_outputs[1:] + encoder_outputs + return ((total_loss,) + output) if total_loss is not None else output + + return Seq2SeqQuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + 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, + ) + + +__all__ = [ + "MT5EncoderModel", + "MT5ForConditionalGeneration", + "MT5ForQuestionAnswering", + "MT5ForSequenceClassification", + "MT5ForTokenClassification", + "MT5Model", + "MT5PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/nanochat/__init__.py b/third_party/transformers/src/transformers/models/nanochat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6610a71fa701d0f8854c96ed5b63ef389666cc90 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nanochat/__init__.py @@ -0,0 +1,14 @@ +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_nanochat import * + from .modeling_nanochat 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/nanochat/configuration_nanochat.py b/third_party/transformers/src/transformers/models/nanochat/configuration_nanochat.py new file mode 100644 index 0000000000000000000000000000000000000000..bd915e06dab4027213197a1dc6cefc8e7e8bab47 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nanochat/configuration_nanochat.py @@ -0,0 +1,81 @@ +# 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 huggingface_hub.dataclasses import strict + +from ...configuration_utils import PretrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="karpathy/nanochat-d32") +@strict +class NanoChatConfig(PretrainedConfig): + r""" + Example: + + ```python + >>> from transformers import NanoChatModel, NanoChatConfig + + >>> # Initializing a NanoChat style configuration + >>> configuration = NanoChatConfig() + + >>> # Initializing a model from the NanoChat style configuration + >>> model = NanoChatModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "nanochat" + 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.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise", + } + + vocab_size: int = 50304 + hidden_size: int = 768 + intermediate_size: int = 8192 + num_hidden_layers: int = 12 + num_attention_heads: int = 6 + num_key_value_heads: int | None = None + max_position_embeddings: int = 2048 + hidden_act: str = "relu2" + attention_dropout: float | int = 0.0 + rms_norm_eps: float = 1e-6 + initializer_range: float = 0.02 + rope_parameters: RopeParameters | dict | None = None + use_cache: bool = True + final_logit_softcapping: float | None = 15.0 + attention_bias: bool = False + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + pad_token_id: int | None = 1 + tie_word_embeddings: bool = False + + 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__ = ["NanoChatConfig"] diff --git a/third_party/transformers/src/transformers/models/nanochat/convert_nanochat_checkpoints.py b/third_party/transformers/src/transformers/models/nanochat/convert_nanochat_checkpoints.py new file mode 100644 index 0000000000000000000000000000000000000000..977e857d6aceaf2706a425015f1dcfd88b32bbfe --- /dev/null +++ b/third_party/transformers/src/transformers/models/nanochat/convert_nanochat_checkpoints.py @@ -0,0 +1,306 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import gc +import json +import os +from pathlib import Path + +import torch + +from transformers import AutoTokenizer, NanoChatConfig, NanoChatForCausalLM + + +def infer_kv_heads(config: NanoChatConfig, state_dict: dict[str, torch.Tensor]) -> int: + key_weight = state_dict.get("transformer.h.0.attn.c_k.weight") + if key_weight is None: + return config.num_key_value_heads + rows = key_weight.shape[0] + head_dim = config.hidden_size // config.num_attention_heads + if rows % head_dim != 0: + return config.num_key_value_heads + inferred = rows // head_dim + print(f"Inferred {inferred} key_value heads from checkpoint") + return max(inferred, 1) + + +def convert_layer(old_prefix: str, new_prefix: str) -> dict[str, str]: + return { + f"{old_prefix}.attn.c_q.weight": f"{new_prefix}.self_attn.q_proj.weight", + f"{old_prefix}.attn.c_k.weight": f"{new_prefix}.self_attn.k_proj.weight", + f"{old_prefix}.attn.c_v.weight": f"{new_prefix}.self_attn.v_proj.weight", + f"{old_prefix}.attn.c_proj.weight": f"{new_prefix}.self_attn.o_proj.weight", + f"{old_prefix}.mlp.c_fc.weight": f"{new_prefix}.mlp.fc1.weight", + f"{old_prefix}.mlp.c_proj.weight": f"{new_prefix}.mlp.fc2.weight", + } + + +def load_config_from_checkpoint(input_path: Path) -> NanoChatConfig: + """Load config from either meta_*.json or config.json in the checkpoint directory.""" + # Try to find meta_*.json first + meta_files = list(input_path.glob("meta_*.json")) + + if meta_files: + meta_file = meta_files[0] + print(f"Loading config from {meta_file.name}") + with open(meta_file, "r") as f: + meta_config = json.load(f) + + # Extract model config from meta file + if "model_config" in meta_config: + model_config = meta_config["model_config"] + else: + model_config = meta_config + + # Map to NanoChat config parameters + config_kwargs = { + "vocab_size": model_config.get("vocab_size", 50304), + "hidden_size": model_config.get("n_embd", 768), + "num_hidden_layers": model_config.get("n_layer", 12), + "num_attention_heads": model_config.get("n_head", 6), + "num_key_value_heads": model_config.get("n_kv_head"), + "max_position_embeddings": model_config.get("sequence_len", 2048), + "intermediate_size": model_config.get("intermediate_size", model_config.get("n_embd", 768) * 4), + } + + # Try to load existing config.json for additional parameters + config_file = input_path / "config.json" + if config_file.exists(): + print("Loading additional config from config.json") + with open(config_file, "r") as f: + extra_config = json.load(f) + + # Add additional parameters from config.json + for key in [ + "hidden_act", + "attention_dropout", + "rms_norm_eps", + "initializer_range", + "logits_soft_cap", + "attention_bias", + "intermediate_size", + "bos_token_id", + "eos_token_id", + "pad_token_id", + ]: + if key in extra_config: + config_kwargs[key] = extra_config[key] + # Handle legacy qkv_bias -> attention_bias conversion + elif key == "attention_bias" and "qkv_bias" in extra_config: + config_kwargs[key] = extra_config["qkv_bias"] + + # Handle rope_theta as a direct kwarg for the rope_parameters processing + if "rope_theta" in extra_config: + config_kwargs["rope_theta"] = extra_config["rope_theta"] + + # Handle rope_parameters or rope_scaling if present + if "rope_parameters" in extra_config: + config_kwargs["rope_parameters"] = extra_config["rope_parameters"] + elif "rope_scaling" in extra_config and extra_config["rope_scaling"] is not None: + config_kwargs["rope_parameters"] = extra_config["rope_scaling"] + + config = NanoChatConfig(**config_kwargs) + else: + # Fallback to loading from config.json if it exists + config_file = input_path / "config.json" + if config_file.exists(): + print("Loading config from config.json") + config = NanoChatConfig.from_pretrained(input_path) + # Handle legacy qkv_bias -> attention_bias conversion + if hasattr(config, "qkv_bias") and not hasattr(config, "attention_bias"): + config.attention_bias = config.qkv_bias + else: + raise ValueError(f"No config file found in {input_path}. Expected meta_*.json or config.json") + + return config + + +def write_model(input_dir, output_dir): + """Convert NanoChat model from original checkpoint format to HuggingFace format.""" + print("Converting the model.") + os.makedirs(output_dir, exist_ok=True) + + input_path = Path(input_dir) + + # Load config + config = load_config_from_checkpoint(input_path) + print(f"Loaded config hidden_size={config.hidden_size} num_layers={config.num_hidden_layers}") + + # Load checkpoint - try model_*.pt first, then pytorch_model.bin + checkpoint_files = list(input_path.glob("model_*.pt")) + if checkpoint_files: + checkpoint_path = checkpoint_files[0] + else: + checkpoint_path = input_path / "pytorch_model.bin" + + print(f"Fetching all parameters from the checkpoint at {checkpoint_path}...") + old_state = torch.load(checkpoint_path, map_location="cpu") + + # Original nanochat weights are in bfloat16 + for key in old_state: + if old_state[key].dtype == torch.float32: + old_state[key] = old_state[key].to(torch.bfloat16) + + # Infer key-value heads from checkpoint + inferred_kv = infer_kv_heads(config, old_state) + config.num_key_value_heads = inferred_kv + if config.num_attention_heads % config.num_key_value_heads != 0: + print(f"Adjusting num_attention_heads from {config.num_attention_heads} to {config.num_key_value_heads}") + config.num_attention_heads = config.num_key_value_heads + + print("Converting model...") + state_dict = {} + rename_map = {} + + def assign( + old_key: str, + new_key: str, + old_state: dict[str, torch.Tensor], + state_dict: dict[str, torch.Tensor], + rename_map: dict[str, str], + ) -> None: + tensor = old_state.get(old_key) + if tensor is None: + return + state_dict[new_key] = tensor.clone() + rename_map[old_key] = new_key + + # Convert embeddings and head + assign("transformer.wte.weight", "model.embed_tokens.weight", old_state, state_dict, rename_map) + assign("lm_head.weight", "lm_head.weight", old_state, state_dict, rename_map) + + # Convert layers + for layer_idx in range(config.num_hidden_layers): + old_prefix = f"transformer.h.{layer_idx}" + new_prefix = f"model.layers.{layer_idx}" + mapping = convert_layer(old_prefix, new_prefix) + for old_key, new_key in mapping.items(): + assign(old_key, new_key, old_state, state_dict, rename_map) + + missing = [key for key in old_state.keys() if key not in rename_map] + if missing: + print(f"Skipped {len(missing)} legacy entries that have no equivalent in the shared implementation") + + del old_state + gc.collect() + + # Update config + config.torch_dtype = torch.bfloat16 + config.tie_word_embeddings = False + + # Load the checkpoint into the model + print("Loading the checkpoint in a NanoChat model.") + with torch.device("meta"): + model = NanoChatForCausalLM(config) + model.load_state_dict(state_dict, strict=True, assign=True) + print("Checkpoint loaded successfully.") + + if hasattr(model.config, "_name_or_path"): + del model.config._name_or_path + + print("Saving the model.") + model.save_pretrained(output_dir) + del state_dict, model + + # Safety check: reload the converted model + gc.collect() + print("Reloading the model to check if it's saved correctly.") + NanoChatForCausalLM.from_pretrained(output_dir, torch_dtype=torch.bfloat16, device_map="auto") + print("Model reloaded successfully.") + + +def write_tokenizer(input_dir, output_dir): + """Convert and save the tokenizer.""" + input_path = Path(input_dir) + + # Convert the pickle tokenizer to HF format + tokenizer_pkl = input_path / "tokenizer.pkl" + if tokenizer_pkl.exists(): + try: + import pickle + + from transformers.integrations.tiktoken import convert_tiktoken_to_fast + + with open(tokenizer_pkl, "rb") as f: + tok_pkl = pickle.load(f) + convert_tiktoken_to_fast(tok_pkl, output_dir) + print("Converted tokenizer.pkl to HuggingFace format") + except Exception as e: + print(f"Warning: Failed to convert tokenizer.pkl: {e}") + # Fallback: copy tokenizer files if they exist + for filename in ("tokenizer.json", "tokenizer_config.json"): + src = input_path / filename + if src.exists(): + (Path(output_dir) / filename).write_bytes(src.read_bytes()) + else: + # No pickle tokenizer, copy JSON files + for filename in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"): + src = input_path / filename + if src.exists(): + (Path(output_dir) / filename).write_bytes(src.read_bytes()) + + print("Tokenizer saved successfully.") + + +def run_test(output_dir: str, prompt: str, max_new_tokens: int = 64) -> None: + """Run a quick generation test to verify the converted model works correctly.""" + print(f"Running quick generation test with prompt: {prompt}") + tokenizer = AutoTokenizer.from_pretrained(output_dir) + model = NanoChatForCausalLM.from_pretrained(output_dir, torch_dtype=torch.bfloat16) + model.eval() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + inputs = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + output = model.generate(**inputs, max_new_tokens=max_new_tokens) + generated = tokenizer.decode(output[0, inputs.input_ids.shape[1] :], skip_special_tokens=True) + print(f"Generated text: {generated}") + + +def main(): + parser = argparse.ArgumentParser(description="Convert NanoChat checkpoints to HuggingFace format") + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="Path to the original checkpoint directory", + ) + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--test_prompt", + type=str, + default=None, + help="Optional prompt for a quick generation test", + ) + args = parser.parse_args() + + write_model( + args.input_dir, + args.output_dir, + ) + + write_tokenizer(args.input_dir, args.output_dir) + + if args.test_prompt: + run_test(args.output_dir, args.test_prompt) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/nanochat/modeling_nanochat.py b/third_party/transformers/src/transformers/models/nanochat/modeling_nanochat.py new file mode 100644 index 0000000000000000000000000000000000000000..9205b89cd360e1363fcda18dc7826adf01dd2724 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nanochat/modeling_nanochat.py @@ -0,0 +1,518 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/nanochat/modular_nanochat.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_nanochat.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. + +import math +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_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_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_nanochat import NanoChatConfig + + +class NanoChatRMSNorm(torch.nn.Module): + def __init__(self, eps: float = 1e-6): + super().__init__() + self.eps = eps + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + return self._norm(x.float()).type_as(x) + + def extra_repr(self): + return f"eps={self.eps}" + + +class NanoChatRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: NanoChatConfig, 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: NanoChatConfig | 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) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +def rotate_half(x): + """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" + 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 NanoChatAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: NanoChatConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + self.q_norm = NanoChatRMSNorm(eps=config.rms_norm_eps) + self.k_norm = NanoChatRMSNorm(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, + 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) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # RoPE -> Norm (instead of usual Norm -> RoPE) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class NanoChatMLP(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, bias=False) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + 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 NanoChatDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: NanoChatConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = NanoChatAttention(config=config, layer_idx=layer_idx) + + self.mlp = NanoChatMLP(config) + + self.input_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps) + self.post_attention_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class NanoChatPreTrainedModel(PreTrainedModel): + config: NanoChatConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["NanoChatDecoderLayer"] + _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": NanoChatDecoderLayer, + "attentions": NanoChatAttention, + } + + def _init_weights(self, module: nn.Module) -> None: + super()._init_weights(module) + if isinstance(module, NanoChatAttention): + init.normal_( + module.o_proj.weight, + mean=0.0, + std=self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers), + ) + + +@auto_docstring +class NanoChatModel(NanoChatPreTrainedModel): + def __init__(self, config: NanoChatConfig): + 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( + [NanoChatDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + self.norm = NanoChatRMSNorm(eps=config.rms_norm_eps) + self.rotary_emb = NanoChatRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + hidden_states = self.norm(hidden_states) # Additional norm before the layers + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class NanoChatForCausalLM(NanoChatPreTrainedModel, 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 = NanoChatModel(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, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("karpathy/nanochat-d32") + + >>> tokenizer = AutoTokenizer.from_pretrained("karpathy/nanochat-d32") + + >>> conversation = [ + {"role": "user", "content": "What is the capital of France?"}, + ] + + >>> inputs = tokenizer.apply_chat_template( + conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ).to(device) + + >>> with torch.no_grad(): + >>> outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False) + + >>> generated_tokens = outputs[0, inputs["input_ids"].shape[1] :] + >>> output = tokenizer.decode(generated_tokens, skip_special_tokens=True) + ```""" + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + 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, :]) + if self.config.final_logit_softcapping is not None: + logits = logits / self.config.final_logit_softcapping + logits = torch.tanh(logits) + logits = logits * self.config.final_logit_softcapping + + 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__ = ["NanoChatPreTrainedModel", "NanoChatModel", "NanoChatForCausalLM"] diff --git a/third_party/transformers/src/transformers/models/nanochat/modular_nanochat.py b/third_party/transformers/src/transformers/models/nanochat/modular_nanochat.py new file mode 100644 index 0000000000000000000000000000000000000000..713cc29b81eb304e2d92f506294021858bb58c92 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nanochat/modular_nanochat.py @@ -0,0 +1,235 @@ +# 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 collections.abc import Callable + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ..clip.modeling_clip import CLIPMLP +from ..gemma2.modeling_gemma2 import Gemma2ForCausalLM +from ..llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaModel, + LlamaPreTrainedModel, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from ..llama4.modeling_llama4 import Llama4TextL2Norm +from ..qwen3.modeling_qwen3 import Qwen3Attention +from .configuration_nanochat import NanoChatConfig + + +class NanoChatRMSNorm(Llama4TextL2Norm): + pass + + +class NanoChatRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +def rotate_half(x): + """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((x2, -x1), dim=-1) + + +class NanoChatAttention(Qwen3Attention): + def __init__(self, config: NanoChatConfig, layer_idx: int): + super().__init__(config, layer_idx) + del self.sliding_window + del self.layer_type + + self.q_norm = NanoChatRMSNorm(eps=config.rms_norm_eps) + self.k_norm = NanoChatRMSNorm(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, + 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) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # RoPE -> Norm (instead of usual Norm -> RoPE) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class NanoChatMLP(CLIPMLP): + def __init__(self, config): + super().__init__(config) + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + +class NanoChatDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: NanoChatConfig, layer_idx: int): + super().__init__() + + self.input_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps) + self.post_attention_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps) + + +@auto_docstring +class NanoChatPreTrainedModel(LlamaPreTrainedModel): + def _init_weights(self, module: nn.Module) -> None: + PreTrainedModel._init_weights(self, module) + if isinstance(module, NanoChatAttention): + init.normal_( + module.o_proj.weight, + mean=0.0, + std=self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers), + ) + + +@auto_docstring +class NanoChatModel(LlamaModel): + def __init__(self, config: NanoChatConfig): + super().__init__(config) + + self.norm = NanoChatRMSNorm(eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + hidden_states = self.norm(hidden_states) # Additional norm before the layers + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class NanoChatForCausalLM(Gemma2ForCausalLM): + _tp_plan = {"lm_head": "colwise_gather_output"} + + def forward(self, **super_kwargs) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("karpathy/nanochat-d32") + + >>> tokenizer = AutoTokenizer.from_pretrained("karpathy/nanochat-d32") + + >>> conversation = [ + {"role": "user", "content": "What is the capital of France?"}, + ] + + >>> inputs = tokenizer.apply_chat_template( + conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ).to(device) + + >>> with torch.no_grad(): + >>> outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False) + + >>> generated_tokens = outputs[0, inputs["input_ids"].shape[1] :] + >>> output = tokenizer.decode(generated_tokens, skip_special_tokens=True) + ```""" + super().forward(**super_kwargs) + + +__all__ = [ + "NanoChatPreTrainedModel", + "NanoChatModel", + "NanoChatForCausalLM", +] diff --git a/third_party/transformers/src/transformers/models/nemotron/__init__.py b/third_party/transformers/src/transformers/models/nemotron/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2bd66ddd622e067fc6b2863ffdc745a9e09f2307 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron/__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_nemotron import * + from .modeling_nemotron 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/configuration_nemotron.py b/third_party/transformers/src/transformers/models/nemotron/configuration_nemotron.py new file mode 100644 index 0000000000000000000000000000000000000000..8d4cf07c1f05f1928486aacc672065f6e50f6404 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron/configuration_nemotron.py @@ -0,0 +1,73 @@ +# Copyright 2024 HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024, 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. +"""Nemotron 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="thhaus/nemotron3-8b") +@strict +class NemotronConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import NemotronModel, NemotronConfig + + >>> # Initializing a Nemotron nemotron-15b style configuration + >>> configuration = NemotronConfig() + + >>> # Initializing a model from the nemotron-15b style configuration + >>> model = NemotronModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "nemotron" + keys_to_ignore_at_inference = ["past_key_values"] + + vocab_size: int = 256000 + hidden_size: int = 6144 + intermediate_size: int = 24576 + num_hidden_layers: int = 32 + num_attention_heads: int = 48 + head_dim: int | None = None + num_key_value_heads: int | None = None + hidden_act: str = "relu2" + max_position_embeddings: int = 4096 + initializer_range: float = 0.0134 + norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 2 + eos_token_id: int | list[int] | None = 3 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + mlp_bias: bool = False + + def __post_init__(self, **kwargs): + self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads + kwargs.setdefault("partial_rotary_factor", 0.5) # assign default for BC + super().__post_init__(**kwargs) + + +__all__ = ["NemotronConfig"] diff --git a/third_party/transformers/src/transformers/models/nemotron/convert_nemotron_nemo_to_hf.py b/third_party/transformers/src/transformers/models/nemotron/convert_nemotron_nemo_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..b0856d14166af9d9c89470869f3a06d6312b6352 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron/convert_nemotron_nemo_to_hf.py @@ -0,0 +1,346 @@ +# Copyright (c) 2024, 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 json +import os +import shutil +from argparse import ArgumentParser +from collections import OrderedDict + +import torch +from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer +from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel +from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy +from nemo.utils import logging +from pytorch_lightning import Trainer + +from transformers import LlamaTokenizer, PreTrainedTokenizerFast +from transformers.convert_slow_tokenizer import LlamaConverter + + +""" +Script to convert a nemotron checkpoint in nemo (mcore path) into a HuggingFace checkpoint. +This script can be used to 1) generate only the HF weights, or 2) generate an entire HF model folder. + +1) Generate only HF weights from a nemo file: + + python convert_nemotron_nemo_to_hf.py \ + --input_name_or_path /path/to/file.nemo or /path/to/extracted_folder \ + --output_path /path/to/pytorch_model.bin + +2) Generate the full HF model folder + + python convert_nemotron_nemo_to_hf.py \ + --input_name_or_path /path/to/file.nemo or /path/to/extracted_folder \ + --hf_input_path /path/to/input_hf_folder \ + --hf_output_path /path/to/output_hf_folder \ + + Use the --cpu-only flag if the model cannot fit in the GPU (e.g. Nemotron4 340b). + However this option makes the conversion script significantly slower. +""" + + +def get_args(): + parser = ArgumentParser() + parser.add_argument( + "--input_name_or_path", + type=str, + default=None, + required=True, + help="Path to .nemo file or extracted folder", + ) + parser.add_argument("--output_path", type=str, default=None, required=False, help="Path to HF .bin file") + parser.add_argument( + "--hf_input_path", + type=str, + default=None, + help="A HF model path, e.g. a folder containing https://huggingface.co/nvidia/Minitron-8B-Base", + ) + parser.add_argument( + "--hf_output_path", + type=str, + default=None, + help="Output HF model path, with the same format as above but user's own weights", + ) + parser.add_argument( + "--precision", + type=str, + default=None, + help="Precision of output weights." + "Defaults to precision of the input nemo weights (model.cfg.trainer.precision)", + ) + parser.add_argument( + "--cpu-only", + action="store_true", + help="Load model in cpu only. Useful if the model cannot fit in GPU memory, " + "but this option makes the conversion script significantly slower.", + ) + args = parser.parse_args() + return args + + +def convert_hf_config(nemo_config, tokenizer, vocab_size, dtype, hf_output_path, hf_url="nvidia/Minitron-8B-Base"): + """ + Convert NeMo config to HF config + """ + NEMO_ACT2HF = { + "squared-relu": "relu2", + "fast-swiglu": "silu", + } + DTYPE2HF = { + torch.bfloat16: "bfloat16", + torch.float16: "float16", + torch.float32: "float32", + } + hf_config = { + "_name_or_path": hf_url, + "architectures": ["NemotronForCausalLM"], + "bos_token_id": tokenizer.bos_id, + "eos_token_id": tokenizer.eos_id, + "hidden_act": NEMO_ACT2HF[nemo_config.activation], + "hidden_size": nemo_config.hidden_size, + "initializer_range": nemo_config.init_method_std, + "intermediate_size": nemo_config.ffn_hidden_size, + "max_position_embeddings": nemo_config.max_position_embeddings, + "model_type": "nemotron", + "num_attention_heads": nemo_config.num_attention_heads, + "num_hidden_layers": nemo_config.num_layers, + "num_key_value_heads": nemo_config.get("num_query_groups", nemo_config.num_attention_heads), + "norm_eps": nemo_config.layernorm_epsilon, + "rope_theta": nemo_config.get("rotary_base", 10000), + "partial_rotary_factor": nemo_config.get("rotary_percentage", 1.0), + "tie_word_embeddings": False, + "dtype": DTYPE2HF[dtype], + "transformers_version": "4.32.0.dev0", # TODO + "use_cache": True, + "vocab_size": vocab_size, + } + if nemo_config.kv_channels is not None: + hf_config["kv_channels"] = nemo_config.kv_channels + json.dump(hf_config, open(f"{hf_output_path}/config.json", "w"), indent=2) + + +def convert(input_nemo_file, output_hf_file, precision=None, cpu_only=False) -> None: + """ + Convert NeMo weights to HF weights + """ + dummy_trainer = Trainer(devices=1, accelerator="cpu", strategy=NLPDDPStrategy()) + model_config = MegatronGPTModel.restore_from(input_nemo_file, trainer=dummy_trainer, return_config=True) + model_config.tensor_model_parallel_size = 1 + model_config.pipeline_model_parallel_size = 1 + model_config.sequence_parallel = False + model_config.transformer_engine = True + if cpu_only: + map_location = torch.device("cpu") + model_config.use_cpu_initialization = True + model_config.dist_ckpt_load_on_device = False + else: + map_location = None + + if cpu_only: + logging.info("******** Loading model on CPU. This will take a significant amount of time.") + + model = MegatronGPTModel.restore_from( + input_nemo_file, trainer=dummy_trainer, override_config_path=model_config, map_location=map_location + ) + + vocab_size = model.padded_vocab_size + + if precision is None: + precision = model.cfg.precision + if precision in [32, "32"]: + dtype = torch.float32 + elif precision in [16, "16", "16-mixed"]: + dtype = torch.float16 + elif precision in ["bf16", "bf16-mixed"]: + dtype = torch.bfloat16 + else: + logging.warning(f"Precision string {precision} is not recognized, falling back to fp32") + dtype = torch.float32 # fallback + logging.info(f"Using precision {dtype}") + + def param_to_weights(param): + return param.to(dtype) + + checkpoint = OrderedDict() + + hidden_size = model.cfg.hidden_size + head_num = model.cfg.num_attention_heads + num_layers = model.cfg.num_layers + ffn_hidden_size = model.cfg.ffn_hidden_size + num_query_groups = model.cfg.get("num_query_groups", head_num) # different num_query_groups for 70B + if num_query_groups is None: + num_query_groups = head_num + heads_per_group = head_num // num_query_groups + qkv_total_dim = head_num + 2 * num_query_groups + + # Embedding + embed_weight = model.state_dict()["model.embedding.word_embeddings.weight"] + embed_weights_base_name = "model.embed_tokens.weight" + checkpoint[embed_weights_base_name] = param_to_weights(embed_weight) + + for l in range(int(num_layers)): + print(f"converting layer {l}") + + qkv_weights = model.state_dict()[f"model.decoder.layers.{l}.self_attention.linear_qkv.weight"] + qkv_weights = qkv_weights.reshape([qkv_total_dim, -1, hidden_size]) + + q_slice = torch.cat( + [ + torch.arange((heads_per_group + 2) * i, (heads_per_group + 2) * i + heads_per_group) + for i in range(num_query_groups) + ] + ) + k_slice = torch.arange(heads_per_group, qkv_total_dim, (heads_per_group + 2)) + v_slice = torch.arange(heads_per_group + 1, qkv_total_dim, (heads_per_group + 2)) + ## Example of slices + ## (without GQA): num_query_groups = head_num = 32, + ## q_slice = [0, 3, 6, 9 , ... 90, 93] + ## k_slice = [1, 4, 7, 10, ... 91, 94] + ## v_slice = [2, 5, 8, 11, ... 92, 95] + ## (with GQA): num_query_groups = 8, head_num = 64 + ## q_slice = [0, 1, .. 6, 7, 10, 11, .. 16, 17, 20, 21, .. 67, 70, ... 76, 77] + ## k_slice = [8, 18, 28, ... 68, 78] + ## v_slice = [9, 19, 29, ... 69, 79] + + q_weights_base_name = f"model.layers.{l}.self_attn.q_proj.weight" + k_weights_base_name = f"model.layers.{l}.self_attn.k_proj.weight" + v_weights_base_name = f"model.layers.{l}.self_attn.v_proj.weight" + + checkpoint[q_weights_base_name] = param_to_weights(qkv_weights[q_slice].reshape(-1, hidden_size)) + checkpoint[k_weights_base_name] = param_to_weights(qkv_weights[k_slice].reshape(-1, hidden_size)) + checkpoint[v_weights_base_name] = param_to_weights(qkv_weights[v_slice].reshape(-1, hidden_size)) + + # attention dense + o_weight = model.state_dict()[f"model.decoder.layers.{l}.self_attention.linear_proj.weight"] + o_weight_base_name = f"model.layers.{l}.self_attn.o_proj.weight" + checkpoint[o_weight_base_name] = param_to_weights(o_weight) + + # mlp + mlp_weights = model.state_dict()[f"model.decoder.layers.{l}.mlp.linear_fc1.weight"] + mlp_up_proj_weight = model.state_dict()[f"model.decoder.layers.{l}.mlp.linear_fc2.weight"] + + if mlp_weights.shape[0] != mlp_up_proj_weight.shape[1]: + # Has projection (used for swi-glu) + logging.warning( + "Gated projection layers detected in NeMo checkpoint. Currently Nemotron HF does not support gated MLP." + ) + assert mlp_weights.shape[0] == 2 * mlp_up_proj_weight.shape[1] + + mlp_down_proj_weight = mlp_weights[:ffn_hidden_size, :] + mlp_gate_proj_weight = mlp_weights[ffn_hidden_size:, :] + + mlp_down_proj_base_name = f"model.layers.{l}.mlp.gate_proj.weight" + mlp_gate_proj_base_name = f"model.layers.{l}.mlp.up_proj.weight" + + checkpoint[mlp_down_proj_base_name] = param_to_weights(mlp_down_proj_weight) + checkpoint[mlp_gate_proj_base_name] = param_to_weights(mlp_gate_proj_weight) + else: + mlp_down_proj_weight = mlp_weights + mlp_down_proj_base_name = f"model.layers.{l}.mlp.up_proj.weight" + checkpoint[mlp_down_proj_base_name] = param_to_weights(mlp_down_proj_weight) + + mlp_up_proj_base_name = f"model.layers.{l}.mlp.down_proj.weight" + checkpoint[mlp_up_proj_base_name] = param_to_weights(mlp_up_proj_weight) + + # layernorm + input_ln_weight = model.state_dict()[f"model.decoder.layers.{l}.self_attention.linear_qkv.layer_norm_weight"] + input_ln_base_name = f"model.layers.{l}.input_layernorm.weight" + checkpoint[input_ln_base_name] = param_to_weights(input_ln_weight) + if ( + model.state_dict().get(f"model.decoder.layers.{l}.self_attention.linear_qkv.layer_norm_bias", None) + is not None + ): + input_ln_bias = model.state_dict()[f"model.decoder.layers.{l}.self_attention.linear_qkv.layer_norm_bias"] + input_ln_bias_name = f"model.layers.{l}.input_layernorm.bias" + checkpoint[input_ln_bias_name] = param_to_weights(input_ln_bias) + + post_attn_ln_weight = model.state_dict()[f"model.decoder.layers.{l}.mlp.linear_fc1.layer_norm_weight"] + post_attn_ln_base_name = f"model.layers.{l}.post_attention_layernorm.weight" + checkpoint[post_attn_ln_base_name] = param_to_weights(post_attn_ln_weight) + if model.state_dict().get(f"model.decoder.layers.{l}.mlp.linear_fc1.layer_norm_bias", None) is not None: + post_attn_ln_bias = model.state_dict()[f"model.decoder.layers.{l}.mlp.linear_fc1.layer_norm_bias"] + post_attn_ln_bias_name = f"model.layers.{l}.post_attention_layernorm.bias" + checkpoint[post_attn_ln_bias_name] = param_to_weights(post_attn_ln_bias) + + print(f"done layer {l}") + + final_ln_weight = model.state_dict()["model.decoder.final_layernorm.weight"] + final_ln_base_name = "model.norm.weight" + checkpoint[final_ln_base_name] = param_to_weights(final_ln_weight) + if model.state_dict().get("model.decoder.final_layernorm.bias", None) is not None: + final_ln_bias = model.state_dict()["model.decoder.final_layernorm.bias"] + final_ln_bias_name = "model.norm.bias" + checkpoint[final_ln_bias_name] = param_to_weights(final_ln_bias) + + output_layer_weight = model.state_dict()["model.output_layer.weight"] + output_layer_base_name = "lm_head.weight" + checkpoint[output_layer_base_name] = param_to_weights(output_layer_weight) + + os.makedirs(os.path.dirname(output_hf_file), exist_ok=True) + torch.save(checkpoint, output_hf_file) + logging.info(f"Weights saved to {output_hf_file}") + + return model_config, model.tokenizer, dtype, vocab_size + + +def extract_nemotron_tokenizer(nemo_file, model_config, output_hf_path, nemo_tokenizer): + tokenizer_cfg = model_config.tokenizer + if tokenizer_cfg.library == "sentencepiece": + # For sentencepiece tokenizer, we are wrapping with HF's LlamaTokenizer + # and convert it to a PreTrainedTokenizerFast + tokenizer_fn = tokenizer_cfg.model[5:] + output_tokenizer = f"{output_hf_path}/tokenizer.model" + if nemo_file.endswith(".nemo"): + import tarfile + + archive = tarfile.open(nemo_file, "r") + tokenizer_filename = "./" + tokenizer_fn # exclude 'nemo:' prefix + archive.extract(tokenizer_filename, output_hf_path) + archive.close() + os.rename(f"{output_hf_path}/{tokenizer_fn}", output_tokenizer) + elif os.path.isdir(nemo_file): + shutil.copy(f"{nemo_file}/{tokenizer_fn}", output_tokenizer) + # We use LlamaTokenizer for sentencepiece based tokenizer + tokenizer = LlamaTokenizer.from_pretrained(output_hf_path, legacy=False) + # Convert the LlamaTokenizer to a PreTrainedTokenizerFast instance + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=LlamaConverter(tokenizer).converted(), model_input_names=["input_ids", "token_type_ids"] + ) + tokenizer.save_pretrained(output_hf_path) + logging.info(f"Setencepiece tokenizer has been saved to {output_tokenizer}") + elif isinstance(nemo_tokenizer, AutoTokenizer): + nemo_tokenizer.tokenizer.save_pretrained(output_hf_path) + logging.info(f"HF AutoTokenizer has been saved to {output_hf_path}") + else: + raise ValueError(f"Unsupported tokenizer type: library: {tokenizer_cfg.library}, type: {tokenizer_cfg.type}") + + +if __name__ == "__main__": + args = get_args() + if not args.hf_output_path: + assert args.output_path is not None, "Need to provide either output_path or hf_output_path" + else: + args.output_path = f"{args.hf_output_path}/pytorch_model.bin" + logging.info(f"weight will be saved to {args.output_path}") + + nemo_config, nemo_tokenizer, dtype, vocab_size = convert( + args.input_name_or_path, args.output_path, precision=args.precision, cpu_only=args.cpu_only + ) + if args.hf_input_path and args.hf_output_path: + convert_hf_config(nemo_config, nemo_tokenizer, vocab_size, dtype, args.hf_output_path, args.hf_input_path) + extract_nemotron_tokenizer(args.input_name_or_path, nemo_config, args.hf_output_path, nemo_tokenizer) + else: + logging.info("`hf_input_path` and/or `hf_output_path` not provided, not generating full HF model.") + logging.info(f".bin file is saved to {args.output_path}") diff --git a/third_party/transformers/src/transformers/models/nemotron/modeling_nemotron.py b/third_party/transformers/src/transformers/models/nemotron/modeling_nemotron.py new file mode 100644 index 0000000000000000000000000000000000000000..b91b45ffd18364836de3b1cf2a16bbe3511ec82d --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron/modeling_nemotron.py @@ -0,0 +1,731 @@ +# Copyright 2024 HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024, 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 Nemotron model.""" + +import math +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Size, Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, StaticCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask +from ...modeling_layers import ( + GenericForQuestionAnswering, + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from ...modeling_rope_utils import ( + ROPE_INIT_FUNCTIONS, + dynamic_rope_update, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_nemotron import NemotronConfig + + +logger = logging.get_logger(__name__) + + +def _cast_if_autocast_enabled(device_type, *args): + if not torch.is_autocast_enabled(): + return args + else: + target_dtype = torch.get_autocast_dtype(device_type) + return torch.amp.autocast_mode._cast(args, device_type, target_dtype) + + +class NemotronLayerNorm1P(nn.LayerNorm): + def __init__( + self, + normalized_shape: int | list[int] | Size, + eps: float = 1e-5, + elementwise_affine: bool = True, + bias: bool = True, + device=None, + dtype=None, + ): + super().__init__(normalized_shape, eps, elementwise_affine, bias, device, dtype) + + def forward(self, input: Tensor) -> Tensor: + device_type = input.device.type if input.device.type != "mps" else "cpu" + args = _cast_if_autocast_enabled( + device_type, input, self.normalized_shape, self.weight + 1, self.bias, self.eps + ) + with maybe_autocast(device_type=input.device.type, enabled=False): + return F.layer_norm(*args) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron +class NemotronRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: NemotronConfig, 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 + # Ignore copy + def compute_default_rope_parameters( + config: NemotronConfig | 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) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +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) + + rot_dim = cos.shape[-1] + # If q_pass/k_pass is empty, rotary pos embedding is applied to all tensor q/k + q, q_pass = q[..., :rot_dim], q[..., rot_dim:] + k, k_pass = k[..., :rot_dim], k[..., rot_dim:] + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return torch.cat((q_embed, q_pass), dim=-1), torch.cat((k_embed, k_pass), dim=-1) + + +class NemotronMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = 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.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class NemotronAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: NemotronConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + + self.partial_rotary_factor = config.rope_parameters["partial_rotary_factor"] + self.is_causal = True + self.rotary_emb = NemotronRotaryEmbedding(config=config) + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.head_dim * self.num_heads, self.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 = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + 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) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +# NO LONGER EXIST Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron +# TODO cyril: modular +class NemotronFlashAttention2(NemotronAttention): + """ + Nemotron flash attention module. This module inherits from `NemotronAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask() + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if isinstance(past_key_values, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + 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) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (NemotronRMSNorm handles it correctly) + + input_dtype = query_states.dtype + device_type = query_states.device.type if query_states.device.type != "mps" else "cpu" + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_dtype(device_type) + # Handle the case where the model is quantized + elif hasattr(self.config, "_is_quantized"): + target_dtype = self.config.dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, None + + +# NO LONGER EXIST Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron +# TODO cyril: modular +class NemotronSdpaAttention(NemotronAttention): + """ + Nemotron attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `NemotronAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if output_attentions: + logger.warning_once( + f"{self.__class__.__name__} does not support `output_attentions=True`. The returned attention weights will " + "be `None`. If you want to get attention weights, please set `attn_implementation='eager'` when loading the model." + ) + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + 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) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: + causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + is_causal = causal_mask is None and q_len > 1 + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, None + + +NEMOTRON_ATTENTION_CLASSES = { + "eager": NemotronAttention, + "flash_attention_2": NemotronFlashAttention2, + "sdpa": NemotronSdpaAttention, +} + + +# copied from transformers.models.llama.modeling_llama.LlamaDecoderLayer with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron +# no longer copied after attention refactors +class NemotronDecoderLayer(GradientCheckpointingLayer): + # Ignore copy + def __init__(self, config: NemotronConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = NEMOTRON_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + + self.mlp = NemotronMLP(config) + self.input_layernorm = NemotronLayerNorm1P(config.hidden_size, eps=config.norm_eps) + self.post_attention_layernorm = NemotronLayerNorm1P(config.hidden_size, eps=config.norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring +class NemotronPreTrainedModel(PreTrainedModel): + config: NemotronConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["NemotronDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _can_record_outputs = { + "hidden_states": NemotronDecoderLayer, + "attentions": NemotronAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, NemotronLayerNorm1P): + init.ones_(module.weight) + init.zeros_(module.bias) + + +@auto_docstring +class NemotronModel(NemotronPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`NemotronDecoderLayer`] + + Args: + config: NemotronConfig + """ + + def __init__(self, config: NemotronConfig): + 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( + [NemotronDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = NemotronLayerNorm1P(config.hidden_size, eps=config.norm_eps) + self.rotary_emb = NemotronRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + 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: + 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 BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +# TODO: re-enable check: Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron +class NemotronForCausalLM(NemotronPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = NemotronModel(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""" + 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, NemotronForCausalLM + + >>> model = NemotronForCausalLM.from_pretrained("thhaus/nemotron3-8b") + >>> tokenizer = AutoTokenizer.from_pretrained("thhaus/nemotron3-8b") + + >>> 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, :]) + + 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, + ) + + +class NemotronForSequenceClassification(GenericForSequenceClassification, NemotronPreTrainedModel): ... + + +class NemotronForQuestionAnswering(GenericForQuestionAnswering, NemotronPreTrainedModel): + base_model_prefix = "transformer" + + +class NemotronForTokenClassification(GenericForTokenClassification, NemotronPreTrainedModel): ... + + +__all__ = [ + "NemotronForQuestionAnswering", + "NemotronForCausalLM", + "NemotronModel", + "NemotronPreTrainedModel", + "NemotronForSequenceClassification", + "NemotronForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/nomic_bert/__init__.py b/third_party/transformers/src/transformers/models/nomic_bert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70984a23cf9de8eab531cf017870301061d98938 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nomic_bert/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_nomic_bert import * + from .modeling_nomic_bert import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/nomic_bert/configuration_nomic_bert.py b/third_party/transformers/src/transformers/models/nomic_bert/configuration_nomic_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..8e13aa47e3f1d6073885acc3c95bc6a6354c8cec --- /dev/null +++ b/third_party/transformers/src/transformers/models/nomic_bert/configuration_nomic_bert.py @@ -0,0 +1,77 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/nomic_bert/modular_nomic_bert.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_nomic_bert.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="nomic-ai/nomic-embed-text-v1.5") +@strict +class NomicBertConfig(PreTrainedConfig): + r""" + Examples: + + ```python + >>> from transformers import NomicBertConfig, NomicBertModel + + >>> # Initializing a Nomic BERT nomic-ai/nomic-embed-text-v1.5 style configuration + >>> configuration = NomicBertConfig() + + >>> # Initializing a model (with random weights) from the nomic-ai/nomic-embed-text-v1.5 style configuration + >>> model = NomicBertModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "nomic_bert" + + vocab_size: int = 30528 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "silu" + hidden_dropout_prob: float = 0.0 + attention_probs_dropout_prob: float = 0.0 + max_position_embeddings: int = 2048 + type_vocab_size: int = 2 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + pad_token_id: int = 0 + classifier_dropout: float | None = None + bos_token_id: int | None = None + eos_token_id: int | None = None + tie_word_embeddings = True + default_theta = 1000.0 + rope_parameters: RopeParameters | dict | None = None + head_dim: int | None = None + + def __post_init__(self, **kwargs): + super().__post_init__(**kwargs) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + + +__all__ = ["NomicBertConfig"] diff --git a/third_party/transformers/src/transformers/models/nomic_bert/modeling_nomic_bert.py b/third_party/transformers/src/transformers/models/nomic_bert/modeling_nomic_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..f2836a1ec0f6e65d592ce96b2c28ea1a04e9c5bf --- /dev/null +++ b/third_party/transformers/src/transformers/models/nomic_bert/modeling_nomic_bert.py @@ -0,0 +1,721 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/nomic_bert/modular_nomic_bert.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_nomic_bert.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn as nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPooling, + MaskedLMOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_nomic_bert import NomicBertConfig + + +class NomicBertEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config: NomicBertConfig): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + self.register_buffer( + "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + embeddings = inputs_embeds + if inputs_embeds is None: + embeddings = self.word_embeddings(input_ids) + + input_shape = embeddings.shape[:-1] + device = embeddings.device + + if token_type_ids is None: + if hasattr(self, "token_type_ids"): + # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0]) + buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1) + buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids) + token_type_ids = buffered_token_type_ids.expand(*input_shape) + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = embeddings + token_type_embeddings + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + + return embeddings + + +class NomicBertRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: NomicBertConfig, 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: NomicBertConfig | 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 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 + + +@use_kernelized_func(apply_rotary_pos_emb) +class NomicBertAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_probs_dropout_prob + self.is_causal = False + 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_attention_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_attention_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, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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 NomicBertMLP(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 NomicBertLayer(GradientCheckpointingLayer): + def __init__(self, config: NomicBertConfig): + super().__init__() + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + self.post_attention_dropout = nn.Dropout(config.hidden_dropout_prob) + self.post_mlp_dropout = nn.Dropout(config.hidden_dropout_prob) + self.mlp = NomicBertMLP(config) + self.self_attn = NomicBertAttention(config=config) + self.post_mlp_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + attention_output, _ = self.self_attn( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + self.post_attention_dropout(attention_output) + hidden_states = self.post_attention_layernorm(hidden_states) + + residual = hidden_states + hidden_states = self.mlp(hidden_states) + hidden_states = residual + self.post_mlp_dropout(hidden_states) + hidden_states = self.post_mlp_layernorm(hidden_states) + return hidden_states + + +class NomicBertLMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = NomicBertPredictionHeadTransform(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 + + +@auto_docstring +class NomicBertPreTrainedModel(PreTrainedModel): + config_class = NomicBertConfig + base_model_prefix = "nomic_bert" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": NomicBertLayer, + "attentions": NomicBertAttention, + } + + # Are kept as non-persistent buffers to avoid being saved in the state dict + # and causing mismatch when loading from a checkpoint that doesn't have them + _keys_to_ignore_on_load_unexpected = ["inv_freq", "original_inv_freq"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, NomicBertLMPredictionHead): + init.zeros_(module.bias) + elif isinstance(module, NomicBertEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + init.zeros_(module.token_type_ids) + + +class NomicBertPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +@auto_docstring +class NomicBertModel(NomicBertPreTrainedModel): + _no_split_modules = ["NomicBertEmbeddings", "NomicBertLayer"] + + def __init__(self, config, add_pooling_layer=False): + r""" + add_pooling_layer (`bool`, *optional*, defaults to `False`): + Whether to add a pooling layer. + """ + super().__init__(config) + self.config = config + self.gradient_checkpointing = False + + self.embeddings = NomicBertEmbeddings(config) + + self.pooler = NomicBertPooler(config) if add_pooling_layer else None + self.rotary_emb = NomicBertRotaryEmbedding(config) + self.layers = nn.ModuleList([NomicBertLayer(config) for _ in range(config.num_hidden_layers)]) + + # 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, + attention_mask: torch.Tensor | 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[torch.Tensor] | BaseModelOutputWithPooling: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if input_ids is not None: + seq_length = input_ids.shape[1] + device = input_ids.device + else: + seq_length = inputs_embeds.shape[1] + device = inputs_embeds.device + + if position_ids is None: + position_ids = torch.arange(seq_length, dtype=torch.long, device=device)[None, :] + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + ) + + hidden_states = embedding_output + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for encoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + **kwargs, + ) + + sequence_output = hidden_states + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=pooled_output, + ) + + +class NomicBertPredictionHeadTransform(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 + # Use layer_norm rather than LayerNorm to avoid bert legacy mappings weights and bias to gamma and beta + self.layer_norm = 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.layer_norm(hidden_states) + return hidden_states + + +class NomicBertOnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = NomicBertLMPredictionHead(config) + + def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +@auto_docstring +class NomicBertForMaskedLM(NomicBertPreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.weight": "nomic_bert.embeddings.word_embeddings.weight", + "cls.predictions.decoder.bias": "cls.predictions.bias", + } + + def __init__(self, config): + super().__init__(config) + + self.nomic_bert = NomicBertModel(config) + self.cls = NomicBertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | MaskedLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + outputs = self.nomic_bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_dict=True, + **kwargs, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + NomicBert 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 NomicBertForSequenceClassification(NomicBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + self.nomic_bert = NomicBertModel(config, add_pooling_layer=True) + classifier_dropout = ( + config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | SequenceClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + outputs = self.nomic_bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_dict=True, + **kwargs, + ) + + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class NomicBertForTokenClassification(NomicBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.nomic_bert = NomicBertModel(config, add_pooling_layer=False) + classifier_dropout = ( + config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + outputs = self.nomic_bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_dict=True, + **kwargs, + ) + + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "NomicBertPreTrainedModel", + "NomicBertModel", + "NomicBertForMaskedLM", + "NomicBertForSequenceClassification", + "NomicBertForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/nomic_bert/modular_nomic_bert.py b/third_party/transformers/src/transformers/models/nomic_bert/modular_nomic_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..bd27e8034b626faf55d299919430aa489baa2086 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nomic_bert/modular_nomic_bert.py @@ -0,0 +1,299 @@ +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torch.nn as nn +from huggingface_hub.dataclasses import strict +from torch.nn import CrossEntropyLoss + +from ...configuration_utils import PreTrainedConfig +from ...masking_utils import create_bidirectional_mask +from ...modeling_outputs import ( + BaseModelOutputWithPooling, + MaskedLMOutput, +) +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 can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..bert.configuration_bert import BertConfig +from ..bert.modeling_bert import ( + BertForMaskedLM, + BertForSequenceClassification, + BertForTokenClassification, + BertOnlyMLMHead, + BertPredictionHeadTransform, + BertPreTrainedModel, +) +from ..gemma.modeling_gemma import GemmaMLP +from ..jina_embeddings_v3.modeling_jina_embeddings_v3 import ( + JinaEmbeddingsV3Attention, + JinaEmbeddingsV3Embeddings, + JinaEmbeddingsV3Layer, + JinaEmbeddingsV3Model, +) +from ..llama.modeling_llama import LlamaRotaryEmbedding + + +@auto_docstring(checkpoint="nomic-ai/nomic-embed-text-v1.5") +@strict +class NomicBertConfig(BertConfig): + r""" + Examples: + + ```python + >>> from transformers import NomicBertConfig, NomicBertModel + + >>> # Initializing a Nomic BERT nomic-ai/nomic-embed-text-v1.5 style configuration + >>> configuration = NomicBertConfig() + + >>> # Initializing a model (with random weights) from the nomic-ai/nomic-embed-text-v1.5 style configuration + >>> model = NomicBertModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "nomic_bert" + default_theta = 1000.0 + + vocab_size: int = 30528 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "silu" + hidden_dropout_prob: float = 0.0 + attention_probs_dropout_prob: float = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + classifier_dropout: float | None = None + type_vocab_size: int = 2 + bos_token_id: int | None = None + eos_token_id: int | None = None + pad_token_id: int = 0 + tie_word_embeddings = True + rope_parameters: RopeParameters | dict | None = None + max_position_embeddings: int = 2048 + head_dim: int | None = None + is_decoder = AttributeError() + add_cross_attention = AttributeError() + use_cache = AttributeError() + + def __post_init__(self, **kwargs): + PreTrainedConfig.__post_init__(self, **kwargs) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + + +class NomicBertEmbeddings(JinaEmbeddingsV3Embeddings): + pass + + +class NomicBertRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class NomicBertAttention(JinaEmbeddingsV3Attention): + def __init__(self, config): + super().__init__(config) + 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_attention_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + + +class NomicBertMLP(GemmaMLP): + pass + + +class NomicBertLayer(JinaEmbeddingsV3Layer): + pass + + +class NomicBertPreTrainedModel(BertPreTrainedModel): + config_class = NomicBertConfig + base_model_prefix = "nomic_bert" + + # Are kept as non-persistent buffers to avoid being saved in the state dict + # and causing mismatch when loading from a checkpoint that doesn't have them + _keys_to_ignore_on_load_unexpected = ["inv_freq", "original_inv_freq"] + _can_record_outputs = { + "hidden_states": NomicBertLayer, + "attentions": NomicBertAttention, + } + + +@auto_docstring +class NomicBertModel(JinaEmbeddingsV3Model): + def __init__(self, config, add_pooling_layer=False): + r""" + add_pooling_layer (`bool`, *optional*, defaults to `False`): + Whether to add a pooling layer. + """ + super().__init__(config, add_pooling_layer=add_pooling_layer) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | 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[torch.Tensor] | BaseModelOutputWithPooling: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if input_ids is not None: + seq_length = input_ids.shape[1] + device = input_ids.device + else: + seq_length = inputs_embeds.shape[1] + device = inputs_embeds.device + + if position_ids is None: + position_ids = torch.arange(seq_length, dtype=torch.long, device=device)[None, :] + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + ) + + hidden_states = embedding_output + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for encoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + **kwargs, + ) + + sequence_output = hidden_states + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=pooled_output, + ) + + +class NomicBertPredictionHeadTransform(BertPredictionHeadTransform): + def __init__(self, config): + super().__init__(config) + # Use layer_norm rather than LayerNorm to avoid bert legacy mappings weights and bias to gamma and beta + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + del self.LayerNorm + + 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.layer_norm(hidden_states) + return hidden_states + + +class NomicBertOnlyMLMHead(BertOnlyMLMHead): + pass + + +@auto_docstring +class NomicBertForMaskedLM(BertForMaskedLM): + def __init__(self, config): + PreTrainedModel.__init__(self, config) + + self.nomic_bert = NomicBertModel(config) + self.cls = NomicBertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | MaskedLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + outputs = self.nomic_bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_dict=True, + **kwargs, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class NomicBertForSequenceClassification(BertForSequenceClassification): + def __init__(self, config): + super().__init__(config) + self.nomic_bert = NomicBertModel(config, add_pooling_layer=True) + + +class NomicBertForTokenClassification(BertForTokenClassification): + pass + + +__all__ = [ + "NomicBertConfig", + "NomicBertPreTrainedModel", + "NomicBertModel", + "NomicBertForMaskedLM", + "NomicBertForSequenceClassification", + "NomicBertForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/nougat/__init__.py b/third_party/transformers/src/transformers/models/nougat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9a7fc4fbdcf203f27d6a9e58c82af2bd41a742d4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/__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 .image_processing_nougat import * + from .image_processing_pil_nougat import * + from .processing_nougat import * + from .tokenization_nougat 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/nougat/convert_nougat_to_hf.py b/third_party/transformers/src/transformers/models/nougat/convert_nougat_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..4feb7a9eb154137333f92e0df378bc640a3b6e3f --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/convert_nougat_to_hf.py @@ -0,0 +1,281 @@ +# 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 Nougat checkpoints using the original `nougat` library. URL: +https://github.com/facebookresearch/nougat/tree/main""" + +import argparse + +import torch +from huggingface_hub import hf_hub_download +from nougat import NougatModel +from nougat.dataset.rasterize import rasterize_paper +from nougat.utils.checkpoint import get_checkpoint +from PIL import Image + +from transformers import ( + DonutSwinConfig, + DonutSwinModel, + MBartConfig, + MBartForCausalLM, + NougatImageProcessor, + NougatProcessor, + NougatTokenizerFast, + VisionEncoderDecoderModel, +) + + +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, + tie_word_embeddings=False, + ) + + return encoder_config, decoder_config + + +# Copied from transformers.models.donut.convert_donut_to_pytorch.rename_key +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 + + +# Copied from transformers.models.donut.convert_donut_to_pytorch.convert_state_dict +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_nougat_checkpoint(model_tag, pytorch_dump_folder_path=None, push_to_hub=False): + # load original model + checkpoint_path = get_checkpoint(None, model_tag) + original_model = NougatModel.from_pretrained(checkpoint_path) + original_model.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 PDF + filepath = hf_hub_download(repo_id="ysharma/nougat", filename="input/nougat.pdf", repo_type="space") + images = rasterize_paper(pdf=filepath, return_pil=True) + image = Image.open(images[0]) + + tokenizer_file = checkpoint_path / "tokenizer.json" + tokenizer = NougatTokenizerFast(tokenizer_file=str(tokenizer_file)) + tokenizer.pad_token = "" + tokenizer.bos_token = "" + tokenizer.eos_token = "" + tokenizer.unk_token = "" + tokenizer.model_max_length = original_model.config.max_length + + size = {"height": original_model.config.input_size[0], "width": original_model.config.input_size[1]} + image_processor = NougatImageProcessor( + do_align_long_axis=original_model.config.align_long_axis, + size=size, + ) + processor = NougatProcessor(image_processor=image_processor, tokenizer=tokenizer) + + # verify pixel_values + pixel_values = processor(image, return_tensors="pt").pixel_values + original_pixel_values = original_model.encoder.prepare_input(image).unsqueeze(0) + + assert torch.allclose(original_pixel_values, pixel_values) + + # verify patch embeddings + 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) + + # 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) + + # NOTE original model does not use tied weights for embeddings of decoder + original_embeddings = original_model.decoder.model.model.decoder.embed_tokens + embeddings = model.decoder.model.decoder.embed_tokens + assert torch.allclose(original_embeddings.weight, embeddings.weight, atol=1e-3) + + # verify decoder hidden states + prompt = "hello world" + decoder_input_ids = original_model.decoder.tokenizer( + prompt, add_special_tokens=False, return_tensors="pt" + ).input_ids + decoder_attention_mask = torch.ones_like(decoder_input_ids) + original_logits = original_model( + image_tensors=pixel_values, decoder_input_ids=decoder_input_ids, attention_mask=decoder_attention_mask + ).logits + logits = model( + pixel_values, + decoder_input_ids=decoder_input_ids[:, :-1], + decoder_attention_mask=decoder_attention_mask[:, :-1], + ).logits + assert torch.allclose(original_logits, logits, atol=1e-3) + + # verify generation + outputs = model.generate( + pixel_values, + min_length=1, + max_length=30, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + use_cache=True, + bad_words_ids=[ + [tokenizer.unk_token_id], + ], + return_dict_in_generate=True, + do_sample=False, + ) + generated = tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)[0] + + if model_tag == "0.1.0-base": + expected_generation = "# Nougat: Neural Optical Understanding for Academic Documents\n\nLukas Blecher\n\nCorrespondence to: lblec" + elif model_tag == "0.1.0-small": + expected_generation = ( + "# Nougat: Neural Optical Understanding for Academic Documents\n\nLukas Blecher\n\nCorrespondence to: lble" + ) + else: + raise ValueError(f"Unexpected model tag: {model_tag}") + + assert generated == expected_generation + 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: + tag_to_name = {"0.1.0-base": "nougat-base", "0.1.0-small": "nougat-small"} + model_name = tag_to_name[model_tag] + + model.push_to_hub(f"facebook/{model_name}") + processor.push_to_hub(f"facebook/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_tag", + default="0.1.0-base", + required=False, + type=str, + choices=["0.1.0-base", "0.1.0-small"], + help="Tag 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_nougat_checkpoint(args.model_tag, args.pytorch_dump_folder_path, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/nougat/image_processing_nougat.py b/third_party/transformers/src/transformers/models/nougat/image_processing_nougat.py new file mode 100644 index 0000000000000000000000000000000000000000..8e738dbf18380b24ca7d65dfe6556e79704be38d --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/image_processing_nougat.py @@ -0,0 +1,304 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for Nougat.""" + +import torch +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import ( + get_resize_output_image_size, + group_images_by_shape, + reorder_images, +) +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, +) + + +class NougatImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_crop_margin (`bool`, *optional*, defaults to `self.do_crop_margin`): + Whether to crop the image margins. + 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_crop_margin: bool + do_thumbnail: bool + do_align_long_axis: bool + + +@auto_docstring +class NougatImageProcessor(TorchvisionBackend): + valid_kwargs = NougatImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"height": 896, "width": 672} + do_resize = True + do_normalize = True + do_thumbnail = True + do_align_long_axis = False + do_pad = True + do_rescale = True + do_crop_margin = True + + def __init__(self, **kwargs: Unpack[NougatImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[NougatImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def python_find_non_zero( + self, + image: "torch.Tensor", + ): + """This is a reimplementation of a findNonZero function equivalent to cv2.""" + + non_zero_indices = torch.nonzero(image, as_tuple=False) + idxvec = non_zero_indices[:, [2, 1]] + idxvec = idxvec.reshape(-1, 1, 2) + return idxvec + + def python_bounding_rect(self, coordinates): + """This is a reimplementation of a BoundingRect function equivalent to cv2.""" + + min_values = torch.amin(coordinates, axis=(0, 1)).to(torch.int) + max_values = torch.amax(coordinates, axis=(0, 1)).to(torch.int) + + x_min, y_min = min_values[0], min_values[1] + width = max_values[0] - x_min + 1 + height = max_values[1] - y_min + 1 + return x_min, y_min, width, height + + def crop_margin( + self, + image: "torch.Tensor", + gray_threshold: int = 200, + ) -> "torch.Tensor": + """ + Crops the margin of the image. Gray pixels are considered margin (i.e., pixels with a value below the + threshold). + + Args: + image (`torch.Tensor`): + The image to be cropped. + gray_threshold (`int`, *optional*, defaults to `200`) + Value below which pixels are considered to be gray. + """ + data = tvF.rgb_to_grayscale(image, num_output_channels=1) + + max_val = torch.max(data) + min_val = torch.min(data) + + if max_val == min_val: + return image + data = (data - min_val) / (max_val - min_val) * 255 + gray = data < gray_threshold + coords = self.python_find_non_zero(gray) + x_min, y_min, width, height = self.python_bounding_rect(coords) + image = image[:, y_min : y_min + height, x_min : x_min + width] + + return image + + 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. + + Args: + image (`torch.Tensor`): + The image to be aligned. + size (`SizeDict`): + The size to align the long axis to. + Returns: + `torch.Tensor`: The aligned image. + """ + 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 + ): + image = torch.rot90(image, 3, dims=[1, 2]) + + return image + + def thumbnail( + self, + image: "torch.Tensor", + size: SizeDict, + ) -> "torch.Tensor": + """ + Resize the image to make a thumbnail. The image is resized so that no dimension is larger than any + corresponding dimension of the specified size. + + Args: + image (`torch.tensor`): + The image to be resized. + size (`SizeDict`): + The size to resize the image to. + """ + + input_height, input_width = image.shape[-2:] + output_height, output_width = size.height, size.width + + # We always resize to the smallest of either the input or output size. + 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) + + new_size = (height, width) + + return tvF.resize(image, new_size, interpolation=tvF.InterpolationMode.BICUBIC) + + def pad_images( + self, + image: "torch.Tensor", + size: SizeDict, + ) -> "torch.Tensor": + """ + Pads a batch of images to the specified size at the top, bottom, left and right. + + Args: + image (`torch.tensor`): + The image to be padded. + size (`SizeDict`): + The size to pad the image to. + """ + input_height, input_width = image.shape[-2:] + output_height, output_width = size.height, size.width + + delta_width = output_width - input_width + delta_height = output_height - input_height + + 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 resize( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + antialias: bool = True, + **kwargs, + ) -> "torch.Tensor": + """ + Resize an image to `(size.height, size.width)`. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Size of the output image. + resample (`PILImageResampling | tvF.InterpolationMode | int`, *optional*): + Resampling filter to use when resizing the image. + Returns: + `torch.Tensor`: The resized image. + """ + shortest_edge = min(size.height, size.width) + + new_size = get_resize_output_image_size( + image, size=shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST + ) + return super().resize( + image, SizeDict(height=new_size[0], width=new_size[1]), resample=resample, antialias=antialias, **kwargs + ) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + do_align_long_axis: bool = False, + do_thumbnail: bool = True, + do_crop_margin: bool = True, + **kwargs, + ) -> BatchFeature: + # Crop images + if do_crop_margin: + images = [self.crop_margin(image) for image in images] + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_align_long_axis: + stacked_images = self.align_long_axis(image=stacked_images, size=size) + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, resample=resample) + if do_thumbnail: + stacked_images = self.thumbnail(image=stacked_images, size=size) + if do_pad: + stacked_images = self.pad_images(image=stacked_images, size=size) + 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(): + # 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) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + +__all__ = ["NougatImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/nougat/image_processing_pil_nougat.py b/third_party/transformers/src/transformers/models/nougat/image_processing_pil_nougat.py new file mode 100644 index 0000000000000000000000000000000000000000..a4c240229b32b4570746a90312595536ebdc8a8f --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/image_processing_pil_nougat.py @@ -0,0 +1,293 @@ +# 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. +"""Image processor class for Nougat.""" + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import ( + get_resize_output_image_size, + pad, + to_channel_dimension_format, + to_pil_image, +) +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, +) + + +# Adapted from transformers.models.nougat.image_processing_nougat.NougatImageProcessorKwargs +class NougatImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_crop_margin (`bool`, *optional*, defaults to `self.do_crop_margin`): + Whether to crop the image margins. + 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_crop_margin: bool + do_thumbnail: bool + do_align_long_axis: bool + + +@auto_docstring +class NougatImageProcessorPil(PilBackend): + valid_kwargs = NougatImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"height": 896, "width": 672} + do_resize = True + do_normalize = True + do_thumbnail = True + do_align_long_axis = False + do_pad = True + do_rescale = True + do_crop_margin = True + + def __init__(self, **kwargs: Unpack[NougatImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[NougatImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def python_find_non_zero(self, image: np.ndarray): + """This is a reimplementation of a findNonZero function equivalent to cv2.""" + non_zero_indices = np.column_stack(np.nonzero(image)) + idxvec = non_zero_indices[:, [1, 0]] + idxvec = idxvec.reshape(-1, 1, 2) + return idxvec + + def python_bounding_rect(self, coordinates): + """This is a reimplementation of a BoundingRect function equivalent to cv2.""" + min_values = np.min(coordinates, axis=(0, 1)).astype(int) + max_values = np.max(coordinates, axis=(0, 1)).astype(int) + x_min, y_min = min_values[0], min_values[1] + width = max_values[0] - x_min + 1 + height = max_values[1] - y_min + 1 + return x_min, y_min, width, height + + def crop_margin( + self, + image: np.ndarray, + gray_threshold: int = 200, + ) -> np.ndarray: + """ + Crops the margin of the image. Gray pixels are considered margin (i.e., pixels with a value below the + threshold). + + Args: + image (`np.ndarray`): + The image to be cropped. + gray_threshold (`int`, *optional*, defaults to `200`) + Value below which pixels are considered to be gray. + """ + image_pil = to_pil_image(image, input_data_format=ChannelDimension.FIRST) + data = np.array(image_pil.convert("L")).astype(np.uint8) + max_val = data.max() + min_val = data.min() + if max_val == min_val: + return image + data = (data - min_val) / (max_val - min_val) * 255 + gray = data < gray_threshold + coords = self.python_find_non_zero(gray) + x_min, y_min, width, height = self.python_bounding_rect(coords) + image_pil = image_pil.crop((x_min, y_min, x_min + width, y_min + height)) + image = np.array(image_pil).astype(np.uint8) + image = to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=ChannelDimension.LAST) + return image + + 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. + + Args: + image (`np.ndarray`): + The image to be aligned. + size (`SizeDict`): + The size to align the long axis to. + Returns: + `np.ndarray`: The aligned image. + """ + 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 thumbnail( + self, + image: np.ndarray, + size: SizeDict, + ) -> np.ndarray: + """ + Resize the image to make a thumbnail. The image is resized so that no dimension is larger than any + corresponding dimension of the specified size. + + Args: + image (`np.ndarray`): + The image to be resized. + size (`SizeDict`): + The size to resize the image to. + """ + + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = size.height, size.width + + # We always resize to the smallest of either the input or output size. + 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) + + # Use np_resize for exact dimensions; self.resize uses shortest-edge logic and would produce + # different output due to rounding in get_resize_output_image_size. + return super().resize( + image, SizeDict(height=height, width=width), resample=PILImageResampling.BICUBIC, reducing_gap=2.0 + ) + + def pad_images( + self, + image: np.ndarray, + size: SizeDict, + ) -> np.ndarray: + """ + Pads a batch of images to the specified size at the top, bottom, left and right. + + Args: + image (`np.ndarray`): + The image to be padded. + size (`SizeDict`): + The size to pad the image to. + """ + + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = size.height, size.width + + delta_width = output_width - input_width + delta_height = output_height - input_height + + pad_top = delta_height // 2 + pad_left = delta_width // 2 + + pad_bottom = delta_height - pad_top + pad_right = delta_width - pad_left + + padding = ((pad_top, pad_bottom), (pad_left, pad_right)) + return pad(image, padding, input_data_format=ChannelDimension.FIRST, data_format=ChannelDimension.FIRST) + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: "PILImageResampling | None" = None, + reducing_gap: int | None = None, + **kwargs, + ) -> np.ndarray: + """ + Resize an image to `(size.height, size.width)`. + + Args: + image (`np.ndarray`): + Image to resize. + size (`SizeDict`): + Size of the output image. + resample (`PILImageResampling | int`, *optional*): + Resampling filter to use when resizing the image. + Returns: + `np.ndarray`: The resized image. + """ + + shortest_edge = min(size.height, size.width) + + new_size = get_resize_output_image_size( + image, size=shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST + ) + return super().resize( + image, + SizeDict(height=new_size[0], width=new_size[1]), + resample=resample, + reducing_gap=reducing_gap, + **kwargs, + ) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + return_tensors: str | TensorType | None, + do_align_long_axis: bool = False, + do_thumbnail: bool = True, + do_crop_margin: bool = True, + do_pad: bool | None = None, + **kwargs, + ) -> BatchFeature: + processed_images = [] + for image in images: + if do_crop_margin: + image = self.crop_margin(image) + if do_align_long_axis: + image = self.align_long_axis(image, size) + if do_resize: + image = self.resize(image, size, resample) + if do_thumbnail: + image = self.thumbnail(image, size) + if do_pad: + image = self.pad_images(image, 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__ = ["NougatImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/nougat/processing_nougat.py b/third_party/transformers/src/transformers/models/nougat/processing_nougat.py new file mode 100644 index 0000000000000000000000000000000000000000..64b9d122e100a1f419b95482e07c3586d1d01b39 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/processing_nougat.py @@ -0,0 +1,142 @@ +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for Nougat. +""" + +from typing import Optional, Union + +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput, TruncationStrategy + +from ...processing_utils import ProcessorMixin +from ...utils import PaddingStrategy, TensorType, auto_docstring + + +@auto_docstring +class NougatProcessor(ProcessorMixin): + def __init__(self, image_processor, tokenizer): + super().__init__(image_processor, tokenizer) + + @auto_docstring + def __call__( + self, + images=None, + text=None, + do_crop_margin: bool | None = None, + do_resize: bool | None = None, + size: dict[str, int] | None = None, + resample: "PILImageResampling" = None, # noqa: F821 + do_thumbnail: bool | None = None, + do_align_long_axis: bool | None = None, + do_pad: bool | None = None, + do_rescale: bool | None = None, + rescale_factor: int | float | None = None, + do_normalize: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + data_format: Optional["ChannelDimension"] = "channels_first", # noqa: F821 + input_data_format: Union[str, "ChannelDimension"] | None = None, # noqa: F821 + text_pair: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + text_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + text_pair_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: int | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + ): + r""" + do_crop_margin (`bool`, *optional*): + Whether to automatically crop white margins from document images. When enabled, the processor detects + and removes white space around the edges of document pages, which is useful for processing scanned + documents or PDFs with large margins. + do_thumbnail (`bool`, *optional*): + Whether to create a thumbnail version of the image. When enabled, a smaller version of the image is + generated alongside the main processed image, which can be useful for preview or faster processing. + do_align_long_axis (`bool`, *optional*): + Whether to automatically align images so that the longer axis is horizontal. When enabled, portrait + images are rotated to landscape orientation, which is typically better for document processing tasks. + """ + if images is None and text is None: + raise ValueError("You need to specify either an `images` or `text` input to process.") + + if images is not None: + inputs = self.image_processor( + images, + do_crop_margin=do_crop_margin, + do_resize=do_resize, + size=size, + resample=resample, + do_thumbnail=do_thumbnail, + do_align_long_axis=do_align_long_axis, + do_pad=do_pad, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + return_tensors=return_tensors, + data_format=data_format, + input_data_format=input_data_format, + ) + if text is not None: + encodings = self.tokenizer( + text, + text_pair=text_pair, + text_target=text_target, + text_pair_target=text_pair_target, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + ) + + if text is None: + return inputs + elif images is None: + return encodings + else: + inputs["labels"] = encodings["input_ids"] + return inputs + + def post_process_generation(self, *args, **kwargs): + """ + This method forwards all its arguments to NougatTokenizer's [`~PreTrainedTokenizer.post_process_generation`]. + Please refer to the docstring of this method for more information. + """ + return self.tokenizer.post_process_generation(*args, **kwargs) + + +__all__ = ["NougatProcessor"] diff --git a/third_party/transformers/src/transformers/models/nougat/tokenization_nougat.py b/third_party/transformers/src/transformers/models/nougat/tokenization_nougat.py new file mode 100644 index 0000000000000000000000000000000000000000..f9ed3d3fc5fdc611fefdc29ccf776bb729447a0f --- /dev/null +++ b/third_party/transformers/src/transformers/models/nougat/tokenization_nougat.py @@ -0,0 +1,660 @@ +# 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. +""" +Tokenizer class for Nougat. +""" + +import re +from functools import partial +from multiprocessing import Pool + +import numpy as np +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import BPE + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import is_levenshtein_available, is_nltk_available, logging, requires_backends + + +if is_levenshtein_available(): + from Levenshtein import ratio + +if is_nltk_available(): + import nltk + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"} + + +def markdown_compatible(text: str) -> str: + """ + Make text compatible with Markdown formatting. + + This function makes various text formatting adjustments to make it compatible with Markdown. + + Args: + text (`str`): + The input text to be made Markdown-compatible. + + Returns: + `str`: The Markdown-compatible text. + """ + # equation tag + # Replace lines that start with a pattern like (decimal) \[some text\] with \[[some text] \tag{decimal}\]. + text = re.sub(r"^\(([\d.]+[a-zA-Z]?)\) \\\[(.+?)\\\]$", r"\[\2 \\tag{\1}\]", text, flags=re.MULTILINE) + # Replace lines that start with a pattern like \[some text\] (decimal) with \[[some text] \tag{decimal}\]. + text = re.sub(r"^\\\[(.+?)\\\] \(([\d.]+[a-zA-Z]?)\)$", r"\[\1 \\tag{\2}\]", text, flags=re.MULTILINE) + # Replace lines that start with a pattern like \[some text\] (digits) \[another text\] with \[[some text] \tag{digits}\] [another text]. + text = re.sub( + r"^\\\[(.+?)\\\] \(([\d.]+[a-zA-Z]?)\) (\\\[.+?\\\])$", + r"\[\1 \\tag{\2}\] \3", + text, + flags=re.MULTILINE, + ) + # multi line + text = text.replace(r"\. ", ". ") + # bold formatting + text = text.replace(r"\bm{", r"\mathbf{").replace(r"{\\bm ", r"\mathbf{") + text = re.sub(r"\\mbox{ ?\\boldmath\$(.*?)\$}", r"\\mathbf{\1}", text) + # Reformat urls (http, ftp and https only) to markdown [url](url) clickable format + text = re.sub( + r"((?:http|ftp|https):\/\/(?:[\w_-]+(?:(?:\.[\w_-]+)+))(?:[\w.,@?^=%&:\/~+#-]*[\w@?^=%&\/~+#-]))", + r"[\1](\1)", + text, + ) + # algorithms + text = re.sub(r"```\s*(.+?)\s*```", r"```\n\1\n```", text, flags=re.DOTALL) + + return text + + +def normalize_list_like_lines(generation): + """ + Normalize lines in the given text that resemble list items. The function looks for lines that start optionally with + '-' or '*', possibly followed by Roman numerals or digits indicating nesting levels. The function reformats such + lines to make them more structured. + + Args: + generation (str): The input text containing lines that need to be normalized. + + Returns: + str: The input text with the list-like lines normalized. + + Note: + The function uses regular expressions to identify and reformat the list-like lines. The patterns capture + optional bullet points, nesting levels indicated by numerals, and the actual list item content. The + normalization adjusts the bullet point style and nesting levels based on the captured patterns. + """ + + lines = generation.split("\n") + output_lines = [] + for line_no, line in enumerate(lines): + match = re.search(r". ([-*]) ", line) + if not match or line[0] not in ("-", "*"): + output_lines.append(line) + continue # Doesn't fit the pattern we want, no changes + delim = match.group(1) + " " + splits = line.split(delim)[1:] + replacement = "" + delim1 = line[0] + " " + + for i, item in enumerate(splits): + level = 0 + potential_numeral, _, rest = item.strip().partition(" ") + if not rest: + continue + # Infer current nesting level based on detected numbering + if re.match(r"^[\dixv]+((?:\.[\dixv])?)+$", potential_numeral, flags=re.IGNORECASE | re.MULTILINE): + level = potential_numeral.count(".") + + replacement += ( + ("\n" if i > 0 else "") + ("\t" * level) + (delim if i > 0 or line_no == 0 else delim1) + item.strip() + ) + + if line_no == len(lines) - 1: # If this is the last line in the generation + replacement += "\n" # Add an empty line to the end of the generation + + output_lines.append(replacement) + + return "\n".join(output_lines) + + +def find_next_punctuation(text: str, start_idx=0): + """ + Find the index of the next punctuation mark. + + Args: + text (`str`): + String to examine + start_idx (`int`, *optional*) + Index where to start + """ + + for i in range(start_idx, len(text)): + if text[i] in [".", "?", "!", "\n"]: + return i + + return None + + +def truncate_repetitions(text: str, min_len: int = 30) -> str: + """ + Attempt to truncate repeating segments in the input string. + + This function looks for the longest repeating substring at the end of the input string and truncates it to appear + only once. To be considered for removal, repetitions need to be continuous. + + Args: + text (`str`): + The input raw prediction to be truncated. + min_len (int): + The minimum length of the repeating segment. + + Returns: + `str`: The input string with repeated segments truncated. + """ + text_lower = text.lower() + text_length = len(text_lower) + + if text_length < 2 * min_len: + return text + + # try to find a length at which the tail is repeating + max_repetition_length = None + for repetition_length in range(min_len, int(text_length / 2)): + # check if there is a repetition at the end + same = True + for i in range(0, repetition_length): + if text_lower[text_length - repetition_length - i - 1] != text_lower[text_length - i - 1]: + same = False + break + + if same: + max_repetition_length = repetition_length + + if max_repetition_length is None: + return text + + lcs = text_lower[-max_repetition_length:] + + # remove all but the last repetition + substituted_text = text + substituted_text_lower = text_lower + while substituted_text_lower.endswith(lcs): + substituted_text = substituted_text[:-max_repetition_length] + substituted_text_lower = substituted_text_lower[:-max_repetition_length] + + # this is the tail with the repetitions + repeating_tail = text_lower[len(substituted_text_lower) :] + + # add until next punctuation and make sure last sentence is not repeating + substituted_text_lower_out = substituted_text_lower + while True: + sentence_end = find_next_punctuation(text_lower, len(substituted_text_lower_out)) + sentence_start = find_next_punctuation(text_lower[::-1], len(substituted_text_lower_out)) + if sentence_end and sentence_start: + sentence = text_lower[sentence_start:sentence_end] + substituted_text_lower_out = text_lower[: sentence_end + 1] + if sentence in repeating_tail: + break + else: + break + + text_out = text[: len(substituted_text_lower_out)] + + return text_out + + +def remove_numbers(lines): + def _clean(s): + return re.sub(r"(?:[\d_]|\*\*)", "", s).strip() + + if isinstance(lines, str): + return _clean(lines) + out = [] + for l in lines: + out.append(_clean(l)) + return out + + +def get_slices(lines, clean_lines): + """ + Get slices of text based on specific criteria within the lines. + + This function identifies and returns slices of text from the input lines based on certain conditions. + + These conditions were chosen by the Nougat authors: + - The slice is less than 200 characters long. + - The slice is more than 3 characters long. + - The slice does not start with "[MISSING_PAGE". + - The slice is either the same as the next slice or the ratio of the two in terms of Levenshtein distance is + greater than 0.9. + + Args: + lines (`list[str]`): + The list of lines containing the text. + clean_lines (`list[str]`): + A cleaned version of the text (without numbers). + + Returns: + `list[tuple]`: A list of tuples representing the start and end indices of text slices. + """ + indices = np.zeros(len(lines)) + for i in range(len(lines) - 1): + j = i + 1 + while not clean_lines[j] and j < len(lines) - 1: + j += 1 + if ( + len(clean_lines[i]) < 200 + and len(clean_lines[i]) > 3 + and len(clean_lines[j]) < 200 + and len(clean_lines[j]) > 3 + and not clean_lines[i].startswith("[MISSING_PAGE") + and (clean_lines[i] == clean_lines[j] or ratio(clean_lines[i], clean_lines[j]) > 0.9) + ): + indices[i:j] = 1 + ids = np.where(indices)[0] + slices = [] + if len(ids) == 0: + return slices + j0 = 0 + for j, x in enumerate(np.diff(ids) > 3): + if x: + slices.append((ids[j0], ids[j] + 2)) + j0 = j + 1 + slices.append((ids[j0], ids[-1] + 2)) + return [sli for sli in slices if sli[1] - sli[0] > 15] + + +def remove_slice_from_lines(lines, clean_text, slice) -> str: + """ + Remove a slice of text from the lines based on specific criteria. + + This function identifies a slice of text within the lines and removes it based on certain conditions. + + Args: + lines (list of str): The list of lines containing the text. + clean_text (list of str): A cleaned version of the text (without numbers). + slice (tuple): A tuple representing the start and end indices of the slice to be removed. + + Returns: + str: The removed slice of text as a single string. + """ + base = clean_text[slice[0]] + section = list(slice) + check_start_flag = False + # backwards pass, at most 5 lines + for line_idx in range(max(0, slice[0] - 1), max(0, slice[0] - 5), -1): + if not lines[line_idx]: + continue + if lines[line_idx] == "## References": + section[0] = line_idx + break + elif ratio(base, remove_numbers(lines[line_idx])) < 0.9: + section[0] = line_idx + 1 + potential_ref = remove_numbers(lines[max(0, line_idx - 1)].partition("* [")[-1]) + if len(potential_ref) >= 0.75 * len(base) and ratio(base, potential_ref) < 0.9: + section[0] = line_idx + check_start_flag = True + break + # forward pass, at most 5 lines + for line_idx in range(min(len(lines), slice[1]), min(len(lines), slice[1] + 5)): + if ratio(base, remove_numbers(lines[line_idx])) < 0.9: + section[1] = line_idx + break + if len(lines) <= section[1]: + section[1] = len(lines) - 1 + to_delete = "\n".join(lines[section[0] : section[1] + 1]) + # cut off next page content + itera, iterb = enumerate(lines[section[1] - 1]), enumerate(lines[section[1]]) + while True: + try: + (ia, a) = next(itera) + while a.isnumeric(): + (ia, a) = next(itera) + (ib, b) = next(iterb) + while b.isnumeric(): + (ib, b) = next(iterb) + if a != b: + break + except StopIteration: + break + if check_start_flag and "* [" in to_delete: + to_delete = "* [" + to_delete.partition("* [")[-1] + try: + delta = len(lines[section[1]]) - ib - 1 + if delta > 0: + to_delete = to_delete[:-delta] + except UnboundLocalError: + pass + + return to_delete.strip() + + +class NougatTokenizer(TokenizersBackend): + """ + Tokenizer for Nougat (backed by HuggingFace tokenizers library). + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. This class mainly adds Nougat-specific + methods for postprocessing the generated text. + + Args: + vocab_file (`str`, *optional*): + Path to the vocabulary file. + merges_file (`str`, *optional*): + Path to the merges file. + tokenizer_file (`str`, *optional*): + [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that + contains everything needed to load the tokenizer. + + clean_up_tokenization_spaces (`str`, *optional*, defaults to `False`): + Whether to cleanup spaces after decoding, cleanup consists in removing potential artifacts like extra + spaces. + + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + + vocab (`str`, `dict` or `list`, *optional*): + Custom vocabulary dictionary. If not provided, vocabulary is loaded from vocab_file. + + merges (`str` or `list`, *optional*): + Custom merges list. If not provided, merges are loaded from merges_file. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + errors: str = "replace", + unk_token: str = "", + bos_token: str = "", + eos_token: str = "", + pad_token: str = "", + vocab: str | dict | list | None = None, + merges: str | list | None = None, + **kwargs, + ): + self._vocab = ( + vocab + if vocab is not None + else { + str(bos_token): 0, + str(pad_token): 1, + str(eos_token): 2, + str(unk_token): 3, + "[START_REF]": 4, + } + ) + self._merges = merges or [] + self._tokenizer = Tokenizer( + BPE( + vocab=self._vocab, + merges=self._merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + self._tokenizer.normalizer = normalizers.NFKC() + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split(pattern="SPL1T-TH1S-Pl3A5E", behavior="removed", invert=False), + pre_tokenizers.Digits(individual_digits=True), + pre_tokenizers.Split( + pattern=r"[\(\)\[\]\{\}]|([!\"#\$%\&'\*\+,\-\./:;<=>\?\\\^_`\|\~])\1*", + behavior="isolated", + invert=False, + ), + pre_tokenizers.Split(pattern="\n", behavior="isolated", invert=False), + pre_tokenizers.ByteLevel(add_prefix_space=False, trim_offsets=True, use_regex=True), + ] + ) + self._tokenizer.decoder = decoders.ByteLevel(add_prefix_space=True, trim_offsets=True, use_regex=True) + + super().__init__( + errors=errors, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + **kwargs, + ) + self._tokenizer.post_processor = processors.TemplateProcessing( + single=f"{bos_token}:0 $A:0 {eos_token}:0", + pair="$A:0 $B:1", + special_tokens=[ + (str(eos_token), self.eos_token_id), + (str(bos_token), self.bos_token_id), + ], + ) + + # Enable truncation and padding + self._tokenizer.enable_truncation(max_length=4096) + self._tokenizer.enable_padding(length=4096, pad_id=self.pad_token_id, pad_token=str(pad_token)) + + def remove_hallucinated_references(self, text: str) -> str: + """ + Remove hallucinated or missing references from the text. + + This function identifies and removes references that are marked as missing or hallucinated from the input text. + + Args: + text (`str`): + The input text containing references. + + Returns: + `str`: The text with hallucinated references removed. + """ + lines = text.split("\n") + if len(lines) == 0: + return "" + clean_lines = remove_numbers(lines) + slices = get_slices(lines, clean_lines) + to_delete = [] + for slice in slices: + to_delete.append(remove_slice_from_lines(lines, clean_lines, slice)) + for to_delete in reversed(to_delete): + text = text.replace(to_delete, "\n\n[MISSING_PAGE_POST]\n\n") + text = re.sub( + r"## References\n+\[MISSING_PAGE_POST(:\d+)?\]", + "\n\n[MISSING_PAGE_POST\\1]", + text, + ) + return text + + def correct_tables(self, generation: str) -> str: + """ + Takes a generated string and fixes tables/tabulars to make them match the markdown format needed. + + Args: + generation (str): The generated text to be postprocessed. + + Returns: + str: The postprocessed text. + + Example: + + ```python + correct_tables("\\begin{table} \\begin{tabular}{l l} & \\ \\end{tabular} \\end{table}") + "\\begin{table}\n\\begin{tabular}{l l} & \\ \\end{tabular}\n\\end{table}" + ``` + """ + # remove obvious wrong tables + for l in generation.split("\n"): + if l.count("\\begin{tabular}") > 15 or l.count("\\multicolumn") > 60 or l.count("&") > 400: + generation = generation.replace(l, "") + # whitespace corrections + + generation = generation.replace("\\begin{table} \\begin{tabular}", "\\begin{table}\n\\begin{tabular}") + generation = generation.replace("\\end{tabular} \\end{table}", "\\end{tabular}\n\\end{table}") + generation = generation.replace("\\end{table} Tab", "\\end{table}\nTab") + + generation = re.sub(r"(^.+)\\begin{tab", r"\1\n\\begin{tab", generation, flags=re.MULTILINE) + + # Remove left-aligned empty LaTeX tabular blocks. + generation = generation.replace(r"\begin{tabular}{l l} & \\ \end{tabular}", "") + # Remove tabulars with just 2 newline characters. + generation = generation.replace("\\begin{tabular}{}\n\n\\end{tabular}", "") + return generation + + def post_process_single(self, generation: str, fix_markdown: bool = True) -> str: + """ + Postprocess a single generated text. Regular expressions used here are taken directly from the Nougat article + authors. These expressions are commented for clarity and tested end-to-end in most cases. + + Args: + generation (str): The generated text to be postprocessed. + fix_markdown (bool, optional): Whether to perform Markdown formatting fixes. Default is True. + + Returns: + str: The postprocessed text. + """ + generation = re.sub( + r"(?:\n|^)#+ \d*\W? ?(.{100,})", r"\n\1", generation + ) # too long section titles probably are none + generation = generation.strip() + # Remove LaTeX left margin tag + generation = generation.replace("\n* [leftmargin=*]\n", "\n") + # Remove lines with markdown headings starting with #, with numerals, + # and possibly roman numerals with trailing spaces and newlines + generation = re.sub(r"^#+ (?:[\d+\.]+|[ixv\.]+)?\s*(?:$|\n\s*)", "", generation, flags=re.MULTILINE) + # most likely hallucinated titles + lines = generation.split("\n") + if lines[-1].startswith("#") and lines[-1].lstrip("#").startswith(" ") and len(lines) > 1: + logger.info("Likely hallucinated title at the end of the page: " + lines[-1]) + generation = "\n".join(lines[:-1]) + # obvious repetition detection + generation = truncate_repetitions(generation) + # Reference corrections + generation = self.remove_hallucinated_references(generation) + # Remove lines starting with asterisks and numbers like "*[1]" and followed by capital letters and periods (ie too long references) + generation = re.sub(r"^\* \[\d+\](\s?[A-W]\.+\s?){10,}.*$", "", generation, flags=re.MULTILINE) + # Remove empty brackets after a reference number in brackets. *[12][]ABC will become *[12]ABC + generation = re.sub(r"^(\* \[\d+\])\[\](.*)$", r"\1\2", generation, flags=re.MULTILINE) + # Remove single characters before or after 2 new lines + generation = re.sub(r"(^\w\n\n|\n\n\w$)", "", generation) + # pmc math artifact correction + generation = re.sub( + r"([\s.,()])_([a-zA-Z0-9])__([a-zA-Z0-9]){1,3}_([\s.,:()])", + r"\1\(\2_{\3}\)\4", + generation, + ) + generation = re.sub(r"([\s.,\d])_([a-zA-Z0-9])_([\s.,\d;])", r"\1\(\2\)\3", generation) + # footnote mistakes + generation = re.sub( + r"(\nFootnote .*?:) (?:footnotetext|thanks):\W*(.*(?:\n\n|$))", + r"\1 \2", + generation, + ) + # TODO Come up with footnote formatting inside a table + generation = re.sub(r"\[FOOTNOTE:.+?\](.*?)\[ENDFOOTNOTE\]", "", generation) + # itemize post processing + generation = normalize_list_like_lines(generation) + + if generation.endswith((".", "}")): + generation += "\n\n" + if re.match(r"[A-Z0-9,;:]$", generation): + # add space in case it there is a comma or word ending + generation += " " + elif generation.startswith(("#", "**", "\\begin")): + generation = "\n\n" + generation + elif generation.split("\n")[-1].startswith(("#", "Figure", "Table")): + generation = generation + "\n\n" + else: + try: + last_word = generation.split(" ")[-1] + if last_word in nltk.corpus.words.words(): + generation += " " + except LookupError: + # add space just in case. Will split words but better than concatenating them + generation += " " + + # table corrections + generation = self.correct_tables(generation) + # Remove optional, empty square brackets after begin{array} + generation = generation.replace("\\begin{array}[]{", "\\begin{array}{") + # Remove empty or malformed LaTeX tabular blocks with 2 or more columns specified, with spaces and ampersands. + generation = re.sub( + r"\\begin{tabular}{([clr ]){2,}}\s*[& ]*\s*(\\\\)? \\end{tabular}", + "", + generation, + ) + # Remove lines containing "S.A.B." one or more times. Was included in Nougat's code. + generation = re.sub(r"(\*\*S\. A\. B\.\*\*\n+){2,}", "", generation) + # Remove markdown-style headers that are incomplete or empty on multiple lines. + generation = re.sub(r"^#+( [\[\d\w])?$", "", generation, flags=re.MULTILINE) + # Remove lines with just one period. + generation = re.sub(r"^\.\s*$", "", generation, flags=re.MULTILINE) + # Replace instances of three or more newlines with just two newlines. + generation = re.sub(r"\n{3,}", "\n\n", generation) + if fix_markdown: + return markdown_compatible(generation) + else: + return generation + + def post_process_generation( + self, + generation: str | list[str], + fix_markdown: bool = True, + num_workers: int | None = None, + ) -> str | list[str]: + """ + Postprocess a generated text or a list of generated texts. + + This function can be used to perform postprocessing on generated text, such as fixing Markdown formatting. + + Postprocessing is quite slow so it is recommended to use multiprocessing to speed up the process. + + Args: + generation (Union[str, list[str]]): + The generated text or a list of generated texts. + fix_markdown (`bool`, *optional*, defaults to `True`): + Whether to perform Markdown formatting fixes. + num_workers (`int`, *optional*): + Optional number of workers to pass to leverage multiprocessing (postprocessing several texts in + parallel). + + Returns: + Union[str, list[str]]: The postprocessed text or list of postprocessed texts. + """ + requires_backends(self, ["nltk", "levenshtein"]) + + if isinstance(generation, list): + if num_workers is not None and isinstance(num_workers, int): + with Pool(num_workers) as p: + return p.map(partial(self.post_process_single, fix_markdown=fix_markdown), generation) + else: + return [self.post_process_single(s, fix_markdown=fix_markdown) for s in generation] + else: + return self.post_process_single(generation, fix_markdown=fix_markdown) + + +__all__ = ["NougatTokenizer"] diff --git a/third_party/transformers/src/transformers/models/olmo/__init__.py b/third_party/transformers/src/transformers/models/olmo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..139af5473e41fabbe978ee4b9c956dfd9b6e4453 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmo/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2024 EleutherAI and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_olmo import * + from .modeling_olmo 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/olmo/configuration_olmo.py b/third_party/transformers/src/transformers/models/olmo/configuration_olmo.py new file mode 100644 index 0000000000000000000000000000000000000000..186cc3a704fb2e79b2a2a6c6a66568788dbdfbe6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmo/configuration_olmo.py @@ -0,0 +1,92 @@ +# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OLMo 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="allenai/OLMo-7B-hf") +@strict +class OlmoConfig(PreTrainedConfig): + r""" + clip_qkv (`float`, *optional*): + If not `None`, elements of query, key and value attention states are clipped so that their + absolute value does not exceed this value. + + ```python + >>> from transformers import OlmoModel, OlmoConfig + + >>> # Initializing a OLMo 7B style configuration + >>> configuration = OlmoConfig() + + >>> # Initializing a model from the OLMo 7B style configuration + >>> model = OlmoModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "olmo" + 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 = 50304 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + use_cache: bool = True + pad_token_id: int | None = 1 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = 50279 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + clip_qkv: float | 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__ = ["OlmoConfig"] diff --git a/third_party/transformers/src/transformers/models/olmo/convert_olmo_weights_to_hf.py b/third_party/transformers/src/transformers/models/olmo/convert_olmo_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..6cb626e143f8d4acf9b7c873a60fb78e29c29fa3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmo/convert_olmo_weights_to_hf.py @@ -0,0 +1,246 @@ +# Copyright 2024 EleutherAI and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import gc +import json +import os +import shutil +from pathlib import Path + +import torch +import yaml +from tokenizers import Tokenizer + +from transformers import OlmoConfig, OlmoForCausalLM +from transformers.models.gpt_neox.tokenization_gpt_neox_fast import GPTNeoXTokenizerFast + + +""" +Sample usage: + +``` +python src/transformers/models/olmo/convert_olmo_weights_to_hf.py \ + --input_dir /path/to/downloaded/olmo/weights --model_size 7B --output_dir /output/path +``` + +Thereafter, models can be loaded via: + +```py +from transformers import OlmoForCausalLM, AutoTokenizer + +model = OlmoForCausalLM.from_pretrained("/output/path") +tokenizer = AutoTokenizer.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, tokenizer_path=None, fix_eos_token_id=True): + os.makedirs(model_path, exist_ok=True) + tmp_model_path = os.path.join(model_path, "tmp") + os.makedirs(tmp_model_path, exist_ok=True) + + config_path = Path(input_base_path) / "config.yaml" + olmo_config = yaml.safe_load(config_path.read_text())["model"] + + n_layers = olmo_config["n_layers"] + n_heads = olmo_config["n_heads"] + dim = olmo_config["d_model"] + dims_per_head = dim // n_heads + base = 10000.0 + inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)) + max_position_embeddings = olmo_config["max_sequence_length"] + + vocab_size = olmo_config.get("embedding_size", olmo_config["vocab_size"]) + + if olmo_config.get("n_kv_heads", None) is not None: + num_key_value_heads = olmo_config["n_kv_heads"] # for GQA / MQA + elif olmo_config["multi_query_attention"]: # compatibility with other checkpoints + num_key_value_heads = 1 + else: + num_key_value_heads = n_heads + + print(f"Fetching all parameters from the checkpoint at {input_base_path}.") + + # Not sharded + # (The sharded implementation would also work, but this is simpler.) + loaded = torch.load(os.path.join(input_base_path, "model.pt"), map_location="cpu", weights_only=True) + + param_count = 0 + index_dict = {"weight_map": {}} + for layer_i in range(n_layers): + filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin" + # Unsharded + # TODO: Layernorm stuff + # TODO: multi query attention + fused_dims = [dim, dims_per_head * num_key_value_heads, dims_per_head * num_key_value_heads] + q_proj_weight, k_proj_weight, v_proj_weight = torch.split( + loaded[f"transformer.blocks.{layer_i}.att_proj.weight"], fused_dims, dim=0 + ) + up_proj_weight, gate_proj_weight = torch.chunk( + loaded[f"transformer.blocks.{layer_i}.ff_proj.weight"], 2, dim=0 + ) + state_dict = { + f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj_weight, + f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj_weight, + f"model.layers.{layer_i}.self_attn.v_proj.weight": v_proj_weight, + f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"transformer.blocks.{layer_i}.attn_out.weight"], + f"model.layers.{layer_i}.mlp.gate_proj.weight": gate_proj_weight, + f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"transformer.blocks.{layer_i}.ff_out.weight"], + f"model.layers.{layer_i}.mlp.up_proj.weight": up_proj_weight, + } + + state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq + + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin" + + # Unsharded + # TODO: Deal with weight-tying + state_dict = { + "model.embed_tokens.weight": loaded["transformer.wte.weight"], + "lm_head.weight": loaded["transformer.ff_out.weight"] + if "transformer.ff_out.weight" in loaded + else loaded["transformer.wte.weight"], + } + + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + # Write configs + index_dict["metadata"] = {"total_size": param_count * 2} + write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json")) + + if olmo_config.get("mlp_hidden_size", None) is not None: + intermediate_size = olmo_config["mlp_hidden_size"] // 2 + else: + intermediate_size = (dim * olmo_config["mlp_ratio"]) // 2 + + config = OlmoConfig( + vocab_size=vocab_size, + hidden_size=dim, + intermediate_size=intermediate_size, + num_hidden_layers=n_layers, + num_attention_heads=n_heads, + num_key_value_heads=num_key_value_heads, + max_position_embeddings=max_position_embeddings, + pad_token_id=olmo_config["pad_token_id"], + bos_token_id=None, + eos_token_id=olmo_config["eos_token_id"], + tie_word_embeddings=olmo_config["weight_tying"], + rope_theta=base, + clip_qkv=olmo_config.get("clip_qkv"), + ) + config.save_pretrained(tmp_model_path) + + # Make space so we can load the model properly now. + del state_dict + del loaded + gc.collect() + + if tokenizer_path is not None: + _write_tokenizer(model_path, config, tokenizer_path, fix_eos_token_id) + + print("Loading the checkpoint in a OLMo model.") + model = OlmoForCausalLM.from_pretrained(tmp_model_path, dtype=torch.float32) + # Avoid saving this as part of the config. + del model.config._name_or_path + print("Saving in the Transformers format.") + model.save_pretrained(model_path) + shutil.rmtree(tmp_model_path) + + +def _write_tokenizer( + output_path: Path, config: OlmoConfig, input_tokenizer_path: Path, fix_eos_token_id: bool = True +) -> None: + print(f"Saving a {GPTNeoXTokenizerFast.__name__} to {output_path}.") + + base_tokenizer = Tokenizer.from_file(str(input_tokenizer_path)) + + eos_token_id = config.eos_token_id if config.eos_token_id is not None else base_tokenizer.get_vocab_size() - 1 + pad_token_id = config.pad_token_id if config.pad_token_id is not None else eos_token_id + + if fix_eos_token_id and eos_token_id == 0: + # Fixing a bug in OLMo where eos token id was incorrectly set + print("Changing eos_token_id from 0 to 50279.") + eos_token_id = 50279 + + tokenizer = GPTNeoXTokenizerFast( + tokenizer_object=base_tokenizer, + eos_token=base_tokenizer.decode([eos_token_id], skip_special_tokens=False), + pad_token=base_tokenizer.decode([pad_token_id], skip_special_tokens=False), + unk_token=None, + bos_token=None, + ) + + tokenizer.save_pretrained(output_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + required=True, + help="Location of OLMo weights, which contains config.yaml and model.pt.", + ) + parser.add_argument( + "--tokenizer_json_path", + default=None, + help="Location of OLMo tokenizer json file.", + ) + parser.add_argument( + "--output_dir", + required=True, + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--no_fix_eos_token_id", + action="store_false", + dest="fix_eos_token_id", + help="If set, does not change eos token id from 0 to 50279 if it is 0. Changing 0 to 50279 is a bug fix, so use this option with care.", + ) + # Different OLMo versions used different default values for max_position_embeddings, hence the need to be able to specify which version is being used. + args = parser.parse_args() + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + tokenizer_path=args.tokenizer_json_path, + fix_eos_token_id=args.fix_eos_token_id, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/olmo/modeling_olmo.py b/third_party/transformers/src/transformers/models/olmo/modeling_olmo.py new file mode 100644 index 0000000000000000000000000000000000000000..a0949e9570579a7ce339d91554faad74c2d8d574 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmo/modeling_olmo.py @@ -0,0 +1,499 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/olmo/modular_olmo.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_olmo.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 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 as nn +import torch.nn.functional as F + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernelized_func +from ...masking_utils import create_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_olmo import OlmoConfig + + +class OlmoLayerNorm(nn.Module): + """LayerNorm but with no learnable weight or bias.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.normalized_shape = (hidden_size,) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_dtype = hidden_states.dtype + return F.layer_norm(hidden_states.to(dtype=torch.float32), self.normalized_shape, None, None, eps=1e-5).to( + orig_dtype + ) + + +class OlmoMLP(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 OlmoRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: OlmoConfig, 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: OlmoConfig | 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, sin + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def 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. + """ + q_type, k_type = q.dtype, k.dtype + 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(q_type), k_embed.to(k_type) + + +@use_kernelized_func(apply_rotary_pos_emb) +class OlmoAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: OlmoConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + 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_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + if self.config.clip_qkv is not None: + query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + + 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 OlmoDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: OlmoConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = OlmoAttention(config=config, layer_idx=layer_idx) + + self.mlp = OlmoMLP(config) + self.input_layernorm = OlmoLayerNorm(config.hidden_size) + self.post_attention_layernorm = OlmoLayerNorm(config.hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class OlmoPreTrainedModel(PreTrainedModel): + config: OlmoConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["OlmoDecoderLayer"] + _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": OlmoDecoderLayer, + "attentions": OlmoAttention, + } + + +@auto_docstring +class OlmoModel(OlmoPreTrainedModel): + def __init__(self, config: OlmoConfig): + 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( + [OlmoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = OlmoLayerNorm(config.hidden_size) + self.rotary_emb = OlmoRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class OlmoForCausalLM(OlmoPreTrainedModel, 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 = OlmoModel(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, OlmoForCausalLM + + >>> model = OlmoForCausalLM.from_pretrained("meta-olmo/Olmo-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-olmo/Olmo-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__ = ["OlmoForCausalLM", "OlmoModel", "OlmoPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/olmo/modular_olmo.py b/third_party/transformers/src/transformers/models/olmo/modular_olmo.py new file mode 100644 index 0000000000000000000000000000000000000000..6d1588b03e2e4367e016f3b4637253be3ce18b05 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmo/modular_olmo.py @@ -0,0 +1,189 @@ +# Copyright 2024 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 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...cache_utils import Cache +from ...modeling_rope_utils import dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...utils import logging +from ...utils.generic import maybe_autocast +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaMLP, + LlamaModel, + LlamaRotaryEmbedding, + eager_attention_forward, + rotate_half, +) +from .configuration_olmo import OlmoConfig + + +logger = logging.get_logger(__name__) + + +class OlmoLayerNorm(nn.Module): + """LayerNorm but with no learnable weight or bias.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.normalized_shape = (hidden_size,) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_dtype = hidden_states.dtype + return F.layer_norm(hidden_states.to(dtype=torch.float32), self.normalized_shape, None, None, eps=1e-5).to( + orig_dtype + ) + + +class OlmoMLP(LlamaMLP): + def __init__(self, config): + super().__init__(config) + 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) + + +# This is identical to LlamaRotaryEmbedding except the output cos and sin are returned +# as float32 rather than the input type. +class OlmoRotaryEmbedding(LlamaRotaryEmbedding): + @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, sin + + +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. + """ + q_type, k_type = q.dtype, k.dtype + 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(q_type), k_embed.to(k_type) + + +class OlmoAttention(LlamaAttention): + 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_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + if self.config.clip_qkv is not None: + query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + + 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 OlmoDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: OlmoConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.input_layernorm = OlmoLayerNorm(config.hidden_size) + self.post_attention_layernorm = OlmoLayerNorm(config.hidden_size) + self.self_attn = OlmoAttention(config=config, layer_idx=layer_idx) + + +class OlmoModel(LlamaModel): + def __init__(self, config: OlmoConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [OlmoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = OlmoLayerNorm(config.hidden_size) + + +class OlmoForCausalLM(LlamaForCausalLM): + pass + + +__all__ = [ + "OlmoForCausalLM", + "OlmoModel", + "OlmoPreTrainedModel", # noqa: F822 +] diff --git a/third_party/transformers/src/transformers/models/phi/__init__.py b/third_party/transformers/src/transformers/models/phi/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cffe33da73ee42eb20a2e7f33ea9351bb7da75c2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phi/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2023 Microsoft 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_phi import * + from .modeling_phi 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/phi/configuration_phi.py b/third_party/transformers/src/transformers/models/phi/configuration_phi.py new file mode 100644 index 0000000000000000000000000000000000000000..2a65f0f16aecb1d6008a0eac053c87101ee76cae --- /dev/null +++ b/third_party/transformers/src/transformers/models/phi/configuration_phi.py @@ -0,0 +1,92 @@ +# Copyright 2023 Microsoft 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. + +"""Phi 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="microsoft/phi-1") +@strict +class PhiConfig(PreTrainedConfig): + r""" + qk_layernorm (`bool`, *optional*, defaults to `False`): + Whether or not to normalize the Queries and Keys after projecting the hidden states. + + Example: + + ```python + >>> from transformers import PhiModel, PhiConfig + + >>> # Initializing a Phi-1 style configuration + >>> configuration = PhiConfig.from_pretrained("microsoft/phi-1") + + >>> # Initializing a model from the configuration + >>> model = PhiModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "phi" + 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.dense": "rowwise", + "layers.*.mlp.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "embed_dropout": (["inputs_embeds"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "final_layernorm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 51200 + hidden_size: int = 2048 + intermediate_size: int = 8192 + num_hidden_layers: int = 24 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + resid_pdrop: float | int = 0.0 + embd_pdrop: float | int = 0.0 + attention_dropout: float | int | None = 0.0 + hidden_act: str = "gelu_new" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + use_cache: bool = True + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + qk_layernorm: bool = False + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + pad_token_id: int | None = None + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + kwargs.setdefault("partial_rotary_factor", 0.5) # assign default for BC + super().__post_init__(**kwargs) + + +__all__ = ["PhiConfig"] diff --git a/third_party/transformers/src/transformers/models/phi/convert_phi_weights_to_hf.py b/third_party/transformers/src/transformers/models/phi/convert_phi_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..016aaac7fe1260eeb7bbfd154cfaf90c67488195 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phi/convert_phi_weights_to_hf.py @@ -0,0 +1,206 @@ +# Copyright 2023 Microsoft 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. + +""" +Weights conversion script for Phi + +This script downloads both Phi-1 and Phi-1.5 checkpoints to "checkpoint_path" and then converts the weights to +HugfgingFace model's format and saves them in "pytorch_dump_folder_path". + +Example : $python ./convert_phi_weights_to_hf.py --model_name "microsoft/phi-2" --pytorch_dump_folder ./dump_folder/ --checkpoint_path ./ckpt_path/ +""" + +import argparse +import gc +import os + +import safetensors +import torch +from huggingface_hub import hf_hub_download + +from transformers import PhiConfig, PhiForCausalLM + + +_MODELS = { + "microsoft/phi-1": ["https://huggingface.co/microsoft/phi-1/blob/main/pytorch_model.bin"], + "microsoft/phi-1_5": ["https://huggingface.co/microsoft/phi-1_5/blob/main/pytorch_model.bin"], + "microsoft/phi-2": [ + "https://huggingface.co/microsoft/phi-2/blob/main/model-00001-of-00002.safetensors", + "https://huggingface.co/microsoft/phi-2/blob/main/model-00002-of-00002.safetensors", + ], +} + +PHI_MAPPING = { + "transformer.embd.wte.weight": "model.embed_tokens.weight", + "lm_head.linear": "lm_head", + "lm_head.ln": "model.final_layernorm", + "layers": "model.layers", + "transformer": "model", + ".h.": ".layers.", + "ln": "input_layernorm", + "mixer": "self_attn", + "Wqkv": "query_key_value", + "out_proj": "dense", +} + + +def convert_weights(original_weights, mapping, config): + converted_weights = {} + original_weights_keys = sorted(original_weights.keys()) + + for original_weights_key in original_weights_keys: + new_key = original_weights_key + + if "rotary_emb" in new_key: + continue + + if "Wqkv" in new_key: + if "weight" in new_key: + weight = original_weights[new_key] + weights_shape = weight.shape + weight = ( + weight.view(3, config.num_attention_heads, -1, config.hidden_size) + .transpose(0, 1) + .reshape(*weights_shape) + ) + original_weights[new_key] = weight + elif "bias" in new_key: + bias = original_weights[new_key] + bias_shape = bias.shape + bias = bias.view(3, config.num_attention_heads, -1).transpose(0, 1).reshape(*bias_shape) + original_weights[new_key] = bias + + for k, v in mapping.items(): + if k in new_key: + new_key = new_key.replace(k, v) + + converted_weights[new_key] = original_weights.pop(original_weights_key) + + return converted_weights + + +def _download(url: str, root: str): + repo_id = f"{url.split('/')[3]}/{url.split('/')[4]}" + filename = f"{url.split('/')[-1]}" + hf_hub_download( + repo_id=repo_id, + filename=filename, + force_filename=root, + local_dir_use_symlinks=False, + ) + + +def convert_phi_weights( + model_name, checkpoint_path, pytorch_dump_folder_path, use_cuda, save_weights_directly, _MODELS +): + _MODELS = _MODELS if model_name not in _MODELS else {model_name: _MODELS.get(model_name)} + device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" + for model_name, model_url in _MODELS.items(): + converted_checkpoint = {} + model_checkpoint = {} + + # for phi-2 the weights are stored in 2 different safetensors file so we need to iterate over that list and download one at a time + for model_each_url in model_url: + model_path = os.path.join(checkpoint_path, model_name + "_" + model_each_url.split("/")[-1]) + if not os.path.exists(model_path): + print(f"\n{model_name} was not found! Downloading it to {model_path}") + _download(url=model_each_url, root=model_path) + + if model_path.endswith("safetensors"): + loaded_weights = safetensors.torch.load_file(model_path, device=device) + else: + loaded_weights = torch.load(model_path, map_location=device, weights_only=True) + model_checkpoint.update(**loaded_weights) + + model_type = model_name.split("/")[1] # phi-1 or phi-1_5 or phi-2 + + # init the config for phi-1 and phi-1.5 + config = PhiConfig() + # if we are dealing with phi-2 then update the config + if model_type == "phi-2": + config.hidden_size = 2560 + config.intermediate_size = 10240 + config.num_hidden_layers = 32 + config.resid_pdrop = 0.1 + config.partial_rotary_factor = 0.4 + config.num_hidden_layers = 32 + config.dtype = "float16" + + # Converting the weights + converted_checkpoint.update(**convert_weights(model_checkpoint, PHI_MAPPING, config)) + + # Save either the whole model or the converted weights + if save_weights_directly: + save_weights_path = os.path.join(pytorch_dump_folder_path, model_type + "_pytorch_model.bin") + torch.save(converted_checkpoint, save_weights_path) + print(f"Model weights saved at {save_weights_path}!") + + else: + model = PhiForCausalLM(config).to(device) + model.load_state_dict(converted_checkpoint, strict=True) + save_model_path = os.path.join(pytorch_dump_folder_path, model_type) + model.save_pretrained(save_model_path) + print(f"Model saved at {save_model_path}!") + + # release GPU memory for the 2nd model if cuda was used. + del config, model + + # release GPU memory for the 2nd model if cuda was used. + del model_checkpoint, converted_checkpoint + if use_cuda: + torch.cuda.empty_cache() + gc.collect() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # # Required parameters + parser.add_argument( + "--model_name", + type=str, + help="Name of the model to convert. (Please enter one of the following: phi-1, phi-1_5, phi-2). If nothing is provided, all models will be converted.", + default=None, + ) + parser.add_argument( + "--checkpoint_path", type=str, help="Path to the folder of downloaded checkpoints. (Please enter full path)" + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model. (Please enter full path)", + ) + parser.add_argument( + "--use_cuda", + default=False, + type=bool, + help="Whether to load the weights on GPU during conversion or not, False by default", + ) + parser.add_argument( + "--save_weights_directly", + default=True, + type=bool, + help="Whether to save the weights directly after conversion or load the weight to the Phi model and then save " + "the Phi model along with weights. True by default", + ) + + args = parser.parse_args() + convert_phi_weights( + args.model_name, + args.checkpoint_path, + args.pytorch_dump_folder_path, + args.use_cuda, + args.save_weights_directly, + _MODELS, + ) diff --git a/third_party/transformers/src/transformers/models/phi/modeling_phi.py b/third_party/transformers/src/transformers/models/phi/modeling_phi.py new file mode 100644 index 0000000000000000000000000000000000000000..e3f97a01ee4c01d07b6b4d8101877fea51a7ba79 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phi/modeling_phi.py @@ -0,0 +1,494 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/phi/modular_phi.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_phi.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn as nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import ( + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_phi import PhiConfig + + +class PhiRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: PhiConfig, 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: PhiConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernelized_func(apply_rotary_pos_emb) +class PhiAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: PhiConfig, 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=True) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) + self.dense = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True) + self.rotary_ndims = int(self.head_dim * config.rope_parameters["partial_rotary_factor"]) + self.qk_layernorm = config.qk_layernorm + if self.qk_layernorm: + self.q_layernorm = nn.LayerNorm( + config.hidden_size // config.num_attention_heads, eps=config.layer_norm_eps, elementwise_affine=True + ) + self.k_layernorm = nn.LayerNorm( + config.hidden_size // config.num_attention_heads, eps=config.layer_norm_eps, elementwise_affine=True + ) + + 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_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 self.qk_layernorm: + query_states = self.q_layernorm(query_states) + key_states = self.k_layernorm(key_states) + + cos, sin = position_embeddings + # Partial rotary embedding + query_rot, query_pass = ( + query_states[..., : self.rotary_ndims], + query_states[..., self.rotary_ndims :], + ) + key_rot, key_pass = ( + key_states[..., : self.rotary_ndims], + key_states[..., self.rotary_ndims :], + ) + # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] + query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin) + + # [batch_size, seq_length, num_heads, head_dim] + query_states = torch.cat((query_rot, query_pass), dim=-1) + key_states = torch.cat((key_rot, key_pass), dim=-1) + + 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.dense(attn_output) + return attn_output, attn_weights + + +class PhiMLP(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 PhiDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: PhiConfig, layer_idx: int): + super().__init__() + self.self_attn = PhiAttention(config, layer_idx=layer_idx) + self.mlp = PhiMLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + attn_outputs, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + attn_outputs = self.resid_dropout(attn_outputs) + + feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states)) + hidden_states = attn_outputs + feed_forward_hidden_states + residual + + return hidden_states + + +@auto_docstring +class PhiPreTrainedModel(PreTrainedModel): + config: PhiConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["PhiDecoderLayer"] + _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": PhiDecoderLayer, + "attentions": PhiAttention, + } + + +@auto_docstring +class PhiModel(PhiPreTrainedModel): + def __init__(self, config: PhiConfig): + 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( + [PhiDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.rotary_emb = PhiRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.embed_dropout = nn.Dropout(config.embd_pdrop) + self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_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], + ) -> 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) + + 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, + ) + + inputs_embeds = self.embed_dropout(inputs_embeds) + 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.final_layernorm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class PhiForCausalLM(PhiPreTrainedModel, 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 = PhiModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True) + + # 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, PhiForCausalLM + + >>> model = PhiForCausalLM.from_pretrained("meta-phi/Phi-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-phi/Phi-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class PhiForSequenceClassification(GenericForSequenceClassification, PhiPreTrainedModel): + pass + + +class PhiForTokenClassification(GenericForTokenClassification, PhiPreTrainedModel): + pass + + +__all__ = [ + "PhiPreTrainedModel", + "PhiModel", + "PhiForCausalLM", + "PhiForSequenceClassification", + "PhiForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/phi/modular_phi.py b/third_party/transformers/src/transformers/models/phi/modular_phi.py new file mode 100644 index 0000000000000000000000000000000000000000..b0afe5712e2afc8caa4da5ec5a24d18449bfeb49 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phi/modular_phi.py @@ -0,0 +1,288 @@ +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn as nn + +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPast, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..clip.modeling_clip import CLIPMLP +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaForTokenClassification, + LlamaModel, + LlamaPreTrainedModel, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from .configuration_phi import PhiConfig + + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "microsoft/phi-1" +_CONFIG_FOR_DOC = "PhiConfig" + + +class PhiRotaryEmbedding(LlamaRotaryEmbedding): + @staticmethod + def compute_default_rope_parameters( + config: PhiConfig | 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 + + +class PhiAttention(LlamaAttention): + def __init__(self, config: PhiConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) + self.dense = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True) + del self.o_proj + self.rotary_ndims = int(self.head_dim * config.rope_parameters["partial_rotary_factor"]) + self.qk_layernorm = config.qk_layernorm + if self.qk_layernorm: + self.q_layernorm = nn.LayerNorm( + config.hidden_size // config.num_attention_heads, eps=config.layer_norm_eps, elementwise_affine=True + ) + self.k_layernorm = nn.LayerNorm( + config.hidden_size // config.num_attention_heads, eps=config.layer_norm_eps, elementwise_affine=True + ) + + 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_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 self.qk_layernorm: + query_states = self.q_layernorm(query_states) + key_states = self.k_layernorm(key_states) + + cos, sin = position_embeddings + # Partial rotary embedding + query_rot, query_pass = ( + query_states[..., : self.rotary_ndims], + query_states[..., self.rotary_ndims :], + ) + key_rot, key_pass = ( + key_states[..., : self.rotary_ndims], + key_states[..., self.rotary_ndims :], + ) + # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] + query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin) + + # [batch_size, seq_length, num_heads, head_dim] + query_states = torch.cat((query_rot, query_pass), dim=-1) + key_states = torch.cat((key_rot, key_pass), dim=-1) + + 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.dense(attn_output) + return attn_output, attn_weights + + +class PhiMLP(CLIPMLP): + pass + + +class PhiDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: PhiConfig, layer_idx: int): + super().__init__() + self.self_attn = PhiAttention(config, layer_idx=layer_idx) + self.mlp = PhiMLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + attn_outputs, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + attn_outputs = self.resid_dropout(attn_outputs) + + feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states)) + hidden_states = attn_outputs + feed_forward_hidden_states + residual + + return hidden_states + + +class PhiPreTrainedModel(LlamaPreTrainedModel): + _can_record_outputs = { + "hidden_states": PhiDecoderLayer, + "attentions": PhiAttention, + } + + +class PhiModel(LlamaModel): + def __init__(self, config: PhiConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [PhiDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.embed_dropout = nn.Dropout(config.embd_pdrop) + self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + del self.norm + + @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) + + 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, + ) + + inputs_embeds = self.embed_dropout(inputs_embeds) + 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.final_layernorm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class PhiForCausalLM(LlamaForCausalLM): + def __init__(self, config): + super().__init__(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True) + + +class PhiForSequenceClassification(LlamaForSequenceClassification): + pass + + +class PhiForTokenClassification(LlamaForTokenClassification): + pass + + +__all__ = [ + "PhiPreTrainedModel", + "PhiModel", + "PhiForCausalLM", + "PhiForSequenceClassification", + "PhiForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/pp_doclayout_v2/__init__.py b/third_party/transformers/src/transformers/models/pp_doclayout_v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5c971aba9111bf143a16c7bbdaa27dd6cec2bdb8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_doclayout_v2/__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_doclayout_v2 import * + from .image_processing_pp_doclayout_v2 import * + from .modeling_pp_doclayout_v2 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_doclayout_v2/configuration_pp_doclayout_v2.py b/third_party/transformers/src/transformers/models/pp_doclayout_v2/configuration_pp_doclayout_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..08c28c0e3de78bb21980b93b28043d93c6c15cf5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_doclayout_v2/configuration_pp_doclayout_v2.py @@ -0,0 +1,266 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.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_doclayout_v2.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-DocLayoutV2_safetensors") +@strict +class PPDocLayoutV2ReadingOrderConfig(PreTrainedConfig): + r""" + has_relative_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a relative attention bias in the self-attention mechanism. + has_spatial_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a spatial attention bias in the self-attention mechanism. + rel_pos_bins (`int`, *optional*, defaults to 32): + The number of relative position bins to be used in the self-attention mechanism. + max_rel_pos (`int`, *optional*, defaults to 128): + The maximum number of relative positions to be used in the self-attention mechanism. + rel_2d_pos_bins (`int`, *optional*, defaults to 64): + The number of 2D relative position bins in the self-attention mechanism. + max_rel_2d_pos (`int`, *optional*, defaults to 256): + The maximum number of relative 2D positions in the self-attention mechanism. + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum value that the 2D position embedding might ever be used with. Typically set this to something + large just in case (e.g., 1024). + start_token_id (`int`, *optional*, defaults to 0): + Token id representing the start of a sequence. + end_token_id (`int`, *optional*, defaults to 2): + Token id representing the end of a sequence. + pred_token_id (`int`, *optional*, defaults to 3): + Token id representing valid prediction positions (placeholders) in the sequence. + coordinate_size (`int`, *optional*, defaults to 171): + Dimension of the coordinate embeddings. + shape_size (`int`, *optional*, defaults to 170): + Dimension of the width and height embeddings. + num_classes (`int`, *optional*, defaults to 20): + Number of labels or classes for the layout elements. + relation_bias_embed_dim (`int`, *optional*, defaults to 16): + Embedding dimension for the relation bias. + relation_bias_theta (`float`, *optional*, defaults to 10000): + Temperature parameter used for relation bias scaling. + relation_bias_scale (`float`, *optional*, defaults to 100): + Scale parameter for the relation bias. + global_pointer_head_size (`int`, *optional*, defaults to 64): + The size of the global pointer head. + gp_dropout_value (`float`, *optional*, defaults to 0.0): + The dropout probability in the global pointer head. + """ + + hidden_size: int = 512 + num_attention_heads: int = 8 + attention_probs_dropout_prob: float | int = 0.1 + has_relative_attention_bias: bool = False + has_spatial_attention_bias: bool = True + layer_norm_eps: float = 1e-5 + hidden_dropout_prob: float | int = 0.1 + intermediate_size: int = 2048 + hidden_act: str = "gelu" + num_hidden_layers: int = 6 + rel_pos_bins: int = 32 + max_rel_pos: int = 128 + rel_2d_pos_bins: int = 64 + max_rel_2d_pos: int = 256 + max_position_embeddings: int = 514 + max_2d_position_embeddings: int = 1024 + type_vocab_size: int = 1 + vocab_size: int = 4 + initializer_range: float = 0.01 + start_token_id: int = 0 + pad_token_id: int | None = 1 + end_token_id: int = 2 + pred_token_id: int = 3 + coordinate_size: int = 171 + shape_size: int = 170 + num_classes: int = 20 + relation_bias_embed_dim: int = 16 + relation_bias_theta: int = 10000 + relation_bias_scale: int = 100 + global_pointer_head_size: int = 64 + gp_dropout_value: float | int = 0.0 + + +@auto_docstring(checkpoint="PaddlePaddle/PP-DocLayoutV2_safetensors") +@strict +class PPDocLayoutV2Config(PreTrainedConfig): + r""" + initializer_bias_prior_prob (`float`, *optional*): + The prior probability used by the bias initializer to initialize biases for `enc_score_head` and `class_embed`. + If `None`, `prior_prob` computed as `prior_prob = 1 / (num_labels + 1)` while initializing model weights. + freeze_backbone_batch_norms (`bool`, *optional*, defaults to `True`): + Whether to freeze the batch normalization layers in the backbone. + encoder_hidden_dim (`int`, *optional*, defaults to 256): + Dimension of the layers in hybrid encoder. + encoder_in_channels (`list`, *optional*, defaults to `[512, 1024, 2048]`): + Multi level features input for encoder. + feat_strides (`list[int]`, *optional*, defaults to `[8, 16, 32]`): + Strides used in each feature map. + encoder_layers (`int`, *optional*, defaults to 1): + Total of layers to be used by the encoder. + encode_proj_layers (`list[int]`, *optional*, defaults to `[2]`): + Indexes of the projected layers to be used in the encoder. + positional_encoding_temperature (`int`, *optional*, defaults to 10000): + The temperature parameter used to create the positional encodings. + encoder_activation_function (`str`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"silu"` and `"gelu_new"` are supported. + eval_size (`tuple[int, int]`, *optional*): + Height and width used to computes the effective height and width of the position embeddings after taking + into account the stride. + normalize_before (`bool`, *optional*, defaults to `False`): + Determine whether to apply layer normalization in the transformer encoder layer before self-attention and + feed-forward modules. + hidden_expansion (`float`, *optional*, defaults to 1.0): + Expansion ratio to enlarge the dimension size of RepVGGBlock and CSPRepLayer. + num_queries (`int`, *optional*, defaults to 300): + Number of object queries. + decoder_in_channels (`list`, *optional*, defaults to `[256, 256, 256]`): + Multi level features dimension for decoder + decoder_ffn_dim (`int`, *optional*, defaults to 1024): + Dimension of the "intermediate" (often named feed-forward) layer in decoder. + num_feature_levels (`int`, *optional*, defaults to 3): + The number of input feature levels. + decoder_n_points (`int`, *optional*, defaults to 4): + The number of sampled keys in each feature level for each attention head in the decoder. + decoder_activation_function (`str`, *optional*, defaults to `"relu"`): + The non-linear activation function (function or string) in the decoder. If string, `"gelu"`, + `"relu"`, `"silu"` and `"gelu_new"` are supported. + num_denoising (`int`, *optional*, defaults to 100): + The total number of denoising tasks or queries to be used for contrastive denoising. + label_noise_ratio (`float`, *optional*, defaults to 0.5): + The fraction of denoising labels to which random noise should be added. + box_noise_scale (`float`, *optional*, defaults to 1.0): + Scale or magnitude of noise to be added to the bounding boxes. + learn_initial_query (`bool`, *optional*, defaults to `False`): + Indicates whether the initial query embeddings for the decoder should be learned during training + anchor_image_size (`tuple[int, int]`, *optional*): + Height and width of the input image used during evaluation to generate the bounding box anchors. If None, automatic generate anchor is applied. + disable_custom_kernels (`bool`, *optional*, defaults to `True`): + Whether to disable custom kernels. + class_thresholds (`list[float]`, *optional*): + The thresholds for each label. + class_order (`list[int]`, *optional*): + The priority for each label. + reading_order_config (`dict`, *optional*): + The configuration of a `PPDocLayoutV2ReadingOrder`. + + Examples: + + ```python + >>> from transformers import PPDocLayoutV2Config, PPDocLayoutV2ForObjectDetection + + >>> # Initializing a PP-DocLayoutV2 configuration + >>> configuration = PPDocLayoutV2Config() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = PPDocLayoutV2ForObjectDetection(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pp_doclayout_v2" + sub_configs = {"backbone_config": AutoConfig, "reading_order_config": PPDocLayoutV2ReadingOrderConfig} + + layer_types = ("basic", "bottleneck") + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + } + + initializer_range: float = 0.01 + initializer_bias_prior_prob: float | None = None + layer_norm_eps: float = 1e-5 + batch_norm_eps: float = 1e-5 + backbone_config: PreTrainedConfig | dict | None = None + freeze_backbone_batch_norms: bool = True + encoder_hidden_dim: int = 256 + encoder_in_channels: list[int] | tuple[int, ...] | None = (512, 1024, 2048) + feat_strides: list[int] | tuple[int, ...] | None = (8, 16, 32) + encoder_layers: int = 1 + encoder_ffn_dim: int = 1024 + encoder_attention_heads: int = 8 + dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + encode_proj_layers: list[int] | tuple[int, ...] | None = (2,) + positional_encoding_temperature: int = 10000 + encoder_activation_function: str = "gelu" + activation_function: str = "silu" + eval_size: list[int] | None = None + normalize_before: bool = False + hidden_expansion: float = 1.0 + d_model: int = 256 + num_queries: int = 300 + decoder_in_channels: list[int] | tuple[int, ...] | None = (256, 256, 256) + decoder_ffn_dim: int = 1024 + num_feature_levels: int = 3 + decoder_n_points: int = 4 + decoder_layers: int = 6 + decoder_attention_heads: int = 8 + decoder_activation_function: str = "relu" + attention_dropout: float | int = 0.0 + num_denoising: int = 100 + label_noise_ratio: float = 0.5 + box_noise_scale: float = 1.0 + learn_initial_query: bool = False + anchor_image_size: list[int] | None = None + disable_custom_kernels: bool = True + is_encoder_decoder: bool = True + class_thresholds: list[float] | None = None + class_order: list[int] | None = None + reading_order_config: PreTrainedConfig | dict | None = None + + def __post_init__(self, **kwargs): + if isinstance(self.reading_order_config, dict): + self.reading_order_config = self.sub_configs["reading_order_config"](**self.reading_order_config) + elif self.reading_order_config is None: + self.reading_order_config = self.sub_configs["reading_order_config"]() + + 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": [1, 2, 3], + "freeze_stem_only": True, + "freeze_at": 0, + "freeze_norm": True, + "lr_mult_list": [0, 0.05, 0.05, 0.05, 0.05], + "out_features": ["stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + self.encoder_in_channels = list(self.encoder_in_channels) + self.feat_strides = list(self.feat_strides) + self.encode_proj_layers = list(self.encode_proj_layers) + self.eval_size = list(self.eval_size) if self.eval_size is not None else None + self.decoder_in_channels = list(self.decoder_in_channels) + self.anchor_image_size = list(self.anchor_image_size) if self.anchor_image_size is not None else None + + super().__post_init__(**kwargs) + + +__all__ = ["PPDocLayoutV2Config"] diff --git a/third_party/transformers/src/transformers/models/pp_doclayout_v2/image_processing_pp_doclayout_v2.py b/third_party/transformers/src/transformers/models/pp_doclayout_v2/image_processing_pp_doclayout_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..ec8dce9c8b2538c094d19e8ee1382cf997cc80c3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_doclayout_v2/image_processing_pp_doclayout_v2.py @@ -0,0 +1,193 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.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_doclayout_v2.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 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 PILImageResampling, SizeDict +from ...utils import auto_docstring +from ...utils.generic import TensorType + + +@auto_docstring +class PPDocLayoutV2ImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BICUBIC + image_mean = [0, 0, 0] + image_std = [1, 1, 1] + size = {"height": 800, "width": 800} + do_resize = True + do_rescale = True + do_normalize = True + + # We require `self.resize(..., antialias=False)` to approximate the output of `cv2.resize` + 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 = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, resample=resample, antialias=False) + 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: + 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_order_seqs(self, order_logits): + """ + Computes the order sequences for a batch of inputs based on logits. + + This function takes in the order logits, calculates order scores using a sigmoid activation, + and determines the order sequences by ranking the votes derived from the scores. + + Args: + order_logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_queries)`): + Stacked order logits. + + Returns: + torch.Tensor: A tensor of shape `(batch_size, num_queries)`: + Containing the computed order sequences for each input in the batch. Each row represents the ranked order of elements for the corresponding input in the batch. + """ + order_scores = torch.sigmoid(order_logits) + batch_size, sequence_length, _ = order_scores.shape + + order_votes = order_scores.triu(diagonal=1).sum(dim=1) + (1.0 - order_scores.transpose(1, 2)).tril( + diagonal=-1 + ).sum(dim=1) + + order_pointers = torch.argsort(order_votes, dim=1) + order_seq = torch.empty_like(order_pointers) + ranks = torch.arange(sequence_length, device=order_pointers.device, dtype=order_pointers.dtype).expand( + batch_size, -1 + ) + order_seq.scatter_(1, order_pointers, ranks) + + return order_seq + + def extract_custom_vertices(self): + raise AttributeError("Not needed for PPDocLayoutV2") + + def post_process_object_detection( + self, + outputs, + threshold: float = 0.5, + target_sizes: TensorType | list[tuple] | None = None, + ): + """ + Converts the raw output of [`DetrForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y, + bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + PPDocLayoutV2 is identical to PPDocLayoutV3, except that it does not return `polygon_points`. + + Args: + outputs ([`DetrObjectDetectionOutput`]): + Raw outputs of the model. + 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. + """ + boxes = outputs.pred_boxes + logits = outputs.logits + order_logits = outputs.order_logits + + order_seqs = self._get_order_seqs(order_logits) + + box_centers, box_dims = torch.split(boxes, 2, dim=-1) + top_left_coords = box_centers - 0.5 * box_dims + bottom_right_coords = box_centers + 0.5 * box_dims + boxes = torch.cat([top_left_coords, bottom_right_coords], dim=-1) + + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + if isinstance(target_sizes, list): + img_height, img_width = torch.as_tensor(target_sizes).unbind(1) + else: + img_height, img_width = target_sizes.unbind(1) + scale_factor = torch.stack([img_width, img_height, img_width, img_height], dim=1).to(boxes.device) + boxes = boxes * scale_factor[:, None, :] + + num_top_queries = logits.shape[1] + num_classes = logits.shape[2] + + scores = torch.nn.functional.sigmoid(logits) + scores, index = torch.topk(scores.flatten(1), num_top_queries, dim=-1) + labels = index % num_classes + index = index // num_classes + boxes = boxes.gather(dim=1, index=index.unsqueeze(-1).repeat(1, 1, boxes.shape[-1])) + order_seqs = order_seqs.gather(dim=1, index=index) + + results = [] + for score, label, box, order_seq in zip(scores, labels, boxes, order_seqs): + order_seq = order_seq[score >= threshold] + order_seq, indices = torch.sort(order_seq) + results.append( + { + "scores": score[score >= threshold][indices], + "labels": label[score >= threshold][indices], + "boxes": box[score >= threshold][indices], + "order_seq": order_seq, + } + ) + + return results + + +__all__ = ["PPDocLayoutV2ImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/pp_doclayout_v2/modeling_pp_doclayout_v2.py b/third_party/transformers/src/transformers/models/pp_doclayout_v2/modeling_pp_doclayout_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..7310bd55942e6c4c3ad62b9199800cd01592efaf --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_doclayout_v2/modeling_pp_doclayout_v2.py @@ -0,0 +1,2480 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.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_doclayout_v2.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 warnings +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2CLS, ACT2FN +from ...backbone_utils import load_backbone +from ...image_transforms import center_to_corners_format, corners_to_center_format +from ...integrations import use_kernel_forward_from_hub +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import apply_chunking_to_forward, compile_compatible_method_lru_cache +from ...utils import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + can_return_tuple, + torch_compilable_check, + torch_int, +) +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_pp_doclayout_v2 import PPDocLayoutV2Config + + +class PPDocLayoutV2GlobalPointer(nn.Module): + def __init__(self, config): + super().__init__() + self.head_size = config.global_pointer_head_size + self.dense = nn.Linear(config.hidden_size, self.head_size * 2) + self.dropout = nn.Dropout(config.gp_dropout_value) + + def forward(self, inputs): + batch_size, sequence_length, _ = inputs.shape + query_key_projection = self.dense(inputs).reshape(batch_size, sequence_length, 2, self.head_size) + query_key_projection = self.dropout(query_key_projection) + queries, keys = torch.unbind(query_key_projection, dim=2) + + logits = (queries @ keys.transpose(-2, -1)) / (self.head_size**0.5) + mask = torch.tril(torch.ones(sequence_length, sequence_length, device=logits.device)).bool() + logits = logits.masked_fill(mask.unsqueeze(0), -1e4) + + return logits + + +class PPDocLayoutV2PositionRelationEmbedding(nn.Module): + inv_freq: torch.Tensor + + def __init__(self, config, device=None): + super().__init__() + self.config = config + self.embed_dim = config.relation_bias_embed_dim + self.scale = config.relation_bias_scale + self.pos_proj = nn.Conv2d( + in_channels=self.embed_dim * 4, out_channels=config.num_attention_heads, kernel_size=1 + ) + inv_freq, self.attention_scaling = self.compute_default_rope_parameters(config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @staticmethod + def compute_default_rope_parameters( + config: PPDocLayoutV2Config | 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.relation_bias_theta + dim = config.relation_bias_embed_dim + half_dim = dim // 2 + + 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) / half_dim) + ) + return inv_freq, attention_factor + + def box_relative_encoding( + self, source_boxes: torch.Tensor, target_boxes: torch.Tensor = None, epsilon: float = 1e-5 + ): + source_boxes, target_boxes = source_boxes.unsqueeze(-2), target_boxes.unsqueeze(-3) + source_coordinates, source_dim = source_boxes[..., :2], source_boxes[..., 2:] + target_coordinates, target_dim = target_boxes[..., :2], target_boxes[..., 2:] + + coordinate_difference = torch.abs(source_coordinates - target_coordinates) + relative_coordinates = torch.log(coordinate_difference / (source_dim + epsilon) + 1.0) + relative_dim = torch.log((source_dim + epsilon) / (target_dim + epsilon)) + + relative_encoding = torch.cat([relative_coordinates, relative_dim], dim=-1) + + return relative_encoding + + def get_position_embedding(self, x: torch.Tensor, scale: float = 100.0): + embedding = (x * scale).unsqueeze(-1) * self.inv_freq + embedding = torch.cat((embedding.sin(), embedding.cos()), dim=-1).flatten(start_dim=-2).to(x.dtype) + + return embedding + + def forward(self, source_boxes: torch.Tensor, target_boxes: torch.Tensor = None): + if target_boxes is None: + target_boxes = source_boxes + with torch.no_grad(): + relative_encoding = self.box_relative_encoding(source_boxes, target_boxes) + position_embedding = self.get_position_embedding(relative_encoding, self.scale) + position_embedding = position_embedding.permute(0, 3, 1, 2) + out = self.pos_proj(position_embedding) + return out + + +class PPDocLayoutV2ReadingOrderSelfAttention(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.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.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + def cogview_attention(self, attention_scores, alpha=32): + """ + https://huggingface.co/papers/2105.13290 Section 2.4 Stabilization of training: Precision Bottleneck Relaxation + (PB-Relax). A replacement of the original nn.Softmax(dim=-1)(attention_scores). Seems the new attention_probs + will result in a slower speed and a little bias. Can use torch.allclose(standard_attention_probs, + cogview_attention_probs, atol=1e-08) for comparison. The smaller atol (e.g., 1e-08), the better. + """ + scaled_attention_scores = attention_scores / alpha + max_value = scaled_attention_scores.amax(dim=(-1)).unsqueeze(-1) + new_attention_scores = (scaled_attention_scores - max_value) * alpha + return nn.Softmax(dim=-1)(new_attention_scores) + + def forward( + self, hidden_states, attention_mask=None, rel_pos=None, rel_2d_pos=None, **kwargs: Unpack[TransformersKwargs] + ): + batch_size, seq_length, _ = hidden_states.shape + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + + # Take the dot product between "query" and "key" to get the raw attention scores. + # The attention scores QT K/√d could be significantly larger than input elements, and result in overflow. + # Changing the computational order into QT(K/√d) alleviates the problem. (https://huggingface.co/papers/2105.13290) + attention_scores = torch.matmul(query_layer / math.sqrt(self.attention_head_size), key_layer.transpose(-1, -2)) + + # NOTE: + # PPDocLayoutV2ReadingOrder applies unscaled rel_2d_pos + # unlike LayoutLMv3 which uses (rel_pos + rel_2d_pos) / math.sqrt(self.attention_head_size) + if rel_2d_pos is not None: + attention_scores += rel_2d_pos + elif self.has_relative_attention_bias: + attention_scores += rel_pos / math.sqrt(self.attention_head_size) + + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + # Use the trick of the CogView paper to stabilize training + attention_probs = self.cogview_attention(attention_scores) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer, attention_probs + + +class PPDocLayoutV2ReadingOrderSelfOutput(nn.Module): + 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.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + 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.norm(hidden_states + input_tensor) + return hidden_states + + +class PPDocLayoutV2ReadingOrderIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +class PPDocLayoutV2ReadingOrderOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + 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.norm(hidden_states + input_tensor) + return hidden_states + + +class PPDocLayoutV2ReadingOrderAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = PPDocLayoutV2ReadingOrderSelfAttention(config) + self.output = PPDocLayoutV2ReadingOrderSelfOutput(config) + + def forward( + self, + hidden_states, + attention_mask=None, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + attention_output, _ = self.self( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + attention_output = self.output(attention_output, residual) + return attention_output + + +class PPDocLayoutV2ReadingOrderLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = PPDocLayoutV2ReadingOrderAttention(config) + self.intermediate = PPDocLayoutV2ReadingOrderIntermediate(config) + self.output = PPDocLayoutV2ReadingOrderOutput(config) + + def forward( + self, + hidden_states, + attention_mask=None, + output_attentions=False, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + attention_output = self.attention( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + ) + + layer_output = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output + ) + + return layer_output + + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class PPDocLayoutV2ReadingOrderEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([PPDocLayoutV2ReadingOrderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + self.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + if self.has_relative_attention_bias: + self.rel_pos_bins = config.rel_pos_bins + self.max_rel_pos = config.max_rel_pos + self.rel_pos_bias = nn.Linear(self.rel_pos_bins, config.num_attention_heads, bias=False) + + if self.has_spatial_attention_bias: + self.max_rel_2d_pos = config.max_rel_2d_pos + self.rel_2d_pos_bins = config.rel_2d_pos_bins + self.rel_pos_x_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + self.rel_pos_y_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + self.rel_bias_module = PPDocLayoutV2PositionRelationEmbedding(config) + + def relative_position_bucket(self, relative_position, bidirectional=True, num_buckets=32, max_distance=128): + ret = 0 + if bidirectional: + num_buckets //= 2 + ret += (relative_position > 0).long() * num_buckets + n = torch.abs(relative_position) + else: + n = torch.max(-relative_position, torch.zeros_like(relative_position)) + # now n is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = n < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + val_if_large = max_exact + ( + torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) + ).to(torch.long) + val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1)) + + ret += torch.where(is_small, n, val_if_large) + return ret + + def _cal_1d_pos_emb(self, position_ids): + rel_pos_mat = position_ids.unsqueeze(-2) - position_ids.unsqueeze(-1) + + rel_pos = self.relative_position_bucket( + rel_pos_mat, + num_buckets=self.rel_pos_bins, + max_distance=self.max_rel_pos, + ) + # Since this is a simple indexing operation that is independent of the input, + # no need to track gradients for this operation + # + # Without this no_grad context, training speed slows down significantly + with torch.no_grad(): + rel_pos = self.rel_pos_bias.weight.t()[rel_pos].permute(0, 3, 1, 2) + rel_pos = rel_pos.contiguous() + return rel_pos + + def _cal_2d_pos_emb(self, bbox): + x_min, y_min, x_max, y_max = ( + bbox[..., 0], + bbox[..., 1], + bbox[..., 2], + bbox[..., 3], + ) + + width = (x_max - x_min).clamp(min=1e-3) + height = (y_max - y_min).clamp(min=1e-3) + + center_x = (x_min + x_max) * 0.5 + center_y = (y_min + y_max) * 0.5 + + center_width_height_bbox = torch.stack([center_x, center_y, width, height], dim=-1) + + result = self.rel_bias_module(center_width_height_bbox) + + return result + + def forward( + self, + hidden_states, + bbox=None, + attention_mask=None, + position_ids=None, + patch_height=None, + patch_width=None, + **kwargs: Unpack[TransformersKwargs], + ): + rel_pos = self._cal_1d_pos_emb(position_ids) if self.has_relative_attention_bias else None + rel_2d_pos = self._cal_2d_pos_emb(bbox) if self.has_spatial_attention_bias else None + + for layer_module in self.layer: + hidden_states = layer_module( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +class PPDocLayoutV2TextEmbeddings(nn.Module): + """ + PPDocLayoutV2 text embeddings. Same as `RobertaEmbeddings` but with added spatial (layout) embeddings. + """ + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + self.padding_idx = config.pad_token_id + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx + ) + + self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + spatial_embed_dim = 4 * config.coordinate_size + 2 * config.shape_size + self.spatial_proj = nn.Linear(spatial_embed_dim, config.hidden_size) + + def calculate_spatial_position_embeddings(self, bbox): + 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(torch.clip(bbox[:, :, 3] - bbox[:, :, 1], 0, 1023)) + w_position_embeddings = self.w_position_embeddings(torch.clip(bbox[:, :, 2] - bbox[:, :, 0], 0, 1023)) + + # below is the difference between LayoutLMEmbeddingsV2 (torch.cat) and LayoutLMEmbeddingsV1 (add) + spatial_position_embeddings = torch.cat( + [ + left_position_embeddings, + upper_position_embeddings, + right_position_embeddings, + lower_position_embeddings, + h_position_embeddings, + w_position_embeddings, + ], + dim=-1, + ) + return spatial_position_embeddings + + def create_position_ids_from_input_ids(self, input_ids, padding_idx): + """ + 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`. + """ + # 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)) * mask + return incremental_indices.long() + padding_idx + + 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. + """ + 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) + + def forward( + self, + input_ids=None, + bbox=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + ): + if position_ids is None: + if input_ids is not None: + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx).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 = 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 + + # custom new spatial embeddings + spatial_position_embeddings = self.calculate_spatial_position_embeddings(bbox) + spatial_position_embeddings = self.spatial_proj(spatial_position_embeddings) + embeddings += spatial_position_embeddings + return embeddings + + +@use_kernel_forward_from_hub("MultiScaleDeformableAttention") +class MultiScaleDeformableAttention(nn.Module): + def forward( + self, + value: Tensor, + value_spatial_shapes: Tensor, + value_spatial_shapes_list: list[tuple], + level_start_index: Tensor, + sampling_locations: Tensor, + attention_weights: Tensor, + im2col_step: int, + ): + batch_size, _, num_heads, hidden_dim = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([height * width for height, width in value_spatial_shapes_list], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level_id, (height, width) in enumerate(value_spatial_shapes_list): + # batch_size, height*width, num_heads, hidden_dim + # -> batch_size, height*width, num_heads*hidden_dim + # -> batch_size, num_heads*hidden_dim, height*width + # -> batch_size*num_heads, hidden_dim, height, width + value_l_ = ( + value_list[level_id] + .flatten(2) + .transpose(1, 2) + .reshape(batch_size * num_heads, hidden_dim, height, width) + ) + # batch_size, num_queries, num_heads, num_points, 2 + # -> batch_size, num_heads, num_queries, num_points, 2 + # -> batch_size*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1) + # batch_size*num_heads, hidden_dim, num_queries, num_points + sampling_value_l_ = nn.functional.grid_sample( + value_l_, + sampling_grid_l_, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ) + sampling_value_list.append(sampling_value_l_) + # (batch_size, num_queries, num_heads, num_levels, num_points) + # -> (batch_size, num_heads, num_queries, num_levels, num_points) + # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + batch_size * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(batch_size, num_heads * hidden_dim, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +class PPDocLayoutV2MultiscaleDeformableAttention(nn.Module): + """ + Multiscale deformable attention as proposed in Deformable DETR. + """ + + def __init__(self, config: PPDocLayoutV2Config, num_heads: int, n_points: int): + super().__init__() + + self.attn = MultiScaleDeformableAttention() + + if config.d_model % num_heads != 0: + raise ValueError( + f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}" + ) + dim_per_head = config.d_model // num_heads + # check if dim_per_head is power of 2 + if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0): + warnings.warn( + "You'd better set embed_dim (d_model) in PPDocLayoutV2MultiscaleDeformableAttention to make the" + " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA" + " implementation." + ) + + self.im2col_step = 64 + + self.d_model = config.d_model + self.n_levels = config.num_feature_levels + self.n_heads = num_heads + self.n_points = n_points + + self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 2) + self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points) + self.value_proj = nn.Linear(config.d_model, config.d_model) + self.output_proj = nn.Linear(config.d_model, config.d_model) + + self.disable_custom_kernels = config.disable_custom_kernels + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states=None, + encoder_attention_mask=None, + position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes=None, + spatial_shapes_list=None, + level_start_index=None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + # add position embeddings to the hidden states before projecting to queries and keys + if position_embeddings is not None: + hidden_states = hidden_states + position_embeddings + + batch_size, num_queries, _ = hidden_states.shape + batch_size, sequence_length, _ = encoder_hidden_states.shape + total_elements = sum(height * width for height, width in spatial_shapes_list) + torch_compilable_check( + total_elements == sequence_length, + "Make sure to align the spatial shapes with the sequence length of the encoder hidden states", + ) + + value = self.value_proj(encoder_hidden_states) + if attention_mask is not None: + # we invert the attention_mask + value = value.masked_fill(~attention_mask[..., None], float(0)) + value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads) + sampling_offsets = self.sampling_offsets(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2 + ) + attention_weights = self.attention_weights(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels * self.n_points + ) + attention_weights = F.softmax(attention_weights, -1).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points + ) + # batch_size, num_queries, n_heads, n_levels, n_points, 2 + num_coordinates = reference_points.shape[-1] + if num_coordinates == 2: + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif num_coordinates == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) + else: + raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}") + + output = self.attn( + value, + spatial_shapes, + spatial_shapes_list, + level_start_index, + sampling_locations, + attention_weights, + self.im2col_step, + ) + + output = self.output_proj(output) + + return output, attention_weights + + +@auto_docstring +class PPDocLayoutV2PreTrainedModel(PreTrainedModel): + config: PPDocLayoutV2Config + base_model_prefix = "pp_doclayout_v2" + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = [r"PPDocLayoutV2HybridEncoder", r"PPDocLayoutV2DecoderLayer"] + _supports_sdpa = True + _supports_flash_attn = True + _supports_attention_backend = True + _supports_flex_attn = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, PPDocLayoutV2ForObjectDetection): + if module.model.decoder.class_embed is not None: + for layer in module.model.decoder.class_embed: + prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1) + bias = float(-math.log((1 - prior_prob) / prior_prob)) + init.xavier_uniform_(layer.weight) + init.constant_(layer.bias, bias) + + if module.model.decoder.bbox_embed is not None: + for layer in module.model.decoder.bbox_embed: + init.constant_(layer.layers[-1].weight, 0) + init.constant_(layer.layers[-1].bias, 0) + + elif isinstance(module, PPDocLayoutV2MultiscaleDeformableAttention): + init.constant_(module.sampling_offsets.weight, 0.0) + default_dtype = torch.get_default_dtype() + thetas = torch.arange(module.n_heads, dtype=torch.int64).to(default_dtype) * ( + 2.0 * math.pi / module.n_heads + ) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(module.n_heads, 1, 1, 2) + .repeat(1, module.n_levels, module.n_points, 1) + ) + for i in range(module.n_points): + grid_init[:, :, i, :] *= i + 1 + + init.copy_(module.sampling_offsets.bias, grid_init.view(-1)) + init.constant_(module.attention_weights.weight, 0.0) + init.constant_(module.attention_weights.bias, 0.0) + init.xavier_uniform_(module.value_proj.weight) + init.constant_(module.value_proj.bias, 0.0) + init.xavier_uniform_(module.output_proj.weight) + init.constant_(module.output_proj.bias, 0.0) + + elif isinstance(module, PPDocLayoutV2Model): + prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1) + bias = float(-math.log((1 - prior_prob) / prior_prob)) + init.xavier_uniform_(module.enc_score_head.weight) + init.constant_(module.enc_score_head.bias, bias) + + elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)): + 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, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + + elif isinstance(module, nn.LayerNorm): + init.ones_(module.weight) + init.zeros_(module.bias) + + if hasattr(module, "weight_embedding") and self.config.learn_initial_query: + init.xavier_uniform_(module.weight_embedding.weight) + if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0: + init.xavier_uniform_(module.denoising_class_embed.weight) + if isinstance(module, PPDocLayoutV2TextEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + if isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + init.zeros_(module.weight.data[module.padding_idx]) + if isinstance(module, PPDocLayoutV2PositionRelationEmbedding): + inv_freq, _ = module.compute_default_rope_parameters(module.config, module.inv_freq.device) + module.register_buffer("inv_freq", inv_freq, persistent=False) + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 ReadingOrder Model. This model consists of an encoder and a GlobalPointer head. + It takes layout features as input and outputs logits representing the relative ordering relationships + between elements, which are used to determine the final reading sequence. + """ +) +class PPDocLayoutV2ReadingOrder(PPDocLayoutV2PreTrainedModel): + # Attention is based on LayoutLMv3 (no interface) + _supports_sdpa = False + _supports_flash_attn = False + _supports_attention_backend = False + _supports_flex_attn = False + + def __init__(self, config): + super().__init__(config) + self.embeddings = PPDocLayoutV2TextEmbeddings(config) + self.label_embeddings = nn.Embedding(config.num_classes, config.hidden_size) + self.label_features_projection = nn.Linear(config.hidden_size, config.hidden_size) + self.encoder = PPDocLayoutV2ReadingOrderEncoder(config) + self.relative_head = PPDocLayoutV2GlobalPointer(config) + self.config = config + + self.post_init() + + @auto_docstring + def forward(self, boxes, labels=None, mask=None, **kwargs: Unpack[TransformersKwargs]): + r""" + boxes (`torch.Tensor` of shape `(batch_size, sequence_length, 4)`): + Bounding box coordinates of the detected layout elements **in [0, 1000] scale**. + Format is `[x_min, y_min, x_max, y_max]`. + The tensor usually contains sorted valid boxes followed by zero-padding. + labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + The **remapped** class indices for each layout element. + These are not necessarily the raw detection class IDs, but indices mapped via + `config.class_order` (e.g., mapping text/title/figure to specific reading-order category IDs). + mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Boolean or Binary mask indicating valid detected elements after threshold filtering. + - True: Valid layout element. + - False: Padding/Empty element. + Used to determine the sequence length (`num_pred`) for the pointer mechanism. + """ + device = mask.device + batch_size, seq_len = mask.shape + num_pred = mask.sum(dim=1) + + input_ids = torch.full( + (batch_size, seq_len + 2), self.config.pad_token_id, dtype=torch.long, device=boxes.device + ) + input_ids[:, 0] = self.config.start_token_id + + pred_col_idx = torch.arange(seq_len + 2, device=device).unsqueeze(0) + pred_mask = (pred_col_idx >= 1) & (pred_col_idx <= num_pred.unsqueeze(1)) + input_ids[pred_mask] = self.config.pred_token_id + end_col_indices = num_pred + 1 + input_ids[torch.arange(batch_size, device=device), end_col_indices] = self.config.end_token_id + + pad_box = torch.zeros(size=[boxes.shape[0], 1, boxes.shape[-1]], dtype=boxes.dtype, device=boxes.device) + pad_boxes = torch.cat([pad_box, boxes, pad_box], dim=1) + bbox_embedding = self.embeddings(input_ids=input_ids, bbox=pad_boxes.long()) + + if labels is not None: + label_embs = self.label_embeddings(labels) + label_proj = self.label_features_projection(label_embs) + pad = torch.zeros( + size=[label_proj.shape[0], 1, label_proj.shape[-1]], dtype=label_proj.dtype, device=labels.device + ) + label_proj = torch.cat([pad, label_proj, pad], dim=1) + else: + label_proj = torch.zeros_like(bbox_embedding) + + final_embeddings = bbox_embedding + label_proj + final_embeddings = self.embeddings.norm(final_embeddings) + final_embeddings = self.embeddings.dropout(final_embeddings) + + attention_mask = pred_col_idx < (num_pred + 2).unsqueeze(1) + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=final_embeddings, + attention_mask=attention_mask, + ) + encoder_output = self.encoder(hidden_states=final_embeddings, bbox=pad_boxes, attention_mask=attention_mask) + encoder_output = encoder_output.last_hidden_state + token = encoder_output[:, 1 : 1 + seq_len, :] + read_order_logits = self.relative_head(token) + return read_order_logits + + +@dataclass +@auto_docstring +class PPDocLayoutV2ForObjectDetectionOutput(ModelOutput): + r""" + 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 [`~RTDetrImageProcessor.post_process_object_detection`] to retrieve the + unnormalized (absolute) bounding boxes. + order_logits (`tuple` of `torch.FloatTensor` of shape `(batch_size, num_queries, num_queries)`): + Order logits for all queries. The first dimension of each tensor is the batch size. The second dimension is the number of queries. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, 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 `(batch_size, config.decoder_layers, num_queries, hidden_size)`): + Stacked intermediate hidden states (output of each layer of the decoder). + intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, config.num_labels)`): + Stacked intermediate logits (logits of each layer of the decoder). + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate predicted corners (predicted corners of each layer of the decoder). + initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked initial reference points (initial reference points of each layer of the decoder). + init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Initial reference points sent through the Transformer decoder. + enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the encoder. + enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the encoder. + enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are + picked as region proposals in the first stage. Output of bounding box binary classification (i.e. + foreground and background). + enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the first stage. + denoising_meta_values (`dict`): + Extra dictionary for the denoising related values + """ + + logits: torch.FloatTensor | None = None + pred_boxes: torch.FloatTensor | None = None + order_logits: tuple[torch.FloatTensor] | None = None + last_hidden_state: torch.FloatTensor | None = None + intermediate_hidden_states: torch.FloatTensor | None = None + intermediate_logits: torch.FloatTensor | None = None + intermediate_reference_points: torch.FloatTensor | None = None + intermediate_predicted_corners: torch.FloatTensor | None = None + initial_reference_points: 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 + init_reference_points: tuple[torch.FloatTensor] | None = None + enc_topk_logits: torch.FloatTensor | None = None + enc_topk_bboxes: torch.FloatTensor | None = None + enc_outputs_class: torch.FloatTensor | None = None + enc_outputs_coord_logits: torch.FloatTensor | None = None + denoising_meta_values: dict | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the PP-DocLayoutV2 encoder-decoder model. + """ +) +class PPDocLayoutV2ModelOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, 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 `(batch_size, config.decoder_layers, num_queries, hidden_size)`): + Stacked intermediate hidden states (output of each layer of the decoder). + intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`): + Stacked intermediate logits (logits of each layer of the decoder). + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate predicted corners (predicted corners of each layer of the decoder). + initial_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Initial reference points used for the first decoder layer. + init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Initial reference points sent through the Transformer decoder. + enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): + Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are + picked as region proposals in the encoder stage. Output of bounding box binary classification (i.e. + foreground and background). + enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`): + Logits of predicted bounding boxes coordinates in the encoder stage. + enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are + picked as region proposals in the first stage. Output of bounding box binary classification (i.e. + foreground and background). + enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the first stage. + denoising_meta_values (`dict`): + Extra dictionary for the denoising related values. + """ + + last_hidden_state: torch.FloatTensor | None = None + intermediate_hidden_states: torch.FloatTensor | None = None + intermediate_logits: torch.FloatTensor | None = None + intermediate_reference_points: torch.FloatTensor | None = None + intermediate_predicted_corners: torch.FloatTensor | None = None + initial_reference_points: 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 + init_reference_points: torch.FloatTensor | None = None + enc_topk_logits: torch.FloatTensor | None = None + enc_topk_bboxes: torch.FloatTensor | None = None + enc_outputs_class: torch.FloatTensor | None = None + enc_outputs_coord_logits: torch.FloatTensor | None = None + denoising_meta_values: dict | None = None + + +class PPDocLayoutV2MLPPredictionHead(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 + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the PPDocLayoutV2Decoder. This class adds two attributes to + BaseModelOutputWithCrossAttentions, namely: + - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer) + - a stacked tensor of intermediate reference points. + """ +) +class PPDocLayoutV2DecoderOutput(ModelOutput): + r""" + intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`): + Stacked intermediate hidden states (output of each layer of the decoder). + intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`): + Stacked intermediate logits (logits of each layer of the decoder). + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate predicted corners (predicted corners of each layer of the decoder). + initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked initial reference points (initial reference points of each layer of the decoder). + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, + used to compute the weighted average in the cross-attention heads. + """ + + last_hidden_state: torch.FloatTensor | None = None + intermediate_hidden_states: torch.FloatTensor | None = None + intermediate_logits: torch.FloatTensor | None = None + intermediate_reference_points: torch.FloatTensor | None = None + intermediate_predicted_corners: torch.FloatTensor | None = None + initial_reference_points: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + cross_attentions: tuple[torch.FloatTensor] | None = None + + +class PPDocLayoutV2MLP(nn.Module): + def __init__( + self, config: PPDocLayoutV2Config, hidden_size: int, intermediate_size: int, activation_function: str + ): + super().__init__() + self.fc1 = nn.Linear(hidden_size, intermediate_size) + self.fc2 = nn.Linear(intermediate_size, hidden_size) + self.activation_fn = ACT2FN[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 PPDocLayoutV2FrozenBatchNorm2d(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 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 PPDocLayoutV2SelfAttention(nn.Module): + """ + Multi-headed self-attention from 'Attention Is All You Need' paper. + + In PP_DOCLAYOUT_V2, position embeddings are added to both queries and keys (but not values) in self-attention. + """ + + def __init__( + self, + config: PPDocLayoutV2Config, + 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 + + +def replace_batch_norm(model): + r""" + Recursively replace all `torch.nn.BatchNorm2d` with `PPDocLayoutV2FrozenBatchNorm2d`. + + Args: + model (torch.nn.Module): + input model + """ + for name, module in model.named_children(): + if isinstance(module, nn.BatchNorm2d): + new_module = PPDocLayoutV2FrozenBatchNorm2d(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 PPDocLayoutV2ConvEncoder(nn.Module): + """ + Convolutional backbone using the modeling_pp_doclayout_v2_resnet.py. + + nn.BatchNorm2d layers are replaced by PPDocLayoutV2FrozenBatchNorm2d as defined above. + https://github.com/lyuwenyu/RT-DETR/blob/main/PPDocLayoutV2_pytorch/src/nn/backbone/presnet.py#L142 + """ + + def __init__(self, config): + super().__init__() + + backbone = load_backbone(config) + + if config.freeze_backbone_batch_norms: + # replace batch norm by frozen batch norm + with torch.no_grad(): + replace_batch_norm(backbone) + self.model = backbone + self.intermediate_channel_sizes = self.model.channels + + 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).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 PPDocLayoutV2ConvNormLayer(nn.Module): + def __init__(self, config, in_channels, out_channels, kernel_size, stride, padding=None, activation=None): + super().__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding=(kernel_size - 1) // 2 if padding is None else padding, + bias=False, + ) + self.norm = nn.BatchNorm2d(out_channels, config.batch_norm_eps) + self.activation = nn.Identity() if activation is None else ACT2CLS[activation]() + + def forward(self, hidden_state): + hidden_state = self.conv(hidden_state) + hidden_state = self.norm(hidden_state) + hidden_state = self.activation(hidden_state) + return hidden_state + + +class PPDocLayoutV2EncoderLayer(nn.Module): + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + self.normalize_before = config.normalize_before + self.hidden_size = config.encoder_hidden_dim + + # self-attention + self.self_attn = PPDocLayoutV2SelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.num_attention_heads, + dropout=config.dropout, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + self.dropout = config.dropout + self.mlp = PPDocLayoutV2MLP( + config, self.hidden_size, config.encoder_ffn_dim, config.encoder_activation_function + ) + self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + + 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 + if self.normalize_before: + hidden_states = self.self_attn_layer_norm(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 + if not self.normalize_before: + hidden_states = self.self_attn_layer_norm(hidden_states) + + if self.normalize_before: + hidden_states = self.final_layer_norm(hidden_states) + residual = hidden_states + + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + if not self.normalize_before: + 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 PPDocLayoutV2RepVggBlock(nn.Module): + """ + RepVGG architecture block introduced by the work "RepVGG: Making VGG-style ConvNets Great Again". + """ + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + + activation = config.activation_function + hidden_channels = int(config.encoder_hidden_dim * config.hidden_expansion) + self.conv1 = PPDocLayoutV2ConvNormLayer(config, hidden_channels, hidden_channels, 3, 1, padding=1) + self.conv2 = PPDocLayoutV2ConvNormLayer(config, hidden_channels, hidden_channels, 1, 1, padding=0) + self.activation = nn.Identity() if activation is None else ACT2CLS[activation]() + + def forward(self, x): + y = self.conv1(x) + self.conv2(x) + return self.activation(y) + + +class PPDocLayoutV2CSPRepLayer(nn.Module): + """ + Cross Stage Partial (CSP) network layer with RepVGG blocks. + """ + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + + in_channels = config.encoder_hidden_dim * 2 + out_channels = config.encoder_hidden_dim + num_blocks = 3 + activation = config.activation_function + + hidden_channels = int(out_channels * config.hidden_expansion) + self.conv1 = PPDocLayoutV2ConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation) + self.conv2 = PPDocLayoutV2ConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation) + self.bottlenecks = nn.Sequential(*[PPDocLayoutV2RepVggBlock(config) for _ in range(num_blocks)]) + if hidden_channels != out_channels: + self.conv3 = PPDocLayoutV2ConvNormLayer(config, hidden_channels, out_channels, 1, 1, activation=activation) + else: + self.conv3 = nn.Identity() + + def forward(self, hidden_state): + hidden_state_1 = self.conv1(hidden_state) + hidden_state_1 = self.bottlenecks(hidden_state_1) + hidden_state_2 = self.conv2(hidden_state) + return self.conv3(hidden_state_1 + hidden_state_2) + + +class PPDocLayoutV2DecoderLayer(nn.Module): + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + self.hidden_size = config.d_model + + # self-attention + self.self_attn = PPDocLayoutV2SelfAttention( + 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, eps=config.layer_norm_eps) + # cross-attention + self.encoder_attn = PPDocLayoutV2MultiscaleDeformableAttention( + config, + num_heads=config.decoder_attention_heads, + n_points=config.decoder_n_points, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + # feedforward neural networks + self.mlp = PPDocLayoutV2MLP( + config, self.hidden_size, config.decoder_ffn_dim, config.decoder_activation_function + ) + self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + object_queries_position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes=None, + spatial_shapes_list=None, + level_start_index=None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: 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)`. + object_queries_position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings for the object query slots. These are added to both queries and keys + in the self-attention layer (not values). + reference_points (`torch.FloatTensor`, *optional*): + Reference points. + spatial_shapes (`torch.LongTensor`, *optional*): + Spatial shapes. + level_start_index (`torch.LongTensor`, *optional*): + Level start index. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(batch, seq_len, hidden_size)` + 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. + """ + residual = hidden_states + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=encoder_attention_mask, + position_embeddings=object_queries_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 + + # Cross-Attention + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + position_embeddings=object_queries_position_embeddings, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + **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 PPDocLayoutV2SinePositionEmbedding(nn.Module): + """ + 2D sinusoidal position embedding used in RT-DETR hybrid encoder. + """ + + def __init__(self, embed_dim: int = 256, temperature: int = 10000): + super().__init__() + self.embed_dim = embed_dim + self.temperature = temperature + + @compile_compatible_method_lru_cache(maxsize=32) + def forward( + self, + width: int, + height: int, + device: torch.device | str, + dtype: torch.dtype, + ) -> torch.Tensor: + """ + Generate 2D sinusoidal position embeddings. + + Returns: + Position embeddings of shape (1, height*width, embed_dim) + """ + grid_w = torch.arange(torch_int(width), device=device).to(dtype) + grid_h = torch.arange(torch_int(height), device=device).to(dtype) + grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing="xy") + if self.embed_dim % 4 != 0: + raise ValueError("Embed dimension must be divisible by 4 for 2D sin-cos position embedding") + pos_dim = self.embed_dim // 4 + omega = torch.arange(pos_dim, device=device).to(dtype) / pos_dim + omega = 1.0 / (self.temperature**omega) + + out_w = grid_w.flatten()[..., None] @ omega[None] + out_h = grid_h.flatten()[..., None] @ omega[None] + + return torch.concat([out_h.sin(), out_h.cos(), out_w.sin(), out_w.cos()], dim=1)[None, :, :] + + +class PPDocLayoutV2AIFILayer(nn.Module): + """ + AIFI (Attention-based Intra-scale Feature Interaction) layer used in RT-DETR hybrid encoder. + """ + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + self.config = config + self.encoder_hidden_dim = config.encoder_hidden_dim + self.eval_size = config.eval_size + + self.position_embedding = PPDocLayoutV2SinePositionEmbedding( + embed_dim=self.encoder_hidden_dim, + temperature=config.positional_encoding_temperature, + ) + self.layers = nn.ModuleList([PPDocLayoutV2EncoderLayer(config) for _ in range(config.encoder_layers)]) + + def forward( + self, + hidden_states: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`): + Feature map to process. + """ + batch_size = hidden_states.shape[0] + height, width = hidden_states.shape[2:] + + hidden_states = hidden_states.flatten(2).permute(0, 2, 1) + + if self.training or self.eval_size is None: + pos_embed = self.position_embedding( + width=width, + height=height, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + else: + pos_embed = None + + for layer in self.layers: + hidden_states = layer( + hidden_states, + attention_mask=None, + spatial_position_embeddings=pos_embed, + **kwargs, + ) + + hidden_states = ( + hidden_states.permute(0, 2, 1).reshape(batch_size, self.encoder_hidden_dim, height, width).contiguous() + ) + + return hidden_states + + +class PPDocLayoutV2HybridEncoder(PPDocLayoutV2PreTrainedModel): + """ + Hybrid encoder consisting of AIFI (Attention-based Intra-scale Feature Interaction) layers, + a top-down Feature Pyramid Network (FPN) and a bottom-up Path Aggregation Network (PAN). + More details on the paper: https://huggingface.co/papers/2304.08069 + + Args: + config: PPDocLayoutV2Config + """ + + _can_record_outputs = { + "hidden_states": PPDocLayoutV2AIFILayer, + "attentions": PPDocLayoutV2SelfAttention, + } + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__(config) + self.config = config + self.in_channels = config.encoder_in_channels + self.feat_strides = config.feat_strides + self.encoder_hidden_dim = config.encoder_hidden_dim + self.encode_proj_layers = config.encode_proj_layers + self.positional_encoding_temperature = config.positional_encoding_temperature + self.eval_size = config.eval_size + self.out_channels = [self.encoder_hidden_dim for _ in self.in_channels] + self.out_strides = self.feat_strides + self.num_fpn_stages = len(self.in_channels) - 1 + self.num_pan_stages = len(self.in_channels) - 1 + + # AIFI (Attention-based Intra-scale Feature Interaction) layers + self.aifi = nn.ModuleList([PPDocLayoutV2AIFILayer(config) for _ in range(len(self.encode_proj_layers))]) + + # top-down FPN + self.lateral_convs = nn.ModuleList() + self.fpn_blocks = nn.ModuleList() + for _ in range(self.num_fpn_stages): + lateral_conv = PPDocLayoutV2ConvNormLayer( + config, + in_channels=self.encoder_hidden_dim, + out_channels=self.encoder_hidden_dim, + kernel_size=1, + stride=1, + activation=config.activation_function, + ) + fpn_block = PPDocLayoutV2CSPRepLayer(config) + self.lateral_convs.append(lateral_conv) + self.fpn_blocks.append(fpn_block) + + # bottom-up PAN + self.downsample_convs = nn.ModuleList() + self.pan_blocks = nn.ModuleList() + for _ in range(self.num_pan_stages): + downsample_conv = PPDocLayoutV2ConvNormLayer( + config, + in_channels=self.encoder_hidden_dim, + out_channels=self.encoder_hidden_dim, + kernel_size=3, + stride=2, + activation=config.activation_function, + ) + pan_block = PPDocLayoutV2CSPRepLayer(config) + self.downsample_convs.append(downsample_conv) + self.pan_blocks.append(pan_block) + + self.post_init() + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + def forward( + self, + inputs_embeds=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. + """ + feature_maps = inputs_embeds + + # AIFI: Apply transformer encoder to specified feature levels + if self.config.encoder_layers > 0: + for i, enc_ind in enumerate(self.encode_proj_layers): + feature_maps[enc_ind] = self.aifi[i](feature_maps[enc_ind], **kwargs) + + # top-down FPN + fpn_feature_maps = [feature_maps[-1]] + for idx, (lateral_conv, fpn_block) in enumerate(zip(self.lateral_convs, self.fpn_blocks)): + backbone_feature_map = feature_maps[self.num_fpn_stages - idx - 1] + top_fpn_feature_map = fpn_feature_maps[-1] + # apply lateral block + top_fpn_feature_map = lateral_conv(top_fpn_feature_map) + fpn_feature_maps[-1] = top_fpn_feature_map + # apply fpn block + top_fpn_feature_map = F.interpolate(top_fpn_feature_map, scale_factor=2.0, mode="nearest") + fused_feature_map = torch.concat([top_fpn_feature_map, backbone_feature_map], dim=1) + new_fpn_feature_map = fpn_block(fused_feature_map) + fpn_feature_maps.append(new_fpn_feature_map) + + fpn_feature_maps.reverse() + + # bottom-up PAN + pan_feature_maps = [fpn_feature_maps[0]] + for idx, (downsample_conv, pan_block) in enumerate(zip(self.downsample_convs, self.pan_blocks)): + top_pan_feature_map = pan_feature_maps[-1] + fpn_feature_map = fpn_feature_maps[idx + 1] + downsampled_feature_map = downsample_conv(top_pan_feature_map) + fused_feature_map = torch.concat([downsampled_feature_map, fpn_feature_map], dim=1) + new_pan_feature_map = pan_block(fused_feature_map) + pan_feature_maps.append(new_pan_feature_map) + + return BaseModelOutput(last_hidden_state=pan_feature_maps) + + +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) + + +class PPDocLayoutV2Decoder(PPDocLayoutV2PreTrainedModel): + _can_record_outputs = { + "hidden_states": PPDocLayoutV2DecoderLayer, + "attentions": PPDocLayoutV2SelfAttention, + "cross_attentions": PPDocLayoutV2MultiscaleDeformableAttention, + } + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__(config) + + self.dropout = config.dropout + self.layers = nn.ModuleList([PPDocLayoutV2DecoderLayer(config) for _ in range(config.decoder_layers)]) + self.query_pos_head = PPDocLayoutV2MLPPredictionHead(4, 2 * config.d_model, config.d_model, num_layers=2) + + # hack implementation for iterative bounding box refinement and two-stage Deformable DETR + self.bbox_embed = None + self.class_embed = None + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + inputs_embeds=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + reference_points=None, + spatial_shapes=None, + spatial_shapes_list=None, + level_start_index=None, + **kwargs: Unpack[TransformersKwargs], + ): + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): + The query embeddings that are passed into the decoder. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, 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, 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**). + reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*): + Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area. + spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`): + Spatial shapes of the feature maps. + level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`, *optional*): + Indexes for the start of each feature level. In range `[0, sequence_length]`. + """ + if inputs_embeds is not None: + hidden_states = inputs_embeds + + # decoder layers + intermediate = () + intermediate_reference_points = () + intermediate_logits = () + + reference_points = F.sigmoid(reference_points) + + # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/PPDocLayoutV2_pytorch/src/zoo/PPDocLayoutV2/PPDocLayoutV2_decoder.py#L252 + for idx, decoder_layer in enumerate(self.layers): + reference_points_input = reference_points.unsqueeze(2) + object_queries_position_embeddings = self.query_pos_head(reference_points) + + hidden_states = decoder_layer( + hidden_states, + object_queries_position_embeddings=object_queries_position_embeddings, + encoder_hidden_states=encoder_hidden_states, + reference_points=reference_points_input, + spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + encoder_attention_mask=encoder_attention_mask, + **kwargs, + ) + + # hack implementation for iterative bounding box refinement + if self.bbox_embed is not None: + predicted_corners = self.bbox_embed[idx](hidden_states) + new_reference_points = F.sigmoid(predicted_corners + inverse_sigmoid(reference_points)) + reference_points = new_reference_points.detach() + + intermediate += (hidden_states,) + intermediate_reference_points += ( + (new_reference_points,) if self.bbox_embed is not None else (reference_points,) + ) + + if self.class_embed is not None: + logits = self.class_embed[idx](hidden_states) + intermediate_logits += (logits,) + + # Keep batch_size as first dimension + intermediate = torch.stack(intermediate, dim=1) + intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1) + if self.class_embed is not None: + intermediate_logits = torch.stack(intermediate_logits, dim=1) + + return PPDocLayoutV2DecoderOutput( + last_hidden_state=hidden_states, + intermediate_hidden_states=intermediate, + intermediate_logits=intermediate_logits, + intermediate_reference_points=intermediate_reference_points, + ) + + +def get_contrastive_denoising_training_group( + targets, + num_classes, + num_queries, + class_embed, + num_denoising_queries=100, + label_noise_ratio=0.5, + box_noise_scale=1.0, +): + """ + Creates a contrastive denoising training group using ground-truth samples. It adds noise to labels and boxes. + + Args: + targets (`list[dict]`): + The target objects, each containing 'class_labels' and 'boxes' for objects in an image. + num_classes (`int`): + Total number of classes in the dataset. + num_queries (`int`): + Number of query slots in the transformer. + class_embed (`callable`): + A function or a model layer to embed class labels. + num_denoising_queries (`int`, *optional*, defaults to 100): + Number of denoising queries. + label_noise_ratio (`float`, *optional*, defaults to 0.5): + Ratio of noise applied to labels. + box_noise_scale (`float`, *optional*, defaults to 1.0): + Scale of noise applied to bounding boxes. + Returns: + `tuple` comprising various elements: + - **input_query_class** (`torch.FloatTensor`) -- + Class queries with applied label noise. + - **input_query_bbox** (`torch.FloatTensor`) -- + Bounding box queries with applied box noise. + - **attn_mask** (`torch.FloatTensor`) -- + Attention mask for separating denoising and reconstruction queries. + - **denoising_meta_values** (`dict`) -- + Metadata including denoising positive indices, number of groups, and split sizes. + """ + + if num_denoising_queries <= 0: + return None, None, None, None + + num_ground_truths = [len(t["class_labels"]) for t in targets] + device = targets[0]["class_labels"].device + + max_gt_num = max(num_ground_truths) + if max_gt_num == 0: + return None, None, None, None + + num_groups_denoising_queries = num_denoising_queries // max_gt_num + num_groups_denoising_queries = 1 if num_groups_denoising_queries == 0 else num_groups_denoising_queries + # pad gt to max_num of a batch + batch_size = len(num_ground_truths) + + input_query_class = torch.full([batch_size, max_gt_num], num_classes, dtype=torch.int32, device=device) + input_query_bbox = torch.zeros([batch_size, max_gt_num, 4], device=device) + pad_gt_mask = torch.zeros([batch_size, max_gt_num], dtype=torch.bool, device=device) + + for i in range(batch_size): + num_gt = num_ground_truths[i] + if num_gt > 0: + input_query_class[i, :num_gt] = targets[i]["class_labels"] + input_query_bbox[i, :num_gt] = targets[i]["boxes"] + pad_gt_mask[i, :num_gt] = 1 + # each group has positive and negative queries. + input_query_class = input_query_class.tile([1, 2 * num_groups_denoising_queries]) + input_query_bbox = input_query_bbox.tile([1, 2 * num_groups_denoising_queries, 1]) + pad_gt_mask = pad_gt_mask.tile([1, 2 * num_groups_denoising_queries]) + # positive and negative mask + negative_gt_mask = torch.zeros([batch_size, max_gt_num * 2, 1], device=device) + negative_gt_mask[:, max_gt_num:] = 1 + negative_gt_mask = negative_gt_mask.tile([1, num_groups_denoising_queries, 1]) + positive_gt_mask = 1 - negative_gt_mask + # contrastive denoising training positive index + positive_gt_mask = positive_gt_mask.squeeze(-1) * pad_gt_mask + denoise_positive_idx = torch.nonzero(positive_gt_mask)[:, 1] + denoise_positive_idx = torch.split( + denoise_positive_idx, [n * num_groups_denoising_queries for n in num_ground_truths] + ) + # total denoising queries + num_denoising_queries = torch_int(max_gt_num * 2 * num_groups_denoising_queries) + + if label_noise_ratio > 0: + mask = torch.rand_like(input_query_class, dtype=torch.float) < (label_noise_ratio * 0.5) + # randomly put a new one here + new_label = torch.randint_like(mask, 0, num_classes, dtype=input_query_class.dtype) + input_query_class = torch.where(mask & pad_gt_mask, new_label, input_query_class) + + if box_noise_scale > 0: + known_bbox = center_to_corners_format(input_query_bbox) + diff = torch.tile(input_query_bbox[..., 2:] * 0.5, [1, 1, 2]) * box_noise_scale + rand_sign = torch.randint_like(input_query_bbox, 0, 2) * 2.0 - 1.0 + rand_part = torch.rand_like(input_query_bbox) + rand_part = (rand_part + 1.0) * negative_gt_mask + rand_part * (1 - negative_gt_mask) + rand_part *= rand_sign + known_bbox += rand_part * diff + known_bbox.clip_(min=0.0, max=1.0) + input_query_bbox = corners_to_center_format(known_bbox) + input_query_bbox = inverse_sigmoid(input_query_bbox) + + input_query_class = class_embed(input_query_class) + + target_size = num_denoising_queries + num_queries + attn_mask = torch.full([target_size, target_size], 0, dtype=torch.float, device=device) + # match query cannot see the reconstruction + attn_mask[num_denoising_queries:, :num_denoising_queries] = -torch.inf + + # reconstructions cannot see each other + for i in range(num_groups_denoising_queries): + idx_block_start = max_gt_num * 2 * i + idx_block_end = max_gt_num * 2 * (i + 1) + attn_mask[idx_block_start:idx_block_end, :idx_block_start] = -torch.inf + attn_mask[idx_block_start:idx_block_end, idx_block_end:num_denoising_queries] = -torch.inf + + denoising_meta_values = { + "dn_positive_idx": denoise_positive_idx, + "dn_num_group": num_groups_denoising_queries, + "dn_num_split": [num_denoising_queries, num_queries], + } + + return input_query_class, input_query_bbox, attn_mask, denoising_meta_values + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 Model (consisting of a backbone and encoder-decoder) outputting raw hidden states without any head on top. + """ +) +class PPDocLayoutV2Model(PPDocLayoutV2PreTrainedModel): + def __init__(self, config: PPDocLayoutV2Config): + super().__init__(config) + + # Create backbone + self.backbone = PPDocLayoutV2ConvEncoder(config) + intermediate_channel_sizes = self.backbone.intermediate_channel_sizes + + # Create encoder input projection layers + # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/PPDocLayoutV2_pytorch/src/zoo/PPDocLayoutV2/hybrid_encoder.py#L212 + num_backbone_outs = len(intermediate_channel_sizes) + encoder_input_proj_list = [] + for i in range(num_backbone_outs): + in_channels = intermediate_channel_sizes[i] + encoder_input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, config.encoder_hidden_dim, kernel_size=1, bias=False), + nn.BatchNorm2d(config.encoder_hidden_dim), + ) + ) + self.encoder_input_proj = nn.ModuleList(encoder_input_proj_list) + + # Create encoder + self.encoder = PPDocLayoutV2HybridEncoder(config) + + # denoising part + if config.num_denoising > 0: + self.denoising_class_embed = nn.Embedding( + config.num_labels + 1, config.d_model, padding_idx=config.num_labels + ) + + # decoder embedding + if config.learn_initial_query: + self.weight_embedding = nn.Embedding(config.num_queries, config.d_model) + + # encoder head + self.enc_output = nn.Sequential( + nn.Linear(config.d_model, config.d_model), + nn.LayerNorm(config.d_model, eps=config.layer_norm_eps), + ) + self.enc_score_head = nn.Linear(config.d_model, config.num_labels) + self.enc_bbox_head = PPDocLayoutV2MLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) + + # init encoder output anchors and valid_mask + if config.anchor_image_size: + self.anchors, self.valid_mask = self.generate_anchors(dtype=self.dtype) + + # Create decoder input projection layers + # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/PPDocLayoutV2_pytorch/src/zoo/PPDocLayoutV2/PPDocLayoutV2_decoder.py#L412 + num_backbone_outs = len(config.decoder_in_channels) + decoder_input_proj_list = [] + for i in range(num_backbone_outs): + in_channels = config.decoder_in_channels[i] + decoder_input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, config.d_model, kernel_size=1, bias=False), + nn.BatchNorm2d(config.d_model, config.batch_norm_eps), + ) + ) + for _ in range(config.num_feature_levels - num_backbone_outs): + decoder_input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(config.d_model, config.batch_norm_eps), + ) + ) + in_channels = config.d_model + self.decoder_input_proj = nn.ModuleList(decoder_input_proj_list) + + # decoder + self.decoder = PPDocLayoutV2Decoder(config) + self.denoising_class_embed = nn.Embedding(config.num_labels, config.d_model) + + self.post_init() + + def freeze_backbone(self): + for param in self.backbone.parameters(): + param.requires_grad_(False) + + def unfreeze_backbone(self): + for param in self.backbone.parameters(): + param.requires_grad_(True) + + @compile_compatible_method_lru_cache(maxsize=32) + def generate_anchors(self, spatial_shapes=None, grid_size=0.05, device="cpu", dtype=torch.float32): + if spatial_shapes is None: + spatial_shapes = [ + [int(self.config.anchor_image_size[0] / s), int(self.config.anchor_image_size[1] / s)] + for s in self.config.feat_strides + ] + anchors = [] + for level, (height, width) in enumerate(spatial_shapes): + grid_y, grid_x = torch.meshgrid( + torch.arange(end=height, device=device).to(dtype), + torch.arange(end=width, device=device).to(dtype), + indexing="ij", + ) + grid_xy = torch.stack([grid_x, grid_y], -1) + grid_xy = grid_xy.unsqueeze(0) + 0.5 + grid_xy[..., 0] /= width + grid_xy[..., 1] /= height + wh = torch.ones_like(grid_xy) * grid_size * (2.0**level) + anchors.append(torch.concat([grid_xy, wh], -1).reshape(-1, height * width, 4)) + # define the valid range for anchor coordinates + eps = 1e-2 + anchors = torch.concat(anchors, 1) + valid_mask = ((anchors > eps) * (anchors < 1 - eps)).all(-1, keepdim=True) + anchors = torch.log(anchors / (1 - anchors)) + anchors = torch.where(valid_mask, anchors, torch.tensor(torch.finfo(dtype).max, dtype=dtype, device=device)) + + return anchors, valid_mask + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: list[dict] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | PPDocLayoutV2ModelOutput: + r""" + 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. + 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, PPDocLayoutV2Model + >>> 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("PekingU/PPDocLayoutV2_r50vd") + >>> model = PPDocLayoutV2Model.from_pretrained("PekingU/PPDocLayoutV2_r50vd") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 300, 256] + ```""" + if pixel_values is None and inputs_embeds is None: + raise ValueError("You have to specify either pixel_values or inputs_embeds") + + if inputs_embeds is None: + 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) + features = self.backbone(pixel_values, pixel_mask) + proj_feats = [self.encoder_input_proj[level](source) for level, (source, mask) in enumerate(features)] + else: + batch_size = inputs_embeds.shape[0] + device = inputs_embeds.device + proj_feats = inputs_embeds + + if encoder_outputs is None: + encoder_outputs = self.encoder( + proj_feats, + **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, + ) + + # Equivalent to def _get_encoder_input + # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/PPDocLayoutV2_pytorch/src/zoo/PPDocLayoutV2/PPDocLayoutV2_decoder.py#L412 + sources = [] + for level, source in enumerate(encoder_outputs.last_hidden_state): + sources.append(self.decoder_input_proj[level](source)) + + # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage + if self.config.num_feature_levels > len(sources): + _len_sources = len(sources) + sources.append(self.decoder_input_proj[_len_sources](encoder_outputs.last_hidden_state)[-1]) + for i in range(_len_sources + 1, self.config.num_feature_levels): + sources.append(self.decoder_input_proj[i](encoder_outputs.last_hidden_state[-1])) + + # Prepare encoder inputs (by flattening) + source_flatten = [] + spatial_shapes_list = [] + spatial_shapes = torch.empty((len(sources), 2), device=device, dtype=torch.long) + for level, source in enumerate(sources): + height, width = source.shape[-2:] + spatial_shapes[level, 0] = height + spatial_shapes[level, 1] = width + spatial_shapes_list.append((height, width)) + source = source.flatten(2).transpose(1, 2) + source_flatten.append(source) + source_flatten = torch.cat(source_flatten, 1) + level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) + + # prepare denoising training + if self.training and self.config.num_denoising > 0 and labels is not None: + ( + denoising_class, + denoising_bbox_unact, + attention_mask, + denoising_meta_values, + ) = get_contrastive_denoising_training_group( + targets=labels, + num_classes=self.config.num_labels, + num_queries=self.config.num_queries, + class_embed=self.denoising_class_embed, + num_denoising_queries=self.config.num_denoising, + label_noise_ratio=self.config.label_noise_ratio, + box_noise_scale=self.config.box_noise_scale, + ) + else: + denoising_class, denoising_bbox_unact, attention_mask, denoising_meta_values = None, None, None, None + + batch_size = len(source_flatten) + device = source_flatten.device + dtype = source_flatten.dtype + + # prepare input for decoder + if self.training or self.config.anchor_image_size is None: + # Pass spatial_shapes as tuple to make it hashable and make sure + # lru_cache is working for generate_anchors() + spatial_shapes_tuple = tuple(spatial_shapes_list) + anchors, valid_mask = self.generate_anchors(spatial_shapes_tuple, device=device, dtype=dtype) + else: + anchors, valid_mask = self.anchors, self.valid_mask + anchors, valid_mask = anchors.to(device, dtype), valid_mask.to(device, dtype) + + # use the valid_mask to selectively retain values in the feature map where the mask is `True` + memory = valid_mask.to(source_flatten.dtype) * source_flatten + + output_memory = self.enc_output(memory) + + enc_outputs_class = self.enc_score_head(output_memory) + enc_outputs_coord_logits = self.enc_bbox_head(output_memory) + anchors + + _, topk_ind = torch.topk(enc_outputs_class.max(-1).values, self.config.num_queries, dim=1) + + reference_points_unact = enc_outputs_coord_logits.gather( + dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_coord_logits.shape[-1]) + ) + + enc_topk_bboxes = F.sigmoid(reference_points_unact) + if denoising_bbox_unact is not None: + reference_points_unact = torch.concat([denoising_bbox_unact, reference_points_unact], 1) + + enc_topk_logits = enc_outputs_class.gather( + dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_class.shape[-1]) + ) + + # extract region features + if self.config.learn_initial_query: + target = self.weight_embedding.tile([batch_size, 1, 1]) + else: + target = output_memory.gather(dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1])) + target = target.detach() + + if denoising_class is not None: + target = torch.concat([denoising_class, target], 1) + + init_reference_points = reference_points_unact.detach() + + # decoder + decoder_outputs = self.decoder( + inputs_embeds=target, + encoder_hidden_states=source_flatten, + encoder_attention_mask=attention_mask, + reference_points=init_reference_points, + spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, + level_start_index=level_start_index, + **kwargs, + ) + + return PPDocLayoutV2ModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, + intermediate_logits=decoder_outputs.intermediate_logits, + intermediate_reference_points=decoder_outputs.intermediate_reference_points, + intermediate_predicted_corners=decoder_outputs.intermediate_predicted_corners, + initial_reference_points=decoder_outputs.initial_reference_points, + 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, + init_reference_points=init_reference_points, + enc_topk_logits=enc_topk_logits, + enc_topk_bboxes=enc_topk_bboxes, + enc_outputs_class=enc_outputs_class, + enc_outputs_coord_logits=enc_outputs_coord_logits, + denoising_meta_values=denoising_meta_values, + ) + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 Model (consisting of a backbone and encoder-decoder) outputting bounding boxes, logits and order_logits to be further + decoded into scores, classes and their reading order. + """ +) +class PPDocLayoutV2ForObjectDetection(PPDocLayoutV2PreTrainedModel): + # When using clones, all layers > 0 will be clones, but layer 0 *is* required + # We can't initialize the model on meta device as some weights are modified during the initialization + _no_split_modules = None + _keys_to_ignore_on_load_missing = ["num_batches_tracked", "rel_pos_y_bias", "rel_pos_x_bias"] + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__(config) + self.model = PPDocLayoutV2Model(config) + num_pred = config.decoder_layers + self.model.decoder.class_embed = nn.ModuleList( + [torch.nn.Linear(config.d_model, config.num_labels) for _ in range(num_pred)] + ) + self.model.decoder.bbox_embed = nn.ModuleList( + [PPDocLayoutV2MLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) for _ in range(num_pred)] + ) + + self.reading_order = PPDocLayoutV2ReadingOrder(config.reading_order_config) + self.num_queries = config.num_queries + self.config = config + # if two-stage, the last class_embed and bbox_embed is for region proposal generation + self.post_init() + + def _set_aux_loss(self, outputs_class, outputs_coord): + return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class, outputs_coord)] + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_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], + ) -> tuple[torch.FloatTensor] | PPDocLayoutV2ForObjectDetectionOutput: + r""" + 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 AutoModelForObjectDetection, AutoImageProcessor + >>> from PIL import Image + >>> import requests + >>> import torch + + >>> url = "https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/layout_demo.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> model_path = "PaddlePaddle/PP-DocLayoutV2_safetensors" + >>> image_processor = AutoImageProcessor.from_pretrained(model_path) + >>> model = AutoModelForObjectDetection.from_pretrained(model_path) + + >>> # prepare image for the model + >>> inputs = image_processor(images=[image], return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax) + >>> results = image_processor.post_process_object_detection(outputs, target_sizes=torch.tensor([image.size[::-1]])) + + >>> # print outputs + >>> for result in results: + ... for idx, (score, label_id, box) in enumerate(zip(result["scores"], result["labels"], result["boxes"])): + ... score, label = score.item(), label_id.item() + ... box = [round(i, 2) for i in box.tolist()] + ... print(f"Order {idx + 1}: {model.config.id2label[label]}: {score:.2f} {box}") + Order 1: text: 0.99 [335.39, 184.26, 896.49, 654.48] + Order 2: paragraph_title: 0.97 [337.14, 683.49, 869.42, 798.27] + Order 3: text: 0.99 [335.71, 843.04, 891.17, 1454.15] + Order 4: text: 0.99 [920.42, 185.53, 1476.39, 464.25] + Order 5: text: 0.98 [920.62, 483.75, 1480.52, 765.34] + Order 6: text: 0.98 [920.58, 846.75, 1481.94, 1220.53] + Order 7: text: 0.97 [921.12, 1239.27, 1468.87, 1377.33] + Order 8: footnote: 0.82 [334.58, 1614.67, 1483.84, 1731.61] + Order 9: text: 0.51 [334.58, 1614.67, 1483.84, 1731.61] + Order 10: footnote: 0.83 [334.7, 1757.26, 1471.07, 1845.33] + Order 11: text: 0.87 [336.65, 1910.28, 661.33, 1939.92] + Order 12: footnote: 0.95 [336.16, 2114.52, 1450.28, 2171.74] + Order 13: number: 0.87 [106.04, 2257.37, 136.05, 2281.98] + Order 14: footer: 0.93 [338.6, 2255.94, 985.67, 2283.57] + ```""" + outputs = self.model( + pixel_values, + pixel_mask=pixel_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + labels=labels, + **kwargs, + ) + + intermediate_reference_points = outputs.intermediate_reference_points + intermediate_logits = outputs.intermediate_logits + raw_bboxes = intermediate_reference_points[:, -1] + logits = intermediate_logits[:, -1] + + box_centers, box_sizes = raw_bboxes.split(2, dim=-1) + bboxes = torch.cat([box_centers - 0.5 * box_sizes, box_centers + 0.5 * box_sizes], dim=-1) * 1000 + bboxes = bboxes.clamp_(0.0, 1000.0) + + max_logits, class_ids = logits.max(dim=-1) + max_probs = max_logits.sigmoid() + + class_thresholds = torch.tensor(self.config.class_thresholds, dtype=torch.float32, device=logits.device) + thresholds = class_thresholds[class_ids] + mask = max_probs >= thresholds + + indices = torch.argsort(mask.to(torch.int8), dim=1, descending=True) + + sorted_class_ids = torch.take_along_dim(class_ids, indices, dim=1) + sorted_boxes = torch.take_along_dim(bboxes, indices[..., None].expand(-1, -1, 4), dim=1) + pred_boxes = torch.take_along_dim(raw_bboxes, indices[..., None].expand(-1, -1, 4), dim=1) + logits = torch.take_along_dim(logits, indices[..., None].expand(-1, -1, logits.size(-1)), dim=1) + + sorted_mask = torch.take_along_dim(mask, indices, dim=1) + + pad_boxes = torch.where(sorted_mask[..., None], sorted_boxes, torch.zeros_like(sorted_boxes)) + pad_class_ids = torch.where(sorted_mask, sorted_class_ids, torch.zeros_like(sorted_class_ids)) + + class_order = torch.tensor(self.config.class_order, dtype=torch.int32, device=logits.device) + pad_class_ids = class_order[pad_class_ids] + + order_logits = self.reading_order( + boxes=pad_boxes, + labels=pad_class_ids, + mask=mask, + ) + order_logits = order_logits[:, :, : self.num_queries] + + if labels is not None: + raise ValueError("PPDocLayoutV2ForObjectDetection does not support training") + + return PPDocLayoutV2ForObjectDetectionOutput( + logits=logits, + pred_boxes=pred_boxes, + order_logits=order_logits, + last_hidden_state=outputs.last_hidden_state, + intermediate_hidden_states=outputs.intermediate_hidden_states, + intermediate_logits=outputs.intermediate_logits, + intermediate_reference_points=outputs.intermediate_reference_points, + intermediate_predicted_corners=outputs.intermediate_predicted_corners, + initial_reference_points=outputs.initial_reference_points, + 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, + init_reference_points=outputs.init_reference_points, + enc_topk_logits=outputs.enc_topk_logits, + enc_topk_bboxes=outputs.enc_topk_bboxes, + enc_outputs_class=outputs.enc_outputs_class, + enc_outputs_coord_logits=outputs.enc_outputs_coord_logits, + denoising_meta_values=outputs.denoising_meta_values, + ) + + +__all__ = [ + "PPDocLayoutV2ForObjectDetection", + "PPDocLayoutV2Model", + "PPDocLayoutV2PreTrainedModel", + "PPDocLayoutV2ReadingOrder", +] diff --git a/third_party/transformers/src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.py b/third_party/transformers/src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..16fd7c9f552270caa8635686960672d1b8afddf7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_doclayout_v2/modular_pp_doclayout_v2.py @@ -0,0 +1,1001 @@ +# 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 +from dataclasses import dataclass +from typing import Optional + +import torch +from huggingface_hub.dataclasses import strict +from torch import nn + +from ... import initialization as init +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...masking_utils import create_bidirectional_mask +from ...processing_utils import Unpack +from ...utils import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, +) +from ...utils.generic import TensorType +from ..auto import AutoConfig +from ..layoutlmv3.modeling_layoutlmv3 import ( + LayoutLMv3Attention, + LayoutLMv3Encoder, + LayoutLMv3Intermediate, + LayoutLMv3Layer, + LayoutLMv3Output, + LayoutLMv3SelfAttention, + LayoutLMv3SelfOutput, + LayoutLMv3TextEmbeddings, +) +from ..pp_doclayout_v3.image_processing_pp_doclayout_v3 import PPDocLayoutV3ImageProcessor +from ..pp_doclayout_v3.modeling_pp_doclayout_v3 import PPDocLayoutV3GlobalPointer +from ..rt_detr.modeling_rt_detr import ( + RTDetrForObjectDetection, + RTDetrMLPPredictionHead, + RTDetrModel, + RTDetrModelOutput, + RTDetrPreTrainedModel, +) + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="PaddlePaddle/PP-DocLayoutV2_safetensors") +@strict +class PPDocLayoutV2ReadingOrderConfig(PreTrainedConfig): + r""" + has_relative_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a relative attention bias in the self-attention mechanism. + has_spatial_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a spatial attention bias in the self-attention mechanism. + rel_pos_bins (`int`, *optional*, defaults to 32): + The number of relative position bins to be used in the self-attention mechanism. + max_rel_pos (`int`, *optional*, defaults to 128): + The maximum number of relative positions to be used in the self-attention mechanism. + rel_2d_pos_bins (`int`, *optional*, defaults to 64): + The number of 2D relative position bins in the self-attention mechanism. + max_rel_2d_pos (`int`, *optional*, defaults to 256): + The maximum number of relative 2D positions in the self-attention mechanism. + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum value that the 2D position embedding might ever be used with. Typically set this to something + large just in case (e.g., 1024). + start_token_id (`int`, *optional*, defaults to 0): + Token id representing the start of a sequence. + end_token_id (`int`, *optional*, defaults to 2): + Token id representing the end of a sequence. + pred_token_id (`int`, *optional*, defaults to 3): + Token id representing valid prediction positions (placeholders) in the sequence. + coordinate_size (`int`, *optional*, defaults to 171): + Dimension of the coordinate embeddings. + shape_size (`int`, *optional*, defaults to 170): + Dimension of the width and height embeddings. + num_classes (`int`, *optional*, defaults to 20): + Number of labels or classes for the layout elements. + relation_bias_embed_dim (`int`, *optional*, defaults to 16): + Embedding dimension for the relation bias. + relation_bias_theta (`float`, *optional*, defaults to 10000): + Temperature parameter used for relation bias scaling. + relation_bias_scale (`float`, *optional*, defaults to 100): + Scale parameter for the relation bias. + global_pointer_head_size (`int`, *optional*, defaults to 64): + The size of the global pointer head. + gp_dropout_value (`float`, *optional*, defaults to 0.0): + The dropout probability in the global pointer head. + """ + + hidden_size: int = 512 + num_attention_heads: int = 8 + attention_probs_dropout_prob: float | int = 0.1 + has_relative_attention_bias: bool = False + has_spatial_attention_bias: bool = True + layer_norm_eps: float = 1e-5 + hidden_dropout_prob: float | int = 0.1 + intermediate_size: int = 2048 + hidden_act: str = "gelu" + num_hidden_layers: int = 6 + rel_pos_bins: int = 32 + max_rel_pos: int = 128 + rel_2d_pos_bins: int = 64 + max_rel_2d_pos: int = 256 + max_position_embeddings: int = 514 + max_2d_position_embeddings: int = 1024 + type_vocab_size: int = 1 + vocab_size: int = 4 + initializer_range: float = 0.01 + start_token_id: int = 0 + pad_token_id: int | None = 1 + end_token_id: int = 2 + pred_token_id: int = 3 + coordinate_size: int = 171 + shape_size: int = 170 + num_classes: int = 20 + relation_bias_embed_dim: int = 16 + relation_bias_theta: int = 10000 + relation_bias_scale: int = 100 + global_pointer_head_size: int = 64 + gp_dropout_value: float | int = 0.0 + + +@auto_docstring(checkpoint="PaddlePaddle/PP-DocLayoutV2_safetensors") +@strict +class PPDocLayoutV2Config(PreTrainedConfig): + r""" + initializer_bias_prior_prob (`float`, *optional*): + The prior probability used by the bias initializer to initialize biases for `enc_score_head` and `class_embed`. + If `None`, `prior_prob` computed as `prior_prob = 1 / (num_labels + 1)` while initializing model weights. + freeze_backbone_batch_norms (`bool`, *optional*, defaults to `True`): + Whether to freeze the batch normalization layers in the backbone. + encoder_hidden_dim (`int`, *optional*, defaults to 256): + Dimension of the layers in hybrid encoder. + encoder_in_channels (`list`, *optional*, defaults to `[512, 1024, 2048]`): + Multi level features input for encoder. + feat_strides (`list[int]`, *optional*, defaults to `[8, 16, 32]`): + Strides used in each feature map. + encoder_layers (`int`, *optional*, defaults to 1): + Total of layers to be used by the encoder. + encode_proj_layers (`list[int]`, *optional*, defaults to `[2]`): + Indexes of the projected layers to be used in the encoder. + positional_encoding_temperature (`int`, *optional*, defaults to 10000): + The temperature parameter used to create the positional encodings. + encoder_activation_function (`str`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"silu"` and `"gelu_new"` are supported. + eval_size (`tuple[int, int]`, *optional*): + Height and width used to computes the effective height and width of the position embeddings after taking + into account the stride. + normalize_before (`bool`, *optional*, defaults to `False`): + Determine whether to apply layer normalization in the transformer encoder layer before self-attention and + feed-forward modules. + hidden_expansion (`float`, *optional*, defaults to 1.0): + Expansion ratio to enlarge the dimension size of RepVGGBlock and CSPRepLayer. + num_queries (`int`, *optional*, defaults to 300): + Number of object queries. + decoder_in_channels (`list`, *optional*, defaults to `[256, 256, 256]`): + Multi level features dimension for decoder + decoder_ffn_dim (`int`, *optional*, defaults to 1024): + Dimension of the "intermediate" (often named feed-forward) layer in decoder. + num_feature_levels (`int`, *optional*, defaults to 3): + The number of input feature levels. + decoder_n_points (`int`, *optional*, defaults to 4): + The number of sampled keys in each feature level for each attention head in the decoder. + decoder_activation_function (`str`, *optional*, defaults to `"relu"`): + The non-linear activation function (function or string) in the decoder. If string, `"gelu"`, + `"relu"`, `"silu"` and `"gelu_new"` are supported. + num_denoising (`int`, *optional*, defaults to 100): + The total number of denoising tasks or queries to be used for contrastive denoising. + label_noise_ratio (`float`, *optional*, defaults to 0.5): + The fraction of denoising labels to which random noise should be added. + box_noise_scale (`float`, *optional*, defaults to 1.0): + Scale or magnitude of noise to be added to the bounding boxes. + learn_initial_query (`bool`, *optional*, defaults to `False`): + Indicates whether the initial query embeddings for the decoder should be learned during training + anchor_image_size (`tuple[int, int]`, *optional*): + Height and width of the input image used during evaluation to generate the bounding box anchors. If None, automatic generate anchor is applied. + disable_custom_kernels (`bool`, *optional*, defaults to `True`): + Whether to disable custom kernels. + class_thresholds (`list[float]`, *optional*): + The thresholds for each label. + class_order (`list[int]`, *optional*): + The priority for each label. + reading_order_config (`dict`, *optional*): + The configuration of a `PPDocLayoutV2ReadingOrder`. + + Examples: + + ```python + >>> from transformers import PPDocLayoutV2Config, PPDocLayoutV2ForObjectDetection + + >>> # Initializing a PP-DocLayoutV2 configuration + >>> configuration = PPDocLayoutV2Config() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = PPDocLayoutV2ForObjectDetection(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pp_doclayout_v2" + sub_configs = {"backbone_config": AutoConfig, "reading_order_config": PPDocLayoutV2ReadingOrderConfig} + + layer_types = ("basic", "bottleneck") + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + } + + initializer_range: float = 0.01 + initializer_bias_prior_prob: float | None = None + layer_norm_eps: float = 1e-5 + batch_norm_eps: float = 1e-5 + backbone_config: PreTrainedConfig | dict | None = None + freeze_backbone_batch_norms: bool = True + encoder_hidden_dim: int = 256 + encoder_in_channels: list[int] | tuple[int, ...] | None = (512, 1024, 2048) + feat_strides: list[int] | tuple[int, ...] | None = (8, 16, 32) + encoder_layers: int = 1 + encoder_ffn_dim: int = 1024 + encoder_attention_heads: int = 8 + dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + encode_proj_layers: list[int] | tuple[int, ...] | None = (2,) + positional_encoding_temperature: int = 10000 + encoder_activation_function: str = "gelu" + activation_function: str = "silu" + eval_size: list[int] | None = None + normalize_before: bool = False + hidden_expansion: float = 1.0 + d_model: int = 256 + num_queries: int = 300 + decoder_in_channels: list[int] | tuple[int, ...] | None = (256, 256, 256) + decoder_ffn_dim: int = 1024 + num_feature_levels: int = 3 + decoder_n_points: int = 4 + decoder_layers: int = 6 + decoder_attention_heads: int = 8 + decoder_activation_function: str = "relu" + attention_dropout: float | int = 0.0 + num_denoising: int = 100 + label_noise_ratio: float = 0.5 + box_noise_scale: float = 1.0 + learn_initial_query: bool = False + anchor_image_size: list[int] | None = None + disable_custom_kernels: bool = True + is_encoder_decoder: bool = True + class_thresholds: list[float] | None = None + class_order: list[int] | None = None + reading_order_config: PreTrainedConfig | dict | None = None + + def __post_init__(self, **kwargs): + if isinstance(self.reading_order_config, dict): + self.reading_order_config = self.sub_configs["reading_order_config"](**self.reading_order_config) + elif self.reading_order_config is None: + self.reading_order_config = self.sub_configs["reading_order_config"]() + + 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": [1, 2, 3], + "freeze_stem_only": True, + "freeze_at": 0, + "freeze_norm": True, + "lr_mult_list": [0, 0.05, 0.05, 0.05, 0.05], + "out_features": ["stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + self.encoder_in_channels = list(self.encoder_in_channels) + self.feat_strides = list(self.feat_strides) + self.encode_proj_layers = list(self.encode_proj_layers) + self.eval_size = list(self.eval_size) if self.eval_size is not None else None + self.decoder_in_channels = list(self.decoder_in_channels) + self.anchor_image_size = list(self.anchor_image_size) if self.anchor_image_size is not None else None + + super().__post_init__(**kwargs) + + +class PPDocLayoutV2ImageProcessor(PPDocLayoutV3ImageProcessor): + def extract_custom_vertices(self): + raise AttributeError("Not needed for PPDocLayoutV2") + + def _mask2polygon(self): + raise AttributeError("Not needed for PPDocLayoutV2") + + def _extract_polygon_points_by_masks(self): + raise AttributeError("Not needed for PPDocLayoutV2") + + def post_process_object_detection( + self, + outputs, + threshold: float = 0.5, + target_sizes: TensorType | list[tuple] | None = None, + ): + """ + Converts the raw output of [`DetrForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y, + bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + PPDocLayoutV2 is identical to PPDocLayoutV3, except that it does not return `polygon_points`. + + Args: + outputs ([`DetrObjectDetectionOutput`]): + Raw outputs of the model. + 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. + """ + boxes = outputs.pred_boxes + logits = outputs.logits + order_logits = outputs.order_logits + + order_seqs = self._get_order_seqs(order_logits) + + box_centers, box_dims = torch.split(boxes, 2, dim=-1) + top_left_coords = box_centers - 0.5 * box_dims + bottom_right_coords = box_centers + 0.5 * box_dims + boxes = torch.cat([top_left_coords, bottom_right_coords], dim=-1) + + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + if isinstance(target_sizes, list): + img_height, img_width = torch.as_tensor(target_sizes).unbind(1) + else: + img_height, img_width = target_sizes.unbind(1) + scale_factor = torch.stack([img_width, img_height, img_width, img_height], dim=1).to(boxes.device) + boxes = boxes * scale_factor[:, None, :] + + num_top_queries = logits.shape[1] + num_classes = logits.shape[2] + + scores = torch.nn.functional.sigmoid(logits) + scores, index = torch.topk(scores.flatten(1), num_top_queries, dim=-1) + labels = index % num_classes + index = index // num_classes + boxes = boxes.gather(dim=1, index=index.unsqueeze(-1).repeat(1, 1, boxes.shape[-1])) + order_seqs = order_seqs.gather(dim=1, index=index) + + results = [] + for score, label, box, order_seq in zip(scores, labels, boxes, order_seqs): + order_seq = order_seq[score >= threshold] + order_seq, indices = torch.sort(order_seq) + results.append( + { + "scores": score[score >= threshold][indices], + "labels": label[score >= threshold][indices], + "boxes": box[score >= threshold][indices], + "order_seq": order_seq, + } + ) + + return results + + +class PPDocLayoutV2GlobalPointer(PPDocLayoutV3GlobalPointer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, self.head_size * 2) + + +class PPDocLayoutV2PositionRelationEmbedding(nn.Module): + inv_freq: torch.Tensor + + def __init__(self, config, device=None): + super().__init__() + self.config = config + self.embed_dim = config.relation_bias_embed_dim + self.scale = config.relation_bias_scale + self.pos_proj = nn.Conv2d( + in_channels=self.embed_dim * 4, out_channels=config.num_attention_heads, kernel_size=1 + ) + inv_freq, self.attention_scaling = self.compute_default_rope_parameters(config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @staticmethod + def compute_default_rope_parameters( + config: PPDocLayoutV2Config | 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.relation_bias_theta + dim = config.relation_bias_embed_dim + half_dim = dim // 2 + + 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) / half_dim) + ) + return inv_freq, attention_factor + + def box_relative_encoding( + self, source_boxes: torch.Tensor, target_boxes: torch.Tensor = None, epsilon: float = 1e-5 + ): + source_boxes, target_boxes = source_boxes.unsqueeze(-2), target_boxes.unsqueeze(-3) + source_coordinates, source_dim = source_boxes[..., :2], source_boxes[..., 2:] + target_coordinates, target_dim = target_boxes[..., :2], target_boxes[..., 2:] + + coordinate_difference = torch.abs(source_coordinates - target_coordinates) + relative_coordinates = torch.log(coordinate_difference / (source_dim + epsilon) + 1.0) + relative_dim = torch.log((source_dim + epsilon) / (target_dim + epsilon)) + + relative_encoding = torch.cat([relative_coordinates, relative_dim], dim=-1) + + return relative_encoding + + def get_position_embedding(self, x: torch.Tensor, scale: float = 100.0): + embedding = (x * scale).unsqueeze(-1) * self.inv_freq + embedding = torch.cat((embedding.sin(), embedding.cos()), dim=-1).flatten(start_dim=-2).to(x.dtype) + + return embedding + + def forward(self, source_boxes: torch.Tensor, target_boxes: torch.Tensor = None): + if target_boxes is None: + target_boxes = source_boxes + with torch.no_grad(): + relative_encoding = self.box_relative_encoding(source_boxes, target_boxes) + position_embedding = self.get_position_embedding(relative_encoding, self.scale) + position_embedding = position_embedding.permute(0, 3, 1, 2) + out = self.pos_proj(position_embedding) + return out + + +class PPDocLayoutV2ReadingOrderSelfAttention(LayoutLMv3SelfAttention): + def forward( + self, hidden_states, attention_mask=None, rel_pos=None, rel_2d_pos=None, **kwargs: Unpack[TransformersKwargs] + ): + batch_size, seq_length, _ = hidden_states.shape + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + + # Take the dot product between "query" and "key" to get the raw attention scores. + # The attention scores QT K/√d could be significantly larger than input elements, and result in overflow. + # Changing the computational order into QT(K/√d) alleviates the problem. (https://huggingface.co/papers/2105.13290) + attention_scores = torch.matmul(query_layer / math.sqrt(self.attention_head_size), key_layer.transpose(-1, -2)) + + # NOTE: + # PPDocLayoutV2ReadingOrder applies unscaled rel_2d_pos + # unlike LayoutLMv3 which uses (rel_pos + rel_2d_pos) / math.sqrt(self.attention_head_size) + if rel_2d_pos is not None: + attention_scores += rel_2d_pos + elif self.has_relative_attention_bias: + attention_scores += rel_pos / math.sqrt(self.attention_head_size) + + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + # Use the trick of the CogView paper to stabilize training + attention_probs = self.cogview_attention(attention_scores) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer, attention_probs + + +class PPDocLayoutV2ReadingOrderSelfOutput(LayoutLMv3SelfOutput): + def __init__(self, config): + super().__init__() + del self.LayerNorm + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + 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.norm(hidden_states + input_tensor) + return hidden_states + + +class PPDocLayoutV2ReadingOrderIntermediate(LayoutLMv3Intermediate): + pass + + +class PPDocLayoutV2ReadingOrderOutput(LayoutLMv3Output): + def __init__(self, config): + super().__init__() + del self.LayerNorm + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + 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.norm(hidden_states + input_tensor) + return hidden_states + + +class PPDocLayoutV2ReadingOrderAttention(LayoutLMv3Attention): + def __init__(self, config): + super().__init__() + self.self = PPDocLayoutV2ReadingOrderSelfAttention(config) + self.output = PPDocLayoutV2ReadingOrderSelfOutput(config) + + +class PPDocLayoutV2ReadingOrderLayer(LayoutLMv3Layer): + def __init__(self, config): + super().__init__() + self.attention = PPDocLayoutV2ReadingOrderAttention(config) + self.intermediate = PPDocLayoutV2ReadingOrderIntermediate(config) + self.output = PPDocLayoutV2ReadingOrderOutput(config) + + +class PPDocLayoutV2ReadingOrderEncoder(LayoutLMv3Encoder): + def __init__(self, config): + super().__init__(config) + self.layer = nn.ModuleList([PPDocLayoutV2ReadingOrderLayer(config) for _ in range(config.num_hidden_layers)]) + self.rel_bias_module = PPDocLayoutV2PositionRelationEmbedding(config) + + def _cal_2d_pos_emb(self, bbox): + x_min, y_min, x_max, y_max = ( + bbox[..., 0], + bbox[..., 1], + bbox[..., 2], + bbox[..., 3], + ) + + width = (x_max - x_min).clamp(min=1e-3) + height = (y_max - y_min).clamp(min=1e-3) + + center_x = (x_min + x_max) * 0.5 + center_y = (y_min + y_max) * 0.5 + + center_width_height_bbox = torch.stack([center_x, center_y, width, height], dim=-1) + + result = self.rel_bias_module(center_width_height_bbox) + + return result + + +class PPDocLayoutV2TextEmbeddings(LayoutLMv3TextEmbeddings): + def __init__(self, config): + super().__init__(config) + + del self.LayerNorm + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + spatial_embed_dim = 4 * config.coordinate_size + 2 * config.shape_size + self.spatial_proj = nn.Linear(spatial_embed_dim, config.hidden_size) + + def forward( + self, + input_ids=None, + bbox=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + ): + if position_ids is None: + if input_ids is not None: + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx).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 = 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 + + # custom new spatial embeddings + spatial_position_embeddings = self.calculate_spatial_position_embeddings(bbox) + spatial_position_embeddings = self.spatial_proj(spatial_position_embeddings) + embeddings += spatial_position_embeddings + return embeddings + + +@auto_docstring +class PPDocLayoutV2PreTrainedModel(RTDetrPreTrainedModel): + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, PPDocLayoutV2TextEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + if isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + init.zeros_(module.weight.data[module.padding_idx]) + if isinstance(module, PPDocLayoutV2PositionRelationEmbedding): + inv_freq, _ = module.compute_default_rope_parameters(module.config, module.inv_freq.device) + module.register_buffer("inv_freq", inv_freq, persistent=False) + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 ReadingOrder Model. This model consists of an encoder and a GlobalPointer head. + It takes layout features as input and outputs logits representing the relative ordering relationships + between elements, which are used to determine the final reading sequence. + """ +) +class PPDocLayoutV2ReadingOrder(PPDocLayoutV2PreTrainedModel): + # Attention is based on LayoutLMv3 (no interface) + _supports_sdpa = False + _supports_flash_attn = False + _supports_attention_backend = False + _supports_flex_attn = False + + def __init__(self, config): + super().__init__(config) + self.embeddings = PPDocLayoutV2TextEmbeddings(config) + self.label_embeddings = nn.Embedding(config.num_classes, config.hidden_size) + self.label_features_projection = nn.Linear(config.hidden_size, config.hidden_size) + self.encoder = PPDocLayoutV2ReadingOrderEncoder(config) + self.relative_head = PPDocLayoutV2GlobalPointer(config) + self.config = config + + self.post_init() + + @auto_docstring + def forward(self, boxes, labels=None, mask=None, **kwargs: Unpack[TransformersKwargs]): + r""" + boxes (`torch.Tensor` of shape `(batch_size, sequence_length, 4)`): + Bounding box coordinates of the detected layout elements **in [0, 1000] scale**. + Format is `[x_min, y_min, x_max, y_max]`. + The tensor usually contains sorted valid boxes followed by zero-padding. + labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + The **remapped** class indices for each layout element. + These are not necessarily the raw detection class IDs, but indices mapped via + `config.class_order` (e.g., mapping text/title/figure to specific reading-order category IDs). + mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Boolean or Binary mask indicating valid detected elements after threshold filtering. + - True: Valid layout element. + - False: Padding/Empty element. + Used to determine the sequence length (`num_pred`) for the pointer mechanism. + """ + device = mask.device + batch_size, seq_len = mask.shape + num_pred = mask.sum(dim=1) + + input_ids = torch.full( + (batch_size, seq_len + 2), self.config.pad_token_id, dtype=torch.long, device=boxes.device + ) + input_ids[:, 0] = self.config.start_token_id + + pred_col_idx = torch.arange(seq_len + 2, device=device).unsqueeze(0) + pred_mask = (pred_col_idx >= 1) & (pred_col_idx <= num_pred.unsqueeze(1)) + input_ids[pred_mask] = self.config.pred_token_id + end_col_indices = num_pred + 1 + input_ids[torch.arange(batch_size, device=device), end_col_indices] = self.config.end_token_id + + pad_box = torch.zeros(size=[boxes.shape[0], 1, boxes.shape[-1]], dtype=boxes.dtype, device=boxes.device) + pad_boxes = torch.cat([pad_box, boxes, pad_box], dim=1) + bbox_embedding = self.embeddings(input_ids=input_ids, bbox=pad_boxes.long()) + + if labels is not None: + label_embs = self.label_embeddings(labels) + label_proj = self.label_features_projection(label_embs) + pad = torch.zeros( + size=[label_proj.shape[0], 1, label_proj.shape[-1]], dtype=label_proj.dtype, device=labels.device + ) + label_proj = torch.cat([pad, label_proj, pad], dim=1) + else: + label_proj = torch.zeros_like(bbox_embedding) + + final_embeddings = bbox_embedding + label_proj + final_embeddings = self.embeddings.norm(final_embeddings) + final_embeddings = self.embeddings.dropout(final_embeddings) + + attention_mask = pred_col_idx < (num_pred + 2).unsqueeze(1) + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=final_embeddings, + attention_mask=attention_mask, + ) + encoder_output = self.encoder(hidden_states=final_embeddings, bbox=pad_boxes, attention_mask=attention_mask) + encoder_output = encoder_output.last_hidden_state + token = encoder_output[:, 1 : 1 + seq_len, :] + read_order_logits = self.relative_head(token) + return read_order_logits + + +@dataclass +@auto_docstring +class PPDocLayoutV2ForObjectDetectionOutput(ModelOutput): + r""" + 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 [`~RTDetrImageProcessor.post_process_object_detection`] to retrieve the + unnormalized (absolute) bounding boxes. + order_logits (`tuple` of `torch.FloatTensor` of shape `(batch_size, num_queries, num_queries)`): + Order logits for all queries. The first dimension of each tensor is the batch size. The second dimension is the number of queries. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, 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 `(batch_size, config.decoder_layers, num_queries, hidden_size)`): + Stacked intermediate hidden states (output of each layer of the decoder). + intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, config.num_labels)`): + Stacked intermediate logits (logits of each layer of the decoder). + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate predicted corners (predicted corners of each layer of the decoder). + initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked initial reference points (initial reference points of each layer of the decoder). + init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Initial reference points sent through the Transformer decoder. + enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the encoder. + enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the encoder. + enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are + picked as region proposals in the first stage. Output of bounding box binary classification (i.e. + foreground and background). + enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the first stage. + denoising_meta_values (`dict`): + Extra dictionary for the denoising related values + """ + + logits: torch.FloatTensor | None = None + pred_boxes: torch.FloatTensor | None = None + order_logits: tuple[torch.FloatTensor] | None = None + last_hidden_state: torch.FloatTensor | None = None + intermediate_hidden_states: torch.FloatTensor | None = None + intermediate_logits: torch.FloatTensor | None = None + intermediate_reference_points: torch.FloatTensor | None = None + intermediate_predicted_corners: torch.FloatTensor | None = None + initial_reference_points: 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 + init_reference_points: tuple[torch.FloatTensor] | None = None + enc_topk_logits: torch.FloatTensor | None = None + enc_topk_bboxes: torch.FloatTensor | None = None + enc_outputs_class: torch.FloatTensor | None = None + enc_outputs_coord_logits: torch.FloatTensor | None = None + denoising_meta_values: dict | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the PP-DocLayoutV2 encoder-decoder model. + """ +) +class PPDocLayoutV2ModelOutput(RTDetrModelOutput): + pass + + +class PPDocLayoutV2MLPPredictionHead(RTDetrMLPPredictionHead): + pass + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 Model (consisting of a backbone and encoder-decoder) outputting raw hidden states without any head on top. + """ +) +class PPDocLayoutV2Model(RTDetrModel): + def __init__(self, config: PPDocLayoutV2Config): + super().__init__() + self.denoising_class_embed = nn.Embedding(config.num_labels, config.d_model) + + +@auto_docstring( + custom_intro=""" + PP-DocLayoutV2 Model (consisting of a backbone and encoder-decoder) outputting bounding boxes, logits and order_logits to be further + decoded into scores, classes and their reading order. + """ +) +class PPDocLayoutV2ForObjectDetection(RTDetrForObjectDetection): + _keys_to_ignore_on_load_missing = ["num_batches_tracked", "rel_pos_y_bias", "rel_pos_x_bias"] + + def __init__(self, config: PPDocLayoutV2Config): + super().__init__(config) + + self.reading_order = PPDocLayoutV2ReadingOrder(config.reading_order_config) + self.num_queries = config.num_queries + self.config = config + + self.post_init() + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_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], + ) -> tuple[torch.FloatTensor] | PPDocLayoutV2ForObjectDetectionOutput: + r""" + 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 AutoModelForObjectDetection, AutoImageProcessor + >>> from PIL import Image + >>> import requests + >>> import torch + + >>> url = "https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/layout_demo.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> model_path = "PaddlePaddle/PP-DocLayoutV2_safetensors" + >>> image_processor = AutoImageProcessor.from_pretrained(model_path) + >>> model = AutoModelForObjectDetection.from_pretrained(model_path) + + >>> # prepare image for the model + >>> inputs = image_processor(images=[image], return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax) + >>> results = image_processor.post_process_object_detection(outputs, target_sizes=torch.tensor([image.size[::-1]])) + + >>> # print outputs + >>> for result in results: + ... for idx, (score, label_id, box) in enumerate(zip(result["scores"], result["labels"], result["boxes"])): + ... score, label = score.item(), label_id.item() + ... box = [round(i, 2) for i in box.tolist()] + ... print(f"Order {idx + 1}: {model.config.id2label[label]}: {score:.2f} {box}") + Order 1: text: 0.99 [335.39, 184.26, 896.49, 654.48] + Order 2: paragraph_title: 0.97 [337.14, 683.49, 869.42, 798.27] + Order 3: text: 0.99 [335.71, 843.04, 891.17, 1454.15] + Order 4: text: 0.99 [920.42, 185.53, 1476.39, 464.25] + Order 5: text: 0.98 [920.62, 483.75, 1480.52, 765.34] + Order 6: text: 0.98 [920.58, 846.75, 1481.94, 1220.53] + Order 7: text: 0.97 [921.12, 1239.27, 1468.87, 1377.33] + Order 8: footnote: 0.82 [334.58, 1614.67, 1483.84, 1731.61] + Order 9: text: 0.51 [334.58, 1614.67, 1483.84, 1731.61] + Order 10: footnote: 0.83 [334.7, 1757.26, 1471.07, 1845.33] + Order 11: text: 0.87 [336.65, 1910.28, 661.33, 1939.92] + Order 12: footnote: 0.95 [336.16, 2114.52, 1450.28, 2171.74] + Order 13: number: 0.87 [106.04, 2257.37, 136.05, 2281.98] + Order 14: footer: 0.93 [338.6, 2255.94, 985.67, 2283.57] + ```""" + outputs = self.model( + pixel_values, + pixel_mask=pixel_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + labels=labels, + **kwargs, + ) + + intermediate_reference_points = outputs.intermediate_reference_points + intermediate_logits = outputs.intermediate_logits + raw_bboxes = intermediate_reference_points[:, -1] + logits = intermediate_logits[:, -1] + + box_centers, box_sizes = raw_bboxes.split(2, dim=-1) + bboxes = torch.cat([box_centers - 0.5 * box_sizes, box_centers + 0.5 * box_sizes], dim=-1) * 1000 + bboxes = bboxes.clamp_(0.0, 1000.0) + + max_logits, class_ids = logits.max(dim=-1) + max_probs = max_logits.sigmoid() + + class_thresholds = torch.tensor(self.config.class_thresholds, dtype=torch.float32, device=logits.device) + thresholds = class_thresholds[class_ids] + mask = max_probs >= thresholds + + indices = torch.argsort(mask.to(torch.int8), dim=1, descending=True) + + sorted_class_ids = torch.take_along_dim(class_ids, indices, dim=1) + sorted_boxes = torch.take_along_dim(bboxes, indices[..., None].expand(-1, -1, 4), dim=1) + pred_boxes = torch.take_along_dim(raw_bboxes, indices[..., None].expand(-1, -1, 4), dim=1) + logits = torch.take_along_dim(logits, indices[..., None].expand(-1, -1, logits.size(-1)), dim=1) + + sorted_mask = torch.take_along_dim(mask, indices, dim=1) + + pad_boxes = torch.where(sorted_mask[..., None], sorted_boxes, torch.zeros_like(sorted_boxes)) + pad_class_ids = torch.where(sorted_mask, sorted_class_ids, torch.zeros_like(sorted_class_ids)) + + class_order = torch.tensor(self.config.class_order, dtype=torch.int32, device=logits.device) + pad_class_ids = class_order[pad_class_ids] + + order_logits = self.reading_order( + boxes=pad_boxes, + labels=pad_class_ids, + mask=mask, + ) + order_logits = order_logits[:, :, : self.num_queries] + + if labels is not None: + raise ValueError("PPDocLayoutV2ForObjectDetection does not support training") + + return PPDocLayoutV2ForObjectDetectionOutput( + logits=logits, + pred_boxes=pred_boxes, + order_logits=order_logits, + last_hidden_state=outputs.last_hidden_state, + intermediate_hidden_states=outputs.intermediate_hidden_states, + intermediate_logits=outputs.intermediate_logits, + intermediate_reference_points=outputs.intermediate_reference_points, + intermediate_predicted_corners=outputs.intermediate_predicted_corners, + initial_reference_points=outputs.initial_reference_points, + 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, + init_reference_points=outputs.init_reference_points, + enc_topk_logits=outputs.enc_topk_logits, + enc_topk_bboxes=outputs.enc_topk_bboxes, + enc_outputs_class=outputs.enc_outputs_class, + enc_outputs_coord_logits=outputs.enc_outputs_coord_logits, + denoising_meta_values=outputs.denoising_meta_values, + ) + + +__all__ = [ + "PPDocLayoutV2ForObjectDetection", + "PPDocLayoutV2ImageProcessor", + "PPDocLayoutV2Config", + "PPDocLayoutV2Model", + "PPDocLayoutV2PreTrainedModel", + "PPDocLayoutV2ReadingOrder", +] diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/__init__.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb0b1324b316d8c950a5b379c781f28c67db88f0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/__init__.py @@ -0,0 +1,32 @@ +# 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_prompt_depth_anything import PromptDepthAnythingConfig + from .image_processing_pil_prompt_depth_anything import PromptDepthAnythingImageProcessorPil + from .image_processing_prompt_depth_anything import PromptDepthAnythingImageProcessor + from .modeling_prompt_depth_anything import ( + PromptDepthAnythingForDepthEstimation, + PromptDepthAnythingPreTrainedModel, + ) +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/prompt_depth_anything/configuration_prompt_depth_anything.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/configuration_prompt_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..b6577a84e3b35a17b71cb48ae98cab8d2337c4e5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/configuration_prompt_depth_anything.py @@ -0,0 +1,103 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/prompt_depth_anything/modular_prompt_depth_anything.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_prompt_depth_anything.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto.configuration_auto import AutoConfig + + +@auto_docstring(checkpoint="LiheYoung/depth-anything-small-hf") +@strict +class PromptDepthAnythingConfig(PreTrainedConfig): + r""" + reassemble_hidden_size (`int`, *optional*, defaults to 384): + The number of input channels of the reassemble layers. + reassemble_factors (`list[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`): + The up/downsampling factors of the reassemble layers. + neck_hidden_sizes (`list[str]`, *optional*, defaults to `[48, 96, 192, 384]`): + The hidden sizes to project to for the feature maps of the backbone. + fusion_hidden_size (`int`, *optional*, defaults to 64): + The number of channels before fusion. + head_in_index (`int`, *optional*, defaults to -1): + The index of the features to use in the depth estimation head. + head_hidden_size (`int`, *optional*, defaults to 32): + The number of output channels in the second convolution of the depth estimation head. + depth_estimation_type (`str`, *optional*, defaults to `"relative"`): + The type of depth estimation to use. Can be one of `["relative", "metric"]`. + max_depth (`float`, *optional*): + The maximum depth to use for the "metric" depth estimation head. 20 should be used for indoor models + and 80 for outdoor models. For "relative" depth estimation, this value is ignored. + + Example: + + ```python + >>> from transformers import PromptDepthAnythingConfig, PromptDepthAnythingForDepthEstimation + + >>> # Initializing a PromptDepthAnything small style configuration + >>> configuration = PromptDepthAnythingConfig() + + >>> # Initializing a model from the PromptDepthAnything small style configuration + >>> model = PromptDepthAnythingForDepthEstimation(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "prompt_depth_anything" + sub_configs = {"backbone_config": AutoConfig} + + backbone_config: dict | PreTrainedConfig | None = None + patch_size: int | list[int] | tuple[int, int] = 14 + initializer_range: float = 0.02 + reassemble_hidden_size: int = 384 + reassemble_factors: list[int | float] | tuple[int | float, ...] = (4, 2, 1, 0.5) + neck_hidden_sizes: list[int] | tuple[int, ...] = (48, 96, 192, 384) + fusion_hidden_size: int = 64 + head_in_index: int = -1 + head_hidden_size: int = 32 + depth_estimation_type: str = "relative" + max_depth: int | None = None + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="dinov2", + default_config_kwargs={ + "image_size": 518, + "hidden_size": 384, + "num_attention_heads": 6, + "out_indices": [9, 10, 11, 12], + "reshape_hidden_states": False, + }, + **kwargs, + ) + + self.max_depth = self.max_depth if self.max_depth else 1 + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.depth_estimation_type not in ["relative", "metric"]: + raise ValueError("depth_estimation_type must be one of ['relative', 'metric']") + + +__all__ = ["PromptDepthAnythingConfig"] diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/convert_prompt_depth_anything_to_hf.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/convert_prompt_depth_anything_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c7788ba1916b7ec8782c7839601b5021fd4b03 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/convert_prompt_depth_anything_to_hf.py @@ -0,0 +1,294 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert Prompt Depth Anything checkpoints from the original repository. URL: +https://github.com/DepthAnything/PromptDA""" + +import argparse +import re +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 ( + Dinov2Config, + PromptDepthAnythingConfig, + PromptDepthAnythingForDepthEstimation, + PromptDepthAnythingImageProcessor, +) +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_dpt_config(model_name): + if "small" in model_name or "vits" in model_name: + out_indices = [3, 6, 9, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-small", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 64 + neck_hidden_sizes = [48, 96, 192, 384] + elif "base" in model_name or "vitb" in model_name: + out_indices = [3, 6, 9, 12] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-base", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 128 + neck_hidden_sizes = [96, 192, 384, 768] + elif "large" in model_name or "vitl" in model_name: + out_indices = [5, 12, 18, 24] + backbone_config = Dinov2Config.from_pretrained( + "facebook/dinov2-large", out_indices=out_indices, apply_layernorm=True, reshape_hidden_states=False + ) + fusion_hidden_size = 256 + neck_hidden_sizes = [256, 512, 1024, 1024] + else: + raise NotImplementedError(f"Model not supported: {model_name}") + + depth_estimation_type = "metric" + max_depth = None + + config = PromptDepthAnythingConfig( + reassemble_hidden_size=backbone_config.hidden_size, + patch_size=backbone_config.patch_size, + backbone_config=backbone_config, + fusion_hidden_size=fusion_hidden_size, + neck_hidden_sizes=neck_hidden_sizes, + depth_estimation_type=depth_estimation_type, + max_depth=max_depth, + ) + + return config + + +def transform_qkv_weights(key, value, config): + if not key.startswith("qkv_transform"): + return value + + layer_idx = int(key.split("_")[-1]) + hidden_size = config.backbone_config.hidden_size + + suffix = "bias" if "bias" in key else "weight" + return { + f"backbone.encoder.layer.{layer_idx}.attention.attention.query.{suffix}": value[:hidden_size], + f"backbone.encoder.layer.{layer_idx}.attention.attention.key.{suffix}": value[hidden_size : hidden_size * 2], + f"backbone.encoder.layer.{layer_idx}.attention.attention.value.{suffix}": value[-hidden_size:], + } + + +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + # Stem + r"pretrained.cls_token": r"backbone.embeddings.cls_token", + r"pretrained.mask_token": r"backbone.embeddings.mask_token", + r"pretrained.pos_embed": r"backbone.embeddings.position_embeddings", + r"pretrained.patch_embed.proj.(weight|bias)": r"backbone.embeddings.patch_embeddings.projection.\1", + # Backbone + r"pretrained.norm.(weight|bias)": r"backbone.layernorm.\1", + # Transformer layers + r"pretrained.blocks.(\d+).ls1.gamma": r"backbone.encoder.layer.\1.layer_scale1.lambda1", + r"pretrained.blocks.(\d+).ls2.gamma": r"backbone.encoder.layer.\1.layer_scale2.lambda1", + r"pretrained.blocks.(\d+).norm1.(weight|bias)": r"backbone.encoder.layer.\1.norm1.\2", + r"pretrained.blocks.(\d+).norm2.(weight|bias)": r"backbone.encoder.layer.\1.norm2.\2", + r"pretrained.blocks.(\d+).mlp.fc1.(weight|bias)": r"backbone.encoder.layer.\1.mlp.fc1.\2", + r"pretrained.blocks.(\d+).mlp.fc2.(weight|bias)": r"backbone.encoder.layer.\1.mlp.fc2.\2", + r"pretrained.blocks.(\d+).attn.proj.(weight|bias)": r"backbone.encoder.layer.\1.attention.output.dense.\2", + r"pretrained.blocks.(\d+).attn.qkv.(weight|bias)": r"qkv_transform_\2_\1", + # Neck + r"depth_head.projects.(\d+).(weight|bias)": r"neck.reassemble_stage.layers.\1.projection.\2", + r"depth_head.scratch.layer(\d+)_rn.weight": lambda m: f"neck.convs.{int(m.group(1)) - 1}.weight", + r"depth_head.resize_layers.(\d+).(weight|bias)": r"neck.reassemble_stage.layers.\1.resize.\2", + # Refinenet (with reversed indices) + r"depth_head.scratch.refinenet(\d+).out_conv.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.projection.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit1.conv1.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.residual_layer1.convolution1.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit1.conv2.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.residual_layer1.convolution2.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit2.conv1.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.residual_layer2.convolution1.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit2.conv2.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.residual_layer2.convolution2.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit_depth.0.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.prompt_depth_layer.convolution1.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit_depth.2.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.prompt_depth_layer.convolution2.{m.group(2)}", + r"depth_head.scratch.refinenet(\d+).resConfUnit_depth.4.(weight|bias)": lambda m: f"neck.fusion_stage.layers.{4 - int(m.group(1))}.prompt_depth_layer.convolution3.{m.group(2)}", + # Head + r"depth_head.scratch.output_conv1.(weight|bias)": r"head.conv1.\1", + r"depth_head.scratch.output_conv2.0.(weight|bias)": r"head.conv2.\1", + r"depth_head.scratch.output_conv2.2.(weight|bias)": r"head.conv3.\1", +} + + +def convert_old_keys_to_new_keys(state_dict_keys: dict | None = None): + """ + Convert old state dict keys to new keys using regex patterns. + """ + output_dict = {} + if state_dict_keys is not None: + for old_key in state_dict_keys: + new_key = old_key + for pattern, replacement in ORIGINAL_TO_CONVERTED_KEY_MAPPING.items(): + match = re.match(pattern, old_key) + if match: + if callable(replacement): + new_key = replacement(match) + else: + new_key = re.sub(pattern, replacement, old_key) + break + output_dict[old_key] = new_key + return output_dict + + +@torch.no_grad() +def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits): + """ + Copy/paste/tweak model's weights to our DPT structure. + """ + + # define DPT configuration + config = get_dpt_config(model_name) + + model_name_to_repo = { + "prompt-depth-anything-vits": "depth-anything/prompt-depth-anything-vits", + "prompt-depth-anything-vits-transparent": "depth-anything/prompt-depth-anything-vits-transparent", + "prompt-depth-anything-vitl": "depth-anything/prompt-depth-anything-vitl", + } + + # load original state_dict + repo_id = model_name_to_repo[model_name] + filename = name_to_checkpoint[model_name] + filepath = hf_hub_download( + repo_id=repo_id, + filename=f"{filename}", + ) + + state_dict = torch.load(filepath, map_location="cpu", weights_only=True)["state_dict"] + state_dict = {key[9:]: state_dict[key] for key in state_dict} + + # Convert state dict using mappings + key_mapping = convert_old_keys_to_new_keys(state_dict.keys()) + new_state_dict = {} + for key, value in state_dict.items(): + new_key = key_mapping[key] + transformed_value = transform_qkv_weights(new_key, value, config) + if isinstance(transformed_value, dict): + new_state_dict.update(transformed_value) + else: + new_state_dict[new_key] = transformed_value + + # load HuggingFace model + model = PromptDepthAnythingForDepthEstimation(config) + model.load_state_dict(new_state_dict, strict=False) + model.eval() + + processor = PromptDepthAnythingImageProcessor( + do_resize=True, + size=756, + ensure_multiple_of=14, + keep_aspect_ratio=True, + do_rescale=True, + do_normalize=True, + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/image.jpg?raw=true" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + + prompt_depth_url = ( + "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/arkit_depth.png?raw=true" + ) + with httpx.stream("GET", prompt_depth_url) as response: + prompt_depth = Image.open(BytesIO(response.read())) + + inputs = processor(image, return_tensors="pt", prompt_depth=prompt_depth) + + # Verify forward pass + with torch.no_grad(): + outputs = model(**inputs) + predicted_depth = outputs.predicted_depth + + print("Shape of predicted depth:", predicted_depth.shape) + print("First values:", predicted_depth[0, :3, :3]) + + # assert logits + if verify_logits: + expected_shape = torch.Size([1, 756, 1008]) + if model_name == "prompt-depth-anything-vits": + expected_slice = torch.tensor( + [[3.0100, 3.0016, 3.0219], [3.0046, 3.0137, 3.0275], [3.0083, 3.0191, 3.0292]] + ) + elif model_name == "prompt-depth-anything-vits-transparent": + expected_slice = torch.tensor( + [[3.0058, 3.0397, 3.0460], [3.0314, 3.0393, 3.0504], [3.0326, 3.0465, 3.0545]] + ) + elif model_name == "prompt-depth-anything-vitl": + expected_slice = torch.tensor( + [[3.1336, 3.1358, 3.1363], [3.1368, 3.1267, 3.1414], [3.1397, 3.1385, 3.1448]] + ) + else: + raise ValueError("Not supported") + assert predicted_depth.shape == torch.Size(expected_shape) + assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=5e-3) # 5mm tolerance + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model and processor to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and processor to hub...") + model.push_to_hub(repo_id=f"{model_name.title()}-hf") + processor.push_to_hub(repo_id=f"{model_name.title()}-hf") + + +name_to_checkpoint = { + "prompt-depth-anything-vits": "model.ckpt", + "prompt-depth-anything-vits-transparent": "model.ckpt", + "prompt-depth-anything-vitl": "model.ckpt", +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="prompt_depth_anything_vits", + type=str, + choices=name_to_checkpoint.keys(), + help="Name of the model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether to push the model to the hub after conversion.", + ) + parser.add_argument( + "--verify_logits", + action="store_false", + required=False, + help="Whether to verify the logits after conversion.", + ) + + args = parser.parse_args() + convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits) diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_pil_prompt_depth_anything.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_pil_prompt_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..26f31249462aa3bb4b6c5c6e68aaa1a7bc50d3f9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_pil_prompt_depth_anything.py @@ -0,0 +1,333 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for PromptDepthAnything.""" + +import math +from typing import TYPE_CHECKING + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import PaddingMode +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring +from ...utils.import_utils import requires + + +if TYPE_CHECKING: + from ...modeling_outputs import DepthEstimatorOutput + + +def _constrain_to_multiple_of(val, multiple, min_val=0, max_val=None): + """Constrain a value to be a multiple of another value.""" + x = round(val / multiple) * multiple + + if max_val is not None and x > max_val: + x = math.floor(val / multiple) * multiple + + if x < min_val: + x = math.ceil(val / multiple) * multiple + + return x + + +def _get_resize_output_image_size( + input_image: np.ndarray, output_size: tuple[int, int], keep_aspect_ratio: bool, multiple: int +) -> tuple[int, int]: + """Get the output size for resizing an image.""" + input_height, input_width = get_image_size(input_image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = output_size + + # determine new height and width + scale_height = output_height / input_height + scale_width = output_width / input_width + + if keep_aspect_ratio: + # scale as little as possible + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + + new_height = _constrain_to_multiple_of(scale_height * input_height, multiple=multiple) + new_width = _constrain_to_multiple_of(scale_width * input_width, multiple=multiple) + + return (new_height, new_width) + + +# Adapted from transformers.models.prompt_depth_anything.image_processing_prompt_depth_anything.PromptDepthAnythingImageProcessorKwargs +class PromptDepthAnythingImageProcessorKwargs(ImagesKwargs, total=False): + r""" + keep_aspect_ratio (`bool`, *optional*): + If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. + ensure_multiple_of (`int`, *optional*): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. + prompt_scale_to_meter (`float`, *optional*): + Scale factor to convert the prompt depth to meters. + size_divisor (`int`, *optional*): + If `do_pad` is `True`, pads the image dimensions to be divisible by this value. + """ + + keep_aspect_ratio: bool + ensure_multiple_of: int + size_divisor: int + prompt_scale_to_meter: float + + +@auto_docstring +class PromptDepthAnythingImageProcessorPil(PilBackend): + model_input_names = ["pixel_values", "prompt_depth"] + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 384, "width": 384} + do_resize = True + do_rescale = True + do_normalize = True + keep_aspect_ratio = False + ensure_multiple_of = 1 + do_pad = False + size_divisor = None + prompt_scale_to_meter = 0.001 + valid_kwargs = PromptDepthAnythingImageProcessorKwargs + + def __init__(self, **kwargs: Unpack[PromptDepthAnythingImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + prompt_depth: ImageInput | None = None, + **kwargs: Unpack[PromptDepthAnythingImageProcessorKwargs], + ) -> BatchFeature: + r""" + prompt_depth (`ImageInput`, *optional*): + Prompt depth to preprocess. + """ + return super().preprocess(images, prompt_depth, **kwargs) + + def resize_with_aspect_ratio( + self, + image: np.ndarray, + size: SizeDict, + keep_aspect_ratio: bool = False, + ensure_multiple_of: int = 1, + resample: PILImageResampling | None = None, + ) -> np.ndarray: + """ + Resize an image to target size while optionally maintaining aspect ratio and ensuring dimensions are multiples. + """ + if resample is None: + resample = PILImageResampling.BICUBIC + + # Custom resize with aspect ratio preservation and ensure_multiple_of constraint + output_size = _get_resize_output_image_size( + image, + output_size=(size.height, size.width), + keep_aspect_ratio=keep_aspect_ratio, + multiple=ensure_multiple_of, + ) + + # Standard resize method with calculated output size + return self.resize(image=image, size=SizeDict(height=output_size[0], width=output_size[1]), resample=resample) + + def pad_image(self, image: np.ndarray, size_divisor: int) -> np.ndarray: + """ + Center pad an image to be a multiple of size_divisor. + """ + + def _get_pad(size, size_divisor): + new_size = math.ceil(size / size_divisor) * size_divisor + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + + # Match slow processor and PyTorch convention: width->left/right, height->top/bottom + pad_size_left, pad_size_right = _get_pad(width, size_divisor) + pad_size_top, pad_size_bottom = _get_pad(height, size_divisor) + + # NumPy padding: ((pad_top, pad_bottom), (pad_left, pad_right)) + padding = ((pad_size_top, pad_size_bottom), (pad_size_left, pad_size_right)) + padded_image = np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=0, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + + return padded_image + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + prompt_depth: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + device: str | None = None, + return_tensors: str | TensorType | None = None, + prompt_scale_to_meter: float | None = None, + **kwargs, + ) -> BatchFeature: + """ + Preprocess image-like inputs, including the main images and optional prompt depth. + """ + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=False, input_data_format=input_data_format, device=device + ) # always use do_convert_rgb=False rather than defining it as a param to match slow processor + + # Process images with the standard pipeline + pixel_values = self._preprocess(images, return_tensors=return_tensors, **kwargs) + + data = {"pixel_values": pixel_values} + + # Process prompt depth if provided + if prompt_depth is not None: + processed_prompt_depths = self._prepare_image_like_inputs( + images=prompt_depth, + do_convert_rgb=False, # Depth maps should not be converted + input_data_format=input_data_format, + device=device, + expected_ndims=2, + ) + + # Validate prompt_depths has same length as images as in slow processor + if len(processed_prompt_depths) != len(images): + raise ValueError( + f"Number of prompt depth images ({len(processed_prompt_depths)}) does not match number of input images ({len(images)})" + ) + + if prompt_scale_to_meter is None: + prompt_scale_to_meter = self.prompt_scale_to_meter + + final_prompt_depths = [] + for depth in processed_prompt_depths: + depth = depth * prompt_scale_to_meter + + # Handle case where depth is constant (min == max) + if depth.min() == depth.max(): + depth[0, 0] = depth[0, 0] + 1e-6 # Add small variation to avoid numerical issues + + if depth.ndim == 2: # Add channel dimension if needed + depth = np.expand_dims(depth, 0) # [H, W] -> [1, H, W] (channels first) + + depth = depth.astype(np.float32) # Convert to float32 to match slow processor + final_prompt_depths.append(depth) + + data["prompt_depth"] = final_prompt_depths + + return BatchFeature(data=data, tensor_type=return_tensors) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + return_tensors: str | TensorType | None, + keep_aspect_ratio: bool | None = None, + ensure_multiple_of: int | None = None, + size_divisor: int | None = None, + **kwargs, + ) -> list[np.ndarray]: + """ + Override the base _preprocess method to handle custom PromptDepthAnything parameters. + """ + processed_images = [] + for image in images: + if do_resize: + image = self.resize_with_aspect_ratio( + image=image, + size=size, + keep_aspect_ratio=keep_aspect_ratio, + ensure_multiple_of=ensure_multiple_of, + resample=resample, + ) + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + if do_pad and size_divisor is not None: + image = self.pad_image(image, size_divisor) + processed_images.append(image) + + return processed_images + + @requires(backends=("torch",)) + def post_process_depth_estimation( + self, outputs: "DepthEstimatorOutput", target_sizes: TensorType | list[tuple[int, int]] | None | None = None + ) -> list[dict[str, TensorType]]: + """ + Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions and depth PIL images. + Only supports PyTorch. + + Args: + outputs ([`DepthEstimatorOutput`]): + Raw outputs of the model. + target_sizes (`TensorType` 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. + + Returns: + `list[dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth + predictions. + """ + import torch + + predicted_depth = outputs.predicted_depth + + if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth" + ) + + results = [] + target_sizes = [None] * len(predicted_depth) if target_sizes is None else target_sizes + for depth, target_size in zip(predicted_depth, target_sizes): + if target_size is not None: + depth = torch.nn.functional.interpolate( + depth.unsqueeze(0).unsqueeze(1), size=target_size, mode="bicubic", align_corners=False + ).squeeze() + + results.append({"predicted_depth": depth}) + + return results + + +__all__ = ["PromptDepthAnythingImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..5d8bda90826970b1e5b005bbda2ab0a13a207838 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything.py @@ -0,0 +1,357 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for PromptDepthAnything.""" + +import math +from typing import TYPE_CHECKING, Union + +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, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, + requires_backends, +) + + +if TYPE_CHECKING: + from ...modeling_outputs import DepthEstimatorOutput + +import torch +from torchvision.transforms.v2 import functional as tvF + + +class PromptDepthAnythingImageProcessorKwargs(ImagesKwargs, total=False): + r""" + keep_aspect_ratio (`bool`, *optional*): + If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. + ensure_multiple_of (`int`, *optional*): + If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. + prompt_scale_to_meter (`float`, *optional*): + Scale factor to convert the prompt depth to meters. + size_divisor (`int`, *optional*): + If `do_pad` is `True`, pads the image dimensions to be divisible by this value. + """ + + keep_aspect_ratio: bool + ensure_multiple_of: int + size_divisor: int + prompt_scale_to_meter: float + + +def _constrain_to_multiple_of(val, multiple, min_val=0, max_val=None): + """Constrain a value to be a multiple of another value.""" + x = round(val / multiple) * multiple + + if max_val is not None and x > max_val: + x = math.floor(val / multiple) * multiple + + if x < min_val: + x = math.ceil(val / multiple) * multiple + + return x + + +def _get_resize_output_image_size( + input_image: "torch.Tensor", + output_size: tuple[int, int], + keep_aspect_ratio: bool, + multiple: int, +) -> tuple[int, int]: + """Get the output size for resizing an image.""" + input_height, input_width = input_image.shape[-2:] + output_height, output_width = output_size + + # determine new height and width + scale_height = output_height / input_height + scale_width = output_width / input_width + + if keep_aspect_ratio: + # scale as little as possible + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + + new_height = _constrain_to_multiple_of(scale_height * input_height, multiple=multiple) + new_width = _constrain_to_multiple_of(scale_width * input_width, multiple=multiple) + + return (new_height, new_width) + + +@auto_docstring +class PromptDepthAnythingImageProcessor(TorchvisionBackend): + model_input_names = ["pixel_values", "prompt_depth"] + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 384, "width": 384} + do_resize = True + do_rescale = True + do_normalize = True + keep_aspect_ratio = False + ensure_multiple_of = 1 + do_pad = False + size_divisor = None + prompt_scale_to_meter = 0.001 + valid_kwargs = PromptDepthAnythingImageProcessorKwargs + + def __init__(self, **kwargs: Unpack[PromptDepthAnythingImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + prompt_depth: ImageInput | None = None, + **kwargs: Unpack[PromptDepthAnythingImageProcessorKwargs], + ) -> BatchFeature: + r""" + prompt_depth (`ImageInput`, *optional*): + Prompt depth to preprocess. + """ + return super().preprocess(images, prompt_depth, **kwargs) + + def resize_with_aspect_ratio( + self, + image: "torch.Tensor", + size: SizeDict, + keep_aspect_ratio: bool = False, + ensure_multiple_of: int = 1, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + ) -> "torch.Tensor": + """ + Resize an image to target size while optionally maintaining aspect ratio and ensuring dimensions are multiples. + """ + # Set default resample to BICUBIC to match the slow processor (causes slight numerical differences otherwise) + if resample is None: + resample = PILImageResampling.BICUBIC + + # Custom resize with aspect ratio preservation and ensure_multiple_of constraint + output_size = _get_resize_output_image_size( + image, + output_size=(size.height, size.width), + keep_aspect_ratio=keep_aspect_ratio, + multiple=ensure_multiple_of, + ) + + # Standard resize method with calculated output size + return self.resize( + image=image, + size=SizeDict(height=output_size[0], width=output_size[1]), + resample=resample, + ) + + def pad_image( + self, + image: "torch.Tensor", + size_divisor: int, + ) -> "torch.Tensor": + """ + Center pad an image to be a multiple of size_divisor. + """ + + def _get_pad(size, size_divisor): + new_size = math.ceil(size / size_divisor) * size_divisor + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + height, width = image.shape[-2:] + + # Match slow processor and PyTorch convention: width->left/right, height->top/bottom + pad_size_left, pad_size_right = _get_pad(width, size_divisor) + pad_size_top, pad_size_bottom = _get_pad(height, size_divisor) + + # Use torchvision padding for fast processing + # /!\ NB: torchvision tvF.pad expects (left, top, right, bottom) for the last two dims (W then H) + # Source: https://docs.pytorch.org/vision/main/generated/torchvision.transforms.Pad.html + # So: (left=width_pad, top=height_pad, right=width_pad, bottom=height_pad) + padding = [pad_size_left, pad_size_top, pad_size_right, pad_size_bottom] + padded_image = tvF.pad(image, padding=padding) + + return padded_image + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + prompt_depth: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + device: Union[str, "torch.device"] | None = None, + return_tensors: str | TensorType | None = None, + prompt_scale_to_meter: float | None = None, + **kwargs, + ) -> BatchFeature: + """ + Preprocess image-like inputs, including the main images and optional prompt depth. + """ + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=False, input_data_format=input_data_format, device=device + ) # always use do_convert_rgb=False rather than defining it as a param to match slow processor + + # Process images with the standard pipeline + pixel_values = self._preprocess(images, return_tensors=return_tensors, **kwargs) + + data = {"pixel_values": pixel_values} + + # Process prompt depth if provided + if prompt_depth is not None: + processed_prompt_depths = self._prepare_image_like_inputs( + images=prompt_depth, + do_convert_rgb=False, # Depth maps should not be converted + input_data_format=input_data_format, + device=images[0].device if images else device, + expected_ndims=2, + ) + + # Validate prompt_depths has same length as images as in slow processor + if len(processed_prompt_depths) != len(images): + raise ValueError( + f"Number of prompt depth images ({len(processed_prompt_depths)}) does not match number of input images ({len(images)})" + ) + + if prompt_scale_to_meter is None: + prompt_scale_to_meter = self.prompt_scale_to_meter + + final_prompt_depths = [] + for depth in processed_prompt_depths: + depth = depth * prompt_scale_to_meter + + # Handle case where depth is constant (min == max) + if depth.min() == depth.max(): + depth[0, 0] = depth[0, 0] + 1e-6 # Add small variation to avoid numerical issues + + if depth.ndim == 2: # Add channel dimension if needed + depth = depth.unsqueeze(0) # [H, W] -> [1, H, W] (channels first) + + depth = depth.float() # Convert to float32 to match slow processor + final_prompt_depths.append(depth) + + data["prompt_depth"] = final_prompt_depths + + return BatchFeature(data=data, tensor_type=return_tensors) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + keep_aspect_ratio: bool | None = None, + ensure_multiple_of: int | None = None, + size_divisor: int | None = None, + **kwargs, + ) -> list["torch.Tensor"]: + """ + Override the base _preprocess method to handle custom PromptDepthAnything parameters. + """ + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize_with_aspect_ratio( + image=stacked_images, + size=size, + keep_aspect_ratio=keep_aspect_ratio, + ensure_multiple_of=ensure_multiple_of, + resample=resample, + ) + resized_images_grouped[shape] = stacked_images + + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + + if do_pad and size_divisor is not None: + stacked_images = self.pad_image(stacked_images, size_divisor) + + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + return processed_images + + def post_process_depth_estimation( + self, + outputs: "DepthEstimatorOutput", + target_sizes: TensorType | list[tuple[int, int]] | None | None = None, + ) -> list[dict[str, TensorType]]: + """ + Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions and depth PIL images. + Only supports PyTorch. + + Args: + outputs ([`DepthEstimatorOutput`]): + Raw outputs of the model. + target_sizes (`TensorType` 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. + + Returns: + `list[dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth + predictions. + """ + requires_backends(self, "torch") + + predicted_depth = outputs.predicted_depth + + if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth" + ) + + results = [] + target_sizes = [None] * len(predicted_depth) if target_sizes is None else target_sizes + for depth, target_size in zip(predicted_depth, target_sizes): + if target_size is not None: + depth = torch.nn.functional.interpolate( + depth.unsqueeze(0).unsqueeze(1), size=target_size, mode="bicubic", align_corners=False + ).squeeze() + + results.append({"predicted_depth": depth}) + + return results + + +__all__ = ["PromptDepthAnythingImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/modeling_prompt_depth_anything.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/modeling_prompt_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..9faecd40f0643f1a513d8e73af7374e7b905f210 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/modeling_prompt_depth_anything.py @@ -0,0 +1,480 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/prompt_depth_anything/modular_prompt_depth_anything.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_prompt_depth_anything.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn + +from ...backbone_utils import load_backbone +from ...modeling_outputs import DepthEstimatorOutput +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import merge_with_config_defaults, torch_int +from ...utils.output_capturing import capture_outputs +from .configuration_prompt_depth_anything import PromptDepthAnythingConfig + + +class PromptDepthAnythingLayer(nn.Module): + def __init__(self, config: PromptDepthAnythingConfig): + super().__init__() + self.convolution1 = nn.Conv2d( + 1, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + self.activation1 = nn.ReLU() + + self.convolution2 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + self.activation2 = nn.ReLU() + + self.convolution3 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + def forward(self, prompt_depth: torch.Tensor) -> torch.Tensor: + hidden_state = self.convolution1(prompt_depth) + hidden_state = self.activation1(hidden_state) + hidden_state = self.convolution2(hidden_state) + hidden_state = self.activation2(hidden_state) + hidden_state = self.convolution3(hidden_state) + return hidden_state + + +class PromptDepthAnythingPreActResidualLayer(nn.Module): + """ + ResidualConvUnit, pre-activate residual unit. + + Args: + config (`[PromptDepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.activation1 = nn.ReLU() + self.convolution1 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + self.activation2 = nn.ReLU() + self.convolution2 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + residual = hidden_state + hidden_state = self.activation1(hidden_state) + hidden_state = self.convolution1(hidden_state) + hidden_state = self.activation2(hidden_state) + hidden_state = self.convolution2(hidden_state) + + return hidden_state + residual + + +class PromptDepthAnythingFeatureFusionLayer(nn.Module): + """Feature fusion layer, merges feature maps from different stages. + + Args: + config (`[PromptDepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config: PromptDepthAnythingConfig): + super().__init__() + + self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True) + + self.residual_layer1 = PromptDepthAnythingPreActResidualLayer(config) + self.residual_layer2 = PromptDepthAnythingPreActResidualLayer(config) + self.prompt_depth_layer = PromptDepthAnythingLayer(config) + + def forward(self, hidden_state, residual=None, size=None, prompt_depth=None): + if residual is not None: + if hidden_state.shape != residual.shape: + residual = nn.functional.interpolate( + residual, size=hidden_state.shape[2:], mode="bilinear", align_corners=False + ) + hidden_state = hidden_state + self.residual_layer1(residual) + + hidden_state = self.residual_layer2(hidden_state) + + if prompt_depth is not None: + prompt_depth = nn.functional.interpolate( + prompt_depth, size=hidden_state.shape[2:], mode="bilinear", align_corners=False + ) + res = self.prompt_depth_layer(prompt_depth) + hidden_state = hidden_state + res + + modifier = {"scale_factor": 2} if size is None else {"size": size} + + hidden_state = nn.functional.interpolate( + hidden_state, + **modifier, + mode="bilinear", + align_corners=True, + ) + hidden_state = self.projection(hidden_state) + + return hidden_state + + +class PromptDepthAnythingFeatureFusionStage(nn.Module): + def __init__(self, config: PromptDepthAnythingConfig): + super().__init__() + self.layers = nn.ModuleList() + for _ in range(len(config.neck_hidden_sizes)): + self.layers.append(PromptDepthAnythingFeatureFusionLayer(config)) + + def forward(self, hidden_states, size=None, prompt_depth=None): + # reversing the hidden_states, we start from the last + hidden_states = hidden_states[::-1] + + fused_hidden_states = [] + fused_hidden_state = None + + for idx, (hidden_state, layer) in enumerate(zip(hidden_states, self.layers)): + size = hidden_states[idx + 1].shape[2:] if idx != (len(hidden_states) - 1) else None + + if fused_hidden_state is None: + # first layer only uses the last hidden_state + fused_hidden_state = layer(hidden_state, size=size, prompt_depth=prompt_depth) + else: + fused_hidden_state = layer(fused_hidden_state, hidden_state, size=size, prompt_depth=prompt_depth) + + fused_hidden_states.append(fused_hidden_state) + + return fused_hidden_states + + +class PromptDepthAnythingDepthEstimationHead(nn.Module): + """ + Output head consisting of 3 convolutional layers. It progressively halves the feature dimension and upsamples + the predictions to the input resolution after the first convolutional layer (details can be found in the DPT paper's + supplementary material). The final activation function is either ReLU or Sigmoid, depending on the depth estimation + type (relative or metric). For metric depth estimation, the output is scaled by the maximum depth used during pretraining. + """ + + def __init__(self, config): + super().__init__() + + self.head_in_index = config.head_in_index + self.patch_size = config.patch_size + + features = config.fusion_hidden_size + self.conv1 = nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1) + self.conv2 = nn.Conv2d(features // 2, config.head_hidden_size, kernel_size=3, stride=1, padding=1) + self.activation1 = nn.ReLU() + self.conv3 = nn.Conv2d(config.head_hidden_size, 1, kernel_size=1, stride=1, padding=0) + if config.depth_estimation_type == "relative": + self.activation2 = nn.ReLU() + elif config.depth_estimation_type == "metric": + self.activation2 = nn.Sigmoid() + else: + raise ValueError(f"Unknown depth estimation type: {config.depth_estimation_type}") + self.max_depth = config.max_depth + + def forward(self, hidden_states: list[torch.Tensor], patch_height: int, patch_width: int) -> torch.Tensor: + hidden_states = hidden_states[-1] + + predicted_depth = self.conv1(hidden_states) + target_height = torch_int(patch_height * self.patch_size) + target_width = torch_int(patch_width * self.patch_size) + predicted_depth = nn.functional.interpolate( + predicted_depth, + (target_height, target_width), + mode="bilinear", + align_corners=True, + ) + predicted_depth = self.conv2(predicted_depth) + predicted_depth = self.activation1(predicted_depth) + predicted_depth = self.conv3(predicted_depth) + predicted_depth = self.activation2(predicted_depth) + # (batch_size, 1, height, width) -> (batch_size, height, width), which + # keeps the same behavior as Depth Anything v1 & v2 + predicted_depth = predicted_depth.squeeze(dim=1) + + return predicted_depth + + +@auto_docstring +class PromptDepthAnythingPreTrainedModel(PreTrainedModel): + config: PromptDepthAnythingConfig + base_model_prefix = "prompt_depth_anything" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + + +class PromptDepthAnythingReassembleLayer(nn.Module): + def __init__(self, config: PromptDepthAnythingConfig, channels: int, factor: int): + super().__init__() + self.projection = nn.Conv2d(in_channels=config.reassemble_hidden_size, out_channels=channels, kernel_size=1) + + # up/down sampling depending on factor + if factor > 1: + self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0) + elif factor == 1: + self.resize = nn.Identity() + elif factor < 1: + # so should downsample + stride = torch_int(1 / factor) + self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=stride, padding=1) + + def forward(self, hidden_state): + hidden_state = self.projection(hidden_state) + hidden_state = self.resize(hidden_state) + + return hidden_state + + +class PromptDepthAnythingReassembleStage(nn.Module): + """ + This class reassembles the hidden states of the backbone into image-like feature representations at various + resolutions. + + This happens in 3 stages: + 1. Take the patch embeddings and reshape them to image-like feature representations. + 2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`. + 3. Resizing the spatial dimensions (height, width). + + Args: + config (`[PromptDepthAnythingConfig]`): + Model configuration class defining the model architecture. + """ + + def __init__(self, config): + super().__init__() + + self.config = config + self.layers = nn.ModuleList() + for channels, factor in zip(config.neck_hidden_sizes, config.reassemble_factors): + self.layers.append(PromptDepthAnythingReassembleLayer(config, channels=channels, factor=factor)) + + def forward(self, hidden_states: list[torch.Tensor], patch_height=None, patch_width=None) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length + 1, hidden_size)`): + List of hidden states from the backbone. + """ + out = [] + + for i, hidden_state in enumerate(hidden_states): + # reshape to (batch_size, num_channels, height, width) + hidden_state = hidden_state[:, 1:] + batch_size, _, num_channels = hidden_state.shape + hidden_state = hidden_state.reshape(batch_size, patch_height, patch_width, num_channels) + hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() + hidden_state = self.layers[i](hidden_state) + out.append(hidden_state) + + return out + + +class PromptDepthAnythingNeck(nn.Module): + """ + PromptDepthAnythingNeck. A neck is a module that is normally used between the backbone and the head. It takes a list of tensors as + input and produces another list of tensors as output. For PromptDepthAnything, it includes 2 stages: + + * PromptDepthAnythingReassembleStage + * PromptDepthAnythingFeatureFusionStage. + + Args: + config (dict): config dict. + """ + + def __init__(self, config): + super().__init__() + self.config = config + + self.reassemble_stage = PromptDepthAnythingReassembleStage(config) + + self.convs = nn.ModuleList() + for channel in config.neck_hidden_sizes: + self.convs.append(nn.Conv2d(channel, config.fusion_hidden_size, kernel_size=3, padding=1, bias=False)) + + # fusion + self.fusion_stage = PromptDepthAnythingFeatureFusionStage(config) + + def forward( + self, + hidden_states: list[torch.Tensor], + patch_height: int | None = None, + patch_width: int | None = None, + prompt_depth: torch.Tensor | None = None, + ) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`): + List of hidden states from the backbone. + """ + if not isinstance(hidden_states, (tuple, list)): + raise TypeError("hidden_states should be a tuple or list of tensors") + + if len(hidden_states) != len(self.config.neck_hidden_sizes): + raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.") + + # postprocess hidden states + hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width) + + features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)] + + # fusion blocks + output = self.fusion_stage(features, prompt_depth=prompt_depth) + + return output + + +@auto_docstring( + custom_intro=""" + Prompt Depth Anything Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2. + """ +) +class PromptDepthAnythingForDepthEstimation(PromptDepthAnythingPreTrainedModel): + _no_split_modules = ["DPTViTEmbeddings"] + + def __init__(self, config): + super().__init__(config) + + self.backbone = load_backbone(config) + self.neck = PromptDepthAnythingNeck(config) + self.head = PromptDepthAnythingDepthEstimationHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + prompt_depth: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | DepthEstimatorOutput: + r""" + prompt_depth (`torch.FloatTensor` of shape `(batch_size, 1, height, width)`, *optional*): + Prompt depth is the sparse or low-resolution depth obtained from multi-view geometry or a + low-resolution depth sensor. It generally has shape (height, width), where height + and width can be smaller than those of the images. It is optional and can be None, which means no prompt depth + will be used. If it is None, the output will be a monocular relative depth. + The values are recommended to be in meters, but this is not necessary. + + Example: + + ```python + >>> from transformers import AutoImageProcessor, AutoModelForDepthEstimation + >>> import torch + >>> import numpy as np + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/image.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("depth-anything/prompt-depth-anything-vits-hf") + >>> model = AutoModelForDepthEstimation.from_pretrained("depth-anything/prompt-depth-anything-vits-hf") + + >>> prompt_depth_url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/arkit_depth.png?raw=true" + >>> with httpx.stream("GET", prompt_depth_url) as response: + ... prompt_depth = Image.open(BytesIO(response.read())) + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt", prompt_depth=prompt_depth) + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # interpolate to original size + >>> post_processed_output = image_processor.post_process_depth_estimation( + ... outputs, + ... target_sizes=[(image.height, image.width)], + ... ) + + >>> # visualize the prediction + >>> predicted_depth = post_processed_output[0]["predicted_depth"] + >>> depth = predicted_depth * 1000. + >>> depth = depth.detach().cpu().numpy() + >>> depth = Image.fromarray(depth.astype("uint16")) # mm + ``` + """ + loss = None + if labels is not None: + raise NotImplementedError("Training is not implemented yet") + + outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, **kwargs) + hidden_states = outputs.feature_maps + + _, _, height, width = pixel_values.shape + patch_size = self.config.patch_size + patch_height = height // patch_size + patch_width = width // patch_size + + if prompt_depth is not None: + batch_size = prompt_depth.shape[0] + depth_min = torch.min(prompt_depth.reshape(batch_size, -1), dim=1).values + depth_max = torch.max(prompt_depth.reshape(batch_size, -1), dim=1).values + depth_min, depth_max = depth_min.view(batch_size, 1, 1, 1), depth_max.view(batch_size, 1, 1, 1) + prompt_depth = (prompt_depth - depth_min) / (depth_max - depth_min) + + hidden_states = self.neck(hidden_states, patch_height, patch_width, prompt_depth=prompt_depth) + + predicted_depth = self.head(hidden_states, patch_height, patch_width) + if prompt_depth is not None: + depth_min = depth_min.squeeze(1).to(predicted_depth.device) + depth_max = depth_max.squeeze(1).to(predicted_depth.device) + predicted_depth = predicted_depth * (depth_max - depth_min) + depth_min + + return DepthEstimatorOutput( + loss=loss, + predicted_depth=predicted_depth, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = ["PromptDepthAnythingForDepthEstimation", "PromptDepthAnythingPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/prompt_depth_anything/modular_prompt_depth_anything.py b/third_party/transformers/src/transformers/models/prompt_depth_anything/modular_prompt_depth_anything.py new file mode 100644 index 0000000000000000000000000000000000000000..6658a240230b4d31eb948186377f62815c58696c --- /dev/null +++ b/third_party/transformers/src/transformers/models/prompt_depth_anything/modular_prompt_depth_anything.py @@ -0,0 +1,328 @@ +# 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 as nn + +from ...modeling_outputs import DepthEstimatorOutput +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import merge_with_config_defaults, torch_int +from ...utils.output_capturing import capture_outputs +from ..depth_anything.configuration_depth_anything import DepthAnythingConfig +from ..depth_anything.modeling_depth_anything import ( + DepthAnythingDepthEstimationHead, + DepthAnythingFeatureFusionLayer, + DepthAnythingFeatureFusionStage, + DepthAnythingForDepthEstimation, + DepthAnythingNeck, + DepthAnythingReassembleStage, +) + + +class PromptDepthAnythingConfig(DepthAnythingConfig): + pass + + +class PromptDepthAnythingLayer(nn.Module): + def __init__(self, config: PromptDepthAnythingConfig): + super().__init__() + self.convolution1 = nn.Conv2d( + 1, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + self.activation1 = nn.ReLU() + + self.convolution2 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + self.activation2 = nn.ReLU() + + self.convolution3 = nn.Conv2d( + config.fusion_hidden_size, + config.fusion_hidden_size, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ) + + def forward(self, prompt_depth: torch.Tensor) -> torch.Tensor: + hidden_state = self.convolution1(prompt_depth) + hidden_state = self.activation1(hidden_state) + hidden_state = self.convolution2(hidden_state) + hidden_state = self.activation2(hidden_state) + hidden_state = self.convolution3(hidden_state) + return hidden_state + + +class PromptDepthAnythingFeatureFusionLayer(DepthAnythingFeatureFusionLayer): + def __init__(self, config: PromptDepthAnythingConfig): + super().__init__(config) + self.prompt_depth_layer = PromptDepthAnythingLayer(config) + + def forward(self, hidden_state, residual=None, size=None, prompt_depth=None): + if residual is not None: + if hidden_state.shape != residual.shape: + residual = nn.functional.interpolate( + residual, size=hidden_state.shape[2:], mode="bilinear", align_corners=False + ) + hidden_state = hidden_state + self.residual_layer1(residual) + + hidden_state = self.residual_layer2(hidden_state) + + if prompt_depth is not None: + prompt_depth = nn.functional.interpolate( + prompt_depth, size=hidden_state.shape[2:], mode="bilinear", align_corners=False + ) + res = self.prompt_depth_layer(prompt_depth) + hidden_state = hidden_state + res + + modifier = {"scale_factor": 2} if size is None else {"size": size} + + hidden_state = nn.functional.interpolate( + hidden_state, + **modifier, + mode="bilinear", + align_corners=True, + ) + hidden_state = self.projection(hidden_state) + + return hidden_state + + +class PromptDepthAnythingFeatureFusionStage(DepthAnythingFeatureFusionStage): + def forward(self, hidden_states, size=None, prompt_depth=None): + # reversing the hidden_states, we start from the last + hidden_states = hidden_states[::-1] + + fused_hidden_states = [] + fused_hidden_state = None + + for idx, (hidden_state, layer) in enumerate(zip(hidden_states, self.layers)): + size = hidden_states[idx + 1].shape[2:] if idx != (len(hidden_states) - 1) else None + + if fused_hidden_state is None: + # first layer only uses the last hidden_state + fused_hidden_state = layer(hidden_state, size=size, prompt_depth=prompt_depth) + else: + fused_hidden_state = layer(fused_hidden_state, hidden_state, size=size, prompt_depth=prompt_depth) + + fused_hidden_states.append(fused_hidden_state) + + return fused_hidden_states + + +class PromptDepthAnythingDepthEstimationHead(DepthAnythingDepthEstimationHead): + def forward(self, hidden_states: list[torch.Tensor], patch_height: int, patch_width: int) -> torch.Tensor: + hidden_states = hidden_states[-1] + + predicted_depth = self.conv1(hidden_states) + target_height = torch_int(patch_height * self.patch_size) + target_width = torch_int(patch_width * self.patch_size) + predicted_depth = nn.functional.interpolate( + predicted_depth, + (target_height, target_width), + mode="bilinear", + align_corners=True, + ) + predicted_depth = self.conv2(predicted_depth) + predicted_depth = self.activation1(predicted_depth) + predicted_depth = self.conv3(predicted_depth) + predicted_depth = self.activation2(predicted_depth) + # (batch_size, 1, height, width) -> (batch_size, height, width), which + # keeps the same behavior as Depth Anything v1 & v2 + predicted_depth = predicted_depth.squeeze(dim=1) + + return predicted_depth + + +@auto_docstring +class PromptDepthAnythingPreTrainedModel(PreTrainedModel): + config: PromptDepthAnythingConfig + base_model_prefix = "prompt_depth_anything" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + + +class PromptDepthAnythingReassembleLayer(nn.Module): + def __init__(self, config: PromptDepthAnythingConfig, channels: int, factor: int): + super().__init__() + self.projection = nn.Conv2d(in_channels=config.reassemble_hidden_size, out_channels=channels, kernel_size=1) + + # up/down sampling depending on factor + if factor > 1: + self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0) + elif factor == 1: + self.resize = nn.Identity() + elif factor < 1: + # so should downsample + stride = torch_int(1 / factor) + self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=stride, padding=1) + + def forward(self, hidden_state): + hidden_state = self.projection(hidden_state) + hidden_state = self.resize(hidden_state) + + return hidden_state + + +class PromptDepthAnythingReassembleStage(DepthAnythingReassembleStage): + pass + + +class PromptDepthAnythingNeck(DepthAnythingNeck): + def forward( + self, + hidden_states: list[torch.Tensor], + patch_height: int | None = None, + patch_width: int | None = None, + prompt_depth: torch.Tensor | None = None, + ) -> list[torch.Tensor]: + """ + Args: + hidden_states (`list[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`): + List of hidden states from the backbone. + """ + if not isinstance(hidden_states, (tuple, list)): + raise TypeError("hidden_states should be a tuple or list of tensors") + + if len(hidden_states) != len(self.config.neck_hidden_sizes): + raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.") + + # postprocess hidden states + hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width) + + features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)] + + # fusion blocks + output = self.fusion_stage(features, prompt_depth=prompt_depth) + + return output + + +@auto_docstring( + custom_intro=""" + Prompt Depth Anything Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2. + """ +) +class PromptDepthAnythingForDepthEstimation(DepthAnythingForDepthEstimation): + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + prompt_depth: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | DepthEstimatorOutput: + r""" + prompt_depth (`torch.FloatTensor` of shape `(batch_size, 1, height, width)`, *optional*): + Prompt depth is the sparse or low-resolution depth obtained from multi-view geometry or a + low-resolution depth sensor. It generally has shape (height, width), where height + and width can be smaller than those of the images. It is optional and can be None, which means no prompt depth + will be used. If it is None, the output will be a monocular relative depth. + The values are recommended to be in meters, but this is not necessary. + + Example: + + ```python + >>> from transformers import AutoImageProcessor, AutoModelForDepthEstimation + >>> import torch + >>> import numpy as np + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/image.jpg?raw=true" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("depth-anything/prompt-depth-anything-vits-hf") + >>> model = AutoModelForDepthEstimation.from_pretrained("depth-anything/prompt-depth-anything-vits-hf") + + >>> prompt_depth_url = "https://github.com/DepthAnything/PromptDA/blob/main/assets/example_images/arkit_depth.png?raw=true" + >>> with httpx.stream("GET", prompt_depth_url) as response: + ... prompt_depth = Image.open(BytesIO(response.read())) + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt", prompt_depth=prompt_depth) + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # interpolate to original size + >>> post_processed_output = image_processor.post_process_depth_estimation( + ... outputs, + ... target_sizes=[(image.height, image.width)], + ... ) + + >>> # visualize the prediction + >>> predicted_depth = post_processed_output[0]["predicted_depth"] + >>> depth = predicted_depth * 1000. + >>> depth = depth.detach().cpu().numpy() + >>> depth = Image.fromarray(depth.astype("uint16")) # mm + ``` + """ + loss = None + if labels is not None: + raise NotImplementedError("Training is not implemented yet") + + outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, **kwargs) + hidden_states = outputs.feature_maps + + _, _, height, width = pixel_values.shape + patch_size = self.config.patch_size + patch_height = height // patch_size + patch_width = width // patch_size + + if prompt_depth is not None: + batch_size = prompt_depth.shape[0] + depth_min = torch.min(prompt_depth.reshape(batch_size, -1), dim=1).values + depth_max = torch.max(prompt_depth.reshape(batch_size, -1), dim=1).values + depth_min, depth_max = depth_min.view(batch_size, 1, 1, 1), depth_max.view(batch_size, 1, 1, 1) + prompt_depth = (prompt_depth - depth_min) / (depth_max - depth_min) + + hidden_states = self.neck(hidden_states, patch_height, patch_width, prompt_depth=prompt_depth) + + predicted_depth = self.head(hidden_states, patch_height, patch_width) + if prompt_depth is not None: + depth_min = depth_min.squeeze(1).to(predicted_depth.device) + depth_max = depth_max.squeeze(1).to(predicted_depth.device) + predicted_depth = predicted_depth * (depth_max - depth_min) + depth_min + + return DepthEstimatorOutput( + loss=loss, + predicted_depth=predicted_depth, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "PromptDepthAnythingConfig", + "PromptDepthAnythingForDepthEstimation", + "PromptDepthAnythingPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/prophetnet/__init__.py b/third_party/transformers/src/transformers/models/prophetnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7485dfbb530c3e11683c2b7dceaf35bf9edc9086 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prophetnet/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_prophetnet import * + from .modeling_prophetnet import * + from .tokenization_prophetnet 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/prophetnet/configuration_prophetnet.py b/third_party/transformers/src/transformers/models/prophetnet/configuration_prophetnet.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4d5e48b32c85e73e9dfe8971dfa3fa5768217d --- /dev/null +++ b/third_party/transformers/src/transformers/models/prophetnet/configuration_prophetnet.py @@ -0,0 +1,89 @@ +# Copyright 2020 The Microsoft 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. +"""ProphetNet model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/prophetnet-large-uncased") +@strict +class ProphetNetConfig(PreTrainedConfig): + r""" + ngram (`int`, *optional*, defaults to 2): + Number of future tokens to predict. Set to 1 to be same as traditional Language model to predict next first + token. + num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. This is for relative position calculation. See the + [T5 paper](see https://huggingface.co/papers/1910.10683) for more details. + relative_max_distance (`int`, *optional*, defaults to 128): + Relative distances greater than this number will be put into the last same bucket. This is for relative + position calculation. See the [T5 paper](see https://huggingface.co/papers/1910.10683) for more details. + disable_ngram_loss (`bool`, *optional*, defaults to `False`): + Whether be trained predicting only the next first token. + eps (`float`, *optional*, defaults to 0.0): + Controls the `epsilon` parameter value for label smoothing in the loss calculation. If set to 0, no label + smoothing is performed. + """ + + model_type = "prophetnet" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "num_attention_heads": "num_encoder_attention_heads", + } + + activation_dropout: float | int = 0.1 + activation_function: str = "gelu" + vocab_size: int = 30522 + hidden_size: int = 1024 + encoder_ffn_dim: int = 4096 + num_encoder_layers: int = 12 + num_encoder_attention_heads: int = 16 + decoder_ffn_dim: int = 4096 + num_decoder_layers: int = 12 + num_decoder_attention_heads: int = 16 + attention_dropout: float | int = 0.1 + dropout: float | int = 0.1 + max_position_embeddings: int = 512 + init_std: float = 0.02 + is_encoder_decoder: bool = True + add_cross_attention: bool = True + decoder_start_token_id: int | None = 0 + ngram: int = 2 + num_buckets: int = 32 + relative_max_distance: int = 128 + disable_ngram_loss: bool = False + eps: float = 0.0 + use_cache: bool = True + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + is_decoder: bool = False + tie_word_embeddings: bool = True + + @property + def num_hidden_layers(self) -> int: + return self.num_encoder_layers + + @num_hidden_layers.setter + def num_hidden_layers(self, value): + raise NotImplementedError( + "This model does not support the setting of `num_hidden_layers`. Please set `num_encoder_layers` and" + " `num_decoder_layers`." + ) + + +__all__ = ["ProphetNetConfig"] diff --git a/third_party/transformers/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..f3f624b4090055965205f21e276b11fd27d218a7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,156 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert ProphetNet checkpoint.""" + +import argparse + +from torch import nn + +# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here +# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively +from transformers_old.modeling_prophetnet import ( + ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld, +) +from transformers_old.modeling_xlm_prophetnet import ( + XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld, +) + +from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging + + +logger = logging.get_logger(__name__) +logging.set_verbosity_info() + + +def convert_prophetnet_checkpoint_to_pytorch(prophetnet_checkpoint_path: str, pytorch_dump_folder_path: str): + """ + Copy/paste/tweak prohpetnet's weights to our prophetnet structure. + """ + if "xprophetnet" in prophetnet_checkpoint_path: + prophet_old = XLMProphetNetForConditionalGenerationOld.from_pretrained(prophetnet_checkpoint_path) + prophet, loading_info = XLMProphetNetForConditionalGeneration.from_pretrained( + prophetnet_checkpoint_path, output_loading_info=True + ) + else: + prophet_old = ProphetNetForConditionalGenerationOld.from_pretrained(prophetnet_checkpoint_path) + prophet, loading_info = ProphetNetForConditionalGeneration.from_pretrained( + prophetnet_checkpoint_path, output_loading_info=True + ) + + special_keys = ["key_proj", "value_proj", "query_proj"] + + mapping = { + "self_attn": "ngram_self_attn", + "cross_attn": "encoder_attn", + "cross_attn_layer_norm": "encoder_attn_layer_norm", + "feed_forward_layer_norm": "final_layer_norm", + "feed_forward": "", + "intermediate": "fc1", + "output": "fc2", + "key_proj": "k_proj", + "query_proj": "q_proj", + "value_proj": "v_proj", + "word_embeddings": "embed_tokens", + "embeddings_layer_norm": "emb_layer_norm", + "relative_pos_embeddings": "relative_linear", + "ngram_embeddings": "ngram_input_embed", + "position_embeddings": "embed_positions", + } + + for key in loading_info["missing_keys"]: + attributes = key.split(".") + + if attributes[0] == "lm_head": + model = prophet + old_model = prophet_old + else: + model = prophet.prophetnet + old_model = prophet_old.model + + is_key_init = False + for attribute in attributes: + if attribute in mapping: + old_attribute = mapping[attribute] + if not hasattr(old_model, old_attribute) and len(old_attribute) > 0: + old_attribute = attribute + elif hasattr(old_model, attribute): + old_attribute = attribute + + if attribute == "weight": + assert old_model.weight.shape == model.weight.shape, "Shapes have to match!" + model.weight = old_model.weight + logger.info(f"{attribute} is initialized.") + is_key_init = True + break + elif attribute == "bias": + assert old_model.bias.shape == model.bias.shape, "Shapes have to match!" + model.bias = old_model.bias + logger.info(f"{attribute} is initialized") + is_key_init = True + break + elif attribute in special_keys and hasattr(old_model, "in_proj_weight"): + embed_dim = old_model.in_proj_weight.shape[0] // 3 + param = getattr(model, attribute) + param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match" + param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match" + if attribute == "query_proj": + model.query_proj.weight = nn.Parameter(old_model.in_proj_weight[:embed_dim, :]) + model.query_proj.bias = nn.Parameter(old_model.in_proj_bias[:embed_dim]) + + elif attribute == "key_proj": + model.key_proj.weight = nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]) + model.key_proj.bias = nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]) + elif attribute == "value_proj": + model.value_proj.weight = nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]) + model.value_proj.bias = nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]) + is_key_init = True + break + elif attribute == "position_embeddings": + assert model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1], ( + "Hidden size has to match" + ) + assert model.position_embeddings.weight.shape[0] == 512, "We want 512 position_embeddings." + model.position_embeddings.weight = nn.Parameter(old_model.embed_positions.weight[:512, :]) + is_key_init = True + break + + if attribute.isdigit(): + model = model[int(attribute)] + old_model = old_model[int(old_attribute)] + else: + model = getattr(model, attribute) + + if old_attribute: + if not hasattr(old_model, old_attribute): + raise ValueError(f"{old_model} does not have {old_attribute}") + old_model = getattr(old_model, old_attribute) + + if not is_key_init: + raise ValueError(f"{key} was not correctly initialized!") + + print(f"Saving model to {pytorch_dump_folder_path}") + prophet.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--prophetnet_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump." + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model." + ) + args = parser.parse_args() + convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path) diff --git a/third_party/transformers/src/transformers/models/prophetnet/modeling_prophetnet.py b/third_party/transformers/src/transformers/models/prophetnet/modeling_prophetnet.py new file mode 100644 index 0000000000000000000000000000000000000000..86089cd98914f0cc59e0b61dd4f9c5ce8df128d4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/prophetnet/modeling_prophetnet.py @@ -0,0 +1,1846 @@ +# Copyright 2020 The Microsoft 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 ProphetNet model, ported from ProphetNet repo(fairsequery_states version).""" + +import copy +import math +from dataclasses import dataclass + +import torch +from torch import Tensor, nn +from torch.nn import LayerNorm + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_prophetnet import ProphetNetConfig + + +logger = logging.get_logger(__name__) + + +def softmax(hidden_state, dim, onnx_trace=False): + if onnx_trace: + return nn.functional.softmax(hidden_state.float(), dim=dim) + else: + return nn.functional.softmax(hidden_state, dim=dim, dtype=torch.float32) + + +def ngram_attention_bias(sequence_length, ngram, device, dtype): + """ + This function computes the bias for the predict stream + """ + left_block = ( + torch.ones((ngram, sequence_length, sequence_length), device=device, dtype=dtype) * torch.finfo(dtype).min + ) + right_block = left_block.detach().clone() + # create bias + for stream_idx in range(ngram): + right_block[stream_idx].fill_diagonal_(0, wrap=False) + left_block[stream_idx].triu_(-stream_idx + 1) + + left_block[:, :, 0] = 0 + return torch.cat([left_block, right_block], dim=2) + + +def compute_relative_buckets(num_buckets, max_distance, relative_positions, is_bidirectional=False): + """ + This function computes individual parts of the relative position buckets. For more detail, see paper. + """ + inv_relative_positions = -relative_positions + rel_positions_bucket = 0 + + if is_bidirectional: + num_buckets = num_buckets // 2 + rel_positions_bucket = ( + rel_positions_bucket + + torch.lt(inv_relative_positions, torch.zeros_like(inv_relative_positions)).int() * num_buckets + ) + inv_relative_positions = torch.abs(inv_relative_positions) + else: + inv_relative_positions = torch.max(inv_relative_positions, torch.zeros_like(inv_relative_positions)) + + max_exact = num_buckets // 2 + is_small = torch.lt(inv_relative_positions, max_exact) + val_if_large = max_exact + torch.log(inv_relative_positions.float() / max_exact) / math.log( + max_distance / max_exact + ) * (num_buckets - max_exact) + val_if_large = torch.min(val_if_large, torch.ones_like(val_if_large) * (num_buckets - 1)).int() + rel_positions_bucket = rel_positions_bucket + torch.where(is_small, inv_relative_positions.int(), val_if_large) + return rel_positions_bucket + + +def compute_all_stream_relative_buckets(num_buckets, max_distance, position_ids): + """ + This function computes both main and predict relative position buckets. For more detail, see paper. + """ + # main stream + main_stream_relative_positions = position_ids.unsqueeze(1).repeat(1, position_ids.size(-1), 1) + main_stream_relative_positions = main_stream_relative_positions - position_ids.unsqueeze(-1) + + # predicting stream + predicting_stream_relative_positions = torch.cat((position_ids - 1, position_ids), dim=-1).unsqueeze(1) + predicting_stream_relative_positions = predicting_stream_relative_positions.repeat(1, position_ids.size(-1), 1) + predicting_stream_relative_positions = predicting_stream_relative_positions - position_ids.unsqueeze(-1) + + # get both position buckets + main_relative_position_buckets = compute_relative_buckets( + num_buckets, max_distance, main_stream_relative_positions, is_bidirectional=False + ) + predict_relative_position_buckets = compute_relative_buckets( + num_buckets, max_distance, predicting_stream_relative_positions, is_bidirectional=False + ) + return main_relative_position_buckets, predict_relative_position_buckets + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for sequence-to-sequence language models outputs. + """ +) +class ProphetNetSeq2SeqLMOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + logits (`torch.FloatTensor` of shape `(batch_size, decoder_sequence_length, config.vocab_size)`): + Prediction scores of the main stream language modeling head (scores for each vocabulary token before + SoftMax). + logits_ngram (`torch.FloatTensor` of shape `(batch_size, ngram * decoder_sequence_length, config.vocab_size)`): + Prediction scores of the predict stream 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. + decoder_ngram_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 layer) of + shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`. + + Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding + outputs. + decoder_ngram_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_attn_heads, + decoder_sequence_length, decoder_sequence_length)`. + + Attentions weights of the predict stream of the decoder, after the attention softmax, used to compute the + weighted average in the self-attention heads. + encoder_last_hidden_state (`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 of the model. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + logits_ngram: torch.FloatTensor | None = None + past_key_values: Cache | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_ngram_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_attentions: tuple[torch.FloatTensor] | None = None + decoder_ngram_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=""" + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class ProphetNetSeq2SeqModelOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, decoder_sequence_length, hidden_size)`): + Sequence of main stream 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. + last_hidden_state_ngram (`torch.FloatTensor` of shape `(batch_size,ngram * decoder_sequence_length, config.vocab_size)`, *optional*): + Sequence of predict stream hidden-states at the output of the last layer of the decoder of the model. + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + decoder_ngram_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 layer) of + shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`. + + Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding + outputs. + decoder_ngram_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_attn_heads, + decoder_sequence_length, decoder_sequence_length)`. + + Attentions weights of the predict stream of the decoder, after the attention softmax, used to compute the + weighted average in the + encoder_last_hidden_state (`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 of the model. + """ + + last_hidden_state: torch.FloatTensor + last_hidden_state_ngram: torch.FloatTensor | None = None + past_key_values: Cache | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_ngram_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_attentions: tuple[torch.FloatTensor] | None = None + decoder_ngram_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=""" + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + """ +) +class ProphetNetDecoderModelOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, decoder_sequence_length, hidden_size)`): + Sequence of main stream 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. + last_hidden_state_ngram (`torch.FloatTensor` of shape `(batch_size, ngram * decoder_sequence_length, config.vocab_size)`): + Sequence of predict stream hidden-states at the output of the last layer of the decoder of the model. + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + hidden_states_ngram (`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 layer) of + shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`. + + Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding + outputs. + ngram_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_attn_heads, + decoder_sequence_length, decoder_sequence_length)`. + + Attentions weights of the predict stream of the decoder, after the attention softmax, used to compute the + weighted average in the + """ + + last_hidden_state: torch.FloatTensor + last_hidden_state_ngram: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states_ngram: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + ngram_attentions: tuple[torch.FloatTensor] | None = None + cross_attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + """ +) +class ProphetNetDecoderLMOutput(ModelOutput): + r""" + ngram_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 layer) of + shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`. + + Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding + outputs. + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + logits (`torch.FloatTensor` of shape `(batch_size, decoder_sequence_length, config.vocab_size)`): + Prediction scores of the main stream language modeling head (scores for each vocabulary token before + SoftMax). + logits_ngram (`torch.FloatTensor` of shape `(batch_size, ngram * decoder_sequence_length, config.vocab_size)`): + Prediction scores of the predict stream 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. + hidden_states_ngram (`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 layer) of + shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`. + + Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding + outputs. + ngram_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_attn_heads, + decoder_sequence_length, decoder_sequence_length)`. + + Attentions weights of the predict stream of the decoder, after the attention softmax, used to compute the + weighted average in the + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + logits_ngram: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states_ngram: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + ngram_attentions: tuple[torch.FloatTensor] | None = None + cross_attentions: tuple[torch.FloatTensor] | None = None + + +@auto_docstring +class ProphetNetPreTrainedModel(PreTrainedModel): + config: ProphetNetConfig + base_model_prefix = "prophetnet" + supports_gradient_checkpointing = True + + def _shift_right(self, input_ids): + decoder_start_token_id = self.config.decoder_start_token_id + pad_token_id = self.config.pad_token_id + + assert decoder_start_token_id is not None, ( + "self.model.config.decoder_start_token_id has to be defined. In ProphetNet it is usually set to the" + " pad_token_id. See ProphetNet docs for more information" + ) + + # shift inputs 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 + + assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined." + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + assert torch.all(shifted_input_ids >= 0).item(), "Verify that `shifted_input_ids` has only positive values" + + return shifted_input_ids + + +class ProphetNetPositionalEmbeddings(nn.Embedding): + """ + This module learns positional embeddings up to a fixed maximum size. Padding ids are ignored by either offsetting + based on padding_idx or by setting padding_idx to None and ensuring that the appropriate position ids are passed to + the forward function. + """ + + def __init__(self, config: ProphetNetConfig) -> None: + self.max_length = config.max_position_embeddings + super().__init__(config.max_position_embeddings, config.hidden_size, config.pad_token_id) + + def forward(self, inputs_shape, device, attention_mask=None, past_key_values=None, position_ids=None): + assert (position_ids is None) or (self.padding_idx is None), ( + "If position_ids is pre-computed then padding_idx should not be set." + ) + + if position_ids is None: + if past_key_values is not None and past_key_values.get_seq_length() != 0: + # position_ids is the same for every token when decoding a single step + # Without the int() cast, it doesn't work in some cases when exporting to ONNX + prev_num_input_ids = past_key_values.get_seq_length() + num_input_ids = inputs_shape[1] + prev_num_input_ids + position_ids = torch.ones((1, 1), dtype=torch.long, device=device) * ( + int(self.padding_idx + num_input_ids) + ) + else: + if attention_mask is None: + attention_mask = torch.ones(inputs_shape, dtype=torch.long, device=device) + + # retrieve position_ids from input_ids / attention_mask + position_ids = ( + torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask + ).long() + self.padding_idx + + # make sure position_ids are not bigger then max_length + position_ids = position_ids.clamp(0, self.max_length - 1) + + return super().forward(position_ids), position_ids + + def _forward(self, position_ids): + return super().forward(position_ids) + + +class ProphetNetAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: ProphetNetConfig, num_attn_heads: int, layer_idx: int | None = None): + super().__init__() + hidden_size = config.hidden_size + + self.attention_dropout = config.attention_dropout + self.dropout = config.dropout + self.num_attn_heads = num_attn_heads + self.head_dim = hidden_size // num_attn_heads + self.layer_idx = layer_idx + + assert self.head_dim * num_attn_heads == hidden_size, ( + "`config.hidden_size` must be divisible by `config.num_encoder_attention_heads` and" + " `config.num_decoder_attention_heads`" + ) + + self.key_proj = nn.Linear(hidden_size, hidden_size) + self.value_proj = nn.Linear(hidden_size, hidden_size) + self.query_proj = nn.Linear(hidden_size, hidden_size) + + self.out_proj = nn.Linear(hidden_size, hidden_size) + + def forward( + self, + hidden_states, + key_value_states: Tensor | None = None, + attention_mask: Tensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[Tensor, Tensor | None]: + batch_size, tgt_len, hidden_size = hidden_states.size() + + # 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 + assert list(hidden_states.size()) == [ + batch_size, + tgt_len, + hidden_size, + ], f"Size of hidden states should be {batch_size, tgt_len, hidden_size}, but is {hidden_states.size()}" + + # previous time steps are cached - no need to recompute key and value if they are static + query_states = self.query_proj(hidden_states) / (self.head_dim**0.5) + + 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.key_proj(current_states) + value_states = self.value_proj(current_states) + key_states = key_states.view(batch_size, -1, self.num_attn_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(batch_size, -1, self.num_attn_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 + + query_states = query_states.view(batch_size, tgt_len, self.num_attn_heads, self.head_dim).transpose(1, 2) + src_len = key_states.size(2) + + attn_weights = torch.einsum("bsij,bsjk->bsik", query_states, key_states.transpose(2, 3)) + expected_shape = (batch_size, self.num_attn_heads, tgt_len, src_len) + if attn_weights.size() != expected_shape: + raise ValueError(f"Attention weights should have size {expected_shape}, but is {attn_weights.size()}") + + # This is part of a workaround to get around fork/join parallelism not supporting Optional types. + if attention_mask is not None and attention_mask.dim() == 0: + attention_mask = None + + expected_shape = (batch_size, self.num_attn_heads, 1, src_len) + if attention_mask is not None and attention_mask.size() != expected_shape: + raise ValueError(f"Attention mask should have size {expected_shape}, but is {attention_mask.size()}") + if attention_mask is not None: # don't attend to padding symbols + attn_weights = attn_weights + attention_mask + if output_attentions: + attn_weights_reshaped = attn_weights + else: + attn_weights_reshaped = None + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_probs = nn.functional.dropout( + attn_weights, + p=self.attention_dropout, + training=self.training, + ) + attn_output = torch.einsum("bsij,bsjk->bsik", attn_probs, value_states) + expected_shape = (batch_size, self.num_attn_heads, tgt_len, self.head_dim) + if attn_output.size() != expected_shape: + raise ValueError(f"`attn_output` should have shape {expected_shape}, but is of shape {attn_output.size()}") + + attn_output = attn_output.transpose(1, 2).reshape(batch_size, tgt_len, hidden_size) + attn_output = self.out_proj(attn_output) + + attn_output = nn.functional.dropout(attn_output, p=self.dropout, training=self.training) + return attn_output, attn_weights_reshaped + + +class ProphetNetFeedForward(nn.Module): + """ + This is the residual two feed-forward layer block based on the original Transformer implementation. + """ + + def __init__(self, config: ProphetNetConfig, ffn_dim: int): + super().__init__() + self.activation_fn = ACT2FN[config.activation_function] + self.intermediate = nn.Linear(config.hidden_size, ffn_dim) + self.output = nn.Linear(ffn_dim, config.hidden_size) + self.activation_dropout = config.activation_dropout + self.dropout = config.dropout + + def forward(self, hidden_states): + hidden_states = self.intermediate(hidden_states) + hidden_states = self.activation_fn(hidden_states) + + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.output(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + return hidden_states + + +class ProphetNetNgramSelfAttention(nn.Module): + def __init__(self, config: ProphetNetConfig, layer_idx=None): + super().__init__() + self.hidden_size = config.hidden_size + + self.num_buckets = config.num_buckets + self.relative_max_distance = config.relative_max_distance + self.num_attn_heads = config.num_decoder_attention_heads + self.dropout = config.dropout + self.attention_dropout = config.attention_dropout + self.head_dim = config.hidden_size // self.num_attn_heads + self.ngram = config.ngram + self.layer_idx = layer_idx + + assert self.head_dim * self.num_attn_heads == config.hidden_size, ( + "config.hidden_size must be divisible by num_attn_heads" + ) + # key, value, query projection + self.key_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.value_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.query_proj = nn.Linear(config.hidden_size, config.hidden_size) + + # out projection + self.out_proj = nn.Linear(config.hidden_size, config.hidden_size) + + # rel position embeddings + self.relative_pos_embeddings = nn.Linear(config.hidden_size, self.num_buckets * self.num_attn_heads) + + # for onnx runtime + self.onnx_trace = False + + def _shape(self, tensor, seq_len, batch_size): + return tensor.view(batch_size, seq_len, self.num_attn_heads, self.head_dim).transpose(1, 2).contiguous() + + def prepare_for_onnx_export_(self): + self.onnx_trace = True + + def forward( + self, + hidden_states, + past_key_values: Cache | None = None, + attention_mask=None, + extended_predict_attention_mask=None, + main_relative_position_buckets=None, + predict_relative_position_buckets=None, + position_ids=None, + **kwargs, + ): + batch_size, ngram_sequence_length, hidden_size = hidden_states.size() + assert list(hidden_states.size()) == [batch_size, ngram_sequence_length, hidden_size], ( + f"`hidden_states` should be of shape {batch_size, ngram_sequence_length, hidden_size}, but is of shape" + f" {hidden_states.shape}" + ) + + # project + query_states = self.query_proj(hidden_states) + key_states = self.key_proj(hidden_states) + value_states = self.value_proj(hidden_states) + + # normalize + query_states = query_states / (self.head_dim**0.5) + + # reshape + query_states = self._shape(query_states, ngram_sequence_length, batch_size) + key_states = self._shape(key_states, -1, batch_size) + value_states = self._shape(value_states, -1, batch_size) + proj_shape = (batch_size, self.num_attn_heads, -1, self.head_dim) + + query_states = query_states.reshape(*proj_shape) + key_states = key_states.reshape(*proj_shape) + value_states = value_states.reshape(*proj_shape) + + # chunk into main stream and predict stream + hidden_states_list = hidden_states.chunk(1 + self.ngram, dim=1) + query_states_list = query_states.chunk(1 + self.ngram, dim=2) + key_states_list = key_states.chunk(1 + self.ngram, dim=2) + value_states_list = value_states.chunk(1 + self.ngram, dim=2) + + main_hidden_states, hidden_states_predict_list = hidden_states_list[0], hidden_states_list[1:] + main_query_states, predict_query_states_list = query_states_list[0], query_states_list[1:] + main_key_states, predict_key_states_list = key_states_list[0], key_states_list[1:] + main_value_states, predict_value_states_list = value_states_list[0], value_states_list[1:] + + # ProphetNet has two separate attention layers, one for self and one for cross attention + # We need to obtain the self attention only for this module, if `EncoderDecoderCache` + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + main_key_states, main_value_states = curr_past_key_values.update( + main_key_states, main_value_states, self.layer_idx + ) + + # get seq_length of main stream only + sequence_length = ngram_sequence_length // (1 + self.ngram) + + # MAIN-STREAM + # main attn weights + # [batch_size, number_heads, sequence_length, head_dimesion] + # x [batch_size, number_heads, head_dimesion, sequence_length] + # -> [batch_size, number_heads, sequence_length, sequence_length] + main_attn_weights = torch.einsum("bntc,bncs->bnts", main_query_states, main_key_states.transpose(2, 3)) + + # retrieve relative position embeddings for each layer -> see paper for more details + main_relative_pos_embeddings = self.get_main_relative_pos_embeddings( + main_hidden_states, main_attn_weights, position_ids, main_relative_position_buckets + ) + + main_attn_weights = main_attn_weights + main_relative_pos_embeddings + + if attention_mask is not None: + main_attn_weights = main_attn_weights + attention_mask + + main_attn_probs = softmax( + main_attn_weights, + dim=-1, + onnx_trace=self.onnx_trace, + ).type_as(main_attn_weights) + + main_attn_probs = nn.functional.dropout(main_attn_probs, p=self.attention_dropout, training=self.training) + # project to attn_output + # [batch_size, number_heads, sequence_length, sequence_length] + # x [batch_size, number_heads, sequence_length, head_dimesion] + # -> [batch_size, number_heads, sequence_length, head_dimesion] + main_attn_output = torch.einsum("bntc,bncs->bnts", main_attn_probs, main_value_states) + # reshape so that num_heads dim is merged into last `head_dim` axis + main_attn_output = main_attn_output.transpose(1, 2).reshape(batch_size, 1, sequence_length, hidden_size) + main_attn_output = self.out_proj(main_attn_output) + + # PREDICT-STREAM + # [batch_size, ngram, number_heads, sequence_length, head_dimesion] + predict_query_states = torch.stack(predict_query_states_list, 1).view( + batch_size, self.ngram, self.num_attn_heads, sequence_length, self.head_dim + ) + + # [batch_size, ngram, number_heads, 2*sequence_length, head_dimesion] + predict_key_states = torch.stack([torch.cat([main_key_states, key], 2) for key in predict_key_states_list], 1) + + # [batch_size, sequence_length, ngram, hidden_size] + predict_hidden_states = torch.stack(hidden_states_predict_list, dim=2) + + # [batch_size, number_heads, ngram, 2*sequence_length, head_dimesion] + predict_value_states = torch.cat( + [torch.cat([main_value_states, v_p], 2).unsqueeze(2) for v_p in predict_value_states_list], 2 + ) + + # [batch_size, ngram, number_heads, sequence_length, head_dimesion] + # x [batch_size, ngram, number_heads, 2*sequence_length, head_dimesion] + # -> [batch_size, ngram, number_heads, sequence_length, 2*sequence_length] + predict_attn_weights = torch.einsum("bnhtc,bnhsc->bnhts", (predict_query_states, predict_key_states)) + + # retrieve relative position embeddings for each layer -> see paper for more details + # [batch_size, ngram, number_heads, sequence_length, predict_relative_pos_embeddings] + predict_relative_pos_embeddings = self.get_predict_relative_pos_embeddings( + predict_hidden_states, predict_attn_weights, position_ids, predict_relative_position_buckets + ) + + # [batch_size, ngram, number_heads, sequence_length, 2*sequence_length] + predict_attn_weights = predict_attn_weights + predict_relative_pos_embeddings + + if extended_predict_attention_mask is not None: + # Permuting Predict attention mask to [batch_size, ngram, number_heads, sequence_length, 2*sequence_length] + extended_predict_attention_mask = extended_predict_attention_mask.permute(0, 2, 1, 3, 4) + extended_predict_attention_mask = extended_predict_attention_mask.to(predict_attn_weights.dtype) + predict_attn_weights = predict_attn_weights + extended_predict_attention_mask + + predict_attn_probs = softmax( + predict_attn_weights, + dim=-1, + onnx_trace=self.onnx_trace, + ).type_as(predict_attn_weights) + + predict_attn_probs = nn.functional.dropout( + predict_attn_probs, p=self.attention_dropout, training=self.training + ) + # project to attention output + # [batch_size, ngram, number_heads, sequence_length, 2*sequence_length] + # x [batch_size, ngram, number_heads, 2*sequence_length, head_dimesion] + # -> [batch_size, ngram, number_heads, sequence_length, head_dimesion] + predict_attn_output = torch.einsum( + "bnhts,bnhsc->bnhtc", (predict_attn_probs, predict_value_states.transpose(1, 2)) + ) + + # reshape so that num_heads dim is merged into last `head_dim` axis + # [batch_size, ngram, number_heads, sequence_length, head_dimesion] -> [batch_size, ngram, sequence_length, hidden_size] + predict_attn_output = predict_attn_output.transpose(2, 3) + predict_attn_output = predict_attn_output.reshape(batch_size, self.ngram, sequence_length, hidden_size) + predict_attn_output = self.out_proj(predict_attn_output) + + # concat to single attn output + # [batch_size, (1+ngram)*sequence_length, hidden_size] + attn_output = torch.cat([main_attn_output, predict_attn_output], 1).view(batch_size, -1, hidden_size) + # reshape into better form for `config.output_attentions` + main_attn_probs = main_attn_probs.view(batch_size, self.num_attn_heads, sequence_length, -1) + + attn_output = nn.functional.dropout(attn_output, p=self.dropout, training=self.training) + + return attn_output, main_attn_probs, predict_attn_probs + + def get_main_relative_pos_embeddings( + self, hidden_states, attn_weights, position_ids, main_relative_position_buckets + ): + # input hidden_states [batch_size, sequence_length, hidden_size] + # input attn_weights [batch_size, num_heads, sequence_length, sequence_length] + # input position_ids [batch_size, sequence_length] or [1,1] + batch_size, num_attn_heads, tgt_len, src_len = attn_weights.shape + attn_weights = attn_weights.view(batch_size, num_attn_heads, tgt_len, src_len) + if main_relative_position_buckets is None: + batch_size, sequence_length = hidden_states.shape[:2] + relative_positions = ( + torch.arange(1, attn_weights.shape[-1] + 1) + .unsqueeze(0) + .unsqueeze(0) + .repeat(batch_size, sequence_length, 1) + .to(position_ids.device) + ) + # [batch_size, sequence_length, sequence_length+1] + relative_positions = relative_positions - position_ids.unsqueeze(0).repeat(batch_size, sequence_length, 1) + main_relative_position_buckets = compute_relative_buckets( + self.num_buckets, self.relative_max_distance, relative_positions, False + ) + + # [batch_size, sequence_length, num_buckets * num_heads] + rel_pos_embeddings = self.relative_pos_embeddings(hidden_states) + rel_pos_embeddings = rel_pos_embeddings.view( + rel_pos_embeddings.shape[:2] + (self.num_buckets, self.num_attn_heads) + ) + rel_pos_embeddings = rel_pos_embeddings.permute(0, 3, 1, 2) + # [batch_size, num_heads, sequence_length, num_buckets] + rel_pos_embeddings = rel_pos_embeddings.reshape(attn_weights.shape[:3] + (-1,)) + + main_relative_position_buckets = main_relative_position_buckets.repeat(1, self.num_attn_heads, 1) + # [batch_size * num_heads * sequence_length, sequence_length] + main_relative_position_buckets = main_relative_position_buckets.view( + -1, main_relative_position_buckets.shape[-1] + ) + main_relative_position_buckets = main_relative_position_buckets.long() + # [batch_size * num_heads * sequence_length, sequence_length] + rel_pos_embeddings = rel_pos_embeddings.reshape(-1, rel_pos_embeddings.size(-1)) + + main_relative_pos_embeddings = torch.gather(rel_pos_embeddings, dim=1, index=main_relative_position_buckets) + main_relative_pos_embeddings = main_relative_pos_embeddings.view(batch_size, num_attn_heads, tgt_len, -1) + return main_relative_pos_embeddings + + def get_predict_relative_pos_embeddings( + self, hidden_states, attn_weights, position_ids, predict_relative_position_buckets + ): + # input hidden_states [batch_size, sequence_length, ngram, hidden_size] + # input attn_weights [batch_size, ngram, num_heads, sequence_length, 2*sequence_length] + # input position_ids [batch_size, sequence_length] or [1,1] + # input predict_relative_position_buckets [batch_size, sequence_length, 2*sequence_length] or None + batch_size, sequence_length = hidden_states.shape[0:2] + + if predict_relative_position_buckets is None: + key_sequence_length = attn_weights.shape[-1] + assert position_ids[0][0] == key_sequence_length - 1, ( + "`position_ids` are incorrect. They should be of the format 1 2 3 4 5 ... (key_sequence_length - 1)" + ) + relative_positions = ( + torch.arange(0, key_sequence_length) + .unsqueeze(0) + .unsqueeze(0) + .repeat(batch_size, sequence_length, 1) + .to(position_ids.device) + ) + + relative_positions = relative_positions - position_ids.unsqueeze(0).repeat(batch_size, sequence_length, 1) + predict_relative_position_buckets = compute_relative_buckets( + self.num_buckets, self.relative_max_distance, relative_positions, False + ) + + # [batch_size, ngram, sequence_length, hidden_size] + hidden_states = hidden_states.transpose(1, 2) + rel_pos_embeddings = self.relative_pos_embeddings(hidden_states) + + # [batch_size, ngram, sequence_length, num_buckets, num_heads] + rel_pos_embeddings = rel_pos_embeddings.view( + hidden_states.shape[:-1] + (self.num_buckets, self.num_attn_heads) + ) + rel_pos_embeddings = rel_pos_embeddings.permute(0, 2, 1, 4, 3) + # [batch_size * ngram * sequence_length * num_heads, num_buckets] + rel_pos_embeddings = rel_pos_embeddings.reshape(-1, self.num_buckets) + # [ngram, batch_size, num_heads * sequence_length, -1] + predict_relative_position_buckets = predict_relative_position_buckets.unsqueeze(0) + predict_relative_position_buckets = predict_relative_position_buckets.repeat( + self.ngram, 1, self.num_attn_heads, 1 + ) + # [ngram * batch_size * num_heads * sequence_length, -1] + predict_relative_position_buckets = predict_relative_position_buckets.view( + -1, predict_relative_position_buckets.size(-1) + ).long() + + predict_relative_pos_embeddings = torch.gather( + rel_pos_embeddings, dim=1, index=predict_relative_position_buckets + ) + + # [batch_size, gram, num_heads, sequence_length, -1] + predict_relative_pos_embeddings = predict_relative_pos_embeddings.view( + batch_size, self.ngram, self.num_attn_heads, sequence_length, -1 + ) + + return predict_relative_pos_embeddings + + +class ProphetNetEncoderLayer(GradientCheckpointingLayer): + """ + Encoder block for Prophetnet + """ + + def __init__(self, config: ProphetNetConfig): + super().__init__() + # 1st residual block + self.self_attn = ProphetNetAttention(config, config.num_encoder_attention_heads) + self.self_attn_layer_norm = LayerNorm(config.hidden_size) + + # 2nd residual block + self.feed_forward = ProphetNetFeedForward(config, config.encoder_ffn_dim) + self.feed_forward_layer_norm = LayerNorm(config.hidden_size) + + def forward( + self, + hidden_states, + attention_mask, + output_attentions: bool = False, + ): + # 1st residual block + attention_output, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + hidden_states = self.self_attn_layer_norm(attention_output + hidden_states) + + # 2nd residual block + feed_forward_output = self.feed_forward(hidden_states) + hidden_states = self.feed_forward_layer_norm(feed_forward_output + hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class ProphetNetDecoderLayer(GradientCheckpointingLayer): + """ + Decoder block for Prophetnet + """ + + def __init__(self, config: ProphetNetConfig, layer_idx=None): + super().__init__() + # 1st residual block + self.self_attn = ProphetNetNgramSelfAttention(config, layer_idx=layer_idx) + self.self_attn_layer_norm = LayerNorm(config.hidden_size) + + # 2nd residual block + if config.add_cross_attention: + self.cross_attn = ProphetNetAttention(config, config.num_decoder_attention_heads, layer_idx=layer_idx) + self.cross_attn_layer_norm = LayerNorm(config.hidden_size) + + # 3rd residual block + self.feed_forward = ProphetNetFeedForward(config, config.decoder_ffn_dim) + self.feed_forward_layer_norm = LayerNorm(config.hidden_size) + + def forward( + self, + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attn_mask=None, + extended_predict_attention_mask=None, + main_relative_position_buckets=None, + predict_relative_position_buckets=None, + position_ids=None, + past_key_values=None, + use_cache: bool | None = True, + output_attentions: bool | None = False, + **kwargs, + ): + # 1st residual block + ngram_attention_output, self_attn_weights, self_attn_weights_ngram = self.self_attn( + hidden_states=hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + extended_predict_attention_mask=extended_predict_attention_mask, + main_relative_position_buckets=main_relative_position_buckets, + predict_relative_position_buckets=predict_relative_position_buckets, + position_ids=position_ids, + ) + hidden_states = self.self_attn_layer_norm(hidden_states + ngram_attention_output) + + cross_attn_weights = None + if encoder_hidden_states is not None: + # 2nd residual block + attention_output, cross_attn_weights = self.cross_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attn_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + hidden_states = self.cross_attn_layer_norm(attention_output + hidden_states) + + # 3rd residual block + feed_forward_output = self.feed_forward(hidden_states) + hidden_states = self.feed_forward_layer_norm(feed_forward_output + hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights, self_attn_weights_ngram, cross_attn_weights) + + return outputs + + +@auto_docstring( + custom_intro=""" + The standalone encoder part of the ProphetNetModel. + """ +) +class ProphetNetEncoder(ProphetNetPreTrainedModel): + def __init__(self, config: ProphetNetConfig): + super().__init__(config) + + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = ProphetNetPositionalEmbeddings(config) + self.embeddings_layer_norm = LayerNorm(config.hidden_size) + + self.layers = nn.ModuleList([ProphetNetEncoderLayer(config) for _ in range(config.num_encoder_layers)]) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.word_embeddings + + def set_input_embeddings(self, value): + self.word_embeddings = value + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, ProphetNetEncoder + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = ProphetNetEncoder.from_pretrained("patrickvonplaten/prophetnet-large-uncased-standalone") + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if input_ids is None and inputs_embeds is None: + raise ValueError("Either input_ids or inputs_embeds has to be passed.") + elif input_ids is not None and inputs_embeds is not None: + raise ValueError("Make sure to only pass input_ids or inputs_embeds.") + elif input_ids is not None and inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + # prepare attention mask + if attention_mask is not None: + extended_attention_mask = ( + 1.0 - attention_mask[:, None, None, :].repeat(1, self.config.num_encoder_attention_heads, 1, 1) + ) * torch.finfo(self.dtype).min + extended_attention_mask = extended_attention_mask.to(inputs_embeds.dtype) + else: + extended_attention_mask = None + + position_embeddings, position_ids = self.position_embeddings(inputs_embeds.shape[:2], inputs_embeds.device) + + hidden_states = inputs_embeds + position_embeddings + hidden_states = self.embeddings_layer_norm(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.config.dropout, training=self.training) + + encoder_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_hidden_states = encoder_hidden_states + (hidden_states,) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask=extended_attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_hidden_states = encoder_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_hidden_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_hidden_states, attentions=all_attentions + ) + + +@auto_docstring( + custom_intro=""" + The standalone decoder part of the ProphetNetModel. + """ +) +class ProphetNetDecoder(ProphetNetPreTrainedModel): + def __init__(self, config: ProphetNetConfig): + super().__init__(config) + + self.ngram = config.ngram + self.num_buckets = config.num_buckets + self.relative_max_distance = config.relative_max_distance + self.dropout = config.dropout + self.max_target_positions = config.max_position_embeddings + + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = ProphetNetPositionalEmbeddings(config) + + self.ngram_embeddings = nn.Embedding(self.ngram, config.hidden_size, None) + self.layers = nn.ModuleList( + [ProphetNetDecoderLayer(config, layer_idx=i) for i in range(config.num_decoder_layers)] + ) + self.embeddings_layer_norm = LayerNorm(config.hidden_size) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.word_embeddings + + def set_input_embeddings(self, value): + self.word_embeddings = value + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + 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, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ProphetNetDecoderModelOutput: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, ProphetNetDecoder + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = ProphetNetDecoder.from_pretrained("microsoft/prophetnet-large-uncased", add_cross_attention=False) + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ```""" + use_cache = use_cache if use_cache is not None else self.config.use_cache + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if input_ids is None and inputs_embeds is None: + raise ValueError("Either `decoder_input_ids` or `decoder_inputs_embeds` has to be passed.") + elif input_ids is not None and inputs_embeds is not None: + raise ValueError("Make sure to only pass `decoder_input_ids` or `decoder_inputs_embeds`.") + elif input_ids is not None and inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + batch_size, sequence_length = inputs_embeds.shape[:2] + + 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)) + if encoder_hidden_states is not None or self.config.is_encoder_decoder + else DynamicCache(config=self.config) + ) + + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + + main_stream_pos_embed, position_ids = self.position_embeddings( + (batch_size, sequence_length), + device=inputs_embeds.device, + past_key_values=past_key_values, + ) + + if past_key_values_length != 0: + main_relative_position_buckets, predict_relative_position_buckets = None, None + else: + ( + main_relative_position_buckets, + predict_relative_position_buckets, + ) = self.compute_buffered_relative_buckets(position_ids) + predicting_stream_pos_embed = self.position_embeddings._forward(position_ids + 1) + + # add position embeddings + hidden_states = inputs_embeds + main_stream_pos_embed + + ngram_embeddings = self.ngram_embeddings.weight + + # prepare attention mask + if past_key_values_length != 0: + assert hidden_states.size(1) == 1, ( + "At the moment `use_cache` is only supported for `decoder_input_ids` of length 1" + ) + + ngram_hidden_states = [ + (ngram_embeddings[ngram - 1] + predicting_stream_pos_embed).repeat(batch_size, 1, 1) + for ngram in range(self.ngram) + ] + extended_attention_mask = None + extended_predict_attention_mask = None + else: + ngram_hidden_states = [ + (ngram_embeddings[ngram - 1] + predicting_stream_pos_embed) for ngram in range(self.ngram) + ] + extended_attention_mask = self.prepare_attention_mask(hidden_states, attention_mask) + extended_predict_attention_mask = self.prepare_predict_attention_mask(hidden_states, attention_mask) + + # prepare encoder attention mask + if encoder_attention_mask is not None: + extended_encoder_attention_mask = ( + 1.0 - encoder_attention_mask[:, None, None, :].repeat(1, self.config.num_decoder_attention_heads, 1, 1) + ) * torch.finfo(self.dtype).min + extended_encoder_attention_mask = extended_encoder_attention_mask.to(inputs_embeds.dtype) + else: + extended_encoder_attention_mask = None + + hidden_states = torch.cat([hidden_states] + ngram_hidden_states, 1) + + if self.embeddings_layer_norm: + hidden_states = self.embeddings_layer_norm(hidden_states) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + # init attentions, hidden_states and cache with empty tuples + all_main_stream_hidden_states = () if output_hidden_states else None + all_ngram_stream_hidden_states = () if output_hidden_states and self.config.ngram > 0 else None + + all_main_stream_attns = () if output_attentions else None + all_ngram_stream_attns = () if output_attentions else None + all_cross_attns = () if output_attentions and self.config.add_cross_attention else None + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + # grad cannot be kept because tensor is sliced + all_main_stream_hidden_states += (hidden_states[:, :sequence_length],) + if self.config.ngram > 0: + all_ngram_stream_hidden_states += (hidden_states[:, sequence_length:],) + + layer_outputs = decoder_layer( + hidden_states, + extended_attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attn_mask=extended_encoder_attention_mask, + extended_predict_attention_mask=extended_predict_attention_mask, + main_relative_position_buckets=main_relative_position_buckets, + predict_relative_position_buckets=predict_relative_position_buckets, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + if output_attentions: + all_main_stream_attns += (layer_outputs[1],) + all_ngram_stream_attns += (layer_outputs[2],) + + if self.config.add_cross_attention: + all_cross_attns += (layer_outputs[3],) + + if output_hidden_states: + all_main_stream_hidden_states += (hidden_states[:, :sequence_length],) + if self.config.ngram > 0: + all_ngram_stream_hidden_states += (hidden_states[:, sequence_length:],) + + # split last_hidden_state for return + last_hidden_state = hidden_states[:, :sequence_length] + last_hidden_state_ngram = hidden_states[:, sequence_length:] if self.config.ngram > 0 else None + + if not return_dict: + return tuple( + v + for v in [ + last_hidden_state, + last_hidden_state_ngram, + past_key_values, + all_main_stream_hidden_states, + all_ngram_stream_hidden_states, + all_main_stream_attns, + all_ngram_stream_attns, + all_cross_attns, + ] + if v is not None + ) + return ProphetNetDecoderModelOutput( + last_hidden_state=last_hidden_state, + last_hidden_state_ngram=last_hidden_state_ngram, + past_key_values=past_key_values, + hidden_states=all_main_stream_hidden_states, + hidden_states_ngram=all_ngram_stream_hidden_states, + attentions=all_main_stream_attns, + ngram_attentions=all_ngram_stream_attns, + cross_attentions=all_cross_attns, + ) + + def compute_buffered_relative_buckets(self, position_ids): + batch_size, sequence_length = position_ids.shape + + position_ids = torch.arange(1, self.max_target_positions).to(position_ids.device).repeat(1, 1) + main_relative_buckets, predict_relative_buckets = compute_all_stream_relative_buckets( + self.num_buckets, self.relative_max_distance, position_ids + ) + + # buffer relative buckets + main_relative_buckets = main_relative_buckets[:, :sequence_length, :sequence_length].repeat(batch_size, 1, 1) + predict_relative_buckets = torch.cat( + [ + predict_relative_buckets[:, :sequence_length, :sequence_length], + predict_relative_buckets[ + :, :sequence_length, self.max_target_positions : self.max_target_positions + sequence_length + ], + ], + 2, + ).repeat(batch_size, 1, 1) + + return main_relative_buckets, predict_relative_buckets + + def prepare_attention_mask(self, hidden_states, attention_mask): + batch_size, seq_length = hidden_states.shape[:2] + + # get causal mask + causal_mask = torch.full( + (seq_length, seq_length), + torch.finfo(hidden_states.dtype).min, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + causal_mask = torch.triu(causal_mask, 1) + + extended_causal_mask = causal_mask[:seq_length, :seq_length][None, None, :, :].expand( + (batch_size, self.config.num_decoder_attention_heads) + causal_mask.shape + ) + + # add usual attention mask + if attention_mask is not None: + extended_attention_mask = (1.0 - attention_mask[:, None, None, :]) * torch.finfo(self.dtype).min + extended_attention_mask = extended_causal_mask + extended_attention_mask + else: + extended_attention_mask = extended_causal_mask + return extended_attention_mask.to(hidden_states.dtype) + + def prepare_predict_attention_mask(self, hidden_states, attention_mask): + batch_size, seq_length = hidden_states.shape[:2] + + # get causal mask + predict_causal_mask = ngram_attention_bias( + self.max_target_positions, self.ngram, hidden_states.device, hidden_states.dtype + ) + predict_causal_mask = torch.cat( + [ + predict_causal_mask[:, :seq_length, :seq_length], + predict_causal_mask[ + :, :seq_length, self.max_target_positions : self.max_target_positions + seq_length + ], + ], + dim=-1, + ) + extended_predict_causal_mask = predict_causal_mask[None, None, :, :, :].expand( + (batch_size, self.config.num_decoder_attention_heads) + predict_causal_mask.shape + ) + + # add usual attention mask + if attention_mask is not None: + extended_attention_mask = (1.0 - attention_mask[:, None, None, None, :]) * torch.finfo(self.dtype).min + extended_attention_mask = extended_attention_mask.expand( + (batch_size, self.config.num_decoder_attention_heads, self.ngram, seq_length, seq_length) + ) + # predicted stream attention_mask should always be 0 + extended_attention_mask = torch.cat( + [extended_attention_mask, torch.zeros_like(extended_attention_mask)], dim=-1 + ) + extended_predict_attention_mask = extended_predict_causal_mask + extended_attention_mask + else: + extended_predict_attention_mask = extended_predict_causal_mask + return extended_predict_attention_mask.to(hidden_states.dtype) + + +@auto_docstring +class ProphetNetModel(ProphetNetPreTrainedModel): + _tied_weights_keys = { + "encoder.word_embeddings.weight": "word_embeddings.weight", + "decoder.word_embeddings.weight": "word_embeddings.weight", + } + + def __init__(self, config: ProphetNetConfig): + super().__init__(config) + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + + encoder_config = copy.deepcopy(config) + encoder_config.use_cache = False + self.encoder = ProphetNetEncoder(encoder_config) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + self.decoder = ProphetNetDecoder(decoder_config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.word_embeddings + + def set_input_embeddings(self, value): + self.word_embeddings = value + self.encoder.word_embeddings = self.word_embeddings + self.decoder.word_embeddings = self.word_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.BoolTensor | None = None, + encoder_outputs: tuple | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ProphetNetSeq2SeqModelOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + ProphetNet 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.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + + Example: + + ```python + >>> from transformers import AutoTokenizer, ProphetNetModel + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = ProphetNetModel.from_pretrained("microsoft/prophetnet-large-uncased") + + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> last_hidden_states = outputs.last_hidden_state # main stream hidden states + >>> last_hidden_states_ngram = outputs.last_hidden_state_ngram # predict hidden states + ```""" + use_cache = use_cache if use_cache is not None else self.config.use_cache + 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 encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + # 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, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + use_cache=use_cache, + return_dict=return_dict, + ) + + if not return_dict: + return decoder_outputs + encoder_outputs + return ProphetNetSeq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + last_hidden_state_ngram=decoder_outputs.last_hidden_state_ngram, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_ngram_hidden_states=decoder_outputs.hidden_states_ngram, + decoder_attentions=decoder_outputs.attentions, + decoder_ngram_attentions=decoder_outputs.ngram_attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The ProphetNet Model with a language modeling head. Can be used for sequence generation tasks. + """ +) +class ProphetNetForConditionalGeneration(ProphetNetPreTrainedModel, GenerationMixin): + _tied_weights_keys = { + "lm_head.weight": "prophetnet.word_embeddings.weight", + } + + def __init__(self, config: ProphetNetConfig): + super().__init__(config) + self.prophetnet = ProphetNetModel(config) + self.padding_idx = config.pad_token_id + self.disable_ngram_loss = config.disable_ngram_loss + + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.prophetnet.word_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.BoolTensor | None = None, + encoder_outputs: torch.Tensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ProphetNetSeq2SeqLMOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + ProphetNet 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.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., + config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for + labels in `[0, ..., config.vocab_size]` + + Example: + + ```python + >>> from transformers import AutoTokenizer, ProphetNetForConditionalGeneration + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = ProphetNetForConditionalGeneration.from_pretrained("microsoft/prophetnet-large-uncased") + + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> logits_next_token = outputs.logits # logits to predict next token as usual + >>> logits_ngram_next_tokens = outputs.logits_ngram # logits to predict 2nd, 3rd, ... next tokens + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: + # get decoder inputs from shifting lm labels to the right + decoder_input_ids = self._shift_right(labels) + + outputs = self.prophetnet( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + 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, + ) + batch_size, sequence_length = ( + decoder_input_ids.shape if decoder_input_ids is not None else decoder_inputs_embeds.shape[:2] + ) + + predicting_streams = outputs[1].view(batch_size, self.config.ngram, sequence_length, -1) + predict_logits = self.lm_head(predicting_streams) + + logits = predict_logits[:, 0] + logits_ngram = predict_logits[:, 1:] if self.config.ngram > 1 else None + + # To use .view in loss computation, make sure that logits is contiguous. + if not logits.is_contiguous(): + logits = logits.contiguous() + + loss = None + if labels is not None: + loss = self._compute_loss(predict_logits, labels) + + if not return_dict: + all_logits = tuple(v for v in [logits, logits_ngram] if v is not None) + return (loss,) + all_logits + outputs[2:] if loss is not None else all_logits + outputs[2:] + else: + return ProphetNetSeq2SeqLMOutput( + loss=loss, + logits=logits, + logits_ngram=logits_ngram, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_ngram_hidden_states=outputs.decoder_ngram_hidden_states, + decoder_attentions=outputs.decoder_attentions, + decoder_ngram_attentions=outputs.decoder_ngram_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + def _compute_loss(self, logits, labels, ignore_index=-100): + expend_targets = labels.new_zeros(self.config.ngram, labels.size(0), labels.size(1)).fill_(ignore_index) + + for i in range(self.config.ngram): + if i > 0 and self.disable_ngram_loss: + break + expend_targets[i, :, :] = labels + + logits = logits.transpose(0, 1).contiguous() + lprobs = nn.functional.log_softmax( + logits.view(-1, logits.size(-1)), + dim=-1, + dtype=torch.float32, + ) + + loss = nn.functional.nll_loss(lprobs, expend_targets.view(-1), reduction="mean") + + if self.config.eps > 0.0: + smooth_loss = -lprobs.sum(dim=-1, keepdim=True) + non_masked_tokens = expend_targets.ne(ignore_index).view(-1) + smooth_loss = smooth_loss[non_masked_tokens] + smooth_loss = smooth_loss.mean() + + eps_i = self.config.eps / lprobs.size(-1) + loss = (1.0 - self.config.eps) * loss + eps_i * smooth_loss + + return loss + + def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): + return self._shift_right(labels) + + def get_encoder(self, modality=None): + if modality is None: + return self.prophetnet.encoder + else: + return super().get_encoder(modality=modality) + + +@auto_docstring( + custom_intro=""" + The standalone decoder part of the ProphetNetModel with a lm head on top. The model can be used for causal + """ +) +class ProphetNetForCausalLM(ProphetNetPreTrainedModel, GenerationMixin): + _tied_weights_keys = { + "lm_head.weight": "prophetnet.word_embeddings.weight", + "prophetnet.decoder.word_embeddings.weight": "prophetnet.word_embeddings.weight", + } + + def __init__(self, config: ProphetNetConfig): + # set config for CLM + config = copy.deepcopy(config) + config.is_decoder = True + config.is_encoder_decoder = False + super().__init__(config) + self.prophetnet = ProphetNetDecoderWrapper(config) + + self.padding_idx = config.pad_token_id + self.disable_ngram_loss = config.disable_ngram_loss + + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.prophetnet.decoder.word_embeddings + + def set_input_embeddings(self, value): + self.prophetnet.decoder.word_embeddings = value + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + 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, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ProphetNetDecoderLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` + + Example: + + ```python + >>> from transformers import AutoTokenizer, ProphetNetForCausalLM + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = ProphetNetForCausalLM.from_pretrained("microsoft/prophetnet-large-uncased") + >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder." + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> logits = outputs.logits + + >>> # Model can also be used with EncoderDecoder framework + >>> from transformers import BertTokenizer, EncoderDecoderModel, AutoTokenizer + >>> import torch + + >>> tokenizer_enc = BertTokenizer.from_pretrained("google-bert/bert-large-uncased") + >>> tokenizer_dec = AutoTokenizer.from_pretrained("microsoft/prophetnet-large-uncased") + >>> model = EncoderDecoderModel.from_encoder_decoder_pretrained( + ... "google-bert/bert-large-uncased", "microsoft/prophetnet-large-uncased" + ... ) + + >>> ARTICLE = ( + ... "the us state department said wednesday it had received no " + ... "formal word from bolivia that it was expelling the us ambassador there " + ... "but said the charges made against him are `` baseless ." + ... ) + >>> input_ids = tokenizer_enc(ARTICLE, return_tensors="pt").input_ids + >>> labels = tokenizer_dec( + ... "us rejects charges against its ambassador in bolivia", return_tensors="pt" + ... ).input_ids + >>> outputs = model(input_ids=input_ids, decoder_input_ids=labels[:, :-1], labels=labels[:, 1:]) + + >>> loss = outputs.loss + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn) + outputs = self.prophetnet.decoder( + input_ids=input_ids, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + batch_size, sequence_length = input_ids.shape if input_ids is not None else inputs_embeds.shape[:2] + + predicting_streams = outputs[1].view(batch_size, self.config.ngram, sequence_length, -1) + predict_logits = self.lm_head(predicting_streams) + + logits = predict_logits[:, 0] + logits_ngram = predict_logits[:, 1:] if self.config.ngram > 1 else None + + loss = None + if labels is not None: + loss = self._compute_loss(predict_logits, labels) + + if not return_dict: + all_logits = tuple(v for v in [logits, logits_ngram] if v is not None) + return (loss,) + all_logits + outputs[2:] if loss is not None else all_logits + outputs[2:] + else: + return ProphetNetDecoderLMOutput( + loss=loss, + logits=logits, + logits_ngram=logits_ngram, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + hidden_states_ngram=outputs.hidden_states_ngram, + attentions=outputs.attentions, + ngram_attentions=outputs.ngram_attentions, + cross_attentions=outputs.cross_attentions, + ) + + def _compute_loss(self, logits, labels, ignore_index=-100): + expend_targets = labels.new_zeros(self.config.ngram, labels.size(0), labels.size(1)).fill_(ignore_index) + + for i in range(self.config.ngram): + if i > 0 and self.disable_ngram_loss: + break + expend_targets[i, :, :] = labels + + logits = logits.transpose(0, 1).contiguous() + lprobs = nn.functional.log_softmax( + logits.view(-1, logits.size(-1)), + dim=-1, + dtype=torch.float32, + ) + + loss = nn.functional.nll_loss(lprobs, expend_targets.view(-1), reduction="mean") + + if self.config.eps > 0.0: + smooth_loss = -lprobs.sum(dim=-1, keepdim=True) + non_masked_tokens = expend_targets.ne(ignore_index).view(-1) + smooth_loss = smooth_loss[non_masked_tokens] + smooth_loss = smooth_loss.mean() + + eps_i = self.config.eps / lprobs.size(-1) + loss = (1.0 - self.config.eps) * loss + eps_i * smooth_loss + + return loss + + +class ProphetNetDecoderWrapper(ProphetNetPreTrainedModel): + """ + This is a wrapper class, so that [`ProphetNetForCausalLM`] can correctly be loaded from pretrained prophetnet + classes. + """ + + _tied_weights_keys = { + "decoder.word_embeddings.weight": "word_embeddings.weight", + } + + def __init__(self, config: ProphetNetConfig): + super().__init__(config) + + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.decoder = ProphetNetDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def forward(self, *args, **kwargs): + return self.decoder(*args, **kwargs) + + +__all__ = [ + "ProphetNetDecoder", + "ProphetNetEncoder", + "ProphetNetForCausalLM", + "ProphetNetForConditionalGeneration", + "ProphetNetModel", + "ProphetNetPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/prophetnet/tokenization_prophetnet.py b/third_party/transformers/src/transformers/models/prophetnet/tokenization_prophetnet.py new file mode 100644 index 0000000000000000000000000000000000000000..b19c414e72d11e3b54b248fc2be7d16093eb419d --- /dev/null +++ b/third_party/transformers/src/transformers/models/prophetnet/tokenization_prophetnet.py @@ -0,0 +1,475 @@ +# Copyright 2020 The Microsoft 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. + +import collections +import os +import unicodedata +from collections.abc import Iterable + +from ...tokenization_python import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "prophetnet.tokenizer"} + + +def whitespace_tokenize(text): + """Runs basic whitespace cleaning and splitting on a piece of text.""" + text = text.strip() + if not text: + return [] + tokens = text.split() + return tokens + + +class BasicTokenizer: + """ + Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.). + + Args: + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + never_split (`Iterable`, *optional*): + Collection of tokens which will never be split during tokenization. Only has an effect when + `do_basic_tokenize=True` + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. + + This should likely be deactivated for Japanese (see this + [issue](https://github.com/huggingface/transformers/issues/328)). + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase` (as in the original BERT). + do_split_on_punc (`bool`, *optional*, defaults to `True`): + In some instances we want to skip the basic punctuation splitting so that later tokenization can capture + the full context of the words, such as contractions. + """ + + def __init__( + self, + do_lower_case=True, + never_split=None, + tokenize_chinese_chars=True, + strip_accents=None, + do_split_on_punc=True, + ): + if never_split is None: + never_split = [] + self.do_lower_case = do_lower_case + self.never_split = set(never_split) + self.tokenize_chinese_chars = tokenize_chinese_chars + self.strip_accents = strip_accents + self.do_split_on_punc = do_split_on_punc + + def tokenize(self, text, never_split=None): + """ + Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer. + + Args: + never_split (`List[str]`, *optional*) + Kept for backward compatibility purposes. Now implemented directly at the base class level (see + [`PreTrainedTokenizer.tokenize`]) List of token not to split. + """ + # union() returns a new set by concatenating the two sets. + never_split = self.never_split.union(set(never_split)) if never_split else self.never_split + text = self._clean_text(text) + + # This was added on November 1st, 2018 for the multilingual and Chinese + # models. This is also applied to the English models now, but it doesn't + # matter since the English models were not trained on any Chinese data + # and generally don't have any Chinese data in them (there are Chinese + # characters in the vocabulary because Wikipedia does have some Chinese + # words in the English Wikipedia.). + if self.tokenize_chinese_chars: + text = self._tokenize_chinese_chars(text) + # prevents treating the same character with different unicode codepoints as different characters + unicode_normalized_text = unicodedata.normalize("NFC", text) + orig_tokens = whitespace_tokenize(unicode_normalized_text) + split_tokens = [] + for token in orig_tokens: + if token not in never_split: + if self.do_lower_case: + token = token.lower() + if self.strip_accents is not False: + token = self._run_strip_accents(token) + elif self.strip_accents: + token = self._run_strip_accents(token) + split_tokens.extend(self._run_split_on_punc(token, never_split)) + + output_tokens = whitespace_tokenize(" ".join(split_tokens)) + return output_tokens + + def _run_strip_accents(self, text): + """Strips accents from a piece of text.""" + text = unicodedata.normalize("NFD", text) + output = [] + for char in text: + cat = unicodedata.category(char) + if cat == "Mn": + continue + output.append(char) + return "".join(output) + + def _run_split_on_punc(self, text, never_split=None): + """Splits punctuation on a piece of text.""" + if not self.do_split_on_punc or (never_split is not None and text in never_split): + return [text] + chars = list(text) + i = 0 + start_new_word = True + output = [] + while i < len(chars): + char = chars[i] + if _is_punctuation(char): + output.append([char]) + start_new_word = True + else: + if start_new_word: + output.append([]) + start_new_word = False + output[-1].append(char) + i += 1 + + return ["".join(x) for x in output] + + def _tokenize_chinese_chars(self, text): + """Adds whitespace around any CJK character.""" + output = [] + for char in text: + cp = ord(char) + if self._is_chinese_char(cp): + output.append(" ") + output.append(char) + output.append(" ") + else: + output.append(char) + return "".join(output) + + def _is_chinese_char(self, cp): + """Checks whether CP is the codepoint of a CJK character.""" + # This defines a "chinese character" as anything in the CJK Unicode block: + # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) + # + # Note that the CJK Unicode block is NOT all Japanese and Korean characters, + # despite its name. The modern Korean Hangul alphabet is a different block, + # as is Japanese Hiragana and Katakana. Those alphabets are used to write + # space-separated words, so they are not treated specially and handled + # like the all of the other languages. + if ( + (cp >= 0x4E00 and cp <= 0x9FFF) + or (cp >= 0x3400 and cp <= 0x4DBF) + or (cp >= 0x20000 and cp <= 0x2A6DF) + or (cp >= 0x2A700 and cp <= 0x2B73F) + or (cp >= 0x2B740 and cp <= 0x2B81F) + or (cp >= 0x2B820 and cp <= 0x2CEAF) + or (cp >= 0xF900 and cp <= 0xFAFF) + or (cp >= 0x2F800 and cp <= 0x2FA1F) + ): + return True + + return False + + def _clean_text(self, text): + """Performs invalid character removal and whitespace cleanup on text.""" + output = [] + for char in text: + cp = ord(char) + if cp == 0 or cp == 0xFFFD or _is_control(char): + continue + if _is_whitespace(char): + output.append(" ") + else: + output.append(char) + return "".join(output) + + +class WordpieceTokenizer: + """Runs WordPiece tokenization.""" + + def __init__(self, vocab, unk_token, max_input_chars_per_word=100): + self.vocab = vocab + self.unk_token = unk_token + self.max_input_chars_per_word = max_input_chars_per_word + + def tokenize(self, text): + """ + Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform + tokenization using the given vocabulary. + + For example, `input = "unaffable"` will return as output `["un", "##aff", "##able"]`. + + Args: + text: A single token or whitespace separated tokens. This should have + already been passed through *BasicTokenizer*. + + Returns: + A list of wordpiece tokens. + """ + + output_tokens = [] + for token in whitespace_tokenize(text): + chars = list(token) + if len(chars) > self.max_input_chars_per_word: + output_tokens.append(self.unk_token) + continue + + is_bad = False + start = 0 + sub_tokens = [] + while start < len(chars): + end = len(chars) + cur_substr = None + while start < end: + substr = "".join(chars[start:end]) + if start > 0: + substr = "##" + substr + if substr in self.vocab: + cur_substr = substr + break + end -= 1 + if cur_substr is None: + is_bad = True + break + sub_tokens.append(cur_substr) + start = end + + if is_bad: + output_tokens.append(self.unk_token) + else: + output_tokens.extend(sub_tokens) + return output_tokens + + +def load_vocab(vocab_file): + """Loads a vocabulary file into a dictionary.""" + vocab = collections.OrderedDict() + with open(vocab_file, "r", encoding="utf-8") as reader: + tokens = reader.readlines() + for index, token in enumerate(tokens): + token = token.rstrip("\n") + vocab[token] = index + return vocab + + +class ProphetNetTokenizer(PreTrainedTokenizer): + r""" + Construct a ProphetNetTokenizer. Based on WordPiece. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + File containing the vocabulary. + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + do_basic_tokenize (`bool`, *optional*, defaults to `True`): + Whether or not to do basic tokenization before WordPiece. + never_split (`Iterable`, *optional*): + Collection of tokens which will never be split during tokenization. Only has an effect when + `do_basic_tokenize=True` + unk_token (`str`, *optional*, defaults to `"[UNK]"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `"[SEP]"`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + x_sep_token (`str`, *optional*, defaults to `"[X_SEP]"`): + Special second separator token, which can be generated by [`ProphetNetForConditionalGeneration`]. It is + used to separate bullet-point like sentences in summarization, *e.g.*. + pad_token (`str`, *optional*, defaults to `"[PAD]"`): + The token used for padding, for example when batching sequences of different lengths. + mask_token (`str`, *optional*, defaults to `"[MASK]"`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. + + This should likely be deactivated for Japanese (see this + [issue](https://github.com/huggingface/transformers/issues/328)). + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase` (as in the original BERT). + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`): + Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like + extra spaces. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + # first name has to correspond to main model input name + # to make sure `tokenizer.pad(...)` works correctly + # `ProphetNet` doesn't have `token_type_ids` as argument. + model_input_names: list[str] = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file: str, + do_lower_case: bool | None = True, + do_basic_tokenize: bool | None = True, + never_split: Iterable | None = None, + unk_token: str | None = "[UNK]", + sep_token: str | None = "[SEP]", + x_sep_token: str | None = "[X_SEP]", + pad_token: str | None = "[PAD]", + mask_token: str | None = "[MASK]", + tokenize_chinese_chars: bool | None = True, + strip_accents: bool | None = None, + clean_up_tokenization_spaces: bool = True, + **kwargs, + ): + if not os.path.isfile(vocab_file): + raise ValueError( + f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained" + " model use `tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + self.vocab = load_vocab(vocab_file) + self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()]) + self.do_basic_tokenize = do_basic_tokenize + if do_basic_tokenize: + self.basic_tokenizer = BasicTokenizer( + do_lower_case=do_lower_case, + never_split=never_split, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + ) + self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token)) + + super().__init__( + do_lower_case=do_lower_case, + do_basic_tokenize=do_basic_tokenize, + never_split=never_split, + unk_token=unk_token, + sep_token=sep_token, + x_sep_token=x_sep_token, + pad_token=pad_token, + mask_token=mask_token, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + @property + def vocab_size(self): + return len(self.vocab) + + def get_vocab(self): + return dict(self.vocab, **self.added_tokens_encoder) + + def _tokenize(self, text): + split_tokens = [] + if self.do_basic_tokenize: + for token in self.basic_tokenizer.tokenize(text, never_split=self.all_special_tokens): + # If the token is part of the never_split set + if token in self.basic_tokenizer.never_split: + split_tokens.append(token) + else: + split_tokens += self.wordpiece_tokenizer.tokenize(token) + else: + split_tokens = self.wordpiece_tokenizer.tokenize(text) + return split_tokens + + def _convert_token_to_id(self, token: str): + """Converts a token (str) in an id using the vocab.""" + return self.vocab.get(token, self.vocab.get(self.unk_token)) + + def _convert_id_to_token(self, index: int): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.ids_to_tokens.get(index, self.unk_token) + + def convert_tokens_to_string(self, tokens: str): + """Converts a sequence of tokens (string) in a single string.""" + out_string = " ".join(tokens).replace(" ##", "").strip() + return out_string + + def get_special_tokens_mask( + self, + token_ids_0: list[int], + token_ids_1: list[int] | None = None, + already_has_special_tokens: bool | None = 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 ([0] * len(token_ids_0)) + [1] + return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + index = 0 + if os.path.isdir(save_directory): + vocab_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] + ) + else: + vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory + with open(vocab_file, "w", encoding="utf-8") as writer: + for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]): + if index != token_index: + logger.warning( + f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive." + " Please check that the vocabulary is not corrupted!" + ) + index = token_index + writer.write(token + "\n") + index += 1 + return (vocab_file,) + + def build_inputs_with_special_tokens( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. A BERT sequence has the following format: + + - single sequence: `[CLS] X [SEP]` + - pair of sequences: `[CLS] A [SEP] B [SEP]` + + Args: + token_ids_0 (`List[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + """ + if token_ids_1 is None: + return token_ids_0 + [self.sep_token_id] + sep = [self.sep_token_id] + return token_ids_0 + sep + token_ids_1 + sep + + +__all__ = ["ProphetNetTokenizer"] diff --git a/third_party/transformers/src/transformers/models/pvt_v2/__init__.py b/third_party/transformers/src/transformers/models/pvt_v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3cb83f130dd42eb08943d5d4e08ed9321e6f076 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pvt_v2/__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_pvt_v2 import * + from .modeling_pvt_v2 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/pvt_v2/configuration_pvt_v2.py b/third_party/transformers/src/transformers/models/pvt_v2/configuration_pvt_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..1c73fa158b5d38e702503930c1833f61db8344b8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pvt_v2/configuration_pvt_v2.py @@ -0,0 +1,93 @@ +# Copyright 2024 Authors: Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, +# Kaitao Song, Ding Liang, Tong Lu, Ping Luo, Ling Shao 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. +"""Pvt V2 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import BackboneConfigMixin +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="OpenGVLab/pvt_v2_b0") +@strict +class PvtV2Config(BackboneConfigMixin, PreTrainedConfig): + r""" + num_encoder_blocks (`[int]`, *optional*, defaults to 4): + The number of encoder blocks (i.e. stages in the Mix Transformer encoder). + sr_ratios (`list[int]`, *optional*, defaults to `[8, 4, 2, 1]`): + Spatial reduction ratios in each encoder block. + patch_sizes (`list[int]`, *optional*, defaults to `[7, 3, 3, 3]`): + Patch size for overlapping patch embedding before each encoder block. + strides (`list[int]`, *optional*, defaults to `[4, 2, 2, 2]`): + Stride for overlapping patch embedding before each encoder block. + num_attention_heads (`list[int]`, *optional*, defaults to `[1, 2, 5, 8]`): + Number of attention heads for each attention layer in each block of the Transformer encoder. + mlp_ratios (`list[int]`, *optional*, defaults to `[8, 8, 4, 4]`): + Ratio of the size of the hidden layer compared to the size of the input layer of the Mix FFNs in the + encoder blocks. + linear_attention (`bool`, *optional*, defaults to `False`): + Use linear attention complexity. If set to True, `sr_ratio` is ignored and average pooling is used for + dimensionality reduction in the attention layers rather than strided convolution. + + Example: + + ```python + >>> from transformers import PvtV2Model, PvtV2Config + + >>> # Initializing a pvt_v2_b0 style configuration + >>> configuration = PvtV2Config() + + >>> # Initializing a model from the OpenGVLab/pvt_v2_b0 style configuration + >>> model = PvtV2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pvt_v2" + + image_size: int | list[int] | tuple[int, int] | dict = 224 + num_channels: int = 3 + num_encoder_blocks: int = 4 + depths: list[int] | tuple[int, ...] = (2, 2, 2, 2) + sr_ratios: list[int] | tuple[int, ...] = (8, 4, 2, 1) + hidden_sizes: list[int] | tuple[int, ...] = (32, 64, 160, 256) + patch_sizes: list[int] | tuple[int, ...] = (7, 3, 3, 3) + strides: list[int] | tuple[int, ...] = (4, 2, 2, 2) + num_attention_heads: list[int] | tuple[int, ...] = (1, 2, 5, 8) + mlp_ratios: list[int] | tuple[int, ...] = (8, 8, 4, 4) + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.0 + attention_probs_dropout_prob: float | int = 0.0 + initializer_range: float = 0.02 + drop_path_rate: float | int = 0.0 + layer_norm_eps: float = 1e-6 + qkv_bias: bool = True + linear_attention: bool = False + _out_features: list[str] | None = None + _out_indices: list[int] | None = None + + def __post_init__(self, **kwargs): + self.image_size = (self.image_size, self.image_size) if isinstance(self.image_size, int) else self.image_size + self.stage_names = [f"stage{idx}" for idx in range(1, len(self.depths) + 1)] + self.set_output_features_output_indices( + out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None) + ) + super().__post_init__(**kwargs) + + +__all__ = ["PvtV2Config"] diff --git a/third_party/transformers/src/transformers/models/pvt_v2/convert_pvt_v2_to_pytorch.py b/third_party/transformers/src/transformers/models/pvt_v2/convert_pvt_v2_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..2316ca36b1ae5ac92483ae1841a3777e2d767a9b --- /dev/null +++ b/third_party/transformers/src/transformers/models/pvt_v2/convert_pvt_v2_to_pytorch.py @@ -0,0 +1,295 @@ +# Copyright 2023 Authors: Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, +# Kaitao Song, Ding Liang, Tong Lu, Ping Luo, Ling Shao 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. +"""Convert PvtV2 checkpoints from the original library.""" + +import argparse +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from PIL import Image + +from transformers import PvtImageProcessor, PvtV2Config, PvtV2ForImageClassification +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def create_rename_keys(config): + rename_keys = [] + for i in range(config.num_encoder_blocks): + # Rename embeddings' parameters + rename_keys.append( + (f"patch_embed{i + 1}.proj.weight", f"pvt_v2.encoder.layers.{i}.patch_embedding.proj.weight") + ) + rename_keys.append((f"patch_embed{i + 1}.proj.bias", f"pvt_v2.encoder.layers.{i}.patch_embedding.proj.bias")) + rename_keys.append( + (f"patch_embed{i + 1}.norm.weight", f"pvt_v2.encoder.layers.{i}.patch_embedding.layer_norm.weight") + ) + rename_keys.append( + (f"patch_embed{i + 1}.norm.bias", f"pvt_v2.encoder.layers.{i}.patch_embedding.layer_norm.bias") + ) + rename_keys.append((f"norm{i + 1}.weight", f"pvt_v2.encoder.layers.{i}.layer_norm.weight")) + rename_keys.append((f"norm{i + 1}.bias", f"pvt_v2.encoder.layers.{i}.layer_norm.bias")) + + for j in range(config.depths[i]): + # Rename blocks' parameters + rename_keys.append( + (f"block{i + 1}.{j}.attn.q.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.query.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.attn.q.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.query.bias") + ) + rename_keys.append( + (f"block{i + 1}.{j}.attn.kv.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.kv.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.attn.kv.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.kv.bias") + ) + + if config.linear_attention or config.sr_ratios[i] > 1: + rename_keys.append( + ( + f"block{i + 1}.{j}.attn.norm.weight", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.layer_norm.weight", + ) + ) + rename_keys.append( + ( + f"block{i + 1}.{j}.attn.norm.bias", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.layer_norm.bias", + ) + ) + rename_keys.append( + ( + f"block{i + 1}.{j}.attn.sr.weight", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.spatial_reduction.weight", + ) + ) + rename_keys.append( + ( + f"block{i + 1}.{j}.attn.sr.bias", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.spatial_reduction.bias", + ) + ) + + rename_keys.append( + (f"block{i + 1}.{j}.attn.proj.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.proj.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.attn.proj.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.proj.bias") + ) + + rename_keys.append( + (f"block{i + 1}.{j}.norm1.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.layer_norm_1.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.norm1.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.layer_norm_1.bias") + ) + + rename_keys.append( + (f"block{i + 1}.{j}.norm2.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.layer_norm_2.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.norm2.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.layer_norm_2.bias") + ) + + rename_keys.append( + (f"block{i + 1}.{j}.mlp.fc1.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dense1.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.mlp.fc1.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dense1.bias") + ) + rename_keys.append( + ( + f"block{i + 1}.{j}.mlp.dwconv.dwconv.weight", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dwconv.dwconv.weight", + ) + ) + rename_keys.append( + ( + f"block{i + 1}.{j}.mlp.dwconv.dwconv.bias", + f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dwconv.dwconv.bias", + ) + ) + rename_keys.append( + (f"block{i + 1}.{j}.mlp.fc2.weight", f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dense2.weight") + ) + rename_keys.append( + (f"block{i + 1}.{j}.mlp.fc2.bias", f"pvt_v2.encoder.layers.{i}.blocks.{j}.mlp.dense2.bias") + ) + + rename_keys.extend( + [ + ("head.weight", "classifier.weight"), + ("head.bias", "classifier.bias"), + ] + ) + + return rename_keys + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_k_v(state_dict, config): + # for each of the encoder blocks: + for i in range(config.num_encoder_blocks): + for j in range(config.depths[i]): + # read in weights + bias of keys and values (which is a single matrix in the original implementation) + kv_weight = state_dict.pop(f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.kv.weight") + kv_bias = state_dict.pop(f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.kv.bias") + # next, add keys and values (in that order) to the state dict + state_dict[f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.key.weight"] = kv_weight[ + : config.hidden_sizes[i], : + ] + state_dict[f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.key.bias"] = kv_bias[: config.hidden_sizes[i]] + + state_dict[f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.value.weight"] = kv_weight[ + config.hidden_sizes[i] :, : + ] + state_dict[f"pvt_v2.encoder.layers.{i}.blocks.{j}.attention.value.bias"] = kv_bias[ + config.hidden_sizes[i] : + ] + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_pvt_v2_checkpoint(pvt_v2_size, pvt_v2_checkpoint, pytorch_dump_folder_path, verify_imagenet_weights=False): + """ + Copy/paste/tweak model's weights to our PVT structure. + """ + + # define default PvtV2 configuration + if pvt_v2_size == "b0": + config_path = "OpenGVLab/pvt_v2_b0" + elif pvt_v2_size == "b1": + config_path = "OpenGVLab/pvt_v2_b1" + elif pvt_v2_size == "b2": + config_path = "OpenGVLab/pvt_v2_b2" + elif pvt_v2_size == "b2-linear": + config_path = "OpenGVLab/pvt_v2_b2_linear" + elif pvt_v2_size == "b3": + config_path = "OpenGVLab/pvt_v2_b3" + elif pvt_v2_size == "b4": + config_path = "OpenGVLab/pvt_v2_b4" + elif pvt_v2_size == "b5": + config_path = "OpenGVLab/pvt_v2_b5" + else: + raise ValueError( + f"Available model sizes: 'b0', 'b1', 'b2', 'b2-linear', 'b3', 'b4', 'b5', but '{pvt_v2_size}' was given" + ) + config = PvtV2Config.from_pretrained(config_path) + # load original model from https://github.com/whai362/PVT + state_dict = torch.load(pvt_v2_checkpoint, map_location="cpu", weights_only=True) + + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + read_in_k_v(state_dict, config) + + # load HuggingFace model + model = PvtV2ForImageClassification(config).eval() + model.load_state_dict(state_dict) + image_processor = PvtImageProcessor(size=config.image_size) + + if verify_imagenet_weights: + # Check outputs on an image, prepared by PvtImageProcessor + print("Verifying conversion of pretrained ImageNet weights...") + encoding = image_processor(images=prepare_img(), return_tensors="pt") + pixel_values = encoding["pixel_values"] + outputs = model(pixel_values) + logits = outputs.logits.detach().cpu() + + if pvt_v2_size == "b0": + expected_slice_logits = torch.tensor([-1.1939, -1.4547, -0.1076]) + elif pvt_v2_size == "b1": + expected_slice_logits = torch.tensor([-0.4716, -0.7335, -0.4600]) + elif pvt_v2_size == "b2": + expected_slice_logits = torch.tensor([0.0795, -0.3170, 0.2247]) + elif pvt_v2_size == "b2-linear": + expected_slice_logits = torch.tensor([0.0968, 0.3937, -0.4252]) + elif pvt_v2_size == "b3": + expected_slice_logits = torch.tensor([-0.4595, -0.2870, 0.0940]) + elif pvt_v2_size == "b4": + expected_slice_logits = torch.tensor([-0.1769, -0.1747, -0.0143]) + elif pvt_v2_size == "b5": + expected_slice_logits = torch.tensor([-0.2943, -0.1008, 0.6812]) + else: + raise ValueError( + f"Available model sizes: 'b0', 'b1', 'b2', 'b2-linear', 'b3', 'b4', 'b5', but " + f"'{pvt_v2_size}' was given" + ) + + assert torch.allclose(logits[0, :3], expected_slice_logits, atol=1e-4), ( + "ImageNet weights not converted successfully." + ) + + print("ImageNet weights verified, conversion successful.") + + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model pytorch_model.bin to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--pvt_v2_size", + default="b0", + type=str, + help="Size of the PVTv2 pretrained model you'd like to convert.", + ) + parser.add_argument( + "--pvt_v2_checkpoint", + default="pvt_v2_b0.pth", + type=str, + help="Checkpoint of the PVTv2 pretrained model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory." + ) + parser.add_argument( + "--verify-imagenet-weights", + action="store_true", + default=False, + help="Verifies the correct conversion of author-published pretrained ImageNet weights.", + ) + + args = parser.parse_args() + convert_pvt_v2_checkpoint( + pvt_v2_size=args.pvt_v2_size, + pvt_v2_checkpoint=args.pvt_v2_checkpoint, + pytorch_dump_folder_path=args.pytorch_dump_folder_path, + verify_imagenet_weights=args.verify_imagenet_weights, + ) diff --git a/third_party/transformers/src/transformers/models/pvt_v2/modeling_pvt_v2.py b/third_party/transformers/src/transformers/models/pvt_v2/modeling_pvt_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..4ecf3a92b06d0186f53adbe728719c5380121b0f --- /dev/null +++ b/third_party/transformers/src/transformers/models/pvt_v2/modeling_pvt_v2.py @@ -0,0 +1,588 @@ +# Copyright 2024 Authors: Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, +# Kaitao Song, Ding Liang, Tong Lu, Ping Luo, Ling Shao 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 PVTv2 model.""" + +import math + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import BackboneMixin, filter_output_hidden_states +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BackboneOutput, BaseModelOutput, ImageClassifierOutput +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from ...utils.generic import can_return_tuple +from .configuration_pvt_v2 import PvtV2Config + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.beit.modeling_beit.drop_path +def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +# Copied from transformers.models.convnext.modeling_convnext.ConvNextDropPath with ConvNext->Pvt +class PvtV2DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float | None = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return f"p={self.drop_prob}" + + +class PvtV2OverlapPatchEmbeddings(nn.Module): + """Image to Patch Embedding""" + + def __init__(self, config: PvtV2Config, layer_idx: int): + super().__init__() + patch_size = config.patch_sizes[layer_idx] + patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size + stride = config.strides[layer_idx] + num_channels = config.num_channels if layer_idx == 0 else config.hidden_sizes[layer_idx - 1] + hidden_size = config.hidden_sizes[layer_idx] + self.patch_size = patch_size + self.proj = nn.Conv2d( + num_channels, + hidden_size, + kernel_size=patch_size, + stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2), + ) + self.layer_norm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + + def forward(self, pixel_values): + embeddings = self.proj(pixel_values) + _, _, height, width = embeddings.shape + embeddings = embeddings.flatten(2).transpose(1, 2) + embeddings = self.layer_norm(embeddings) + return embeddings, height, width + + +class PvtV2DepthWiseConv(nn.Module): + """ + Depth-wise (DW) convolution to infuse positional information using zero-padding. Depth-wise convolutions + have an equal number of groups to the number of input channels, meaning one filter per input channel. This + reduces the overall parameters and compute costs since the key purpose of this layer is position encoding. + """ + + def __init__(self, config: PvtV2Config, dim: int = 768): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, hidden_states, height, width): + batch_size, seq_len, num_channels = hidden_states.shape + hidden_states = hidden_states.transpose(1, 2).view(batch_size, num_channels, height, width) + hidden_states = self.dwconv(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + return hidden_states + + +class PvtV2SelfAttention(nn.Module): + """Efficient self-attention mechanism.""" + + def __init__(self, config: PvtV2Config, hidden_size: int, num_attention_heads: int, spatial_reduction_ratio: int): + super().__init__() + self.linear_attention = config.linear_attention + + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})" + ) + + self.attention_head_size = int(self.hidden_size / self.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(self.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.key = nn.Linear(self.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.attn_drop = nn.Dropout(config.attention_probs_dropout_prob) + self.proj = nn.Linear(self.hidden_size, self.hidden_size) + self.proj_drop = nn.Dropout(config.hidden_dropout_prob) + + self.spatial_reduction_ratio = spatial_reduction_ratio + if self.linear_attention: + self.pool = nn.AdaptiveAvgPool2d(7) + self.spatial_reduction = nn.Conv2d(self.hidden_size, self.hidden_size, kernel_size=1, stride=1) + self.layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + self.act = nn.GELU() + elif spatial_reduction_ratio > 1: + self.spatial_reduction = nn.Conv2d( + self.hidden_size, self.hidden_size, kernel_size=spatial_reduction_ratio, stride=spatial_reduction_ratio + ) + self.layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) + + def transpose_for_scores(self, hidden_states) -> torch.Tensor: + new_shape = hidden_states.size()[:-1] + (self.num_attention_heads, self.attention_head_size) + hidden_states = hidden_states.view(new_shape) + return hidden_states.permute(0, 2, 1, 3) + + def forward( + self, + hidden_states: torch.Tensor, + height: int, + width: int, + output_attentions: bool = False, + ) -> tuple[torch.Tensor]: + batch_size, seq_len, num_channels = hidden_states.shape + query_layer = self.transpose_for_scores(self.query(hidden_states)) + + if self.linear_attention: + hidden_states = hidden_states.permute(0, 2, 1).reshape(batch_size, num_channels, height, width) + hidden_states = ( + self.spatial_reduction(self.pool(hidden_states)).reshape(batch_size, num_channels, -1).permute(0, 2, 1) + ) + hidden_states = self.act(self.layer_norm(hidden_states)) + elif self.spatial_reduction_ratio > 1: + hidden_states = hidden_states.permute(0, 2, 1).reshape(batch_size, num_channels, height, width) + hidden_states = ( + self.spatial_reduction(hidden_states).reshape(batch_size, num_channels, -1).permute(0, 2, 1) + ) + hidden_states = self.layer_norm(hidden_states) + + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + + # 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) + + # 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.attn_drop(attention_probs) + context_layer = (attention_probs @ value_layer).transpose(1, 2).reshape(batch_size, seq_len, num_channels) + context_layer = self.proj(context_layer) + context_layer = self.proj_drop(context_layer) + + outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + + return outputs + + +class PvtV2ConvFeedForwardNetwork(nn.Module): + def __init__( + self, + config: PvtV2Config, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + ): + super().__init__() + out_features = out_features if out_features is not None else in_features + self.dense1 = nn.Linear(in_features, hidden_features) + self.dwconv = PvtV2DepthWiseConv(config, hidden_features) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + self.dense2 = nn.Linear(hidden_features, out_features) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.relu = nn.ReLU() if config.linear_attention else nn.Identity() + + def forward(self, hidden_states: torch.Tensor, height, width) -> torch.Tensor: + hidden_states = self.dense1(hidden_states) + hidden_states = self.relu(hidden_states) + hidden_states = self.dwconv(hidden_states, height, width) + hidden_states = self.intermediate_act_fn(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.dense2(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class PvtV2BlockLayer(nn.Module): + def __init__(self, config: PvtV2Config, layer_idx: int, drop_path: float = 0.0): + super().__init__() + hidden_size: int = config.hidden_sizes[layer_idx] + num_attention_heads: int = config.num_attention_heads[layer_idx] + spatial_reduction_ratio: int = config.sr_ratios[layer_idx] + mlp_ratio: float = config.mlp_ratios[layer_idx] + self.layer_norm_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + self.attention = PvtV2SelfAttention( + config=config, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + spatial_reduction_ratio=spatial_reduction_ratio, + ) + self.drop_path = PvtV2DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + mlp_hidden_size = int(hidden_size * mlp_ratio) + self.mlp = PvtV2ConvFeedForwardNetwork(config=config, in_features=hidden_size, hidden_features=mlp_hidden_size) + + def forward(self, hidden_states: torch.Tensor, height: int, width: int, output_attentions: bool = False): + self_attention_outputs = self.attention( + hidden_states=self.layer_norm_1(hidden_states), + height=height, + width=width, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + outputs = self_attention_outputs[1:] + + attention_output = self.drop_path(attention_output) + hidden_states = attention_output + hidden_states + + mlp_output = self.mlp(self.layer_norm_2(hidden_states), height, width) + + mlp_output = self.drop_path(mlp_output) + layer_output = hidden_states + mlp_output + + outputs = (layer_output,) + outputs + + return outputs + + +class PvtV2EncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: PvtV2Config, layer_idx: int): + super().__init__() + self.patch_embedding = PvtV2OverlapPatchEmbeddings( + config=config, + layer_idx=layer_idx, + ) + # Transformer block + # stochastic depth decay rule + drop_path_decays = torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu").tolist() + block_layers = [] + for block_idx in range(config.depths[layer_idx]): + block_layers.append( + PvtV2BlockLayer( + config=config, + layer_idx=layer_idx, + drop_path=drop_path_decays[sum(config.depths[:layer_idx]) + block_idx], + ) + ) + self.blocks = nn.ModuleList(block_layers) + + # Layer norm + self.layer_norm = nn.LayerNorm(config.hidden_sizes[layer_idx], eps=config.layer_norm_eps) + + def forward(self, hidden_states, output_attentions): + all_self_attentions = () if output_attentions else None + # first, obtain patch embeddings + hidden_states, height, width = self.patch_embedding(hidden_states) + # second, send embeddings through blocks + for block in self.blocks: + layer_outputs = block(hidden_states, height, width, output_attentions) + hidden_states = layer_outputs[0] + if output_attentions: + all_self_attentions += (layer_outputs[1],) + # third, apply layer norm + hidden_states = self.layer_norm(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (all_self_attentions,) + + return outputs, height, width + + +class PvtV2Encoder(nn.Module): + def __init__(self, config: PvtV2Config): + super().__init__() + self.config = config + self.gradient_checkpointing = False + + # encoder layers + self.layers = nn.ModuleList([PvtV2EncoderLayer(config, i) for i in range(config.num_encoder_blocks)]) + + def forward( + self, + pixel_values: torch.FloatTensor, + output_attentions: bool | None = False, + output_hidden_states: bool | None = False, + return_dict: bool | None = True, + ) -> tuple | BaseModelOutput: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + batch_size = pixel_values.shape[0] + hidden_states = pixel_values + for idx, layer in enumerate(self.layers): + layer_output = layer(hidden_states, output_attentions) + outputs, height, width = layer_output + hidden_states = outputs[0] + if output_attentions: + all_self_attentions = all_self_attentions + (outputs[1],) + # reshape back to (batch_size, num_channels, height, width) + hidden_states = hidden_states.reshape(batch_size, height, width, -1).permute(0, 3, 1, 2).contiguous() + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring +class PvtV2PreTrainedModel(PreTrainedModel): + config: PvtV2Config + base_model_prefix = "pvt_v2" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module: nn.Linear | nn.Conv2d | nn.LayerNorm) -> None: + """Initialize the weights""" + if isinstance(module, nn.Linear): + init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv2d): + fan_out = module.kernel_size[0] * module.kernel_size[1] * module.out_channels + fan_out //= module.groups + init.normal_(module.weight, 0, math.sqrt(2.0 / fan_out)) + if module.bias is not None: + init.zeros_(module.bias) + + +@auto_docstring +class PvtV2Model(PvtV2PreTrainedModel): + def __init__(self, config: PvtV2Config): + super().__init__(config) + self.config = config + + # hierarchical Transformer encoder + self.encoder = PvtV2Encoder(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + encoder_outputs = self.encoder( + pixel_values=pixel_values, + output_attentions=output_attentions, + 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 BaseModelOutput( + last_hidden_state=sequence_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + Pvt-v2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state + of the [CLS] token) e.g. for ImageNet. + """ +) +class PvtV2ForImageClassification(PvtV2PreTrainedModel): + def __init__(self, config: PvtV2Config) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.pvt_v2 = PvtV2Model(config) + + # Classifier head + self.classifier = ( + nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ImageClassifierOutput: + 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.pvt_v2( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + # convert last hidden states to (batch_size, height*width, hidden_size) + batch_size = sequence_output.shape[0] + # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels) + sequence_output = sequence_output.permute(0, 2, 3, 1) + sequence_output = sequence_output.reshape(batch_size, -1, self.config.hidden_sizes[-1]) + + # global average pooling + sequence_output = sequence_output.mean(dim=1) + + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(labels, logits, self.config) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + PVTv2 backbone, to be used with frameworks like DETR and MaskFormer. + """ +) +class PvtV2Backbone(BackboneMixin, PvtV2Model): + def __init__(self, config: PvtV2Config): + super().__init__(config) + self.num_features = config.hidden_sizes + + @can_return_tuple + @filter_output_hidden_states + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BackboneOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoBackbone + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> processor = AutoImageProcessor.from_pretrained("OpenGVLab/pvt_v2_b0") + >>> model = AutoBackbone.from_pretrained( + ... "OpenGVLab/pvt_v2_b0", out_features=["stage1", "stage2", "stage3", "stage4"] + ... ) + + >>> inputs = processor(image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> feature_maps = outputs.feature_maps + >>> list(feature_maps[-1].shape) + [1, 256, 7, 7] + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + outputs = self.encoder( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=True, + return_dict=return_dict, + ) + + hidden_states = outputs.hidden_states + + feature_maps = () + for idx, stage in enumerate(self.stage_names): + if stage in self.out_features: + feature_maps += (hidden_states[idx],) + + if not return_dict: + output = (feature_maps,) + if output_hidden_states: + output += (outputs.hidden_states,) + return output + + return BackboneOutput( + feature_maps=feature_maps, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=None, + ) + + +__all__ = ["PvtV2ForImageClassification", "PvtV2Model", "PvtV2PreTrainedModel", "PvtV2Backbone"] diff --git a/third_party/transformers/src/transformers/models/qwen2_5_vl/__init__.py b/third_party/transformers/src/transformers/models/qwen2_5_vl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7a9f44a7a05fc65136552948be9acef53e5ea106 --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen2_5_vl/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2025 The Qwen 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_qwen2_5_vl import * + from .modeling_qwen2_5_vl import * + from .processing_qwen2_5_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/qwen2_5_vl/configuration_qwen2_5_vl.py b/third_party/transformers/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..010fb837ffed568348971fb3ab10f6b71fbea0cf --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py @@ -0,0 +1,209 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_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_qwen2_5_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The Qwen Team 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. +import inspect + +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="Qwen/Qwen2-VL-7B-Instruct") +@strict +class Qwen2_5_VLVisionConfig(PreTrainedConfig): + r""" + tokens_per_second (`int`, *optional*, defaults to 41): + Number of tokens to merge for each second of video. + window_size (`int`, *optional*, defaults to 11): + Size of windows. + out_hidden_size (`int`, *optional*, defaults to 3584): + The output hidden size of the vision model. + fullatt_block_indexes (`int`, *optional*, defaults to `[7, 15, 23, 31]`): + Indices of layers with full attention + """ + + model_type = "qwen2_5_vl" + base_config_key = "vision_config" + + depth: int = 32 + hidden_size: int = 3584 + hidden_act: str = "silu" + intermediate_size: int = 3420 + num_heads: int = 16 + in_channels: int = 3 + patch_size: int | list[int] | tuple[int, int] = 14 + spatial_merge_size: int = 2 + temporal_patch_size: int | list[int] | tuple[int, int] = 2 + tokens_per_second: int = 4 + window_size: int = 112 + out_hidden_size: int = 3584 + fullatt_block_indexes: list[int] | tuple[int, ...] = (7, 15, 23, 31) + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="Qwen/Qwen2-VL-7B-Instruct") +@strict +class Qwen2_5_VLTextConfig(PreTrainedConfig): + r""" + ```python + >>> from transformers import Qwen2_5_VLTextModel, Qwen2_5_VLConfig + + >>> # Initializing a Qwen2_5_VL style configuration + >>> configuration = Qwen2_5_VLConfig() + + >>> # Initializing a model from the Qwen2-VL-7B style configuration + >>> model = Qwen2_5_VLTextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "qwen2_5_vl_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 1000000.0 + # Default tensor parallel plan for base model `Qwen2_5_VL` + 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"]), + } + ignore_keys_at_rope_validation = {"mrope_section"} + + vocab_size: int = 152064 + hidden_size: int = 8192 + intermediate_size: int = 29568 + num_hidden_layers: int = 80 + num_attention_heads: int = 64 + num_key_value_heads: int | None = 8 + hidden_act: str = "silu" + max_position_embeddings: int = 32768 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-05 + use_cache: bool = True + use_sliding_window: bool | None = False + sliding_window: int | None = 4096 + max_window_layers: int | None = 80 + layer_types: list[str] | None = None + attention_dropout: float | int | None = 0.0 + rope_parameters: RopeParameters | dict | None = None + bos_token_id: int | None = 151643 + eos_token_id: int | list[int] | None = 151645 + pad_token_id: int | None = None + + def __post_init__(self, **kwargs): + self.sliding_window = self.sliding_window if self.use_sliding_window else None + + # for backward compatibility + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.layer_types is None: + self.layer_types = [ + "sliding_attention" + if self.sliding_window is not None and i >= self.max_window_layers + else "full_attention" + for i in range(self.num_hidden_layers) + ] + + super().__post_init__(**kwargs) + + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + if self.rope_parameters.get("rope_type", self.rope_parameters.get("type")) == "mrope": + self.rope_parameters["rope_type"] = "default" + self.standardize_rope_params() + return kwargs + + +@auto_docstring(checkpoint="Qwen/Qwen2-VL-7B-Instruct") +@strict +class Qwen2_5_VLConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig + + >>> # Initializing a Qwen2_5_VL style configuration + >>> configuration = Qwen2_5_VLConfig() + + >>> # Initializing a model from the Qwen2-VL-7B style configuration + >>> model = Qwen2_5_VLForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen2_5_vl" + sub_configs = {"vision_config": Qwen2_5_VLVisionConfig, "text_config": Qwen2_5_VLTextConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**self.vision_config) + elif self.vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + # Hub configs are saved as flat dicts so we pop some of kwargs to init `TextConfig` + text_params = inspect.signature(self.sub_configs["text_config"].__init__).parameters.keys() + text_params = list(text_params) + ["rope_parameters", "rope_scaling", "rope_theta"] + text_kwargs = {key: kwargs.pop(key) for key in text_params if key in kwargs} + + if isinstance(self.text_config, dict): + self.text_config = self.sub_configs["text_config"](**self.text_config) + elif self.text_config is None: + # Hub configs are saved as flat dicts so we pop some of kwargs to init `TextConfig` + text_kwargs["dtype"] = kwargs.get("torch_dtype", kwargs.get("dtype")) # don't pop the dtype + self.text_config = self.sub_configs["text_config"](**text_kwargs) + + super().__post_init__(**kwargs) + + +__all__ = ["Qwen2_5_VLConfig", "Qwen2_5_VLTextConfig"] diff --git a/third_party/transformers/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/third_party/transformers/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..f666d5f760f61d7bf8a5187e0cf22146d23762d4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -0,0 +1,1754 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_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_qwen2_5_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The Qwen Team 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. + +import itertools +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub +from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS +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 is_flash_attention_requested, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig + + +logger = logging.get_logger(__name__) + + +@use_kernel_forward_from_hub("RMSNorm") +class Qwen2_5_VLRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Qwen2_5_VLRMSNorm 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 Qwen2_5_VLMLP(nn.Module): + def __init__(self, config, bias: bool = False): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +class Qwen2_5_VisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int | list[int] | tuple[int, int] = 14, + temporal_patch_size: int | list[int] | tuple[int, int] = 2, + in_channels: int = 3, + embed_dim: int = 1152, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.embed_dim = embed_dim + + kernel_size = [temporal_patch_size, patch_size, patch_size] + self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class Qwen2_5_VisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.theta = theta + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class Qwen2_5_VLPatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.ln_q = Qwen2_5_VLRMSNorm(context_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) + return x + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + return q_embed, k_embed + + +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, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Qwen2_5_VLVisionAttention(nn.Module): + def __init__(self, config: Qwen2_5_VLVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 # needed for eager attention + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + if is_flash_attention_requested(self.config): + # Flash Attention: Use cu_seqlens for variable length attention + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = Qwen2_5_VLRMSNorm(config.hidden_size, eps=1e-6) + self.norm2 = Qwen2_5_VLRMSNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen2_5_VLVisionAttention(config=config) + self.mlp = Qwen2_5_VLMLP(config, bias=True) + + @auto_docstring + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + r""" + cu_seqlens (`torch.Tensor`): + Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels. + rotary_pos_emb (`torch.Tensor`, *optional*): + Precomputed rotary positional embeddings applied to the vision attention query/key states. + """ + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +@auto_docstring +class Qwen2_5_VLPreTrainedModel(PreTrainedModel): + config: Qwen2_5_VLConfig + base_model_prefix = "model" + input_modalities = ("image", "video", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Qwen2_5_VisionRotaryEmbedding): + inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim)) + init.copy_(module.inv_freq, inv_freq) + + +class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): + config: Qwen2_5_VLVisionConfig + _no_split_modules = ["Qwen2_5_VLVisionBlock"] + _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5_VLVisionBlock, + "attentions": Qwen2_5_VLVisionAttention, + } + + def __init__(self, config, *inputs, **kwargs) -> None: + super().__init__(config, *inputs, **kwargs) + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.fullatt_block_indexes = config.fullatt_block_indexes + self.window_size = config.window_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen2_5_VisionPatchEmbed( + patch_size=config.patch_size, + temporal_patch_size=config.temporal_patch_size, + in_channels=config.in_channels, + embed_dim=config.hidden_size, + ) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([Qwen2_5_VLVisionBlock(config) for _ in range(config.depth)]) + self.merger = Qwen2_5_VLPatchMerger( + dim=config.out_hidden_size, + context_dim=config.hidden_size, + spatial_merge_size=config.spatial_merge_size, + ) + self.gradient_checkpointing = False + + self.post_init() + + def rot_pos_emb(self, grid_thw): + pos_ids = [] + for t, h, w in grid_thw.tolist(): + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def get_window_index(self, grid_thw): + window_index: list = [] + cu_window_seqlens: list = [0] + window_index_id = 0 + vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + grid_thw_list = grid_thw.tolist() + + for grid_t, grid_h, grid_w in grid_thw_list: + llm_grid_h, llm_grid_w = ( + grid_h // self.spatial_merge_size, + grid_w // self.spatial_merge_size, + ) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + seqlens = (index_padded != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += grid_t * llm_grid_h * llm_grid_w + window_index = torch.cat(window_index, dim=0) + + return window_index, cu_window_seqlens + + @merge_with_config_defaults + @capture_outputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + """ + Args: + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + rotary_pos_emb = self.rot_pos_emb(grid_thw) + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=hidden_states.device, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens_now, + position_embeddings=position_embeddings, + **kwargs, + ) + + merged_hidden_states = self.merger(hidden_states) + reverse_indices = torch.argsort(window_index) + merged_hidden_states = merged_hidden_states[reverse_indices, :] + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Llava outputs, with hidden states and attentions. + """ +) +class Qwen2_5_VLModelOutputWithPast(ModelOutput): + r""" + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + last_hidden_state: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + rope_deltas: torch.LongTensor | None = None + + +class Qwen2_5_VLRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Qwen2_5_VLConfig, 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: Qwen2_5_VLConfig | 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 + + # Ignore copy + def forward(self, x, position_ids): + # In contrast to other models, Qwen2_5_VL has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + 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 Qwen2MLP(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 + + +def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/). + + Explanation: + Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding + sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For + vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately. + Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding. + For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal, + height and width) of text embedding is always the same, so the text embedding rotary position embedding has no + difference with modern LLMs. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + mrope_section(`List(int)`): + Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. + 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. + """ + mrope_section = mrope_section * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen2_5_VLAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.is_causal = True + self.attention_dropout = config.attention_dropout + self.rope_parameters = config.rope_parameters + self.scaling = self.head_dim**-0.5 + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.config.rope_parameters["mrope_section"] + ) + + 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, + position_ids=position_ids, # pass positions for FA2 + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen2_5_VLDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and not is_flash_attention_requested(config): + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + self.self_attn = Qwen2_5_VLAttention(config, layer_idx) + + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2_5_VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2_5_VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring +class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel): + config: Qwen2_5_VLTextConfig + input_modalities = ("text",) + _can_record_outputs = { + "hidden_states": Qwen2_5_VLDecoderLayer, + "attentions": Qwen2_5_VLAttention, + } + + def __init__(self, config: Qwen2_5_VLTextConfig): + 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( + [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2_5_VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.has_sliding_layers = "sliding_attention" in self.config.layer_types + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(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[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # the hard coded `3` is for temporal, height and width. + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions + # where each dim indicates visual spatial positions for temporal/height/width grids. + # There are two scenarios when FA2-like packed masking might be activated. + # 1. User specifically passed packed `position_ids` and no attention mask. + # In this case we expect the useer to create correct position ids for all 3 grids + # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len] + # 2. User runs forward with no attention mask and no position ids. In this case, position ids + # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are + # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass + # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation` + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + position_ids = position_ids[1:] + else: + # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids + text_position_ids = None + + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + # Prepare mask arguments + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": text_position_ids, + } + # Create the masks + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + } + # The sliding window alternating layers are not always activated depending on the config + if self.has_sliding_layers: + causal_mask_mapping["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, + position_ids=text_position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): + base_model_prefix = "model" + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.language_model = Qwen2_5_VLTextModel._from_config(config.text_config) + self.rope_deltas = None # cache rope_deltas here + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + + def get_rope_index( + self, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position 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**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + tokens_per_second = self.config.vision_config.tokens_per_second + + mrope_position_deltas = [] + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + second_per_grid_ts = ( + iter(second_per_grid_ts) if second_per_grid_ts is not None else iter([1] * input_ids.shape[1]) + ) + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 + else: + grid_thw = next(grid_iters[modality_type]) + time_interval = tokens_per_second * int(next(second_per_grid_ts)) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, time_interval, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + vision_outputs = self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs) + split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds + + return vision_outputs + + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + pixel_values = pixel_values.type(self.visual.dtype) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs) + split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs + + def get_placeholder_mask( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_features: torch.FloatTensor | None = None, + video_features: torch.FloatTensor | None = None, + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + special_video_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_video_mask = special_video_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + special_video_mask = input_ids == self.config.video_token_id + + n_image_tokens = special_image_mask.sum() + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if image_features is not None: + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}", + ) + + n_video_tokens = special_video_mask.sum() + special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if video_features is not None: + torch_compilable_check( + inputs_embeds[special_video_mask].numel() == video_features.numel(), + f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}", + ) + return special_image_mask, special_video_mask + + def compute_3d_position_ids( + self, + input_ids: torch.Tensor | None, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + attention_mask: torch.Tensor | None, + past_key_values: torch.Tensor | None, + second_per_grid_ts: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + ) -> torch.Tensor | None: + past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) + + if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, + ) + self.rope_deltas = rope_deltas + # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental + # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids + # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape + # mismatches from stale rope_deltas (e.g., training forward pass after generation). + elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None): + batch_size, seq_length, _ = inputs_embeds.shape + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids = position_ids.masked_fill(attention_mask == 0, 0) + position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device) + else: + position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) + position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) + delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) + position_ids = position_ids + delta.to(device=position_ids.device) + else: + # Can't build correct 3D positions. Let the model infer it + position_ids = None + return position_ids + + @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, + use_cache: bool | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen2_5_VLModelOutputWithPast: + r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + """ + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_embeds = self.get_image_features(pixel_values, image_grid_thw).pooler_output + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw).pooler_output + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if position_ids is None: + position_ids = self.compute_3d_position_ids( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + ) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Qwen2_5_VLModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Qwen2_5_VL causal language model (or autoregressive) outputs. + """ +) +class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + rope_deltas: torch.LongTensor | None = None + + +class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + + def __init__(self, config): + super().__init__(config) + self.model = Qwen2_5_VLModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen2_5_VLCausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + + Example: + + ```python + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg", + }, + {"type": "text", "text": "Describe the image."}, + ], + } + ] + + >>> inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt" + ) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=1024) + >>> generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] + >>> output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(output_text) + ``` + """ + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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 Qwen2_5_VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=outputs.rope_deltas, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + is_first_iteration=False, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + use_cache=use_cache, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if not is_first_iteration and use_cache: + model_inputs["pixel_values"] = None + model_inputs["pixel_values_videos"] = None + + return model_inputs + + def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): + # Overwritten -- requires 3D position ids + + text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs) + + # Early exit in case we are continuing generation from past kv + past_length = 0 + if (cache := model_kwargs.get("past_key_values")) is not None: + past_length = cache.get_seq_length() + if past_length != 0 and self.model.rope_deltas is not None: + position_ids = text_positions[None, ...] + self.model.rope_deltas + return position_ids + + # Otherwise compute 3d position ids for vision tokens and concat with text position ids + if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: + inputs_tensor = model_kwargs["input_ids"] + + is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) + ): + model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} + vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) + self.model.rope_deltas = rope_deltas + else: + vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1) + self.model.rope_deltas = torch.zeros( + inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device + ) + + # Concatenate "text + vision" positions into [4, bs, seq-len] + text_positions = text_positions[None, ...] + position_ids = torch.cat([text_positions, vision_positions], dim=0) + + return position_ids + + def _get_image_nums_and_video_nums( + self, + input_ids: torch.LongTensor | None, + inputs_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Returns: + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) + """ + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + + if inputs_embeds is not None: + vision_start_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + )[..., 0] + image_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + )[..., 0] + video_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + )[..., 0] + else: + vision_start_mask = input_ids == vision_start_token_id + image_mask = input_ids == image_token_id + video_mask = input_ids == video_token_id + + vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) + image_nums = torch.sum(vision_first_mask & image_mask, dim=1) + video_nums = torch.sum(vision_first_mask & video_mask, dim=1) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: torch.LongTensor | None = None, + **model_kwargs, + ) -> tuple[torch.LongTensor, dict[str, Any]]: + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + + def _expand_dict_for_generation_visual(dict_to_expand): + image_grid_thw = model_kwargs.get("image_grid_thw", None) + video_grid_thw = model_kwargs.get("video_grid_thw", None) + image_nums, video_nums = self._get_image_nums_and_video_nums( + input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None) + ) + + def _repeat_interleave_samples(x, lengths, repeat_times): + samples = torch.split(x, lengths) + repeat_args = [repeat_times] + [1] * (x.dim() - 1) + result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # split images into samples + samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "image_grid_thw": + # get the num of images for each sample + lengths = list(image_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "pixel_values_videos": + samples = torch.split(video_grid_thw, list(video_nums)) + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "video_grid_thw": + lengths = list(video_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "second_per_grid_ts": + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size + ) + return dict_to_expand + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if key == "position_ids" and dict_to_expand[key].ndim == 3: + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=1) + elif ( + dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + return dict_to_expand + + model_kwargs = _expand_dict_for_generation_visual(model_kwargs) + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs + + +__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel", "Qwen2_5_VLTextModel"] diff --git a/third_party/transformers/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py b/third_party/transformers/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..8a103cefd2250f16425ee121290a139fd8a9b92a --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py @@ -0,0 +1,927 @@ +# Copyright 2025 The Qwen Team 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 Qwen2.5-VL model.""" + +import itertools + +import torch +import torch.nn as nn +import torch.nn.functional as F +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...configuration_utils import PreTrainedConfig +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPooling +from ...modeling_utils import PreTrainedModel +from ...processing_utils import MultiModalData, ProcessingKwargs, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, can_return_tuple, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ...video_utils import VideoInput +from ..llama.modeling_llama import LlamaRMSNorm +from ..qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig, Qwen2VLTextConfig +from ..qwen2_vl.modeling_qwen2_vl import ( + PatchEmbed, + PatchMerger, + Qwen2VLCausalLMOutputWithPast, + Qwen2VLForConditionalGeneration, + Qwen2VLModel, + Qwen2VLModelOutputWithPast, + Qwen2VLPreTrainedModel, + TransformersKwargs, + VisionAttention, + VisionRotaryEmbedding, +) +from ..qwen2_vl.processing_qwen2_vl import Qwen2VLProcessor + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="Qwen/Qwen2-VL-7B-Instruct") +@strict +class Qwen2_5_VLVisionConfig(PreTrainedConfig): + r""" + tokens_per_second (`int`, *optional*, defaults to 41): + Number of tokens to merge for each second of video. + window_size (`int`, *optional*, defaults to 11): + Size of windows. + out_hidden_size (`int`, *optional*, defaults to 3584): + The output hidden size of the vision model. + fullatt_block_indexes (`int`, *optional*, defaults to `[7, 15, 23, 31]`): + Indices of layers with full attention + """ + + model_type = "qwen2_5_vl" + base_config_key = "vision_config" + + depth: int = 32 + hidden_size: int = 3584 + hidden_act: str = "silu" + intermediate_size: int = 3420 + num_heads: int = 16 + in_channels: int = 3 + patch_size: int | list[int] | tuple[int, int] = 14 + spatial_merge_size: int = 2 + temporal_patch_size: int | list[int] | tuple[int, int] = 2 + tokens_per_second: int = 4 + window_size: int = 112 + out_hidden_size: int = 3584 + fullatt_block_indexes: list[int] | tuple[int, ...] = (7, 15, 23, 31) + initializer_range: float = 0.02 + + +class Qwen2_5_VLTextConfig(Qwen2VLTextConfig): + pass + + +class Qwen2_5_VLConfig(Qwen2VLConfig): + pass + + +class Qwen2_5_VLRMSNorm(LlamaRMSNorm): + pass + + +class Qwen2_5_VLMLP(nn.Module): + def __init__(self, config, bias: bool = False): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +class Qwen2_5_VisionPatchEmbed(PatchEmbed): + pass + + +class Qwen2_5_VisionRotaryEmbedding(VisionRotaryEmbedding): + pass + + +class Qwen2_5_VLPatchMerger(PatchMerger): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: + super().__init__(dim, context_dim, spatial_merge_size) + self.ln_q = Qwen2_5_VLRMSNorm(context_dim, eps=1e-6) + + +class Qwen2_5_VLVisionAttention(VisionAttention): + def __init__(self, config: Qwen2_5_VLVisionConfig) -> None: + super().__init__(config) + self.dim = config.hidden_size + + +class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = Qwen2_5_VLRMSNorm(config.hidden_size, eps=1e-6) + self.norm2 = Qwen2_5_VLRMSNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen2_5_VLVisionAttention(config=config) + self.mlp = Qwen2_5_VLMLP(config, bias=True) + + @auto_docstring + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + r""" + cu_seqlens (`torch.Tensor`): + Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels. + rotary_pos_emb (`torch.Tensor`, *optional*): + Precomputed rotary positional embeddings applied to the vision attention query/key states. + """ + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen2_5_VLPreTrainedModel(Qwen2VLPreTrainedModel): + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, Qwen2_5_VisionRotaryEmbedding): + inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim)) + init.copy_(module.inv_freq, inv_freq) + + +class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): + config: Qwen2_5_VLVisionConfig + _no_split_modules = ["Qwen2_5_VLVisionBlock"] + _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5_VLVisionBlock, + "attentions": Qwen2_5_VLVisionAttention, + } + + def __init__(self, config, *inputs, **kwargs) -> None: + super().__init__(config, *inputs, **kwargs) + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.fullatt_block_indexes = config.fullatt_block_indexes + self.window_size = config.window_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen2_5_VisionPatchEmbed( + patch_size=config.patch_size, + temporal_patch_size=config.temporal_patch_size, + in_channels=config.in_channels, + embed_dim=config.hidden_size, + ) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([Qwen2_5_VLVisionBlock(config) for _ in range(config.depth)]) + self.merger = Qwen2_5_VLPatchMerger( + dim=config.out_hidden_size, + context_dim=config.hidden_size, + spatial_merge_size=config.spatial_merge_size, + ) + self.gradient_checkpointing = False + + self.post_init() + + def rot_pos_emb(self, grid_thw): + pos_ids = [] + for t, h, w in grid_thw.tolist(): + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def get_window_index(self, grid_thw): + window_index: list = [] + cu_window_seqlens: list = [0] + window_index_id = 0 + vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + grid_thw_list = grid_thw.tolist() + + for grid_t, grid_h, grid_w in grid_thw_list: + llm_grid_h, llm_grid_w = ( + grid_h // self.spatial_merge_size, + grid_w // self.spatial_merge_size, + ) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + seqlens = (index_padded != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += grid_t * llm_grid_h * llm_grid_w + window_index = torch.cat(window_index, dim=0) + + return window_index, cu_window_seqlens + + @merge_with_config_defaults + @capture_outputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + """ + Args: + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + rotary_pos_emb = self.rot_pos_emb(grid_thw) + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=hidden_states.device, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens_now, + position_embeddings=position_embeddings, + **kwargs, + ) + + merged_hidden_states = self.merger(hidden_states) + reverse_indices = torch.argsort(window_index) + merged_hidden_states = merged_hidden_states[reverse_indices, :] + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) + + +class Qwen2_5_VLModelOutputWithPast(Qwen2VLModelOutputWithPast): + pass + + +class Qwen2_5_VLModel(Qwen2VLModel): + config: Qwen2_5_VLConfig + base_model_prefix = "model" + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + + def get_rope_index( + self, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position 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**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + tokens_per_second = self.config.vision_config.tokens_per_second + + mrope_position_deltas = [] + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + second_per_grid_ts = ( + iter(second_per_grid_ts) if second_per_grid_ts is not None else iter([1] * input_ids.shape[1]) + ) + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 + else: + grid_thw = next(grid_iters[modality_type]) + time_interval = tokens_per_second * int(next(second_per_grid_ts)) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, time_interval, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + + def compute_3d_position_ids( + self, + input_ids: torch.Tensor | None, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + attention_mask: torch.Tensor | None, + past_key_values: torch.Tensor | None, + second_per_grid_ts: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + ) -> torch.Tensor | None: + past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) + + if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, + ) + self.rope_deltas = rope_deltas + # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental + # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids + # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape + # mismatches from stale rope_deltas (e.g., training forward pass after generation). + elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None): + batch_size, seq_length, _ = inputs_embeds.shape + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids = position_ids.masked_fill(attention_mask == 0, 0) + position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device) + else: + position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) + position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) + delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) + position_ids = position_ids + delta.to(device=position_ids.device) + else: + # Can't build correct 3D positions. Let the model infer it + position_ids = None + return position_ids + + @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, + use_cache: bool | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen2_5_VLModelOutputWithPast: + r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + """ + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_embeds = self.get_image_features(pixel_values, image_grid_thw).pooler_output + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw).pooler_output + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if position_ids is None: + position_ids = self.compute_3d_position_ids( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + ) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Qwen2_5_VLModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + +class Qwen2_5_VLCausalLMOutputWithPast(Qwen2VLCausalLMOutputWithPast): + pass + + +class Qwen2_5_VLForConditionalGeneration(Qwen2VLForConditionalGeneration): + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + + @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, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + second_per_grid_ts: torch.Tensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen2_5_VLCausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + + Example: + + ```python + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg", + }, + {"type": "text", "text": "Describe the image."}, + ], + } + ] + + >>> inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt" + ) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=1024) + >>> generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] + >>> output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(output_text) + ``` + """ + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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 Qwen2_5_VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=outputs.rope_deltas, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + is_first_iteration=False, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + use_cache=use_cache, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if not is_first_iteration and use_cache: + model_inputs["pixel_values"] = None + model_inputs["pixel_values_videos"] = None + + return model_inputs + + +class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": True, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +class Qwen2_5_VLProcessor(Qwen2VLProcessor): + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + video_processor_input_names = self.video_processor.model_input_names + names_from_processor = list( + dict.fromkeys(tokenizer_input_names + image_processor_input_names + video_processor_input_names) + ) + return names_from_processor + ["second_per_grid_ts", "mm_token_type_ids"] + + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + videos: VideoInput | None = None, + **kwargs: Unpack[Qwen2_5_VLProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`. + - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. + - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. + - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. + """ + output_kwargs = self._merge_kwargs( + Qwen2_5_VLProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + image_inputs = videos_inputs = {} + if images is not None: + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + image_grid_thw = image_inputs["image_grid_thw"] + + if videos is not None: + videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"]) + video_grid_thw = videos_inputs["video_grid_thw"] + + # Get video metadata + if not kwargs.get("return_metadata"): + video_metadata = videos_inputs.pop("video_metadata") + else: + video_metadata = videos_inputs["video_metadata"] + + fps = [metadata.sampled_fps for metadata in video_metadata] + + if isinstance(fps, (int, float)): + second_per_grid_ts = [self.video_processor.temporal_patch_size / fps] * len(video_grid_thw) + elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw): + second_per_grid_ts = [self.video_processor.temporal_patch_size / tmp for tmp in fps] + else: + raise ValueError( + f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number." + ) + videos_inputs.update({"second_per_grid_ts": second_per_grid_ts}) + + if not isinstance(text, list): + text = [text] + + text = text.copy() # below lines change text in-place + if images is not None: + merge_length = self.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + num_image_tokens = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + + if videos is not None: + merge_length = self.video_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.video_token in text[i]: + num_video_tokens = video_grid_thw[index].prod() // merge_length + text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.video_token) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + video_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (num_frames, height, width) per each video. + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + + vision_data = {} + if image_sizes is not None: + images_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size + + num_image_patches = [ + self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) + for image_size in image_sizes + ] + num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches] + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + if video_sizes is not None: + videos_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("videos_kwargs", {}) + videos_kwargs.update(kwargs) + num_video_patches = [ + self.video_processor.get_number_of_video_patches(*video_size, videos_kwargs) + for video_size in video_sizes + ] + num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches] + vision_data["num_video_tokens"] = num_video_tokens + + return MultiModalData(**vision_data) + + +__all__ = [ + "Qwen2_5_VLConfig", + "Qwen2_5_VLTextConfig", + "Qwen2_5_VLForConditionalGeneration", + "Qwen2_5_VLModel", + "Qwen2_5_VLPreTrainedModel", + "Qwen2_5_VLProcessor", + "Qwen2_5_VLTextModel", # noqa: F822 +] diff --git a/third_party/transformers/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py b/third_party/transformers/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..8873eb82557add4b59065cca77dd72d31b9ffc8b --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py @@ -0,0 +1,224 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_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_qwen2_5_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The Qwen Team 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 ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring +from ...video_utils import VideoInput + + +class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": True, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +@auto_docstring +class Qwen2_5_VLProcessor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs): + self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token + self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token + self.image_token_id = ( + tokenizer.image_token_id + if getattr(tokenizer, "image_token_id", None) + else tokenizer.convert_tokens_to_ids(self.image_token) + ) + self.video_token_id = ( + tokenizer.video_token_id + if getattr(tokenizer, "video_token_id", None) + else tokenizer.convert_tokens_to_ids(self.video_token) + ) + super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + videos: VideoInput | None = None, + **kwargs: Unpack[Qwen2_5_VLProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`. + - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. + - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. + - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. + """ + output_kwargs = self._merge_kwargs( + Qwen2_5_VLProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + image_inputs = videos_inputs = {} + if images is not None: + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + image_grid_thw = image_inputs["image_grid_thw"] + + if videos is not None: + videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"]) + video_grid_thw = videos_inputs["video_grid_thw"] + + # Get video metadata + if not kwargs.get("return_metadata"): + video_metadata = videos_inputs.pop("video_metadata") + else: + video_metadata = videos_inputs["video_metadata"] + + fps = [metadata.sampled_fps for metadata in video_metadata] + + if isinstance(fps, (int, float)): + second_per_grid_ts = [self.video_processor.temporal_patch_size / fps] * len(video_grid_thw) + elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw): + second_per_grid_ts = [self.video_processor.temporal_patch_size / tmp for tmp in fps] + else: + raise ValueError( + f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number." + ) + videos_inputs.update({"second_per_grid_ts": second_per_grid_ts}) + + if not isinstance(text, list): + text = [text] + + text = text.copy() # below lines change text in-place + if images is not None: + merge_length = self.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + num_image_tokens = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + + if videos is not None: + merge_length = self.video_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.video_token in text[i]: + num_video_tokens = video_grid_thw[index].prod() // merge_length + text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.video_token) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"]) + + if return_mm_token_type_ids: + text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"]) + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + video_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (num_frames, height, width) per each video. + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + + vision_data = {} + if image_sizes is not None: + images_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size + + num_image_patches = [ + self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) + for image_size in image_sizes + ] + num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches] + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + if video_sizes is not None: + videos_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("videos_kwargs", {}) + videos_kwargs.update(kwargs) + num_video_patches = [ + self.video_processor.get_number_of_video_patches(*video_size, videos_kwargs) + for video_size in video_sizes + ] + num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches] + vision_data["num_video_tokens"] = num_video_tokens + + return MultiModalData(**vision_data) + + def post_process_image_text_to_text( + self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs + ): + """ + Post-process the output of the model to decode the text. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + or `(sequence_length,)`. + skip_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `list[str]`: The decoded text. + """ + return self.tokenizer.batch_decode( + generated_outputs, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + video_processor_input_names = self.video_processor.model_input_names + names_from_processor = list( + dict.fromkeys(tokenizer_input_names + image_processor_input_names + video_processor_input_names) + ) + return names_from_processor + ["second_per_grid_ts", "mm_token_type_ids"] + + +__all__ = ["Qwen2_5_VLProcessor"] diff --git a/third_party/transformers/src/transformers/models/rwkv/__init__.py b/third_party/transformers/src/transformers/models/rwkv/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cbfdb86827a0f5551fc99ad84edf2c2684444e96 --- /dev/null +++ b/third_party/transformers/src/transformers/models/rwkv/__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_rwkv import * + from .modeling_rwkv 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/rwkv/configuration_rwkv.py b/third_party/transformers/src/transformers/models/rwkv/configuration_rwkv.py new file mode 100644 index 0000000000000000000000000000000000000000..09f1d05a3a083a792758761c3c8b3763b1f7a597 --- /dev/null +++ b/third_party/transformers/src/transformers/models/rwkv/configuration_rwkv.py @@ -0,0 +1,76 @@ +# Copyright 2023 The OpenAI Team Authors and HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""RWKV configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="RWKV/rwkv-4-169m-pile") +@strict +class RwkvConfig(PreTrainedConfig): + r""" + context_length (`int`, *optional*, defaults to 1024): + The maximum sequence length that this model can be used with in a single forward (using it in RNN mode + lets use any sequence length). + attention_hidden_size (`int`, *optional*): + Dimensionality of the attention hidden states. Will default to `hidden_size` if unset. + rescale_every (`int`, *optional*, defaults to 6): + At inference, the hidden states (and weights of the corresponding output layers) are divided by 2 every + `rescale_every` layer. If set to 0 or a negative number, no rescale is done. + + Example: + + ```python + >>> from transformers import RwkvConfig, RwkvModel + + >>> # Initializing a Rwkv configuration + >>> configuration = RwkvConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = RwkvModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "rwkv" + attribute_map = {"max_position_embeddings": "context_length"} + + vocab_size: int = 50277 + context_length: int = 1024 + hidden_size: int = 4096 + num_hidden_layers: int = 32 + attention_hidden_size: int | None = None + intermediate_size: int | None = None + layer_norm_epsilon: float = 1e-5 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 0 + rescale_every: int = 6 + tie_word_embeddings: bool = False + use_cache: bool = True + + def __post_init__(self, **kwargs): + self.attention_hidden_size = ( + self.attention_hidden_size if self.attention_hidden_size is not None else self.hidden_size + ) + self.intermediate_size = self.intermediate_size if self.intermediate_size is not None else 4 * self.hidden_size + + super().__post_init__(**kwargs) + + +__all__ = ["RwkvConfig"] diff --git a/third_party/transformers/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py b/third_party/transformers/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..61ce995615a6f9653f4f69ea1687b4ebdbf2a873 --- /dev/null +++ b/third_party/transformers/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py @@ -0,0 +1,208 @@ +# 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 a RWKV checkpoint from BlinkDL to the Hugging Face format.""" + +import argparse +import gc +import json +import os +import re + +import torch +from huggingface_hub import hf_hub_download, split_torch_state_dict_into_shards + +from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig +from transformers.modeling_utils import WEIGHTS_INDEX_NAME + + +NUM_HIDDEN_LAYERS_MAPPING = { + "169M": 12, + "430M": 24, + "1B5": 24, + "3B": 32, + "7B": 32, + "14B": 40, +} + +HIDDEN_SIZE_MAPPING = { + "169M": 768, + "430M": 1024, + "1B5": 2048, + "3B": 2560, + "7B": 4096, + "14B": 5120, +} + + +def convert_state_dict(state_dict): + state_dict_keys = list(state_dict.keys()) + for name in state_dict_keys: + weight = state_dict.pop(name) + # emb -> embedding + if name.startswith("emb."): + name = name.replace("emb.", "embeddings.") + # ln_0 -> pre_ln (only present at block 0) + if name.startswith("blocks.0.ln0"): + name = name.replace("blocks.0.ln0", "blocks.0.pre_ln") + # att -> attention + name = re.sub(r"blocks\.(\d+)\.att", r"blocks.\1.attention", name) + # ffn -> feed_forward + name = re.sub(r"blocks\.(\d+)\.ffn", r"blocks.\1.feed_forward", name) + # time_mix_k -> time_mix_key and reshape + if name.endswith(".time_mix_k"): + name = name.replace(".time_mix_k", ".time_mix_key") + # time_mix_v -> time_mix_value and reshape + if name.endswith(".time_mix_v"): + name = name.replace(".time_mix_v", ".time_mix_value") + # time_mix_r -> time_mix_key and reshape + if name.endswith(".time_mix_r"): + name = name.replace(".time_mix_r", ".time_mix_receptance") + + if name != "head.weight": + name = "rwkv." + name + + state_dict[name] = weight + return state_dict + + +def convert_rmkv_checkpoint_to_hf_format( + repo_id, checkpoint_file, output_dir, size=None, tokenizer_file=None, push_to_hub=False, model_name=None +): + # 1. If possible, build the tokenizer. + if tokenizer_file is None: + print("No `--tokenizer_file` provided, we will use the default tokenizer.") + vocab_size = 50277 + tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") + else: + tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_file) + vocab_size = len(tokenizer) + tokenizer.save_pretrained(output_dir) + + # 2. Build the config + possible_sizes = list(NUM_HIDDEN_LAYERS_MAPPING.keys()) + if size is None: + # Try to infer size from the checkpoint name + for candidate in possible_sizes: + if candidate in checkpoint_file: + size = candidate + break + if size is None: + raise ValueError("Could not infer the size, please provide it with the `--size` argument.") + if size not in possible_sizes: + raise ValueError(f"`size` should be one of {possible_sizes}, got {size}.") + + config = RwkvConfig( + vocab_size=vocab_size, + num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size], + hidden_size=HIDDEN_SIZE_MAPPING[size], + ) + config.save_pretrained(output_dir) + + # 3. Download model file then convert state_dict + model_file = hf_hub_download(repo_id, checkpoint_file) + state_dict = torch.load(model_file, map_location="cpu", weights_only=True) + state_dict = convert_state_dict(state_dict) + + # 4. Split in shards and save + state_dict_split = split_torch_state_dict_into_shards(state_dict) + shards = index = None + for tensors in state_dict_split.filename_to_tensors.values(): + shards = {tensor: state_dict[tensor] for tensor in tensors} + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + + for shard_file, shard in shards.items(): + torch.save(shard, os.path.join(output_dir, shard_file)) + + if index is not None: + save_index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + + # 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict + print( + "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model." + ) + shard_files = list(shards.keys()) + + del state_dict + del shards + gc.collect() + + for shard_file in shard_files: + state_dict = torch.load(os.path.join(output_dir, shard_file), weights_only=True) + torch.save({k: v.cpu().clone() for k, v in state_dict.items()}, os.path.join(output_dir, shard_file)) + + del state_dict + gc.collect() + + if push_to_hub: + if model_name is None: + raise ValueError("Please provide a `model_name` to push the model to the Hub.") + model = AutoModelForCausalLM.from_pretrained(output_dir) + model.push_to_hub(model_name, max_shard_size="2GB") + tokenizer.push_to_hub(model_name) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--repo_id", default=None, type=str, required=True, help="Repo ID from which to pull the checkpoint." + ) + parser.add_argument( + "--checkpoint_file", default=None, type=str, required=True, help="Name of the checkpoint file in the repo." + ) + parser.add_argument( + "--output_dir", default=None, type=str, required=True, help="Where to save the converted model." + ) + parser.add_argument( + "--tokenizer_file", + default=None, + type=str, + help="Path to the tokenizer file to use (if not provided, only the model is converted).", + ) + parser.add_argument( + "--size", + default=None, + type=str, + help="Size of the model. Will be inferred from the `checkpoint_file` if not passed.", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Push to the Hub the converted model.", + ) + parser.add_argument( + "--model_name", + default=None, + type=str, + help="Name of the pushed model on the Hub, including the username / organization.", + ) + + args = parser.parse_args() + convert_rmkv_checkpoint_to_hf_format( + args.repo_id, + args.checkpoint_file, + args.output_dir, + size=args.size, + tokenizer_file=args.tokenizer_file, + push_to_hub=args.push_to_hub, + model_name=args.model_name, + ) diff --git a/third_party/transformers/src/transformers/models/rwkv/modeling_rwkv.py b/third_party/transformers/src/transformers/models/rwkv/modeling_rwkv.py new file mode 100644 index 0000000000000000000000000000000000000000..5c7f008fc28e7de091d649795674bfe67c22607a --- /dev/null +++ b/third_party/transformers/src/transformers/models/rwkv/modeling_rwkv.py @@ -0,0 +1,758 @@ +# Copyright 2023 Bo Peng and HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch RWKV model.""" + +import math +from dataclasses import dataclass + +import torch +from torch import nn + +from ... import initialization as init +from ...generation import GenerationMixin +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_utils import PreTrainedModel +from ...utils import ( + ModelOutput, + auto_docstring, + is_bitsandbytes_available, + is_kernels_available, + is_ninja_available, + is_torch_cuda_available, + logging, +) +from .configuration_rwkv import RwkvConfig + + +logger = logging.get_logger(__name__) + + +rwkv_cuda_kernel = None + + +def load_wkv_cuda_kernel(context_length): + global rwkv_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 + + rwkv_cuda_kernel = get_kernel("kernels-community/rwkv") + rwkv_cuda_kernel.max_seq_length = context_length + + +class RwkvLinearAttention(torch.autograd.Function): + @staticmethod + def forward(ctx, time_decay, time_first, key, value, state=None, return_state=False): + batch_size, seq_len, hidden_size = key.size() + if seq_len > rwkv_cuda_kernel.max_seq_length: + raise ValueError( + f"Cannot process a batch with {seq_len} tokens at the same time, use a maximum of " + f"{rwkv_cuda_kernel.max_seq_length} with this model." + ) + if batch_size * hidden_size % min(hidden_size, 32) != 0: + raise ValueError( + f"The product of batch size ({batch_size}) and hidden size ({hidden_size}) needs to be a round " + f"multiple of {min(hidden_size, 32)}." + ) + + ctx.input_dtype = key.dtype + + if ( + time_decay.device.type != "cuda" + or time_first.device.type != "cuda" + or key.device.type != "cuda" + or value.device.type != "cuda" + ): + raise ValueError("Calling the CUDA kernel for wkv attention requires all tensors to be on CUDA devices.") + + time_decay = -torch.exp(time_decay.float().contiguous()) + if key.dtype == torch.float16: + time_first = time_first.float() + key = key.float() + value = value.float() + time_first = time_first.contiguous() + key = key.contiguous() + value = value.contiguous() + # The CUDA kernel will fill this tensor. + output = torch.empty_like(key, memory_format=torch.contiguous_format) + if return_state or state is not None: + if state is None: + state = torch.zeros( + batch_size, + hidden_size, + 3, + dtype=torch.float32, + device=key.device, + memory_format=torch.contiguous_format, + ) + state[:, :, 2] -= 1e38 + else: + state = torch.cat([s.unsqueeze(2) for s in state], dim=2).contiguous() + if key.dtype == torch.bfloat16: + forward_func = rwkv_cuda_kernel.forward_with_state_bf16 + else: + forward_func = rwkv_cuda_kernel.forward_with_state + forward_func(time_decay, time_first, key, value, output, state) + else: + forward_func = rwkv_cuda_kernel.forward_bf16 if key.dtype == torch.bfloat16 else rwkv_cuda_kernel.forward + forward_func(time_decay, time_first, key, value, output) + + ctx.save_for_backward(time_decay, time_first, key, value, output) + + if state is not None: + state = [s.squeeze(2) for s in torch.chunk(state, 3, dim=2)] + + return output.to(ctx.input_dtype), state + + @staticmethod + # g stands for grad + def backward(ctx, g_output, g_state=None): + input_dtype = ctx.input_dtype + + time_decay, time_first, key, value, output = ctx.saved_tensors + # The CUDA kernel will fill those tensors. + g_time_decay = torch.empty_like( + time_decay, + memory_format=torch.contiguous_format, + dtype=torch.bfloat16 if input_dtype == torch.bfloat16 else torch.float32, + ) + g_time_first = torch.empty_like(time_first, memory_format=torch.contiguous_format) + g_key = torch.empty_like(key, memory_format=torch.contiguous_format) + g_value = torch.empty_like(value, memory_format=torch.contiguous_format) + + if input_dtype == torch.float16: + g_output = g_output.float() + backward_func = rwkv_cuda_kernel.backward_bf16 if input_dtype == torch.bfloat16 else rwkv_cuda_kernel.backward + backward_func( + time_decay, + time_first, + key, + value, + output, + g_output.contiguous(), + g_time_decay, + g_time_first, + g_key, + g_value, + ) + + return ( + g_time_decay.to(input_dtype), + g_time_first.to(input_dtype), + g_key.to(input_dtype), + g_value.to(input_dtype), + None, + None, + ) + + +def rwkv_linear_attention_cpu(time_decay, time_first, key, value, state=None, return_state=False): + # For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed + # within a torch.no_grad. + _, seq_length, _ = key.size() + output = torch.zeros_like(key) + + if state is None: + num_state = torch.zeros_like(key[:, 0], dtype=torch.float32) + den_state = torch.zeros_like(key[:, 0], dtype=torch.float32) + max_state = torch.zeros_like(key[:, 0], dtype=torch.float32) - 1e38 + else: + num_state, den_state, max_state = state + # For numerical stability + # real_numerator_state = num_state * torch.exp(max_state) + # real_denominator_state = den_state * torch.exp(max_state) + + time_decay = -torch.exp(time_decay) + + for current_index in range(seq_length): + current_key = key[:, current_index].float() + current_value = value[:, current_index] + + # wkv computation at time t + max_for_output = torch.maximum(max_state, current_key + time_first) + e1 = torch.exp(max_state - max_for_output) + e2 = torch.exp(current_key + time_first - max_for_output) + numerator = e1 * num_state + e2 * current_value + denominator = e1 * den_state + e2 + output[:, current_index] = (numerator / denominator).to(output.dtype) + + # Update state for next iteration + max_for_state = torch.maximum(max_state + time_decay, current_key) + e1 = torch.exp(max_state + time_decay - max_for_state) + e2 = torch.exp(current_key - max_for_state) + num_state = e1 * num_state + e2 * current_value + den_state = e1 * den_state + e2 + max_state = max_for_state + + if return_state or state is not None: + state = [num_state, den_state, max_state] + + return output, state + + +def rwkv_linear_attention(time_decay, time_first, key, value, state=None, return_state=False): + no_cuda = any(t.device.type != "cuda" for t in [time_decay, time_first, key, value]) + # Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version + # in this case). + one_token = key.size(1) == 1 + if rwkv_cuda_kernel is None or no_cuda or one_token: + return rwkv_linear_attention_cpu(time_decay, time_first, key, value, state=state, return_state=return_state) + else: + return RwkvLinearAttention.apply(time_decay, time_first, key, value, state, return_state) + + +class RwkvSelfAttention(nn.Module): + def __init__(self, config, layer_id=0): + super().__init__() + self.config = config + kernel_loaded = rwkv_cuda_kernel is not None and rwkv_cuda_kernel.max_seq_length == config.context_length + if is_ninja_available() and is_torch_cuda_available() and not kernel_loaded: + try: + load_wkv_cuda_kernel(config.context_length) + except Exception: + logger.info("Could not load the custom CUDA kernel for RWKV attention.") + self.layer_id = layer_id + hidden_size = config.hidden_size + attention_hidden_size = ( + config.attention_hidden_size if config.attention_hidden_size is not None else hidden_size + ) + self.attention_hidden_size = attention_hidden_size + + self.time_decay = nn.Parameter(torch.empty(attention_hidden_size)) + self.time_first = nn.Parameter(torch.empty(attention_hidden_size)) + + self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size)) + self.time_mix_value = nn.Parameter(torch.empty(1, 1, hidden_size)) + self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size)) + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + self.key = nn.Linear(hidden_size, attention_hidden_size, bias=False) + self.value = nn.Linear(hidden_size, attention_hidden_size, bias=False) + self.receptance = nn.Linear(hidden_size, attention_hidden_size, bias=False) + self.output = nn.Linear(attention_hidden_size, hidden_size, bias=False) + + # TODO: maybe jit, otherwise move inside forward + def extract_key_value(self, hidden, state=None): + # Mix hidden with the previous timestep to produce key, value, receptance + if hidden.size(1) == 1 and state is not None: + shifted = state[1][:, :, self.layer_id] + else: + shifted = self.time_shift(hidden) + if state is not None: + shifted[:, 0] = state[1][:, :, self.layer_id] + key = hidden * self.time_mix_key + shifted * (1 - self.time_mix_key) + value = hidden * self.time_mix_value + shifted * (1 - self.time_mix_value) + receptance = hidden * self.time_mix_receptance + shifted * (1 - self.time_mix_receptance) + + key = self.key(key) + value = self.value(value) + receptance = torch.sigmoid(self.receptance(receptance)) + if state is not None: + state[1][:, :, self.layer_id] = hidden[:, -1] + return receptance, key, value, state + + def forward(self, hidden, state=None, use_cache=False): + receptance, key, value, state = self.extract_key_value(hidden, state=state) + layer_state = tuple(s[:, :, self.layer_id] for s in state[2:]) if state is not None else None + rwkv, layer_state = rwkv_linear_attention( + self.time_decay, + self.time_first, + key, + value, + state=layer_state, + return_state=use_cache, + ) + + if layer_state is not None: + state[2][:, :, self.layer_id] = layer_state[0] + state[3][:, :, self.layer_id] = layer_state[1] + state[4][:, :, self.layer_id] = layer_state[2] + + return self.output(receptance * rwkv), state + + +class RwkvFeedForward(nn.Module): + def __init__(self, config, layer_id=0): + super().__init__() + self.config = config + self.layer_id = layer_id + hidden_size = config.hidden_size + intermediate_size = ( + config.intermediate_size if config.intermediate_size is not None else 4 * config.hidden_size + ) + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size)) + self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size)) + + self.key = nn.Linear(hidden_size, intermediate_size, bias=False) + self.receptance = nn.Linear(hidden_size, hidden_size, bias=False) + self.value = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, hidden, state=None): + if hidden.size(1) == 1 and state is not None: + shifted = state[0][:, :, self.layer_id] + else: + shifted = self.time_shift(hidden) + if state is not None: + shifted[:, 0] = state[0][:, :, self.layer_id] + key = hidden * self.time_mix_key + shifted * (1 - self.time_mix_key) + receptance = hidden * self.time_mix_receptance + shifted * (1 - self.time_mix_receptance) + + key = torch.square(torch.relu(self.key(key))) + value = self.value(key) + receptance = torch.sigmoid(self.receptance(receptance)) + + if state is not None: + state[0][:, :, self.layer_id] = hidden[:, -1] + + return receptance * value, state + + +class RwkvBlock(GradientCheckpointingLayer): + def __init__(self, config, layer_id): + super().__init__() + self.config = config + self.layer_id = layer_id + + if layer_id == 0: + self.pre_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + self.attention = RwkvSelfAttention(config, layer_id) + self.feed_forward = RwkvFeedForward(config, layer_id) + + def forward(self, hidden, state=None, use_cache=False, output_attentions=False): + if self.layer_id == 0: + hidden = self.pre_ln(hidden) + + attention, state = self.attention(self.ln1(hidden), state=state, use_cache=use_cache) + hidden = hidden + attention + + feed_forward, state = self.feed_forward(self.ln2(hidden), state=state) + hidden = hidden + feed_forward + + outputs = (hidden, state) + if output_attentions: + outputs += (attention,) + else: + outputs += (None,) + + return outputs + + +@auto_docstring +class RwkvPreTrainedModel(PreTrainedModel): + config: RwkvConfig + base_model_prefix = "rwkv" + _no_split_modules = ["RwkvBlock"] + _keep_in_fp32_modules = ["time_decay", "time_first"] + supports_gradient_checkpointing = True + _is_stateful = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + if isinstance(module, RwkvSelfAttention): + layer_id = module.layer_id + num_hidden_layers = module.config.num_hidden_layers + hidden_size = module.config.hidden_size + attention_hidden_size = module.attention_hidden_size + + ratio_0_to_1 = layer_id / (num_hidden_layers - 1) # 0 to 1 + ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0 + + time_weight = torch.tensor( + [i / hidden_size for i in range(hidden_size)], + dtype=module.time_mix_key.dtype, + device=module.time_mix_key.device, + ) + time_weight = time_weight[None, None, :] + + decay_speed = [ + -5 + 8 * (h / (attention_hidden_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1) + for h in range(attention_hidden_size) + ] + decay_speed = torch.tensor(decay_speed, dtype=module.time_decay.dtype, device=module.time_decay.device) + zigzag = ( + torch.tensor( + [(i + 1) % 3 - 1 for i in range(attention_hidden_size)], + dtype=module.time_first.dtype, + device=module.time_first.device, + ) + * 0.5 + ) + + init.copy_(module.time_decay, decay_speed) + init.copy_(module.time_first, torch.ones_like(module.time_first * math.log(0.3) + zigzag)) + + init.copy_(module.time_mix_key, torch.pow(time_weight, ratio_1_to_almost0)) + init.copy_(module.time_mix_value, torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1) + init.copy_(module.time_mix_receptance, torch.pow(time_weight, 0.5 * ratio_1_to_almost0)) + elif isinstance(module, RwkvFeedForward): + layer_id = module.layer_id + num_hidden_layers = module.config.num_hidden_layers + hidden_size = module.config.hidden_size + + ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0 + + time_weight = torch.tensor( + [i / hidden_size for i in range(hidden_size)], + dtype=module.time_mix_key.dtype, + device=module.time_mix_key.device, + ) + time_weight = time_weight[None, None, :] + + init.copy_(module.time_mix_key, torch.pow(time_weight, ratio_1_to_almost0)) + init.copy_(module.time_mix_receptance, torch.pow(time_weight, ratio_1_to_almost0)) + elif isinstance(module, nn.Linear): + shape = module.weight.shape + gain = 1.0 + scale = 1.0 # extra scale for gain + if module.bias is not None: + init.zeros_(module.bias) + if shape[0] > shape[1]: + gain = math.sqrt(shape[0] / shape[1]) + if shape[0] == self.config.vocab_size and shape[1] == self.config.hidden_size: # final projection? + scale = 0.5 + + gain *= scale + init.orthogonal_(module.weight, gain=gain) + elif isinstance(module, nn.Embedding): + shape = module.weight.shape + gain = 1e-4 * math.sqrt(max(shape[0], shape[1])) + init.orthogonal_(module.weight, gain=gain) + elif isinstance(module, nn.LayerNorm): + init.ones_(module.weight) + init.zeros_(module.bias) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for the RWKV model outputs. + """ +) +class RwkvOutput(ModelOutput): + r""" + state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): + The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to + avoid providing the old `input_ids`. + """ + + last_hidden_state: torch.FloatTensor | None = None + state: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for causal language model (or autoregressive) outputs. + """ +) +class RwkvCausalLMOutput(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). + state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): + The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to + avoid providing the old `input_ids`. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + state: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@auto_docstring +class RwkvModel(RwkvPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.blocks = nn.ModuleList([RwkvBlock(config, layer_id=idx) for idx in range(config.num_hidden_layers)]) + self.ln_out = nn.LayerNorm(config.hidden_size) + + self.layers_are_rescaled = False + + self.gradient_checkpointing = False + + # 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 + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + state: list[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 | RwkvOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*): + If passed along, the model uses the previous state in all the blocks (which will give the output for the + `input_ids` provided as if the model add `state_input_ids + input_ids` as context). + use_cache (`bool`, *optional*): + If set to `True`, the last state is returned and can be used to quickly generate the next logits. + """ + 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 if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if attention_mask is not None: + logger.warning_once("`attention_mask` was passed, but it is unused in this model.") + + if self.training == self.layers_are_rescaled: + self._rescale_layers() + + 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.embeddings(input_ids) + + if use_cache and state is None: + shape = (inputs_embeds.size(0), self.config.hidden_size, self.config.num_hidden_layers) + state = [ + torch.zeros( + *shape, dtype=inputs_embeds.dtype if i <= 1 else torch.float32, device=inputs_embeds.device + ) + for i in range(5) + ] + state[4] -= 1e30 + + 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 + + hidden_states = inputs_embeds + + all_self_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + for idx, block in enumerate(self.blocks): + hidden_states, state, attentions = block( + hidden_states, state=state, use_cache=use_cache, output_attentions=output_attentions + ) + + if ( + self.layers_are_rescaled + and self.config.rescale_every > 0 + and (idx + 1) % self.config.rescale_every == 0 + ): + hidden_states = hidden_states / 2 + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if output_attentions: + all_self_attentions = all_self_attentions + (attentions,) + + hidden_states = self.ln_out(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(x for x in [hidden_states, state, all_hidden_states, all_self_attentions] if x is not None) + + return RwkvOutput( + last_hidden_state=hidden_states, + state=state, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + def _rescale_layers(self): + # Layers should be rescaled for inference only. + if self.layers_are_rescaled == (not self.training): + return + if self.config.rescale_every > 0: + with torch.no_grad(): + for block_id, block in enumerate(self.blocks): + if self.training: + block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every)) + block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every)) + else: + # Deal with quantization statistics + if hasattr(block.attention.output.weight, "SCB"): + block.attention.output.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every)) + block.feed_forward.value.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every)) + elif hasattr(block.attention.output.weight, "quant_state"): + self._bnb_4bit_dequantize_and_rescale(block.attention.output, block_id) + self._bnb_4bit_dequantize_and_rescale(block.feed_forward.value, block_id) + else: + block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every)) + block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every)) + + self.layers_are_rescaled = not self.training + + def _bnb_4bit_dequantize_and_rescale(self, target_layer, block_id): + r""" + Perform the dequantization and rescaling of the weights of a given layer. After that operation the layer will + be quantized again. + """ + if not is_bitsandbytes_available(): + raise ImportError("Please install bitsandbytes to use this method.") + import bitsandbytes as bnb + + dequant_weights = bnb.functional.dequantize_4bit(target_layer.weight.data, target_layer.weight.quant_state) + + dequant_weights.div_(2 ** int(block_id // self.config.rescale_every)) + + # re-quantize the model: + # we need to put it first on CPU then back to the device + # this will create an overhead :/ + # We set requires_grad=False as we cannot compute gradients on top of 4bit parameters anyway and to avoid + # bugs with bnb + quant_weight = bnb.nn.Params4bit(dequant_weights.to("cpu"), requires_grad=False).to(dequant_weights.device) + setattr(target_layer, "weight", quant_weight) + + +@auto_docstring( + custom_intro=""" + The RWKV Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """ +) +class RwkvForCausalLM(RwkvPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"head.weight": "rwkv.embeddings.weight"} + + def __init__(self, config): + super().__init__(config) + self.rwkv = RwkvModel(config) + self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.head + + def set_output_embeddings(self, new_embeddings): + self.head = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + state: list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> tuple | RwkvCausalLMOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*): + If passed along, the model uses the previous state in all the blocks (which will give the output for the + `input_ids` provided as if the model add `state_input_ids + input_ids` as context). + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + use_cache (`bool`, *optional*): + If set to `True`, the last state is returned and can be used to quickly generate the next logits. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + rwkv_outputs = self.rwkv( + input_ids, + inputs_embeds=inputs_embeds, + state=state, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = rwkv_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.head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + rwkv_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return RwkvCausalLMOutput( + loss=loss, + logits=logits, + state=rwkv_outputs.state, + hidden_states=rwkv_outputs.hidden_states, + attentions=rwkv_outputs.attentions, + ) + + +__all__ = ["RwkvForCausalLM", "RwkvModel", "RwkvPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/sam3_tracker/__init__.py b/third_party/transformers/src/transformers/models/sam3_tracker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dce9c5ca2c60c735e87adf73bce69c996f4a1241 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sam3_tracker/__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_sam3_tracker import * + from .modeling_sam3_tracker import * + from .processing_sam3_tracker 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/sam3_tracker/configuration_sam3_tracker.py b/third_party/transformers/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..237fbc402ce3f4ab8f957ed1de3309b596a68204 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py @@ -0,0 +1,164 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sam3_tracker/modular_sam3_tracker.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_sam3_tracker.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerPromptEncoderConfig(PreTrainedConfig): + r""" + mask_input_channels (`int`, *optional*, defaults to 16): + The number of channels to be fed to the `MaskDecoder` module. + num_point_embeddings (`int`, *optional*, defaults to 4): + The number of point embeddings to be used. + scale (`float`, *optional*, defaults to 1): + The scale factor for the prompt encoder. + """ + + base_config_key = "prompt_encoder_config" + + hidden_size: int = 256 + + image_size: int | list[int] | tuple[int, int] = 1008 + patch_size: int | list[int] | tuple[int, int] = 14 + mask_input_channels: int = 16 + num_point_embeddings: int = 4 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-6 + scale: int = 1 + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerMaskDecoderConfig(PreTrainedConfig): + r""" + mlp_dim (`int`, *optional*, defaults to 2048): + The dimension of the MLP in the two-way transformer. + attention_downsample_rate (`int`, *optional*, defaults to 2): + The downsample rate for the attention layers. + num_multimask_outputs (`int`, *optional*, defaults to 3): + The number of multimask outputs. + iou_head_depth (`int`, *optional*, defaults to 3): + The depth of the IoU head. + iou_head_hidden_dim (`int`, *optional*, defaults to 256): + The hidden dimension of the IoU head. + dynamic_multimask_via_stability (`bool`, *optional*, defaults to `True`): + Whether to use dynamic multimask via stability. + dynamic_multimask_stability_delta (`float`, *optional*, defaults to 0.05): + The stability delta for the dynamic multimask. + dynamic_multimask_stability_thresh (`float`, *optional*, defaults to 0.98): + The stability threshold for the dynamic multimask. + """ + + base_config_key = "mask_decoder_config" + + hidden_size: int = 256 + hidden_act: str = "gelu" + mlp_dim: int = 2048 + num_hidden_layers: int = 2 + num_attention_heads: int = 8 + attention_downsample_rate: int = 2 + num_multimask_outputs: int = 3 + iou_head_depth: int = 3 + iou_head_hidden_dim: int = 256 + dynamic_multimask_via_stability: bool = True + dynamic_multimask_stability_delta: float = 0.05 + dynamic_multimask_stability_thresh: float = 0.98 + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerConfig(PreTrainedConfig): + r""" + prompt_encoder_config (Union[`dict`, `Sam3TrackerPromptEncoderConfig`], *optional*): + Dictionary of configuration options used to initialize [`Sam3TrackerPromptEncoderConfig`]. + mask_decoder_config (Union[`dict`, `Sam3TrackerMaskDecoderConfig`], *optional*): + Dictionary of configuration options used to initialize [`Sam3TrackerMaskDecoderConfig`]. + + Example: + + ```python + >>> from transformers import ( + ... Sam3TrackerVisionConfig, + ... Sam3TrackerPromptEncoderConfig, + ... Sam3TrackerMaskDecoderConfig, + ... Sam3TrackerModel, + ... ) + + >>> # Initializing a Sam3TrackerConfig with `"facebook/sam3_tracker.1_hiera_tiny"` style configuration + >>> configuration = Sam3TrackerConfig() + + >>> # Initializing a Sam3TrackerModel (with random weights) from the `"facebook/sam3_tracker.1_hiera_tiny"` style configuration + >>> model = Sam3TrackerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a Sam3TrackerConfig from a Sam3TrackerVisionConfig, Sam3TrackerPromptEncoderConfig, and Sam3TrackerMaskDecoderConfig + >>> # Initializing SAM3_TRACKER vision encoder, memory attention, and memory encoder configurations + >>> vision_config = Sam3TrackerVisionConfig() + >>> prompt_encoder_config = Sam3TrackerPromptEncoderConfig() + >>> mask_decoder_config = Sam3TrackerMaskDecoderConfig() + + >>> config = Sam3TrackerConfig(vision_config, prompt_encoder_config, mask_decoder_config) + ``` + """ + + model_type = "sam3_tracker" + sub_configs = { + "vision_config": AutoConfig, + "prompt_encoder_config": Sam3TrackerPromptEncoderConfig, + "mask_decoder_config": Sam3TrackerMaskDecoderConfig, + } + + vision_config: dict | PreTrainedConfig | None = None + prompt_encoder_config: dict | PreTrainedConfig | None = None + mask_decoder_config: dict | PreTrainedConfig | None = None + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "sam3_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["sam3_vision_model"]( + backbone_feature_sizes=[[288, 288], [144, 144], [72, 72]] + ) + + if isinstance(self.prompt_encoder_config, dict): + self.prompt_encoder_config = Sam3TrackerPromptEncoderConfig(**self.prompt_encoder_config) + elif self.prompt_encoder_config is None: + self.prompt_encoder_config = Sam3TrackerPromptEncoderConfig() + + if isinstance(self.mask_decoder_config, dict): + self.mask_decoder_config = Sam3TrackerMaskDecoderConfig(**self.mask_decoder_config) + elif self.mask_decoder_config is None: + self.mask_decoder_config = Sam3TrackerMaskDecoderConfig() + + super().__post_init__(**kwargs) + + +__all__ = ["Sam3TrackerConfig", "Sam3TrackerPromptEncoderConfig", "Sam3TrackerMaskDecoderConfig"] diff --git a/third_party/transformers/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py b/third_party/transformers/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..cd7f674683d0efdc054d6aaa24204f83cf4008c9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py @@ -0,0 +1,1106 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sam3_tracker/modular_sam3_tracker.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_sam3_tracker.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging +from ...utils.generic import TransformersKwargs, is_flash_attention_requested, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..auto import AutoModel +from .configuration_sam3_tracker import Sam3TrackerConfig, Sam3TrackerMaskDecoderConfig, Sam3TrackerPromptEncoderConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring(custom_intro="Base class for the Sam3Tracker model's output.") +class Sam3TrackerImageSegmentationOutput(ModelOutput): + r""" + iou_scores (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks)`): + The Intersection over Union (IoU) scores of the predicted masks. + pred_masks (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks, height, width)`): + The predicted low-resolution masks. This is an alias for `low_res_masks`. These masks need to be post-processed + by the processor to be brought to the original image size. + object_score_logits (`torch.FloatTensor` of shape `(batch_size, point_batch_size, 1)`): + Logits for the object score, indicating if an object is present. + image_embeddings (`tuple(torch.FloatTensor)`): + The features from the FPN, which are used by the mask decoder. This is a tuple of `torch.FloatTensor` where each + tensor has shape `(batch_size, channels, height, width)`. + vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. + Hidden-states of the vision model at the output of each stage. + vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `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 vision model. + mask_decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `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 mask decoder. + """ + + iou_scores: torch.FloatTensor | None = None + pred_masks: torch.FloatTensor | None = None + object_score_logits: torch.FloatTensor | None = None + image_embeddings: tuple[torch.FloatTensor, ...] = None + vision_hidden_states: tuple[torch.FloatTensor, ...] | None = None + vision_attentions: tuple[torch.FloatTensor, ...] | None = None + mask_decoder_attentions: tuple[torch.FloatTensor, ...] | None = None + + +class Sam3TrackerFeedForward(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + activation: str = "relu", + sigmoid_output: bool = False, + ): + super().__init__() + self.num_layers = num_layers + self.activation = ACT2FN[activation] + self.proj_in = nn.Linear(input_dim, hidden_dim) + self.proj_out = nn.Linear(hidden_dim, output_dim) + self.layers = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers - 2)]) + self.sigmoid_output = sigmoid_output + + def forward(self, hidden_states): + hidden_states = self.proj_in(hidden_states) + hidden_states = self.activation(hidden_states) + for layer in self.layers: + hidden_states = self.activation(layer(hidden_states)) + + hidden_states = self.proj_out(hidden_states) + if self.sigmoid_output: + hidden_states = F.sigmoid(hidden_states) + return hidden_states + + +@auto_docstring( + custom_intro=""" + Segment Anything Model 3 (SAM 3) for generating segmentation masks, given an input image and + input points and labels, boxes, or masks. + """ +) +class Sam3TrackerPreTrainedModel(PreTrainedModel): + config_class = Sam3TrackerConfig + base_model_prefix = "sam3_tracker" + main_input_name = "pixel_values" + input_modalities = ("image",) + _supports_sdpa = True + _supports_flash_attn = True + _supports_attention_backend = True + _keys_to_ignore_on_load_unexpected = [ + r"^memory_.*", + r"^mask_downsample.*", + r"^object_pointer_proj.*", + r"^temporal_positional_encoding_projection_layer.*", + "no_memory_positional_encoding", + "no_object_pointer", + "occlusion_spatial_embedding_parameter", + ] + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Sam3TrackerModel): + if module.no_memory_embedding is not None: + init.zeros_(module.no_memory_embedding) + elif isinstance(module, Sam3TrackerPositionalEmbedding): + init.normal_(module.positional_embedding, std=module.scale) + + +class Sam3TrackerPositionalEmbedding(nn.Module): + def __init__(self, config: Sam3TrackerPromptEncoderConfig): + super().__init__() + self.scale = config.scale + positional_embedding = self.scale * torch.randn((2, config.hidden_size // 2)) + self.register_buffer("positional_embedding", positional_embedding) + + def forward(self, input_coords, input_shape=None): + """Positionally encode points that are normalized to [0,1].""" + coordinates = input_coords.clone() + + if input_shape is not None: + coordinates[:, :, :, 0] = coordinates[:, :, :, 0] / input_shape[1] + coordinates[:, :, :, 1] = coordinates[:, :, :, 1] / input_shape[0] + coordinates.to(torch.float32) + + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coordinates = 2 * coordinates - 1 + coordinates = coordinates.to(self.positional_embedding.dtype) + coordinates = coordinates @ self.positional_embedding + coordinates = 2 * np.pi * coordinates + # outputs d_1 x ... x d_n x channel shape + return torch.cat([torch.sin(coordinates), torch.cos(coordinates)], dim=-1) + + +class Sam3TrackerMaskEmbedding(nn.Module): + def __init__(self, config: Sam3TrackerPromptEncoderConfig): + super().__init__() + self.mask_input_channels = config.mask_input_channels // 4 + self.activation = ACT2FN[config.hidden_act] + self.conv1 = nn.Conv2d(1, self.mask_input_channels, kernel_size=2, stride=2) + self.conv2 = nn.Conv2d(self.mask_input_channels, config.mask_input_channels, kernel_size=2, stride=2) + self.conv3 = nn.Conv2d(config.mask_input_channels, config.hidden_size, kernel_size=1) + self.layer_norm1 = Sam3TrackerLayerNorm( + self.mask_input_channels, eps=config.layer_norm_eps, data_format="channels_first" + ) + self.layer_norm2 = Sam3TrackerLayerNorm( + self.mask_input_channels * 4, eps=config.layer_norm_eps, data_format="channels_first" + ) + + def forward(self, masks): + hidden_states = self.conv1(masks) + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.activation(hidden_states) + + hidden_states = self.conv2(hidden_states) + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.activation(hidden_states) + dense_embeddings = self.conv3(hidden_states) + return dense_embeddings + + +class Sam3TrackerPromptEncoder(nn.Module): + def __init__(self, config: Sam3TrackerPromptEncoderConfig): + super().__init__() + self.shared_embedding = Sam3TrackerPositionalEmbedding(config) + self.mask_embed = Sam3TrackerMaskEmbedding(config) + self.no_mask_embed = nn.Embedding(1, config.hidden_size) + + self.image_embedding_size = (config.image_size // config.patch_size, config.image_size // config.patch_size) + self.mask_input_size = (4 * config.image_size // config.patch_size, 4 * config.image_size // config.patch_size) + self.input_image_size = config.image_size + + self.point_embed = nn.Embedding(config.num_point_embeddings, config.hidden_size) + self.hidden_size = config.hidden_size + self.not_a_point_embed = nn.Embedding(1, config.hidden_size) + + def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + points = torch.nn.functional.pad(points, (0, 0, 0, 1), mode="constant", value=0) + labels = torch.nn.functional.pad(labels, (0, 1), mode="constant", value=-1) + input_shape = (self.input_image_size, self.input_image_size) + point_embedding = self.shared_embedding(points, input_shape) + + # torch.where and expanding the labels tensor is required by the ONNX export + point_embedding = torch.where(labels[..., None] == -1, self.not_a_point_embed.weight, point_embedding) + + # This is required for the ONNX export. The dtype, device need to be explicitly + # specified as otherwise torch.onnx.export interprets as double + point_embedding = torch.where( + labels[..., None] != -10, + point_embedding, + torch.zeros_like(point_embedding), + ) + + # Add point embeddings for labels >= 0 + point_embedding = point_embedding + self.point_embed(labels.clamp(min=0)) * (labels >= 0).unsqueeze(-1) + + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.view(*boxes.shape[:2], 2, 2) + # add padding point for consistency with the original implementation + coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0) + corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size)) + corner_embedding[:, :, 0, :] += self.point_embed.weight[2] + corner_embedding[:, :, 1, :] += self.point_embed.weight[3] + corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :]) + return corner_embedding + + def forward( + self, + input_points: tuple[torch.Tensor, torch.Tensor] | None, + input_labels: torch.Tensor | None, + input_boxes: torch.Tensor | None, + input_masks: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense embeddings. + + Args: + points (`torch.Tensor`, *optional*): + point coordinates and labels to embed. + boxes (`torch.Tensor`, *optional*): + boxes to embed + masks (`torch.Tensor`, *optional*): + masks to embed + """ + sparse_embeddings = None + batch_size = 1 + if input_points is not None: + batch_size = input_points.shape[0] + if input_labels is None: + raise ValueError("If points are provided, labels must also be provided.") + point_embeddings = self._embed_points(input_points, input_labels, pad=(input_boxes is None)) + sparse_embeddings = point_embeddings + if input_boxes is not None: + batch_size = input_boxes.shape[0] + box_embeddings = self._embed_boxes(input_boxes) + if sparse_embeddings is None: + sparse_embeddings = box_embeddings + else: + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=2) + if input_masks is not None: + dense_embeddings = self.mask_embed(input_masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + batch_size, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_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(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 + + +class Sam3TrackerAttention(nn.Module): + """ + SAM3_TRACKER's attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and + values. + """ + + def __init__(self, config, downsample_rate=None): + super().__init__() + downsample_rate = config.attention_downsample_rate if downsample_rate is None else downsample_rate + self.config = config + self.hidden_size = config.hidden_size + self.internal_dim = config.hidden_size // downsample_rate + self.num_attention_heads = config.num_attention_heads + self.head_dim = self.internal_dim // config.num_attention_heads + self.scaling = self.head_dim**-0.5 + self.is_causal = False + + self.q_proj = nn.Linear(self.hidden_size, self.internal_dim) + self.k_proj = nn.Linear(self.hidden_size, self.internal_dim) + self.v_proj = nn.Linear(self.hidden_size, self.internal_dim) + self.o_proj = nn.Linear(self.internal_dim, self.hidden_size) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_similarity: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + # Input projections + batch_size, point_batch_size = query.shape[:2] + new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim) + + query = self.q_proj(query).view(*new_shape).transpose(1, 2) + key = self.k_proj(key).view(*new_shape).transpose(1, 2) + value = self.v_proj(value).view(*new_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + if is_flash_attention_requested(self.config) and attention_similarity is not None: + # Target guided masks are represented as float masks and are incompatible with Flash Attention + # Fallback to SDPA for this call only so the rest of the model can still benefit from FA + attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"] + logger.warning_once( + "Falling back to SDPA for target-guided attention because " + "Flash Attention does not support additive bias masks." + ) + + attn_output, attn_weights = attention_interface( + self, + query, + key, + value, + attention_mask=attention_similarity, + dropout=0.0, + scaling=self.scaling, + is_causal=self.is_causal, + **kwargs, + ) + + attn_output = attn_output.reshape( + batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class Sam3TrackerTwoWayAttentionBlock(GradientCheckpointingLayer): + def __init__(self, config: Sam3TrackerMaskDecoderConfig, skip_first_layer_pe: bool = False): + """ + A transformer block with four layers: + (1) self-attention of sparse inputs (2) cross attention of sparse inputs -> dense inputs (3) mlp block on + sparse inputs (4) cross attention of dense inputs -> sparse inputs + + Arguments: + config (`Sam3TrackerMaskDecoderConfig`): + The configuration file used to instantiate the block + attention_downsample_rate (*optionalk*, int, defaults to 2): + The downsample ratio of the block used to reduce the inner dim of the attention. + skip_first_layer_pe (*optional*, bool, defaults to `False`): + Whether or not to skip the addition of the query_point_embedding on the first layer. + """ + super().__init__() + self.self_attn = Sam3TrackerAttention(config, downsample_rate=1) + self.layer_norm1 = nn.LayerNorm(config.hidden_size) + + self.cross_attn_token_to_image = Sam3TrackerAttention(config) + self.layer_norm2 = nn.LayerNorm(config.hidden_size) + + self.mlp = Sam3TrackerFeedForward( + config.hidden_size, config.mlp_dim, config.hidden_size, num_layers=config.num_hidden_layers + ) + self.layer_norm3 = nn.LayerNorm(config.hidden_size) + + self.layer_norm4 = nn.LayerNorm(config.hidden_size) + self.cross_attn_image_to_token = Sam3TrackerAttention(config) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, + queries: Tensor, + keys: Tensor, + query_point_embedding: Tensor, + key_point_embedding: Tensor, + attention_similarity: Tensor, + **kwargs: Unpack[TransformersKwargs], + ): + # Self attention block + if self.skip_first_layer_pe: + queries, _ = self.self_attn(query=queries, key=queries, value=queries) + else: + query = queries + query_point_embedding + attn_out, _ = self.self_attn(query=query, key=query, value=queries) + queries = queries + attn_out + queries = self.layer_norm1(queries) + + # Cross attention block, tokens attending to image embedding + query = queries + query_point_embedding + key = keys + key_point_embedding + + attn_out, _ = self.cross_attn_token_to_image( + query=query, key=key, value=keys, attention_similarity=attention_similarity + ) + queries = queries + attn_out + + queries = self.layer_norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.layer_norm3(queries) + + # Cross attention block, image embedding attending to tokens + query = queries + query_point_embedding + key = keys + key_point_embedding + + attn_out, _ = self.cross_attn_image_to_token(query=key, key=query, value=queries) + keys = keys + attn_out + + keys = self.layer_norm4(keys) + return queries, keys, attn_out + + +class Sam3TrackerTwoWayTransformer(nn.Module): + def __init__(self, config: Sam3TrackerMaskDecoderConfig): + super().__init__() + self.config = config + + self.num_hidden_layers = config.num_hidden_layers + self.layers = nn.ModuleList() + + for i in range(self.num_hidden_layers): + self.layers.append(Sam3TrackerTwoWayAttentionBlock(config, skip_first_layer_pe=(i == 0))) + + self.final_attn_token_to_image = Sam3TrackerAttention(config) + self.layer_norm_final_attn = nn.LayerNorm(config.hidden_size) + + def forward( + self, + point_embeddings: Tensor, + image_embeddings: Tensor, + image_positional_embeddings: Tensor, + attention_similarity: Tensor, + target_embedding=None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + if image_embeddings is None: + raise ValueError("You have to specify an image_embedding") + + image_embeddings = image_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1) + image_positional_embeddings = image_positional_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1) + + # Prepare queries + queries = point_embeddings + keys = image_embeddings + + # Apply transformer blocks and final layernorm + for layer in self.layers: + if target_embedding is not None: + queries += target_embedding + + queries, keys, _ = layer( + queries=queries, + keys=keys, + query_point_embedding=point_embeddings, + key_point_embedding=image_positional_embeddings, + attention_similarity=attention_similarity, + **kwargs, + ) + # Apply the final attention layer from the points to the image + query = queries + point_embeddings + key = keys + image_positional_embeddings + + attn_out, _ = self.final_attn_token_to_image(query=query, key=key, value=keys) + + queries = queries + attn_out + queries = self.layer_norm_final_attn(queries) + return queries, keys + + +class Sam3TrackerLayerNorm(nn.LayerNorm): + r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, + width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs): + super().__init__(normalized_shape, eps=eps, **kwargs) + if data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError(f"Unsupported data format: {data_format}") + self.data_format = data_format + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """ + Args: + features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels) + """ + if self.data_format == "channels_first": + features = features.permute(0, 2, 3, 1) + features = super().forward(features) + features = features.permute(0, 3, 1, 2) + else: + features = super().forward(features) + return features + + +class Sam3TrackerMaskDecoder(nn.Module): + def __init__(self, config: Sam3TrackerMaskDecoderConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + + self.num_multimask_outputs = config.num_multimask_outputs + self.num_mask_tokens = config.num_multimask_outputs + 1 + + self.iou_token = nn.Embedding(1, self.hidden_size) + self.mask_tokens = nn.Embedding(self.num_mask_tokens, self.hidden_size) + + self.transformer = Sam3TrackerTwoWayTransformer(config) + + # should we create a new class for this? + self.upscale_conv1 = nn.ConvTranspose2d(self.hidden_size, self.hidden_size // 4, kernel_size=2, stride=2) + self.upscale_conv2 = nn.ConvTranspose2d(self.hidden_size // 4, self.hidden_size // 8, kernel_size=2, stride=2) + self.upscale_layer_norm = Sam3TrackerLayerNorm(self.hidden_size // 4, data_format="channels_first") + self.activation = nn.GELU() + + mlps_list = [] + for _ in range(self.num_mask_tokens): + mlps_list += [Sam3TrackerFeedForward(self.hidden_size, self.hidden_size, self.hidden_size // 8, 3)] + self.output_hypernetworks_mlps = nn.ModuleList(mlps_list) + self.iou_prediction_head = Sam3TrackerFeedForward( + self.hidden_size, + config.iou_head_hidden_dim, + self.num_mask_tokens, + config.iou_head_depth, + sigmoid_output=True, + ) + + self.conv_s0 = nn.Conv2d(config.hidden_size, config.hidden_size // 8, kernel_size=1, stride=1) + self.conv_s1 = nn.Conv2d(config.hidden_size, config.hidden_size // 4, kernel_size=1, stride=1) + + self.obj_score_token = nn.Embedding(1, self.hidden_size) + self.pred_obj_score_head = Sam3TrackerFeedForward(self.hidden_size, self.hidden_size, 1, 3) + + self.dynamic_multimask_via_stability = config.dynamic_multimask_via_stability + self.dynamic_multimask_stability_delta = config.dynamic_multimask_stability_delta + self.dynamic_multimask_stability_thresh = config.dynamic_multimask_stability_thresh + + def forward( + self, + image_embeddings: torch.Tensor, + image_positional_embeddings: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + high_resolution_features: list[torch.Tensor], + attention_similarity: torch.Tensor | None = None, + target_embedding: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Args: + image_embeddings (`torch.Tensor`): + The embeddings from the image encoder. + image_positional_embeddings (`torch.Tensor`): + Positional encoding with the shape of image_embeddings. + sparse_prompt_embeddings (`torch.Tensor`): + The embeddings of the points and boxes. + dense_prompt_embeddings (`torch.Tensor`): + The embeddings of the mask inputs. + multimask_output (`bool`): + Whether to return multiple masks or a single mask. + high_resolution_features (`list[torch.Tensor]`, *optional*): + The high-resolution features from the vision encoder. + attention_similarity (`torch.Tensor`, *optional*): + The attention similarity tensor. + target_embedding (`torch.Tensor`, *optional*): + The target embedding. + """ + batch_size, num_channels, height, width = image_embeddings.shape + point_batch_size = sparse_prompt_embeddings.shape[1] + # Concatenate output tokens + output_tokens = torch.cat( + [ + self.obj_score_token.weight, + self.iou_token.weight, + self.mask_tokens.weight, + ], + dim=0, + ) + output_tokens = output_tokens.repeat(batch_size, point_batch_size, 1, 1) + + if sparse_prompt_embeddings.shape[0] != 0: + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=2) + else: + tokens = output_tokens + point_embeddings = tokens.to(self.iou_token.weight.dtype) + + # Expand per-image data in batch direction to be per-mask + image_embeddings = image_embeddings + dense_prompt_embeddings + image_embeddings = image_embeddings.repeat_interleave(point_batch_size, dim=0) + image_positional_embeddings = image_positional_embeddings.repeat_interleave(point_batch_size, 0) + # Run the transformer + point_embeddings, image_embeddings = self.transformer( + point_embeddings=point_embeddings, + image_embeddings=image_embeddings, + image_positional_embeddings=image_positional_embeddings, + attention_similarity=attention_similarity, + target_embedding=target_embedding, + **kwargs, + ) + iou_token_out = point_embeddings[:, :, 1, :] + mask_tokens_out = point_embeddings[:, :, 2 : (2 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + image_embeddings = image_embeddings.transpose(2, 3).view( + batch_size * point_batch_size, num_channels, height, width + ) + + feat_s0, feat_s1 = high_resolution_features + feat_s0 = feat_s0.repeat_interleave(point_batch_size, dim=0) + feat_s1 = feat_s1.repeat_interleave(point_batch_size, dim=0) + upscaled_embedding = self.upscale_conv1(image_embeddings) + feat_s1 + upscaled_embedding = self.activation(self.upscale_layer_norm(upscaled_embedding)) + upscaled_embedding = self.activation(self.upscale_conv2(upscaled_embedding) + feat_s0) + + hyper_in_list: list[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + current_mlp = self.output_hypernetworks_mlps[i] + hyper_in_list += [current_mlp(mask_tokens_out[:, :, i, :])] + hyper_in = torch.stack(hyper_in_list, dim=2) + + _, num_channels, height, width = upscaled_embedding.shape + upscaled_embedding = upscaled_embedding.view(batch_size, point_batch_size, num_channels, height * width) + masks = (hyper_in @ upscaled_embedding).view(batch_size, point_batch_size, -1, height, width) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + object_score_logits = self.pred_obj_score_head(point_embeddings[:, :, 0, :]) + + # Select the correct mask or masks for output + if multimask_output: + mask_slice = slice(1, None) + masks = masks[:, :, mask_slice, :, :] + iou_pred = iou_pred[:, :, mask_slice] + elif self.dynamic_multimask_via_stability and not self.training: + mask_slice = slice(0, 1) + masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred) + else: + mask_slice = slice(0, 1) + masks = masks[:, :, mask_slice, :, :] + iou_pred = iou_pred[:, :, mask_slice] + + sam_tokens_out = mask_tokens_out[:, :, mask_slice] # [b, 3, c] shape + + return masks, iou_pred, sam_tokens_out, object_score_logits + + def _get_stability_scores(self, mask_logits): + """ + Compute stability scores of the mask logits based on the IoU between upper and + lower thresholds. + """ + mask_logits = mask_logits.flatten(-2) + stability_delta = self.dynamic_multimask_stability_delta + area_i = torch.sum(mask_logits > stability_delta, dim=-1).float() + area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float() + stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0) + return stability_scores + + def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores): + """ + When outputting a single mask, if the stability score from the current single-mask + output (based on output token 0) falls below a threshold, we instead select from + multi-mask outputs (based on output token 1~3) the mask with the highest predicted + IoU score. This is intended to ensure a valid mask for both clicking and tracking. + """ + # The best mask from multimask output tokens (1~3) + multimask_logits = all_mask_logits[:, :, 1:, :, :] + multimask_iou_scores = all_iou_scores[:, :, 1:] + best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) # [B, P] + best_scores_inds_expanded = best_scores_inds.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + best_scores_inds_expanded = best_scores_inds_expanded.expand( + -1, -1, 1, multimask_logits.size(-2), multimask_logits.size(-1) + ) + best_multimask_logits = torch.gather(multimask_logits, 2, best_scores_inds_expanded) # [B, P, 1, H, W] + best_multimask_iou_scores = torch.gather(multimask_iou_scores, 2, best_scores_inds.unsqueeze(-1)) # [B, P, 1] + + # The mask from singlemask output token 0 and its stability score + singlemask_logits = all_mask_logits[:, :, 0:1, :, :] + singlemask_iou_scores = all_iou_scores[:, :, 0:1] + stability_scores = self._get_stability_scores(singlemask_logits) + is_stable = stability_scores >= self.dynamic_multimask_stability_thresh + + # Dynamically fall back to best multimask output upon low stability scores. + mask_logits_out = torch.where( + is_stable[..., None, None].expand_as(singlemask_logits), + singlemask_logits, + best_multimask_logits, + ) + iou_scores_out = torch.where( + is_stable.expand_as(singlemask_iou_scores), + singlemask_iou_scores, + best_multimask_iou_scores, + ) + return mask_logits_out, iou_scores_out + + +@dataclass +@auto_docstring(custom_intro="Base class for the vision encoder's outputs.") +class Sam3TrackerVisionEncoderOutput(BaseModelOutputWithPooling): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the + model at the output of each stage. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in + the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. + """ + + fpn_hidden_states: torch.FloatTensor | None = None + fpn_position_encoding: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + Segment Anything Model 2 (SAM 2) for generating segmentation masks, given an input image and + input points and labels, boxes, or masks. + """ +) +class Sam3TrackerModel(Sam3TrackerPreTrainedModel): + input_modalities = ("image", "text") + _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(Sam3TrackerTwoWayAttentionBlock, index=2)} + _tied_weights_keys = {} + _base_model_prefix = "tracker_model" + _keys_to_ignore_on_load_unexpected = [ + r"^detector_model.", + r"^memory_.*", + r"^mask_downsample.*", + r"^object_pointer_proj.*", + r"^temporal_positional_encoding_projection_layer.*", + "no_memory_positional_encoding", + "no_object_pointer", + "occlusion_spatial_embedding_parameter", + ] + + def __init__(self, config: Sam3TrackerConfig): + # loading from a sam3_video config + if hasattr(config, "tracker_config") and config.tracker_config is not None: + if isinstance(config.tracker_config, dict): + config.tracker_config = Sam3TrackerConfig(**config.tracker_config) + config = config.tracker_config + super().__init__(config) + self.shared_image_embedding = Sam3TrackerPositionalEmbedding(config.prompt_encoder_config) + self.vision_encoder = AutoModel.from_config(config.vision_config) + self.prompt_encoder = Sam3TrackerPromptEncoder(config.prompt_encoder_config) + # The module using it is not a PreTrainedModel subclass so we need this + config.mask_decoder_config._attn_implementation = config._attn_implementation + self.mask_decoder = Sam3TrackerMaskDecoder(config.mask_decoder_config) + + self.backbone_feature_sizes = config.vision_config.backbone_feature_sizes + # a single token to indicate no memory embedding from previous frames + self.hidden_dim = config.vision_config.fpn_hidden_size + self.no_memory_embedding = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + + self.post_init() + + def get_input_embeddings(self): + return self.vision_encoder.get_input_embeddings() + + def get_image_wide_positional_embeddings(self) -> torch.Tensor: + size = self.prompt_encoder.image_embedding_size + target_device = self.shared_image_embedding.positional_embedding.device + target_dtype = self.shared_image_embedding.positional_embedding.dtype + grid = torch.ones(size, device=target_device, dtype=target_dtype) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / size[0] + x_embed = x_embed / size[1] + + positional_embedding = self.shared_image_embedding(torch.stack([x_embed, y_embed], dim=-1)) + return positional_embedding.permute(2, 0, 1).unsqueeze(0) # channel x height x width + + @torch.no_grad() + def get_image_embeddings( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> list[torch.Tensor]: + r""" + Returns the image embeddings by passing the pixel values through the vision encoder. + + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Input pixel values + """ + batch_size = pixel_values.shape[0] + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + + # add no memory embedding to the last feature map + feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding + + # reshape feature maps to the same shape as the backbone feature sizes + image_embeddings = [ + feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) + for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes) + ] + + return image_embeddings + + @torch.no_grad() + def get_prompt_embeddings( + self, + input_points: torch.FloatTensor | None = None, + input_labels: torch.LongTensor | None = None, + input_boxes: torch.FloatTensor | None = None, + input_masks: torch.LongTensor | None = None, + ): + r""" + Returns the prompt embeddings by passing the input points, labels, boxes and masks through the prompt encoder. + + Args: + input_points (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_points_per_image, 2)`): + Optional input points for the prompt encoder. The padding of the point is automatically done by the + processor. `point_batch_size` refers to the number of masks that we want the model to predict per + point. The model will output `point_batch_size` times 3 masks in total. + input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points_per_image)`): + Optional input labels for the prompt encoder. The padding of the labels is automatically done by the + processor, or can be fed by the user. + input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes_per_image, 4)`): + Optional input boxes for the prompt encoder. The padding of the boxes is automatically done by the + processor. users can also pass manually the input boxes. + input_masks (`torch.LongTensor` of shape `(batch_size, image_size, image_size)`): + Optional input masks for the prompt encoder. + """ + prompt_output = self.prompt_encoder( + input_points=input_points, + input_labels=input_labels, + input_boxes=input_boxes, + input_masks=input_masks, + ) + return prompt_output + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + input_points: torch.FloatTensor | None = None, + input_labels: torch.LongTensor | None = None, + input_boxes: torch.FloatTensor | None = None, + input_masks: torch.LongTensor | None = None, + image_embeddings: torch.FloatTensor | None = None, + multimask_output: bool = True, + attention_similarity: torch.FloatTensor | None = None, + target_embedding: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Sam3TrackerImageSegmentationOutput: + r""" + input_points (`torch.FloatTensor` of shape `(batch_size, num_points, 2)`): + Input 2D spatial points, this is used by the prompt encoder to encode the prompt. Generally yields to much + better results. The points can be obtained by passing a list of list of list to the processor that will + create corresponding `torch` tensors of dimension 4. The first dimension is the image batch size, the + second dimension is the point batch size (i.e. how many segmentation masks do we want the model to predict + per input point), the third dimension is the number of points per segmentation mask (it is possible to pass + multiple points for a single mask), and the last dimension is the x (vertical) and y (horizontal) + coordinates of the point. If a different number of points is passed either for each image, or for each + mask, the processor will create "PAD" points that will correspond to the (0, 0) coordinate, and the + computation of the embedding will be skipped for these points using the labels. + input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points)`): + Input labels for the points, this is used by the prompt encoder to encode the prompt. According to the + official implementation, there are 3 types of labels + + - `1`: the point is a point that contains the object of interest + - `0`: the point is a point that does not contain the object of interest + - `-1`: the point corresponds to the background + + We added the label: + + - `-10`: the point is a padding point, thus should be ignored by the prompt encoder + + The padding labels should be automatically done by the processor. + input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes, 4)`): + Input boxes for the points, this is used by the prompt encoder to encode the prompt. Generally yields to + much better generated masks. The boxes can be obtained by passing a list of list of list to the processor, + that will generate a `torch` tensor, with each dimension corresponding respectively to the image batch + size, the number of boxes per image and the coordinates of the top left and bottom right point of the box. + In the order (`x1`, `y1`, `x2`, `y2`): + + - `x1`: the x coordinate of the top left point of the input box + - `y1`: the y coordinate of the top left point of the input box + - `x2`: the x coordinate of the bottom right point of the input box + - `y2`: the y coordinate of the bottom right point of the input box + input_masks (`torch.FloatTensor` of shape `(batch_size, image_size, image_size)`): + SAM model also accepts segmentation masks as input. The mask will be embedded by the prompt encoder to + generate a corresponding embedding, that will be fed later on to the mask decoder. These masks needs to be + manually fed by the user, and they need to be of shape (`batch_size`, `image_size`, `image_size`). + image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_channels, window_size, window_size)`): + Image embeddings, this is used by the mask decoder to generate masks and iou scores. For more memory + efficient computation, users can first retrieve the image embeddings using the `get_image_embeddings` + method, and then feed them to the `forward` method instead of feeding the `pixel_values`. + multimask_output (`bool`, *optional*): + In the original implementation and paper, the model always outputs 3 masks per image (or per point / per + bounding box if relevant). However, it is possible to just output a single mask, that corresponds to the + "best" mask, by specifying `multimask_output=False`. + attention_similarity (`torch.FloatTensor`, *optional*): + Attention similarity tensor, to be provided to the mask decoder for target-guided attention in case the + model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048). + target_embedding (`torch.FloatTensor`, *optional*): + Embedding of the target concept, to be provided to the mask decoder for target-semantic prompting in case + the model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048). + + Example: + + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoModel, AutoProcessor + + >>> model = AutoModel.from_pretrained("danelcsb/sam3_tracker.1_hiera_tiny") + >>> processor = AutoProcessor.from_pretrained("danelcsb/sam3_tracker.1_hiera_tiny") + + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png" + >>> with httpx.stream("GET", url) as response: + ... raw_image = Image.open(BytesIO(response.read())).convert("RGB") + >>> input_points = [[[400, 650]]] # 2D location of a window on the car + >>> inputs = processor(images=raw_image, input_points=input_points, return_tensors="pt") + + >>> # Get segmentation mask + >>> outputs = model(**inputs) + + >>> # Postprocess masks + >>> masks = processor.post_process_masks( + ... outputs.pred_masks, inputs["original_sizes"] + ... ) + ``` + """ + if not ((pixel_values is None) ^ (image_embeddings is None)): + raise ValueError("Exactly one of pixel_values or image_embeddings must be provided.") + if input_points is not None and input_boxes is not None: + if input_points.shape[1] != input_boxes.shape[1]: + raise ValueError( + f"You should provide as many bounding boxes as input points per box. Got {input_points.shape[1]} and {input_boxes.shape[1]}." + ) + + image_positional_embeddings = self.get_image_wide_positional_embeddings() + # repeat with batch size + batch_size = pixel_values.shape[0] if pixel_values is not None else image_embeddings[-1].shape[0] + image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1) + + vision_attentions = None + vision_hidden_states = None + + if pixel_values is not None: + image_outputs: Sam3TrackerVisionEncoderOutput = self.get_image_features( + pixel_values, return_dict=True, **kwargs + ) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions + + # add no memory embedding to the last feature map + feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding + + # reshape feature maps to the same shape as the backbone feature sizes + image_embeddings = [ + feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) + for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes) + ] + + if input_points is not None and input_labels is None: + input_labels = torch.ones_like(input_points[:, :, :, 0], dtype=torch.int, device=input_points.device) + + if input_points is None and input_boxes is None: + # If no points are provide, pad with an empty point (with label -1) + input_points = torch.zeros( + batch_size, 1, 1, 2, dtype=image_embeddings[-1].dtype, device=image_embeddings[-1].device + ) + input_labels = -torch.ones(batch_size, 1, 1, dtype=torch.int32, device=image_embeddings[-1].device) + + if input_masks is not None: + # If mask_inputs is provided, downsize it into low-res mask input if needed + # and feed it as a dense mask prompt into the SAM mask encoder + if input_masks.shape[-2:] != self.prompt_encoder.mask_input_size: + input_masks = F.interpolate( + input_masks.float(), + size=self.prompt_encoder.mask_input_size, + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ).to(input_masks.dtype) + + sparse_embeddings, dense_embeddings = self.prompt_encoder( + input_points=input_points, + input_labels=input_labels, + input_boxes=input_boxes, + input_masks=input_masks, + ) + low_res_multimasks, iou_scores, _, object_score_logits = self.mask_decoder( + image_embeddings=image_embeddings[-1], + image_positional_embeddings=image_positional_embeddings, + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + high_resolution_features=image_embeddings[:-1], + attention_similarity=attention_similarity, + target_embedding=target_embedding, + **kwargs, + ) + + return Sam3TrackerImageSegmentationOutput( + iou_scores=iou_scores, + pred_masks=low_res_multimasks, + object_score_logits=object_score_logits, + image_embeddings=image_embeddings, + vision_hidden_states=vision_hidden_states, + vision_attentions=vision_attentions, + ) + + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Sam3TrackerVisionEncoderOutput: + r""" + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. + """ + vision_outputs: Sam3TrackerVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) + + feature_maps = vision_outputs.fpn_hidden_states + feature_maps_position_embeddings = vision_outputs.fpn_position_encoding + + # precompute projected level 0 and level 1 features in SAM decoder + # to avoid running it again on every SAM click + feature_maps = list(feature_maps) + feature_maps[0] = self.mask_decoder.conv_s0(feature_maps[0]) + feature_maps[1] = self.mask_decoder.conv_s1(feature_maps[1]) + + # flatten NxCxHxW to HWxNxC + feature_maps = [feature_map.flatten(2).permute(2, 0, 1) for feature_map in feature_maps] + feature_maps_position_embeddings = [ + feature_map_position_embedding.flatten(2).permute(2, 0, 1) + for feature_map_position_embedding in feature_maps_position_embeddings + ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings + + return vision_outputs + + +__all__ = ["Sam3TrackerModel", "Sam3TrackerPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/sam3_tracker/modular_sam3_tracker.py b/third_party/transformers/src/transformers/models/sam3_tracker/modular_sam3_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..1da111af856ff0adfd1a35d0cad6e48780f81762 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sam3_tracker/modular_sam3_tracker.py @@ -0,0 +1,233 @@ +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...configuration_utils import PreTrainedConfig +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoModel +from ..sam2.configuration_sam2 import ( + Sam2Config, + Sam2MaskDecoderConfig, + Sam2PromptEncoderConfig, +) +from ..sam2.modeling_sam2 import ( + Sam2Attention, + Sam2FeedForward, + Sam2ImageSegmentationOutput, + Sam2LayerNorm, + Sam2MaskDecoder, + Sam2MaskEmbedding, + Sam2Model, + Sam2PositionalEmbedding, + Sam2PreTrainedModel, + Sam2PromptEncoder, + Sam2TwoWayAttentionBlock, + Sam2TwoWayTransformer, +) +from ..sam2.processing_sam2 import Sam2Processor + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerPromptEncoderConfig(Sam2PromptEncoderConfig): + r""" + mask_input_channels (`int`, *optional*, defaults to 16): + The number of channels to be fed to the `MaskDecoder` module. + num_point_embeddings (`int`, *optional*, defaults to 4): + The number of point embeddings to be used. + scale (`float`, *optional*, defaults to 1): + The scale factor for the prompt encoder. + """ + + base_config_key = "prompt_encoder_config" + + image_size: int | list[int] | tuple[int, int] = 1008 + patch_size: int | list[int] | tuple[int, int] = 14 + + +class Sam3TrackerProcessor(Sam2Processor): + pass + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerMaskDecoderConfig(Sam2MaskDecoderConfig): + pass + + +@auto_docstring(checkpoint="facebook/sam3") +@strict +class Sam3TrackerConfig(Sam2Config): + r""" + prompt_encoder_config (Union[`dict`, `Sam3TrackerPromptEncoderConfig`], *optional*): + Dictionary of configuration options used to initialize [`Sam3TrackerPromptEncoderConfig`]. + mask_decoder_config (Union[`dict`, `Sam3TrackerMaskDecoderConfig`], *optional*): + Dictionary of configuration options used to initialize [`Sam3TrackerMaskDecoderConfig`]. + + Example: + + ```python + >>> from transformers import ( + ... Sam3TrackerVisionConfig, + ... Sam3TrackerPromptEncoderConfig, + ... Sam3TrackerMaskDecoderConfig, + ... Sam3TrackerModel, + ... ) + + >>> # Initializing a Sam3TrackerConfig with `"facebook/sam3_tracker.1_hiera_tiny"` style configuration + >>> configuration = Sam3TrackerConfig() + + >>> # Initializing a Sam3TrackerModel (with random weights) from the `"facebook/sam3_tracker.1_hiera_tiny"` style configuration + >>> model = Sam3TrackerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a Sam3TrackerConfig from a Sam3TrackerVisionConfig, Sam3TrackerPromptEncoderConfig, and Sam3TrackerMaskDecoderConfig + >>> # Initializing SAM3_TRACKER vision encoder, memory attention, and memory encoder configurations + >>> vision_config = Sam3TrackerVisionConfig() + >>> prompt_encoder_config = Sam3TrackerPromptEncoderConfig() + >>> mask_decoder_config = Sam3TrackerMaskDecoderConfig() + + >>> config = Sam3TrackerConfig(vision_config, prompt_encoder_config, mask_decoder_config) + ``` + """ + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "sam3_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["sam3_vision_model"]( + backbone_feature_sizes=[[288, 288], [144, 144], [72, 72]] + ) + + if isinstance(self.prompt_encoder_config, dict): + self.prompt_encoder_config = Sam3TrackerPromptEncoderConfig(**self.prompt_encoder_config) + elif self.prompt_encoder_config is None: + self.prompt_encoder_config = Sam3TrackerPromptEncoderConfig() + + if isinstance(self.mask_decoder_config, dict): + self.mask_decoder_config = Sam3TrackerMaskDecoderConfig(**self.mask_decoder_config) + elif self.mask_decoder_config is None: + self.mask_decoder_config = Sam3TrackerMaskDecoderConfig() + + PreTrainedConfig.__post_init__(**kwargs) + + +class Sam3TrackerImageSegmentationOutput(Sam2ImageSegmentationOutput): + pass + + +class Sam3TrackerFeedForward(Sam2FeedForward): + pass + + +@auto_docstring( + custom_intro=""" + Segment Anything Model 3 (SAM 3) for generating segmentation masks, given an input image and + input points and labels, boxes, or masks. + """ +) +class Sam3TrackerPreTrainedModel(Sam2PreTrainedModel): + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(module) + if isinstance(module, Sam3TrackerModel): + if module.no_memory_embedding is not None: + init.zeros_(module.no_memory_embedding) + elif isinstance(module, Sam3TrackerPositionalEmbedding): + init.normal_(module.positional_embedding, std=module.scale) + + +class Sam3TrackerPositionalEmbedding(Sam2PositionalEmbedding): + pass + + +class Sam3TrackerMaskEmbedding(Sam2MaskEmbedding): + pass + + +class Sam3TrackerPromptEncoder(Sam2PromptEncoder): + pass + + +class Sam3TrackerAttention(Sam2Attention): + pass + + +class Sam3TrackerTwoWayAttentionBlock(Sam2TwoWayAttentionBlock): + pass + + +class Sam3TrackerTwoWayTransformer(Sam2TwoWayTransformer): + pass + + +class Sam3TrackerLayerNorm(Sam2LayerNorm): + pass + + +class Sam3TrackerMaskDecoder(Sam2MaskDecoder): + pass + + +class Sam3TrackerModel(Sam2Model): + _base_model_prefix = "tracker_model" + _keys_to_ignore_on_load_unexpected = [ + r"^detector_model.", + r"^memory_.*", + r"^mask_downsample.*", + r"^object_pointer_proj.*", + r"^temporal_positional_encoding_projection_layer.*", + "no_memory_positional_encoding", + "no_object_pointer", + "occlusion_spatial_embedding_parameter", + ] + + def __init__(self, config: Sam3TrackerConfig): + # loading from a sam3_video config + if hasattr(config, "tracker_config") and config.tracker_config is not None: + if isinstance(config.tracker_config, dict): + config.tracker_config = Sam3TrackerConfig(**config.tracker_config) + config = config.tracker_config + Sam3TrackerPreTrainedModel.__init__(config) + self.shared_image_embedding = Sam3TrackerPositionalEmbedding(config.prompt_encoder_config) + self.vision_encoder = AutoModel.from_config(config.vision_config) + self.prompt_encoder = Sam3TrackerPromptEncoder(config.prompt_encoder_config) + # The module using it is not a PreTrainedModel subclass so we need this + config.mask_decoder_config._attn_implementation = config._attn_implementation + self.mask_decoder = Sam3TrackerMaskDecoder(config.mask_decoder_config) + + self.backbone_feature_sizes = config.vision_config.backbone_feature_sizes + # a single token to indicate no memory embedding from previous frames + self.hidden_dim = config.vision_config.fpn_hidden_size + self.no_memory_embedding = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + + self.post_init() + + +__all__ = [ + "Sam3TrackerConfig", + "Sam3TrackerPromptEncoderConfig", + "Sam3TrackerMaskDecoderConfig", + "Sam3TrackerProcessor", + "Sam3TrackerModel", + "Sam3TrackerPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/sam3_tracker/processing_sam3_tracker.py b/third_party/transformers/src/transformers/models/sam3_tracker/processing_sam3_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..d9beb6ee65da1f07747b8db9544e7838d08924d9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sam3_tracker/processing_sam3_tracker.py @@ -0,0 +1,510 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sam3_tracker/modular_sam3_tracker.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_sam3_tracker.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 copy import deepcopy + +import numpy as np +import torch + +from ...image_utils import ImageInput +from ...processing_utils import ProcessorMixin +from ...tokenization_utils_base import BatchEncoding +from ...utils import TensorType, auto_docstring +from ...utils.import_utils import requires + + +@requires(backends=("torch",)) +@auto_docstring +class Sam3TrackerProcessor(ProcessorMixin): + def __init__(self, image_processor, target_size: int | None = None, point_pad_value: int = -10, **kwargs): + r""" + target_size (`int`, *optional*): + The target size (in pixels) for normalizing input points and bounding boxes. If not provided, defaults + to the image processor's size configuration. All input coordinates (points and boxes) are normalized + to this size before being passed to the model. This ensures consistent coordinate representation + regardless of the original image dimensions. + point_pad_value (`int`, *optional*, defaults to -10): + The value used for padding input points when batching sequences of different lengths. This value is + used to mark padded positions and is preserved during coordinate normalization. + """ + super().__init__(image_processor, **kwargs) + self.point_pad_value = point_pad_value + self.target_size = target_size if target_size is not None else self.image_processor.size["height"] + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + segmentation_maps: ImageInput | None = None, + input_points: list[list[list[list[float]]]] | torch.Tensor | None = None, + input_labels: list[list[list[int]]] | torch.Tensor | None = None, + input_boxes: list[list[list[float]]] | torch.Tensor | None = None, + original_sizes: list[list[float]] | torch.Tensor | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchEncoding: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to process. + input_points (`list[list[list[list[float]]]]`, `torch.Tensor`, *optional*): + The points to add to the frame. + input_labels (`list[list[list[int]]]`, `torch.Tensor`, *optional*): + The labels for the points. + input_boxes (`list[list[list[float]]]`, `torch.Tensor`, *optional*): + The bounding boxes to add to the frame. + original_sizes (`list[list[float]]`, `torch.Tensor`, *optional*): + The original sizes of the images. + + Returns: + A [`BatchEncoding`] with the following fields: + - `pixel_values` (`torch.Tensor`): The processed image(s). + - `original_sizes` (`list[list[float]]`): The original sizes of the images. + - `labels` (`torch.Tensor`): The processed segmentation maps (if provided). + - `input_points` (`torch.Tensor`): The processed points. + - `input_labels` (`torch.Tensor`): The processed labels. + - `input_boxes` (`torch.Tensor`): The processed bounding boxes. + """ + if images is not None: + encoding_image_processor = self.image_processor( + images, + segmentation_maps=segmentation_maps, + return_tensors=return_tensors, + **kwargs, + ) + elif original_sizes is not None: + if isinstance(original_sizes, torch.Tensor): + original_sizes = original_sizes.cpu().tolist() + encoding_image_processor = BatchEncoding({"original_sizes": original_sizes}, tensor_type=return_tensors) + else: + raise ValueError("Either images or original_sizes must be provided") + + # pop arguments that are not used in the forward but used nevertheless + original_sizes = encoding_image_processor["original_sizes"] + # Check original_sizes is of length 1 or len(images) + if images is not None and len(original_sizes) != 1 and len(original_sizes) != len(images): + raise ValueError( + "original_sizes must be of length 1 or len(images). If you are passing a single image, you must pass a single original_size." + ) + + # Process input points, labels, and boxes if provided + if input_points is not None or input_labels is not None or input_boxes is not None: + # Validate and convert inputs to standardized format + processed_points = self._validate_single_input( + input_points, + expected_depth=4, + input_name="points", + expected_format="[image level, object level, point level, point coordinates]", + expected_coord_size=2, + ) + processed_labels = self._validate_single_input( + input_labels, + expected_depth=3, + input_name="labels", + expected_format="[image level, object level, point level]", + ) + processed_boxes = self._validate_single_input( + input_boxes, + expected_depth=3, + input_name="boxes", + expected_format="[image level, box level, box coordinates]", + expected_coord_size=4, + ) + + # Get padding requirements for all inputs + if processed_points is not None: + points_max_dims = self._get_nested_dimensions(processed_points)[:3] + if processed_labels is not None: + labels_max_dims = self._get_nested_dimensions(processed_labels)[:3] + if processed_boxes is not None: + boxes_max_dims = self._get_nested_dimensions(processed_boxes)[:2] + + # Ensure points and labels have consistent dimensions + if processed_points is not None and processed_labels is not None: + if points_max_dims != labels_max_dims: + raise ValueError( + "Input points and labels have inconsistent dimensions. Please ensure they have the same dimensions." + ) + + # Check that boxes don't need padding (model limitation) + if processed_boxes is not None and len(processed_boxes) >= 2: + if any(len(img_boxes) < boxes_max_dims[1] for img_boxes in processed_boxes): + raise ValueError( + "Input boxes have inconsistent dimensions that would require padding, " + "but boxes cannot be padded due to model limitations. " + "Please ensure all images have the same number of boxes." + ) + + # Pad and normalize all inputs to final tensor format + if processed_points is not None: + padded_points = self._pad_nested_list(processed_points, points_max_dims + [2]) + final_points = torch.tensor(padded_points, dtype=torch.float32) + self._normalize_tensor_coordinates(final_points, original_sizes, preserve_padding=True) + encoding_image_processor.update({"input_points": final_points}) + + if processed_labels is not None: + padded_labels = self._pad_nested_list(processed_labels, labels_max_dims) + final_labels = torch.tensor(padded_labels, dtype=torch.int64) + encoding_image_processor.update({"input_labels": final_labels}) + + if processed_boxes is not None: + final_boxes = torch.tensor(processed_boxes, dtype=torch.float32) + self._normalize_tensor_coordinates(final_boxes, original_sizes, is_bounding_box=True) + encoding_image_processor.update({"input_boxes": final_boxes}) + + return encoding_image_processor + + def _normalize_coordinates( + self, target_size: int, coords: "torch.Tensor", original_size, is_bounding_box=False + ) -> "torch.Tensor": + """ + Expects a numpy array of length 2 in the final dimension. Requires the original image size in (H, W) format. + + Args: + target_size (`int`): + The target size of the image. + coords (`torch.Tensor`): + The coordinates to be normalized. + original_size (`tuple`): + The original size of the image. + is_bounding_box (`bool`, *optional*, defaults to `False`): + Whether the coordinates are bounding boxes. + """ + old_h, old_w = original_size + new_h, new_w = target_size, target_size + coords = deepcopy(coords).float() + + if is_bounding_box: + coords = coords.reshape(-1, 2, 2) + coords[..., 0] = coords[..., 0] * (new_w / old_w) + coords[..., 1] = coords[..., 1] * (new_h / old_h) + + if is_bounding_box: + coords = coords.reshape(-1, 4) + + return coords + + def _convert_to_nested_list(self, data, expected_depth, current_depth=0): + """ + Recursively convert various input formats (tensors, numpy arrays, lists) to nested lists. + + Args: + data: Input data in any format + expected_depth: Expected nesting depth + current_depth: Current depth in recursion + + Returns: + Nested list representation of the data + """ + if data is None: + return None + + # Convert tensor/numpy to list if we're at a leaf level or if it's a multi-dimensional array + if isinstance(data, torch.Tensor): # PyTorch tensor + if current_depth == expected_depth - 2 or len(data.shape) <= 2: # At coordinate level or small tensor + return data.numpy().tolist() + else: + return [self._convert_to_nested_list(item, expected_depth, current_depth + 1) for item in data] + elif isinstance(data, np.ndarray): # NumPy array + if current_depth == expected_depth - 2 or len(data.shape) <= 2: # At coordinate level or small array + return data.tolist() + else: + return [self._convert_to_nested_list(item, expected_depth, current_depth + 1) for item in data] + elif isinstance(data, list): + if current_depth == expected_depth: + # We've reached the expected depth, return as is + return data + else: + # Continue recursion + return [self._convert_to_nested_list(item, expected_depth, current_depth + 1) for item in data] + elif isinstance(data, (int, float)): + return data + else: + raise TypeError(f"Unsupported data type: {type(data)}") + + def _get_nested_dimensions(self, nested_list, max_dims=None): + """ + Get the maximum dimensions at each level of nesting. + + Args: + nested_list (`list`): + Nested list structure. + max_dims (`list`, *optional*): + Current maximum dimensions (for recursion). + + Returns: + `list`: A list of maximum dimensions for each nesting level. + """ + if max_dims is None: + max_dims = [] + + if not isinstance(nested_list, list): + return max_dims + + if len(max_dims) == 0: + max_dims.append(len(nested_list)) + else: + max_dims[0] = max(max_dims[0], len(nested_list)) + + if len(nested_list) > 0: + for item in nested_list: + if isinstance(item, list): + sub_dims = self._get_nested_dimensions(item) + # Merge sub_dims into max_dims + for i, dim in enumerate(sub_dims): + if i + 1 >= len(max_dims): + max_dims.append(dim) + else: + max_dims[i + 1] = max(max_dims[i + 1], dim) + + return max_dims + + def _pad_nested_list(self, nested_list, target_dims, current_level=0, pad_value=None): + """ + Recursively pad a nested list to match target dimensions. + + Args: + nested_list (`list`): + Nested list to pad. + target_dims (`list`): + Target dimensions for each level. + current_level (`int`, *optional*, defaults to 0): + Current nesting level. + pad_value (`int`, *optional*): + Value to use for padding. + + Returns: + `list`: The padded nested list. + """ + if pad_value is None: + pad_value = self.point_pad_value + + if current_level >= len(target_dims): + return nested_list + + # Ensure we have a list + if not isinstance(nested_list, list): + nested_list = [nested_list] + + # Pad current level + current_size = len(nested_list) + target_size = target_dims[current_level] + + # Pad with appropriate values + if current_level == len(target_dims) - 1: + # At the coordinate level, pad with pad_value + nested_list.extend([pad_value] * (target_size - current_size)) + else: + # At higher levels, pad with nested structures + if current_size > 0: + # Create appropriately sized template + if current_level < len(target_dims) - 2: + # For non-coordinate levels, create empty nested structure + template_dims = target_dims[current_level + 1 :] + template = self._create_empty_nested_structure(template_dims, pad_value) + else: + # For coordinate level, create list of pad_values + template = [pad_value] * target_dims[current_level + 1] + + nested_list.extend([deepcopy(template) for _ in range(target_size - current_size)]) + else: + # Create from scratch + template_dims = target_dims[current_level + 1 :] + template = self._create_empty_nested_structure(template_dims, pad_value) + nested_list.extend([deepcopy(template) for _ in range(target_size)]) + + # Recursively pad sublists + if current_level < len(target_dims) - 1: + for i in range(len(nested_list)): + if isinstance(nested_list[i], list): + nested_list[i] = self._pad_nested_list(nested_list[i], target_dims, current_level + 1, pad_value) + + return nested_list + + def _create_empty_nested_structure(self, dims, pad_value): + """ + Create an empty nested structure with given dimensions filled with pad_value. + + Args: + dims (`list`): + The dimensions of the nested structure. + pad_value (`int`): + The value to fill the structure with. + """ + if len(dims) == 1: + return [pad_value] * dims[0] + else: + return [self._create_empty_nested_structure(dims[1:], pad_value) for _ in range(dims[0])] + + def _get_nesting_level(self, input_list): + """ + Get the nesting level of a list structure. + + Args: + input_list (`list`): + The list to get the nesting level of. + """ + if isinstance(input_list, list): + if len(input_list) == 0: + return 1 + return 1 + self._get_nesting_level(input_list[0]) + elif isinstance(input_list, (np.ndarray, torch.Tensor)): + # For arrays/tensors, the nesting level is the number of dimensions + return len(input_list.shape) + return 0 + + def _validate_single_input( + self, + data: torch.Tensor | np.ndarray | list, + expected_depth: int, + input_name: str, + expected_format: str, + expected_coord_size: int | None = None, + ) -> list: + """ + Validate a single input by ensuring proper nesting and raising an error if the input is not valid. + + Args: + data (`torch.Tensor`, `np.ndarray`, or `list`): + Input data to process. + expected_depth (`int`): + Expected nesting depth. + input_name (`str`): + Name of the input for error messages. + expected_format (`str`): + The expected format of the input. + expected_coord_size (`int`, *optional*): + Expected coordinate size (2 for points, 4 for boxes, None for labels). + . + """ + if data is None: + return None + + # Handle tensors and numpy arrays first + if isinstance(data, (torch.Tensor, np.ndarray)): + # For tensors/arrays, we can directly check the number of dimensions + if data.ndim != expected_depth: + raise ValueError( + f"Input {input_name} must be a tensor/array with {expected_depth} dimensions. The expected nesting format is {expected_format}. Got {data.ndim} dimensions." + ) + elif expected_coord_size is not None: + if data.shape[-1] != expected_coord_size: + raise ValueError( + f"Input {input_name} must be a tensor/array with {expected_coord_size} as the last dimension, got {data.shape[-1]}." + ) + return self._convert_to_nested_list(data, expected_depth) + + # Handle nested lists + if isinstance(data, list): + current_depth = self._get_nesting_level(data) + if current_depth != expected_depth: + raise ValueError( + f"Input {input_name} must be a nested list with {expected_depth} levels. The expected nesting format is {expected_format}. Got {current_depth} levels." + ) + return self._convert_to_nested_list(data, expected_depth) + + def _normalize_tensor_coordinates(self, tensor, original_sizes, is_bounding_box=False, preserve_padding=False): + """ + Helper method to normalize coordinates in a tensor across multiple images. + + Args: + tensor (`torch.Tensor`): + Input tensor with coordinates. + original_sizes (`list`): + Original image sizes. + is_bounding_box (`bool`, *optional*, defaults to `False`): + Whether coordinates are bounding boxes. + preserve_padding (`bool`, *optional*, defaults to `False`): + Whether to preserve padding values (for points). + """ + if preserve_padding: + # For points: avoid normalizing pad values + mask = tensor != self.point_pad_value + coord_mask = mask.all(dim=-1, keepdim=True) + + for img_idx in range(len(original_sizes)): + if img_idx < tensor.shape[0]: + original_size = original_sizes[img_idx] if img_idx < len(original_sizes) else original_sizes[0] + normalized_coords = self._normalize_coordinates( + self.target_size, tensor[img_idx], original_size, is_bounding_box=is_bounding_box + ) + + if preserve_padding: + # Only update non-padded values + img_mask = coord_mask[img_idx] + tensor[img_idx] = torch.where( + img_mask.expand_as(tensor[img_idx]), normalized_coords, tensor[img_idx] + ) + else: + tensor[img_idx] = normalized_coords + + def post_process_masks( + self, + masks, + original_sizes, + mask_threshold=0.0, + binarize=True, + max_hole_area=0.0, + max_sprinkle_area=0.0, + apply_non_overlapping_constraints=False, + **kwargs, + ): + """ + Remove padding and upscale masks to the original image size. + + Args: + masks (`Union[List[torch.Tensor], List[np.ndarray]]`): + Batched masks from the mask_decoder in (batch_size, num_channels, height, width) format. + original_sizes (`Union[torch.Tensor, List[Tuple[int,int]]]`): + The original sizes of each image before it was resized to the model's expected input shape, in (height, + width) format. + mask_threshold (`float`, *optional*, defaults to 0.0): + Threshold for binarization and post-processing operations. + binarize (`bool`, *optional*, defaults to `True`): + Whether to binarize the masks. + max_hole_area (`float`, *optional*, defaults to 0.0): + The maximum area of a hole to fill. + max_sprinkle_area (`float`, *optional*, defaults to 0.0): + The maximum area of a sprinkle to fill. + apply_non_overlapping_constraints (`bool`, *optional*, defaults to `False`): + Whether to apply non-overlapping constraints to the masks. + + Returns: + (`torch.Tensor`): Batched masks in batch_size, num_channels, height, width) format, where (height, width) + is given by original_size. + """ + return self.image_processor.post_process_masks( + masks, + original_sizes, + mask_threshold, + binarize, + max_hole_area, + max_sprinkle_area, + apply_non_overlapping_constraints, + **kwargs, + ) + + @property + def model_input_names(self): + image_processor_input_names = self.image_processor.model_input_names + return list(image_processor_input_names + ["original_sizes"]) + + +__all__ = ["Sam3TrackerProcessor"] diff --git a/third_party/transformers/src/transformers/models/sew/__init__.py b/third_party/transformers/src/transformers/models/sew/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0bd614f74a68062d842548bdcba4a65ce19fd62c --- /dev/null +++ b/third_party/transformers/src/transformers/models/sew/__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_sew import * + from .modeling_sew 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/sew/configuration_sew.py b/third_party/transformers/src/transformers/models/sew/configuration_sew.py new file mode 100644 index 0000000000000000000000000000000000000000..acb7be58a454ae8c36d820dfac5118951e1fa7b0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sew/configuration_sew.py @@ -0,0 +1,176 @@ +# Copyright 2021 ASAPP Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""SEW model configuration""" + +import functools +import operator + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="BAAI/seggpt-vit-large") +@strict +class SEWConfig(PreTrainedConfig): + r""" + squeeze_factor (`int`, *optional*, defaults to 2): + Sequence length downsampling factor after the encoder and upsampling factor after the transformer. + feat_proj_dropout (`float`, *optional*, defaults to 0.0): + The dropout probability for output of the feature encoder. + final_dropout (`float`, *optional*, defaults to 0.1): + The dropout probability for the final projection layer of [`SEWForCTC`]. + feat_extract_norm (`str`, *optional*, defaults to `"group"`): + The norm to be applied to 1D convolutional layers in feature encoder. One of `"group"` for group + normalization of only the first 1D convolutional layer or `"layer"` for layer normalization of all 1D + convolutional layers. + feat_extract_activation (`str, `optional`, defaults to `"gelu"`): + The non-linear activation function (function or string) in the 1D convolutional layers of the feature + extractor. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` are supported. + conv_dim (`tuple[int]` or `list[int]`, *optional*, defaults to `(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)`): + A tuple of integers defining the number of input and output channels of each 1D convolutional layer in the + feature encoder. The length of *conv_dim* defines the number of 1D convolutional layers. + conv_stride (`tuple[int]` or `list[int]`, *optional*, defaults to `(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)`): + A tuple of integers defining the stride of each 1D convolutional layer in the feature encoder. The length + of *conv_stride* defines the number of convolutional layers and has to match the length of *conv_dim*. + conv_kernel (`tuple[int]` or `list[int]`, *optional*, defaults to `(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)`): + A tuple of integers defining the kernel size of each 1D convolutional layer in the feature encoder. The + length of *conv_kernel* defines the number of convolutional layers and has to match the length of + *conv_dim*. + conv_bias (`bool`, *optional*, defaults to `False`): + Whether the 1D convolutional layers have a bias. + num_conv_pos_embeddings (`int`, *optional*, defaults to 128): + Number of convolutional positional embeddings. Defines the kernel size of 1D convolutional positional + embeddings layer. + num_conv_pos_embedding_groups (`int`, *optional*, defaults to 16): + Number of groups of 1D convolutional positional embeddings layer. + apply_spec_augment (`bool`, *optional*, defaults to `True`): + Whether to apply *SpecAugment* data augmentation to the outputs of the feature encoder. For reference see + [SpecAugment: A Simple Data Augmentation Method for Automatic Speech + Recognition](https://huggingface.co/papers/1904.08779). + mask_time_prob (`float`, *optional*, defaults to 0.05): + Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking + procedure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If + reasoning from the probability of each feature vector to be chosen as the start of the vector span to be + masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the + actual percentage of masked vectors. This is only relevant if `apply_spec_augment is True`. + mask_time_length (`int`, *optional*, defaults to 10): + Length of vector span along the time axis. + mask_time_min_masks (`int`, *optional*, defaults to 2),: + The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step, + irrespectively of `mask_feature_prob`. Only relevant if ''mask_time_prob*len(time_axis)/mask_time_length < + mask_time_min_masks'' + mask_feature_prob (`float`, *optional*, defaults to 0.0): + Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The + masking procedure generates ''mask_feature_prob*len(feature_axis)/mask_time_length'' independent masks over + the axis. If reasoning from the probability of each feature vector to be chosen as the start of the vector + span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap + may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is + True`. + mask_feature_length (`int`, *optional*, defaults to 10): + Length of vector span along the feature axis. + mask_feature_min_masks (`int`, *optional*, defaults to 0): + The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time + step, irrespectively of `mask_feature_prob`. Only relevant if + ''mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks'' + ctc_zero_infinity (`bool`, *optional*, defaults to `False`): + Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly + occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance + of [`SEWForCTC`]. + use_weighted_layer_sum (`bool`, *optional*, defaults to `False`): + Whether to use a weighted average of layer outputs with learned weights. Only relevant when using an + instance of [`Wav2Vec2ForSequenceClassification`]. + classifier_proj_size (`int`, *optional*, defaults to 256): + Dimensionality of the projection before token mean-pooling for classification. + + Example: + + ```python + >>> from transformers import SEWConfig, SEWModel + + >>> # Initializing a SEW asapp/sew-tiny-100k style configuration + >>> configuration = SEWConfig() + + >>> # Initializing a model (with random weights) from the asapp/sew-tiny-100k style configuration + >>> model = SEWModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "sew" + + vocab_size: int = 32 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + squeeze_factor: int = 2 + hidden_act: str = "gelu" + hidden_dropout: float | int = 0.1 + activation_dropout: float | int = 0.1 + attention_dropout: float | int = 0.1 + feat_proj_dropout: float | int = 0.0 + final_dropout: float | int = 0.1 + layerdrop: float | int = 0.1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + feat_extract_norm: str = "group" + feat_extract_activation: str = "gelu" + conv_dim: list[int] | tuple[int, ...] = (64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512) + conv_stride: list[int] | tuple[int, ...] = (5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1) + conv_kernel: list[int] | tuple[int, ...] = (10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1) + conv_bias: bool = False + num_conv_pos_embeddings: int = 128 + num_conv_pos_embedding_groups: int = 16 + apply_spec_augment: bool = True + mask_time_prob: float | int = 0.05 + mask_time_length: int = 10 + mask_time_min_masks: int = 2 + mask_feature_prob: float | int = 0.0 + mask_feature_length: int = 10 + mask_feature_min_masks: int = 0 + ctc_loss_reduction: str = "mean" + ctc_zero_infinity: bool = False + use_weighted_layer_sum: bool = False + classifier_proj_size: int = 256 + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + + def __post_init__(self, **kwargs): + self.num_feat_extract_layers = len(self.conv_dim) + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if ( + (len(self.conv_stride) != self.num_feat_extract_layers) + or (len(self.conv_kernel) != self.num_feat_extract_layers) + or (len(self.conv_dim) != self.num_feat_extract_layers) + ): + raise ValueError( + "Configuration for convolutional layers is incorrect. " + "It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`, " + f"but is `len(config.conv_dim) = {len(self.conv_dim)}`, `len(config.conv_stride) " + f"= {len(self.conv_stride)}`, `len(config.conv_kernel) = {len(self.conv_kernel)}`." + ) + + @property + def inputs_to_logits_ratio(self): + return functools.reduce(operator.mul, self.conv_stride, 1) + + +__all__ = ["SEWConfig"] diff --git a/third_party/transformers/src/transformers/models/sew/convert_sew_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/sew/convert_sew_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..93c9888d61fbae05ee1113cfcbc7478f87ef04c1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sew/convert_sew_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,304 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert SEW checkpoint.""" + +import argparse +import json +import os + +import fairseq +import torch +from fairseq.data import Dictionary + +# Register SEW's fairseq modules +from sew_asapp import tasks # noqa: F401 + +from transformers import ( + SEWConfig, + SEWForCTC, + SEWModel, + Wav2Vec2CTCTokenizer, + Wav2Vec2FeatureExtractor, + Wav2Vec2Processor, + logging, +) + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + +MAPPING = { + "post_extract_proj": "feature_projection", + "encoder.pos_conv.0": "encoder.pos_conv_embed.conv", + "self_attn.k_proj": "encoder.layers.*.attention.k_proj", + "self_attn.v_proj": "encoder.layers.*.attention.v_proj", + "self_attn.q_proj": "encoder.layers.*.attention.q_proj", + "self_attn.out_proj": "encoder.layers.*.attention.out_proj", + "self_attn_layer_norm": "encoder.layers.*.layer_norm", + "fc1": "encoder.layers.*.feed_forward.intermediate_dense", + "fc2": "encoder.layers.*.feed_forward.output_dense", + "final_layer_norm": "encoder.layers.*.final_layer_norm", + "encoder.upsample.0": "encoder.upsample.projection", + "encoder.layer_norm": "encoder.layer_norm", + "w2v_model.layer_norm": "layer_norm", + "w2v_encoder.proj": "lm_head", + "mask_emb": "masked_spec_embed", +} + + +def set_recursively(hf_pointer, key, value, full_name, weight_type): + for attribute in key.split("."): + hf_pointer = getattr(hf_pointer, attribute) + + if weight_type is not None: + hf_shape = getattr(hf_pointer, weight_type).shape + else: + hf_shape = hf_pointer.shape + + assert hf_shape == value.shape, ( + f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be" + f" {value.shape} for {full_name}" + ) + + if weight_type == "weight": + hf_pointer.weight.data = value + elif weight_type == "weight_g": + hf_pointer.weight_g.data = value + elif weight_type == "weight_v": + hf_pointer.weight_v.data = value + elif weight_type == "bias": + hf_pointer.bias.data = value + else: + hf_pointer.data = value + + logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.") + + +def recursively_load_weights(fairseq_model, hf_model, is_finetuned): + unused_weights = [] + fairseq_dict = fairseq_model.state_dict() + + feature_extractor = hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor + + for name, value in fairseq_dict.items(): + is_used = False + if "conv_layers" in name: + load_conv_layer( + name, + value, + feature_extractor, + unused_weights, + hf_model.config.feat_extract_norm == "group", + ) + is_used = True + else: + for key, mapped_key in MAPPING.items(): + mapped_key = "sew." + mapped_key if (is_finetuned and mapped_key != "lm_head") else mapped_key + + if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]: + is_used = True + if "*" in mapped_key: + layer_index = name.split(key)[0].split(".")[-2] + mapped_key = mapped_key.replace("*", layer_index) + if "weight_g" in name: + weight_type = "weight_g" + elif "weight_v" in name: + weight_type = "weight_v" + elif "weight" in name: + weight_type = "weight" + elif "bias" in name: + weight_type = "bias" + else: + weight_type = None + set_recursively(hf_model, mapped_key, value, name, weight_type) + continue + if not is_used: + unused_weights.append(name) + + logger.warning(f"Unused weights: {unused_weights}") + + +def load_conv_layer(full_name, value, feature_extractor, unused_weights, use_group_norm): + name = full_name.split("conv_layers.")[-1] + items = name.split(".") + layer_id = int(items[0]) + type_id = int(items[1]) + + if type_id == 0: + if "bias" in name: + assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, ( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].conv.bias.data = value + logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.") + elif "weight" in name: + assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, ( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].conv.weight.data = value + logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.") + elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm): + if "bias" in name: + assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, ( + f"{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was" + " found." + ) + feature_extractor.conv_layers[layer_id].layer_norm.bias.data = value + logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.") + elif "weight" in name: + assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, ( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].layer_norm.weight.data = value + logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.") + else: + unused_weights.append(full_name) + + +def convert_config(model, is_finetuned): + config = SEWConfig() + if is_finetuned: + fs_config = model.w2v_encoder.w2v_model.cfg + else: + fs_config = model.cfg + + config.conv_bias = fs_config.conv_bias + conv_layers = eval(fs_config.conv_feature_layers) + config.conv_dim = [x[0] for x in conv_layers] + config.conv_kernel = [x[1] for x in conv_layers] + config.conv_stride = [x[2] for x in conv_layers] + config.feat_extract_activation = "gelu" + config.feat_extract_norm = "layer" if fs_config.extractor_mode == "layer_norm" else "group" + config.final_dropout = 0.0 + config.hidden_act = fs_config.activation_fn.name + config.hidden_size = fs_config.encoder_embed_dim + config.initializer_range = 0.02 + config.intermediate_size = fs_config.encoder_ffn_embed_dim + config.layer_norm_eps = 1e-5 + config.layerdrop = fs_config.encoder_layerdrop + config.num_attention_heads = fs_config.encoder_attention_heads + config.num_conv_pos_embedding_groups = fs_config.conv_pos_groups + config.num_conv_pos_embeddings = fs_config.conv_pos + config.num_feat_extract_layers = len(conv_layers) + config.num_hidden_layers = fs_config.encoder_layers + config.squeeze_factor = fs_config.squeeze_factor + + # take care of any params that are overridden by the Wav2VecCtc model + if is_finetuned: + fs_config = model.cfg + config.final_dropout = fs_config.final_dropout + config.layerdrop = fs_config.layerdrop + config.activation_dropout = fs_config.activation_dropout + config.apply_spec_augment = fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0 + config.attention_dropout = fs_config.attention_dropout + config.feat_proj_dropout = fs_config.dropout_input + config.hidden_dropout = fs_config.dropout + config.mask_feature_length = fs_config.mask_channel_length + config.mask_feature_prob = fs_config.mask_channel_prob + config.mask_time_length = fs_config.mask_length + config.mask_time_prob = fs_config.mask_prob + + config.feature_extractor_type = "Wav2Vec2FeatureExtractor" + config.tokenizer_class = "Wav2Vec2CTCTokenizer" + + return config + + +@torch.no_grad() +def convert_sew_checkpoint( + checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True +): + """ + Copy/paste/tweak model's weights to transformers design. + """ + + if is_finetuned: + model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task( + [checkpoint_path], arg_overrides={"data": "/".join(dict_path.split("/")[:-1])} + ) + else: + model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]) + + if config_path is not None: + config = SEWConfig.from_pretrained(config_path) + else: + config = convert_config(model[0], is_finetuned) + model = model[0].eval() + + return_attention_mask = config.feat_extract_norm == "layer" + feature_extractor = Wav2Vec2FeatureExtractor( + feature_size=1, + sampling_rate=16000, + padding_value=0, + do_normalize=True, + return_attention_mask=return_attention_mask, + ) + + if is_finetuned: + if dict_path: + target_dict = Dictionary.load(dict_path) + + # important change bos & pad token id since CTC symbol is and + # not as in fairseq + target_dict.indices[target_dict.bos_word] = target_dict.pad_index + target_dict.indices[target_dict.pad_word] = target_dict.bos_index + config.bos_token_id = target_dict.pad_index + config.pad_token_id = target_dict.bos_index + config.eos_token_id = target_dict.eos_index + config.vocab_size = len(target_dict.symbols) + vocab_path = os.path.join(pytorch_dump_folder_path, "vocab.json") + if not os.path.isdir(pytorch_dump_folder_path): + logger.error(f"--pytorch_dump_folder_path ({pytorch_dump_folder_path}) should be a directory") + return + os.makedirs(pytorch_dump_folder_path, exist_ok=True) + with open(vocab_path, "w", encoding="utf-8") as vocab_handle: + json.dump(target_dict.indices, vocab_handle) + tokenizer = Wav2Vec2CTCTokenizer( + vocab_path, + unk_token=target_dict.unk_word, + pad_token=target_dict.pad_word, + bos_token=target_dict.bos_word, + eos_token=target_dict.eos_word, + word_delimiter_token="|", + do_lower_case=False, + ) + processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer) + processor.save_pretrained(pytorch_dump_folder_path) + + hf_model = SEWForCTC(config) + else: + hf_model = SEWModel(config) + feature_extractor.save_pretrained(pytorch_dump_folder_path) + + recursively_load_weights(model, hf_model, is_finetuned) + + hf_model.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") + parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint") + parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model") + parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") + parser.add_argument( + "--is_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not" + ) + args = parser.parse_args() + convert_sew_checkpoint( + args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned + ) diff --git a/third_party/transformers/src/transformers/models/sew/modeling_sew.py b/third_party/transformers/src/transformers/models/sew/modeling_sew.py new file mode 100644 index 0000000000000000000000000000000000000000..79070d634608043c0d557cb2409281be158cd8a9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/sew/modeling_sew.py @@ -0,0 +1,1081 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sew/modular_sew.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_sew.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2021 ASAPP Inc. and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Callable + +import numpy as np +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations.deepspeed import is_deepspeed_zero3_enabled +from ...integrations.fsdp import is_fsdp_managed_module +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, CausalLMOutput, SequenceClassifierOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel, get_torch_context_manager_or_global_device +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import is_flash_attention_requested +from .configuration_sew import SEWConfig + + +logger = logging.get_logger(__name__) + + +class SEWNoLayerNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.activation(hidden_states) + return hidden_states + + +class SEWLayerNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + + hidden_states = hidden_states.transpose(-2, -1) + hidden_states = self.layer_norm(hidden_states) + hidden_states = hidden_states.transpose(-2, -1) + + hidden_states = self.activation(hidden_states) + return hidden_states + + +class SEWGroupNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.activation = ACT2FN[config.feat_extract_activation] + + self.layer_norm = nn.GroupNorm(num_groups=self.out_conv_dim, num_channels=self.out_conv_dim, affine=True) + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.activation(hidden_states) + return hidden_states + + +class SEWPositionalConvEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + self.conv = nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size=config.num_conv_pos_embeddings, + padding=config.num_conv_pos_embeddings // 2, + groups=config.num_conv_pos_embedding_groups, + stride=config.squeeze_factor, + ) + + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0): + self.conv = weight_norm(self.conv, name="weight", dim=2) + if hasattr(self.conv, "parametrizations"): + weight_g = self.conv.parametrizations.weight.original0 + weight_v = self.conv.parametrizations.weight.original1 + else: + weight_g = self.conv.weight_g + weight_v = self.conv.weight_v + deepspeed.zero.register_external_parameter(self, weight_v) + deepspeed.zero.register_external_parameter(self, weight_g) + else: + self.conv = weight_norm(self.conv, name="weight", dim=2) + + self.padding = SEWSamePadLayer(config.num_conv_pos_embeddings) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.padding(hidden_states) + hidden_states = self.activation(hidden_states) + + return hidden_states + + +class SEWSamePadLayer(nn.Module): + def __init__(self, num_conv_pos_embeddings): + super().__init__() + self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 0 + + def forward(self, hidden_states): + if self.num_pad_remove > 0: + hidden_states = hidden_states[:, :, : -self.num_pad_remove] + return hidden_states + + +class SEWUpsampling(nn.Module): + def __init__(self, config): + super().__init__() + self.projection = nn.Linear(config.hidden_size, config.hidden_size * config.squeeze_factor) + self.activation = ACT2FN[config.feat_extract_activation] + self.squeeze_factor = config.squeeze_factor + + def forward(self, hidden_states): + hidden_states = self.projection(hidden_states) + hidden_states = self.activation(hidden_states) + + if self.squeeze_factor > 1: + # transform embedding channels to sequence length + bsz, src_len, src_embed_dim = hidden_states.size() + tgt_len = src_len * self.squeeze_factor + tgt_embed_dim = src_embed_dim // self.squeeze_factor + hidden_states = hidden_states.reshape(bsz, src_len, self.squeeze_factor, tgt_embed_dim) + hidden_states = hidden_states.reshape(bsz, tgt_len, tgt_embed_dim) + + return hidden_states + + +class SEWFeatureEncoder(nn.Module): + """Construct the features from raw audio waveform""" + + def __init__(self, config): + super().__init__() + + if config.feat_extract_norm == "group": + conv_layers = [SEWGroupNormConvLayer(config, layer_id=0)] + [ + SEWNoLayerNormConvLayer(config, layer_id=i + 1) for i in range(config.num_feat_extract_layers - 1) + ] + elif config.feat_extract_norm == "layer": + conv_layers = [SEWLayerNormConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers)] + else: + raise ValueError( + f"`config.feat_extract_norm` is {config.feat_extract_norm}, but has to be one of ['group', 'layer']" + ) + self.conv_layers = nn.ModuleList(conv_layers) + self.gradient_checkpointing = False + self._requires_grad = True + + def _freeze_parameters(self): + for param in self.parameters(): + param.requires_grad = False + self._requires_grad = False + + def forward(self, input_values): + hidden_states = input_values[:, None] + + # make sure hidden_states require grad for gradient_checkpointing + if self._requires_grad and self.training: + hidden_states.requires_grad = True + + for conv_layer in self.conv_layers: + hidden_states = conv_layer(hidden_states) + + return hidden_states + + +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 SEWAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + is_causal: bool = False, + config: SEWConfig | None = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + + self.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, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2) + + current_states = key_value_states if is_cross_attention else hidden_states + key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v_proj(current_states).view(*kv_input_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.dropout, + scaling=self.scaling, + output_attentions=output_attentions, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights, None + + +class SEWFeedForward(nn.Module): + def __init__(self, config): + super().__init__() + self.intermediate_dropout = nn.Dropout(config.activation_dropout) + + self.intermediate_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 + + self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.output_dropout = nn.Dropout(config.hidden_dropout) + + def forward(self, hidden_states): + hidden_states = self.intermediate_dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + hidden_states = self.intermediate_dropout(hidden_states) + + hidden_states = self.output_dense(hidden_states) + hidden_states = self.output_dropout(hidden_states) + return hidden_states + + +class SEWEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.attention = SEWAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + dropout=config.attention_dropout, + is_decoder=False, + config=config, + ) + + self.dropout = nn.Dropout(config.hidden_dropout) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.feed_forward = SEWFeedForward(config) + self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + attn_residual = hidden_states + hidden_states, attn_weights, _ = self.attention( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = self.dropout(hidden_states) + hidden_states = attn_residual + hidden_states + + hidden_states = self.layer_norm(hidden_states) + hidden_states = hidden_states + self.feed_forward(hidden_states) + hidden_states = self.final_layer_norm(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class SEWEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pos_conv_embed = SEWPositionalConvEmbedding(config) + self.pool = nn.AvgPool1d(config.squeeze_factor, config.squeeze_factor) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout) + self.layers = nn.ModuleList([SEWEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.upsample = SEWUpsampling(config) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + 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 + + if attention_mask is not None: + expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + if is_flash_attention_requested(self.config): + # make sure padded tokens output 0 + hidden_states[~expand_attention_mask] = 0.0 + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # make sure padded tokens output 0 + hidden_states[~expand_attention_mask] = 0.0 + input_lengths = (attention_mask.long()).sum(-1) + # apply pooling formula to get real output_lengths + output_lengths = input_lengths // self.config.squeeze_factor + max_encoder_length = hidden_states.shape[1] // self.config.squeeze_factor + attention_ids = ( + torch.arange(0, max_encoder_length, device=output_lengths.device) + .view(1, -1) + .expand(output_lengths.shape[0], -1) + ) + attention_mask = (attention_ids < output_lengths.view(-1, 1)).long() + + # extend attention_mask + attention_mask = 1.0 - attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) + attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min + attention_mask = attention_mask.expand( + attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1] + ) + + n_input_timesteps = hidden_states.shape[1] + + hidden_states = hidden_states.transpose(1, 2) + position_embeddings = self.pos_conv_embed(hidden_states) + pooled_hidden_states = self.pool(hidden_states) + min_length = min(position_embeddings.size(-1), pooled_hidden_states.size(-1)) + hidden_states = pooled_hidden_states[..., :min_length] + position_embeddings[..., :min_length] + hidden_states = hidden_states.transpose(1, 2) + + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) + + for layer in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + dropout_probability = torch.rand([]) + + skip_the_layer = self.training and dropout_probability < self.config.layerdrop + if not skip_the_layer or synced_gpus: + # under fsdp or deepspeed zero3 all gpus must run in sync + layer_outputs = layer( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = layer_outputs[0] + + if skip_the_layer: + layer_outputs = (None, None) + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states = self.upsample(hidden_states) + if hidden_states.shape[1] < n_input_timesteps: + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, n_input_timesteps - hidden_states.shape[1])) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring +class SEWPreTrainedModel(PreTrainedModel): + config: SEWConfig + base_model_prefix = "sew" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = False # needs a proper look into the mask creation + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, SEWPositionalConvEmbedding): + init.normal_( + module.conv.weight, + mean=0, + std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)), + ) + init.constant_(module.conv.bias, 0) + elif isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv1d): + if is_deepspeed_zero3_enabled(): + import deepspeed + + if hasattr(module, "weight_v") and hasattr(module, "weight_g"): + with deepspeed.zero.GatheredParameters([module.weight_v, module.weight_g], modifier_rank=0): + init.kaiming_normal_(module.weight) + else: + with deepspeed.zero.GatheredParameters(module.weight, modifier_rank=0): + init.kaiming_normal_(module.weight) + else: + init.kaiming_normal_(module.weight) + + if isinstance(module, (nn.Linear, nn.Conv1d)) and module.bias is not None: + init.zeros_(module.bias) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor | int): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + # 1D convolutional layer output length formula taken + # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html + return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1 + + for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): + input_lengths = _conv_out_length(input_lengths, kernel_size, stride) + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor): + output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long) + batch_size = attention_mask.shape[0] + + attention_mask = torch.zeros( + (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device + ) + # these two operations makes sure that all values before the output lengths idxs are attended to + attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1 + attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool() + return attention_mask + + +def _compute_mask_indices( + shape: tuple[int, int], + mask_prob: float, + mask_length: int, + attention_mask: torch.LongTensor | None = None, + min_masks: int = 0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for + ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on + CPU as part of the preprocessing during training. + + Args: + shape: The shape for which to compute masks. This should be of a tuple of size 2 where + the first element is the batch size and the second element is the length of the axis to span. + mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of + independently generated mask spans of length `mask_length` is computed by + `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the + actual percentage will be smaller. + mask_length: size of the mask + min_masks: minimum number of masked spans + attention_mask: A (right-padded) attention mask which independently shortens the feature axis of + each batch dimension. + """ + batch_size, sequence_length = shape + + if mask_length < 1: + raise ValueError("`mask_length` has to be bigger than 0.") + + if mask_length > sequence_length: + raise ValueError( + f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}" + f" and `sequence_length`: {sequence_length}`" + ) + + # epsilon is used for probabilistic rounding + epsilon = np.random.rand(1).item() + + def compute_num_masked_span(input_length): + """Given input length, compute how many spans should be masked""" + num_masked_span = int(mask_prob * input_length / mask_length + epsilon) + num_masked_span = max(num_masked_span, min_masks) + + # make sure num masked span <= sequence_length + if num_masked_span * mask_length > sequence_length: + num_masked_span = sequence_length // mask_length + + # make sure num_masked span is also <= input_length - (mask_length - 1) + if input_length - (mask_length - 1) < num_masked_span: + num_masked_span = max(input_length - (mask_length - 1), 0) + + return num_masked_span + + # compute number of masked spans in batch + input_lengths = ( + attention_mask.detach().sum(-1).tolist() + if attention_mask is not None + else [sequence_length for _ in range(batch_size)] + ) + + # SpecAugment mask to fill + spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool) + spec_aug_mask_idxs = [] + + max_num_masked_span = compute_num_masked_span(sequence_length) + + if max_num_masked_span == 0: + return spec_aug_mask + + for input_length in input_lengths: + # compute num of masked spans for this input + num_masked_span = compute_num_masked_span(input_length) + + # get random indices to mask + spec_aug_mask_idx = np.random.choice( + np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False + ) + + # pick first sampled index that will serve as a dummy index to pad vector + # to ensure same dimension for all batches due to probabilistic rounding + # Picking first sample just pads those vectors twice. + if len(spec_aug_mask_idx) == 0: + # this case can only happen if `input_length` is strictly smaller then + # `sequence_length` in which case the last token has to be a padding + # token which we can use as a dummy mask id + dummy_mask_idx = sequence_length - 1 + else: + dummy_mask_idx = spec_aug_mask_idx[0] + + spec_aug_mask_idx = np.concatenate( + [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx] + ) + spec_aug_mask_idxs.append(spec_aug_mask_idx) + + spec_aug_mask_idxs = np.array(spec_aug_mask_idxs) + + # expand masked indices to masked spans + spec_aug_mask_idxs = np.broadcast_to( + spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length) + ) + spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length) + + # add offset to the starting indexes so that indexes now create a span + offsets = np.arange(mask_length)[None, None, :] + offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape( + batch_size, max_num_masked_span * mask_length + ) + spec_aug_mask_idxs = spec_aug_mask_idxs + offsets + + # ensure that we cannot have indices larger than sequence_length + if spec_aug_mask_idxs.max() > sequence_length - 1: + spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1 + + # scatter indices to mask + np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1) + + return spec_aug_mask + + +@auto_docstring +class SEWModel(SEWPreTrainedModel): + def __init__(self, config: SEWConfig): + super().__init__(config) + self.config = config + self.feature_extractor = SEWFeatureEncoder(config) + self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps) + + self.project_features = config.conv_dim[-1] != config.hidden_size + if self.project_features: + self.feature_projection = nn.Linear(config.conv_dim[-1], config.hidden_size) + self.feature_dropout = nn.Dropout(config.feat_proj_dropout) + + if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0: + self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_()) + + self.encoder = SEWEncoder(config) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model._mask_hidden_states + def _mask_hidden_states( + self, + hidden_states: torch.FloatTensor, + mask_time_indices: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + ): + """ + Masks extracted features along time axis and/or along feature axis according to + [SpecAugment](https://huggingface.co/papers/1904.08779). + """ + + # `config.apply_spec_augment` can set masking to False + if not getattr(self.config, "apply_spec_augment", True): + return hidden_states + + # generate indices & apply SpecAugment along time axis + batch_size, sequence_length, hidden_size = hidden_states.size() + + if mask_time_indices is not None: + # apply SpecAugment along time axis with given mask_time_indices + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + elif self.config.mask_time_prob > 0 and self.training: + mask_time_indices = _compute_mask_indices( + (batch_size, sequence_length), + mask_prob=self.config.mask_time_prob, + mask_length=self.config.mask_time_length, + attention_mask=attention_mask, + min_masks=self.config.mask_time_min_masks, + ) + mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool) + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + + if self.config.mask_feature_prob > 0 and self.training: + # generate indices & apply SpecAugment along feature axis + mask_feature_indices = _compute_mask_indices( + (batch_size, hidden_size), + mask_prob=self.config.mask_feature_prob, + mask_length=self.config.mask_feature_length, + min_masks=self.config.mask_feature_min_masks, + ) + mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool) + mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1) + hidden_states[mask_feature_indices] = 0 + + return hidden_states + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + mask_time_indices: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + r""" + mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict + masked extracted features in *config.proj_codevector_dim* space. + """ + 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 + + extract_features = self.feature_extractor(input_values) + extract_features = extract_features.transpose(1, 2) + extract_features = self.layer_norm(extract_features) + + if self.project_features: + extract_features = self.feature_projection(extract_features) + hidden_states = self.feature_dropout(extract_features) + + if attention_mask is not None: + # compute reduced attention_mask corresponding to feature vectors + attention_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask) + + hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices) + + encoder_outputs = self.encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + + if not return_dict: + return (hidden_states,) + encoder_outputs[1:] + + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +_HIDDEN_STATES_START_POSITION = 1 + + +@auto_docstring( + custom_intro=""" + SEW Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). + """ +) +class SEWForCTC(SEWPreTrainedModel): + def __init__(self, config, target_lang: str | None = None): + r""" + target_lang (`str`, *optional*): + Language id of adapter weights. Adapter weights are stored in the format adapter..safetensors or + adapter..bin. Only relevant when using an instance of [`SEWForCTC`] with adapters. Uses 'eng' by + default. + """ + super().__init__(config) + + self.sew = SEWModel(config) + self.dropout = nn.Dropout(config.final_dropout) + + self.target_lang = target_lang + + if config.vocab_size is None: + raise ValueError( + f"You are trying to instantiate {self.__class__} with a configuration that " + "does not define the vocabulary size of the language model head. Please " + "instantiate the model as follows: `SEWForCTC.from_pretrained(..., vocab_size=vocab_size)`. " + "or define `vocab_size` of your model's configuration." + ) + output_hidden_size = ( + config.output_hidden_size if hasattr(config, "add_adapter") and config.add_adapter else config.hidden_size + ) + self.lm_head = nn.Linear(output_hidden_size, config.vocab_size) + + # Initialize weights and apply final processing + self.post_init() + + def tie_weights(self, **kwargs): + """ + This method overwrites [`~PreTrainedModel.tie_weights`] so that adapter weights can be correctly loaded when + passing `target_lang=...` to `from_pretrained(...)`. + + This method is **not** supposed to be called by the user and is prone to be changed in the future. + """ + + if get_torch_context_manager_or_global_device() == torch.device("meta"): + return + + # Note that `tie_weights` is usually used to tie input and output embedding weights. The method is re-purposed to + # correctly load adapter layers for SEW so that we do not have to introduce a new API to + # [`PreTrainedModel`]. While slightly hacky, SEW never has to tie input and output embeddings, so that it is + # ok to repurpose this function here. + target_lang = self.target_lang + + if target_lang is not None and getattr(self.config, "adapter_attn_dim", None) is None: + raise ValueError(f"Cannot pass `target_lang`: {target_lang} if `config.adapter_attn_dim` is not defined.") + elif target_lang is None and getattr(self.config, "adapter_attn_dim", None) is not None: + logger.info("By default `target_lang` is set to 'eng'.") + elif target_lang is not None: + self.load_adapter(target_lang, force_load=True) + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.sew.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.sew.parameters(): + param.requires_grad = False + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + **kwargs, + ) -> tuple | CausalLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*): + Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to + the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. + All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., + config.vocab_size - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if labels is not None and labels.max() >= self.config.vocab_size: + raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}") + + outputs = self.sew( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.dropout(hidden_states) + + logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # retrieve loss input_lengths from attention_mask + attention_mask = ( + attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long) + ) + input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long) + + # assuming that padded tokens are filled with -100 + # when not being attended to + labels_mask = labels >= 0 + target_lengths = labels_mask.sum(-1) + flattened_targets = labels.masked_select(labels_mask) + + # ctc_loss doesn't support fp16 + log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1) + + with torch.backends.cudnn.flags(enabled=False): + loss = nn.functional.ctc_loss( + log_probs, + flattened_targets, + input_lengths, + target_lengths, + blank=self.config.pad_token_id, + reduction=self.config.ctc_loss_reduction, + zero_infinity=self.config.ctc_zero_infinity, + ) + + if not return_dict: + output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutput( + loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions + ) + + +@auto_docstring( + custom_intro=""" + SEW Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like + SUPERB Keyword Spotting. + """ +) +class SEWForSequenceClassification(SEWPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + if hasattr(config, "add_adapter") and config.add_adapter: + raise ValueError( + "Sequence classification does not support the use of SEW adapters (config.add_adapter=True)" + ) + self.sew = SEWModel(config) + num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings + if config.use_weighted_layer_sum: + self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers) + self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size) + self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.sew.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.sew.parameters(): + param.requires_grad = False + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutput: + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file + into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library + (`pip install torchcodec`) or the soundfile library (`pip install soundfile`). + To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion + into a tensor of type `torch.FloatTensor`. See [`SEWProcessor.__call__`] for details. + 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 + output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states + + outputs = self.sew( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.config.use_weighted_layer_sum: + hidden_states = outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = outputs[0] + + hidden_states = self.projector(hidden_states) + if attention_mask is None: + pooled_output = hidden_states.mean(dim=1) + else: + padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask) + expand_padding_mask = padding_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + hidden_states[~expand_padding_mask] = 0.0 + pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1) + + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = ["SEWForCTC", "SEWForSequenceClassification", "SEWModel", "SEWPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/sew/modular_sew.py b/third_party/transformers/src/transformers/models/sew/modular_sew.py new file mode 100644 index 0000000000000000000000000000000000000000..312419793a3461f00dc5c9d16d83b106294f403d --- /dev/null +++ b/third_party/transformers/src/transformers/models/sew/modular_sew.py @@ -0,0 +1,449 @@ +# Copyright 2021 ASAPP Inc. and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch SEW model.""" + +import math + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations.deepspeed import is_deepspeed_zero3_enabled +from ...integrations.fsdp import is_fsdp_managed_module +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring +from ...utils.generic import is_flash_attention_requested +from ..wav2vec2.modeling_wav2vec2 import ( + Wav2Vec2Attention, + Wav2Vec2EncoderLayer, + Wav2Vec2FeatureEncoder, + Wav2Vec2FeedForward, + Wav2Vec2ForCTC, + Wav2Vec2ForSequenceClassification, + Wav2Vec2GroupNormConvLayer, + Wav2Vec2LayerNormConvLayer, + Wav2Vec2NoLayerNormConvLayer, + Wav2Vec2SamePadLayer, + _compute_mask_indices, +) +from .configuration_sew import SEWConfig + + +_HIDDEN_STATES_START_POSITION = 1 + + +class SEWNoLayerNormConvLayer(Wav2Vec2NoLayerNormConvLayer): + pass + + +class SEWLayerNormConvLayer(Wav2Vec2LayerNormConvLayer): + pass + + +class SEWGroupNormConvLayer(Wav2Vec2GroupNormConvLayer): + pass + + +class SEWPositionalConvEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + self.conv = nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size=config.num_conv_pos_embeddings, + padding=config.num_conv_pos_embeddings // 2, + groups=config.num_conv_pos_embedding_groups, + stride=config.squeeze_factor, + ) + + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0): + self.conv = weight_norm(self.conv, name="weight", dim=2) + if hasattr(self.conv, "parametrizations"): + weight_g = self.conv.parametrizations.weight.original0 + weight_v = self.conv.parametrizations.weight.original1 + else: + weight_g = self.conv.weight_g + weight_v = self.conv.weight_v + deepspeed.zero.register_external_parameter(self, weight_v) + deepspeed.zero.register_external_parameter(self, weight_g) + else: + self.conv = weight_norm(self.conv, name="weight", dim=2) + + self.padding = SEWSamePadLayer(config.num_conv_pos_embeddings) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.padding(hidden_states) + hidden_states = self.activation(hidden_states) + + return hidden_states + + +class SEWSamePadLayer(Wav2Vec2SamePadLayer): + pass + + +class SEWUpsampling(nn.Module): + def __init__(self, config): + super().__init__() + self.projection = nn.Linear(config.hidden_size, config.hidden_size * config.squeeze_factor) + self.activation = ACT2FN[config.feat_extract_activation] + self.squeeze_factor = config.squeeze_factor + + def forward(self, hidden_states): + hidden_states = self.projection(hidden_states) + hidden_states = self.activation(hidden_states) + + if self.squeeze_factor > 1: + # transform embedding channels to sequence length + bsz, src_len, src_embed_dim = hidden_states.size() + tgt_len = src_len * self.squeeze_factor + tgt_embed_dim = src_embed_dim // self.squeeze_factor + hidden_states = hidden_states.reshape(bsz, src_len, self.squeeze_factor, tgt_embed_dim) + hidden_states = hidden_states.reshape(bsz, tgt_len, tgt_embed_dim) + + return hidden_states + + +class SEWFeatureEncoder(Wav2Vec2FeatureEncoder): + pass + + +class SEWAttention(Wav2Vec2Attention): + pass + + +class SEWFeedForward(Wav2Vec2FeedForward): + pass + + +class SEWEncoderLayer(Wav2Vec2EncoderLayer): + pass + + +class SEWEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pos_conv_embed = SEWPositionalConvEmbedding(config) + self.pool = nn.AvgPool1d(config.squeeze_factor, config.squeeze_factor) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout) + self.layers = nn.ModuleList([SEWEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.upsample = SEWUpsampling(config) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + 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 + + if attention_mask is not None: + expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + if is_flash_attention_requested(self.config): + # make sure padded tokens output 0 + hidden_states[~expand_attention_mask] = 0.0 + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # make sure padded tokens output 0 + hidden_states[~expand_attention_mask] = 0.0 + input_lengths = (attention_mask.long()).sum(-1) + # apply pooling formula to get real output_lengths + output_lengths = input_lengths // self.config.squeeze_factor + max_encoder_length = hidden_states.shape[1] // self.config.squeeze_factor + attention_ids = ( + torch.arange(0, max_encoder_length, device=output_lengths.device) + .view(1, -1) + .expand(output_lengths.shape[0], -1) + ) + attention_mask = (attention_ids < output_lengths.view(-1, 1)).long() + + # extend attention_mask + attention_mask = 1.0 - attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) + attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min + attention_mask = attention_mask.expand( + attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1] + ) + + n_input_timesteps = hidden_states.shape[1] + + hidden_states = hidden_states.transpose(1, 2) + position_embeddings = self.pos_conv_embed(hidden_states) + pooled_hidden_states = self.pool(hidden_states) + min_length = min(position_embeddings.size(-1), pooled_hidden_states.size(-1)) + hidden_states = pooled_hidden_states[..., :min_length] + position_embeddings[..., :min_length] + hidden_states = hidden_states.transpose(1, 2) + + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) + + for layer in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + dropout_probability = torch.rand([]) + + skip_the_layer = self.training and dropout_probability < self.config.layerdrop + if not skip_the_layer or synced_gpus: + # under fsdp or deepspeed zero3 all gpus must run in sync + layer_outputs = layer( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = layer_outputs[0] + + if skip_the_layer: + layer_outputs = (None, None) + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states = self.upsample(hidden_states) + if hidden_states.shape[1] < n_input_timesteps: + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, n_input_timesteps - hidden_states.shape[1])) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring +class SEWPreTrainedModel(PreTrainedModel): + config: SEWConfig + base_model_prefix = "sew" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = False # needs a proper look into the mask creation + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, SEWPositionalConvEmbedding): + init.normal_( + module.conv.weight, + mean=0, + std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)), + ) + init.constant_(module.conv.bias, 0) + elif isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv1d): + if is_deepspeed_zero3_enabled(): + import deepspeed + + if hasattr(module, "weight_v") and hasattr(module, "weight_g"): + with deepspeed.zero.GatheredParameters([module.weight_v, module.weight_g], modifier_rank=0): + init.kaiming_normal_(module.weight) + else: + with deepspeed.zero.GatheredParameters(module.weight, modifier_rank=0): + init.kaiming_normal_(module.weight) + else: + init.kaiming_normal_(module.weight) + + if isinstance(module, (nn.Linear, nn.Conv1d)) and module.bias is not None: + init.zeros_(module.bias) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor | int): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + # 1D convolutional layer output length formula taken + # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html + return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1 + + for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): + input_lengths = _conv_out_length(input_lengths, kernel_size, stride) + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor): + output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long) + batch_size = attention_mask.shape[0] + + attention_mask = torch.zeros( + (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device + ) + # these two operations makes sure that all values before the output lengths idxs are attended to + attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1 + attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool() + return attention_mask + + +@auto_docstring +class SEWModel(SEWPreTrainedModel): + def __init__(self, config: SEWConfig): + super().__init__(config) + self.config = config + self.feature_extractor = SEWFeatureEncoder(config) + self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps) + + self.project_features = config.conv_dim[-1] != config.hidden_size + if self.project_features: + self.feature_projection = nn.Linear(config.conv_dim[-1], config.hidden_size) + self.feature_dropout = nn.Dropout(config.feat_proj_dropout) + + if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0: + self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_()) + + self.encoder = SEWEncoder(config) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model._mask_hidden_states + def _mask_hidden_states( + self, + hidden_states: torch.FloatTensor, + mask_time_indices: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + ): + """ + Masks extracted features along time axis and/or along feature axis according to + [SpecAugment](https://huggingface.co/papers/1904.08779). + """ + + # `config.apply_spec_augment` can set masking to False + if not getattr(self.config, "apply_spec_augment", True): + return hidden_states + + # generate indices & apply SpecAugment along time axis + batch_size, sequence_length, hidden_size = hidden_states.size() + + if mask_time_indices is not None: + # apply SpecAugment along time axis with given mask_time_indices + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + elif self.config.mask_time_prob > 0 and self.training: + mask_time_indices = _compute_mask_indices( + (batch_size, sequence_length), + mask_prob=self.config.mask_time_prob, + mask_length=self.config.mask_time_length, + attention_mask=attention_mask, + min_masks=self.config.mask_time_min_masks, + ) + mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool) + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + + if self.config.mask_feature_prob > 0 and self.training: + # generate indices & apply SpecAugment along feature axis + mask_feature_indices = _compute_mask_indices( + (batch_size, hidden_size), + mask_prob=self.config.mask_feature_prob, + mask_length=self.config.mask_feature_length, + min_masks=self.config.mask_feature_min_masks, + ) + mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool) + mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1) + hidden_states[mask_feature_indices] = 0 + + return hidden_states + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + mask_time_indices: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + r""" + mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict + masked extracted features in *config.proj_codevector_dim* space. + """ + 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 + + extract_features = self.feature_extractor(input_values) + extract_features = extract_features.transpose(1, 2) + extract_features = self.layer_norm(extract_features) + + if self.project_features: + extract_features = self.feature_projection(extract_features) + hidden_states = self.feature_dropout(extract_features) + + if attention_mask is not None: + # compute reduced attention_mask corresponding to feature vectors + attention_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask) + + hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices) + + encoder_outputs = self.encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + + if not return_dict: + return (hidden_states,) + encoder_outputs[1:] + + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class SEWForCTC(Wav2Vec2ForCTC): + pass + + +class SEWForSequenceClassification(Wav2Vec2ForSequenceClassification): + pass + + +__all__ = ["SEWForCTC", "SEWForSequenceClassification", "SEWModel", "SEWPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/speech_to_text/__init__.py b/third_party/transformers/src/transformers/models/speech_to_text/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b4dce9e2cc611046a08cb796f1f74ceb68efb0fa --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/__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_speech_to_text import * + from .feature_extraction_speech_to_text import * + from .modeling_speech_to_text import * + from .processing_speech_to_text import * + from .tokenization_speech_to_text 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/speech_to_text/configuration_speech_to_text.py b/third_party/transformers/src/transformers/models/speech_to_text/configuration_speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..6dbc7884338ae67dafe8ee644dfabd5f1a700e74 --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/configuration_speech_to_text.py @@ -0,0 +1,108 @@ +# Copyright 2021 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. +"""Speech2Text model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="facebook/s2t-small-librispeech-asr") +@strict +class Speech2TextConfig(PreTrainedConfig): + r""" + max_source_positions (`int`, *optional*, defaults to 6000): + The maximum sequence length of log-mel filter-bank features that this model might ever be used with. + max_target_positions (`int`, *optional*, defaults to 1024): + The maximum sequence length that this model might ever be used with. Typically, set this to something large + just in case (e.g., 512 or 1024 or 2048). + num_conv_layers (`int`, *optional*, defaults to 2): + Number of 1D convolutional layers in the conv module. + conv_kernel_sizes (`tuple[int]`, *optional*, defaults to `(5, 5)`): + A tuple of integers defining the kernel size of each 1D convolutional layer in the conv module. The length + of `conv_kernel_sizes` has to match `num_conv_layers`. + conv_channels (`int`, *optional*, defaults to 1024): + An integer defining the number of output channels of each convolution layers except the final one in the + conv module. + input_feat_per_channel (`int`, *optional*, defaults to 80): + An integer specifying the size of feature vector. This is also the dimensions of log-mel filter-bank + features. + + Example: + + ```python + >>> from transformers import Speech2TextConfig, Speech2TextModel + + >>> # Initializing a Speech2Text s2t_transformer_s style configuration + >>> configuration = Speech2TextConfig() + + >>> # Initializing a model (with random weights) from the s2t_transformer_s style configuration + >>> model = Speech2TextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "speech_to_text" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "num_attention_heads": "encoder_attention_heads", + "hidden_size": "d_model", + "num_hidden_layers": "encoder_layers", + } + + vocab_size: int = 10000 + encoder_layers: int = 12 + encoder_ffn_dim: int = 2048 + encoder_attention_heads: int = 4 + decoder_layers: int = 6 + decoder_ffn_dim: int = 2048 + decoder_attention_heads: int = 4 + encoder_layerdrop: float | int = 0.0 + decoder_layerdrop: float | int = 0.0 + use_cache: bool = True + 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 + decoder_start_token_id: int = 2 + scale_embedding: bool = True + pad_token_id: int | None = 1 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 2 + max_source_positions: int = 6000 + max_target_positions: int = 1024 + num_conv_layers: int = 2 + conv_kernel_sizes: list[int] | tuple[int, ...] = (5, 5) + conv_channels: int = 1024 + input_feat_per_channel: int = 80 + input_channels: int = 1 + tie_word_embeddings: bool = True + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if len(self.conv_kernel_sizes) != self.num_conv_layers: + raise ValueError( + "Configuration for convolutional module is incorrect. " + "It is required that `len(config.conv_kernel_sizes)` == `config.num_conv_layers` " + f"but is `len(config.conv_kernel_sizes) = {len(self.conv_kernel_sizes)}`, " + f"`config.num_conv_layers = {self.num_conv_layers}`." + ) + + +__all__ = ["Speech2TextConfig"] diff --git a/third_party/transformers/src/transformers/models/speech_to_text/convert_s2t_fairseq_to_tfms.py b/third_party/transformers/src/transformers/models/speech_to_text/convert_s2t_fairseq_to_tfms.py new file mode 100644 index 0000000000000000000000000000000000000000..9286fae776fd857e4f45b77e7c24f55f5e05ac3b --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/convert_s2t_fairseq_to_tfms.py @@ -0,0 +1,121 @@ +# Copyright 2021 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 torch +from torch import nn + +from transformers import Speech2TextConfig, Speech2TextForConditionalGeneration + + +def remove_ignore_keys_(state_dict): + ignore_keys = [ + "encoder.version", + "decoder.version", + "model.encoder.version", + "model.decoder.version", + "decoder.output_projection.weight", + "_float_tensor", + "encoder.embed_positions._float_tensor", + "decoder.embed_positions._float_tensor", + ] + for k in ignore_keys: + state_dict.pop(k, None) + + +def rename_keys(s_dict): + keys = list(s_dict.keys()) + for key in keys: + if "transformer_layers" in key: + s_dict[key.replace("transformer_layers", "layers")] = s_dict.pop(key) + elif "subsample" in key: + s_dict[key.replace("subsample", "conv")] = s_dict.pop(key) + + +def make_linear_from_emb(emb): + vocab_size, emb_size = emb.weight.shape + lin_layer = nn.Linear(vocab_size, emb_size, bias=False) + lin_layer.weight.data = emb.weight.data + return lin_layer + + +def convert_fairseq_s2t_checkpoint_to_tfms(checkpoint_path, pytorch_dump_folder_path): + m2m_100 = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + args = m2m_100["args"] + state_dict = m2m_100["model"] + lm_head_weights = state_dict["decoder.output_projection.weight"] + + remove_ignore_keys_(state_dict) + rename_keys(state_dict) + + vocab_size = state_dict["decoder.embed_tokens.weight"].shape[0] + + tie_embeds = args.share_decoder_input_output_embed + + conv_kernel_sizes = [int(i) for i in args.conv_kernel_sizes.split(",")] + config = Speech2TextConfig( + vocab_size=vocab_size, + max_source_positions=args.max_source_positions, + max_target_positions=args.max_target_positions, + encoder_layers=args.encoder_layers, + decoder_layers=args.decoder_layers, + encoder_attention_heads=args.encoder_attention_heads, + decoder_attention_heads=args.decoder_attention_heads, + encoder_ffn_dim=args.encoder_ffn_embed_dim, + decoder_ffn_dim=args.decoder_ffn_embed_dim, + d_model=args.encoder_embed_dim, + dropout=args.dropout, + attention_dropout=args.attention_dropout, + activation_dropout=args.activation_dropout, + activation_function="relu", + num_conv_layers=len(conv_kernel_sizes), + conv_channels=args.conv_channels, + conv_kernel_sizes=conv_kernel_sizes, + input_feat_per_channel=args.input_feat_per_channel, + input_channels=args.input_channels, + tie_word_embeddings=tie_embeds, + num_beams=5, + max_length=200, + use_cache=True, + decoder_start_token_id=2, + early_stopping=True, + ) + + model = Speech2TextForConditionalGeneration(config) + missing, unexpected = model.model.load_state_dict(state_dict, strict=False) + if len(missing) > 0 and not set(missing) <= { + "encoder.embed_positions.weights", + "decoder.embed_positions.weights", + }: + raise ValueError( + "Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights` are allowed to be missing," + f" but all the following weights are missing {missing}" + ) + + if tie_embeds: + model.lm_head = make_linear_from_emb(model.model.decoder.embed_tokens) + else: + model.lm_head.weight.data = lm_head_weights + + model.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument("--fairseq_path", type=str, help="Path to the fairseq model (.pt) file.") + parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") + args = parser.parse_args() + convert_fairseq_s2t_checkpoint_to_tfms(args.fairseq_path, args.pytorch_dump_folder_path) diff --git a/third_party/transformers/src/transformers/models/speech_to_text/feature_extraction_speech_to_text.py b/third_party/transformers/src/transformers/models/speech_to_text/feature_extraction_speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..9685e9be01348784ea4f7c7a7b84dd71db04a49f --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/feature_extraction_speech_to_text.py @@ -0,0 +1,311 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Feature extractor class for Speech2Text +""" + +import numpy as np + +from ...audio_utils import mel_filter_bank, spectrogram, window_function +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import PaddingStrategy, TensorType, is_speech_available, logging + + +if is_speech_available(): + import torch + import torchaudio.compliance.kaldi as ta_kaldi + +logger = logging.get_logger(__name__) + + +class Speech2TextFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs a Speech2Text feature extractor. + + This feature extractor inherits from [`Speech2TextFeatureExtractor`] which contains most of the main methods. Users + should refer to this superclass for more information regarding those methods. + + This class extracts mel-filter bank features from raw speech using TorchAudio if installed or using numpy + otherwise, and applies utterance-level cepstral mean and variance normalization to the extracted features. + + Args: + feature_size (`int`, *optional*, defaults to 80): + The feature dimension of the extracted features. + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + num_mel_bins (`int`, *optional*, defaults to 80): + Number of Mel-frequency bins. + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding vectors. + dither (`float`, *optional*, defaults to 0.0): + Adds dithering. In other words, adds a small Gaussian noise to each frame. + E.g. use 4.0 to add dithering with a normal distribution centered + around 0.0 with standard deviation 4.0 (assuming [-32k,+32k] range of kaldi waveform). + The value 0.0 means no dithering. + Dithering has similar effect as `mel_floor`. It reduces the high log_mel_fbank + values for signals with hard-zero sections, when VAD cutoff is present in the signal. + do_ceptral_normalize (`bool`, *optional*, defaults to `True`): + Whether or not to apply utterance-level cepstral mean and variance normalization to extracted features. + normalize_means (`bool`, *optional*, defaults to `True`): + Whether or not to zero-mean normalize the extracted features. + normalize_vars (`bool`, *optional*, defaults to `True`): + Whether or not to unit-variance normalize the extracted features. + """ + + model_input_names = ["input_features", "attention_mask"] + + def __init__( + self, + feature_size=80, + sampling_rate=16000, + num_mel_bins=80, + padding_value=0.0, + dither=0.0, + do_ceptral_normalize=True, + normalize_means=True, + normalize_vars=True, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + self.num_mel_bins = num_mel_bins + self.dither = dither + self.do_ceptral_normalize = do_ceptral_normalize + self.normalize_means = normalize_means + self.normalize_vars = normalize_vars + self.return_attention_mask = True + + if not is_speech_available(): + mel_filters = mel_filter_bank( + num_frequency_bins=257, + num_mel_filters=self.num_mel_bins, + min_frequency=20, + max_frequency=sampling_rate // 2, + sampling_rate=sampling_rate, + norm=None, + mel_scale="kaldi", + triangularize_in_mel_space=True, + ) + + self.mel_filters = mel_filters + self.window = window_function(400, "povey", periodic=False) + + def _extract_fbank_features( + self, + waveform: np.ndarray, + ) -> np.ndarray: + """ + Get mel-filter bank features using TorchAudio. Note that TorchAudio requires 16-bit signed integers as inputs + and hence the waveform should not be normalized before feature extraction. + """ + waveform = waveform * (2**15) # Kaldi compliance: 16-bit signed integers + if is_speech_available(): + waveform = torch.from_numpy(waveform).unsqueeze(0) + features = ta_kaldi.fbank( + waveform, + dither=self.dither, + num_mel_bins=self.num_mel_bins, + sample_frequency=self.sampling_rate, + ) + features = features.numpy() + else: + waveform = np.squeeze(waveform) + features = spectrogram( + waveform, + self.window, + frame_length=400, + hop_length=160, + fft_length=512, + power=2.0, + center=False, + dither=self.dither, + preemphasis=0.97, + mel_filters=self.mel_filters, + log_mel="log", + mel_floor=1.192092955078125e-07, + remove_dc_offset=True, + ).T + return features + + @staticmethod + def utterance_cmvn( + x: np.ndarray, + input_length: int, + normalize_means: bool | None = True, + normalize_vars: bool | None = True, + padding_value: float = 0.0, + ) -> np.ndarray: + # make sure we normalize float32 arrays + if normalize_means: + mean = x[:input_length].mean(axis=0) + x = np.subtract(x, mean) + if normalize_vars: + std = x[:input_length].std(axis=0) + x = np.divide(x, std) + + if input_length < x.shape[0]: + x[input_length:] = padding_value + + # make sure array is in float32 + x = x.astype(np.float32) + + return x + + def normalize( + self, input_features: list[np.ndarray], attention_mask: np.ndarray | None = None + ) -> list[np.ndarray]: + lengths = attention_mask.sum(-1) if attention_mask is not None else [x.shape[0] for x in input_features] + return [ + self.utterance_cmvn(x, n, self.normalize_means, self.normalize_vars, self.padding_value) + for x, n in zip(input_features, lengths) + ] + + def __call__( + self, + raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]], + padding: bool | str | PaddingStrategy = False, + max_length: int | None = None, + truncation: bool = False, + pad_to_multiple_of: int | None = None, + return_tensors: str | TensorType | None = None, + sampling_rate: int | None = None, + return_attention_mask: bool | None = None, + **kwargs, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). + + Args: + raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`): + The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not + stereo, i.e. single float per timestep. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + truncation (`bool`): + Activates truncation to cut input sequences longer than *max_length* to *max_length*. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific feature_extractor's default. + + [What are attention masks?](../glossary#attention-mask) + + + + For Speech2TextTransformer models, `attention_mask` should always be passed for batched inference, to + avoid subtle bugs. + + + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + sampling_rate (`int`, *optional*): + The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors. + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values / vectors. + """ + + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of" + f" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with" + f" {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1 + if is_batched_numpy and len(raw_speech.shape) > 2: + raise ValueError(f"Only mono-channel audio is supported for input to {self}") + is_batched = is_batched_numpy or ( + isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list))) + ) + + if is_batched: + raw_speech = [np.asarray(speech, dtype=np.float32) for speech in raw_speech] + elif not is_batched and not isinstance(raw_speech, np.ndarray): + raw_speech = np.asarray(raw_speech, dtype=np.float32) + elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64): + raw_speech = raw_speech.astype(np.float32) + + # always return batch + if not is_batched: + raw_speech = [raw_speech] + + # extract fbank features + features = [self._extract_fbank_features(waveform) for waveform in raw_speech] + + # convert into correct format for padding + encoded_inputs = BatchFeature({"input_features": features}) + + padded_inputs = self.pad( + encoded_inputs, + padding=padding, + max_length=max_length, + truncation=truncation, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + **kwargs, + ) + + # make sure list is in array format + input_features = padded_inputs.get("input_features") + if isinstance(input_features[0], list): + padded_inputs["input_features"] = [np.asarray(feature, dtype=np.float32) for feature in input_features] + + attention_mask = padded_inputs.get("attention_mask") + if attention_mask is not None: + padded_inputs["attention_mask"] = [np.asarray(array, dtype=np.int32) for array in attention_mask] + + # Utterance-level cepstral mean and variance normalization + if self.do_ceptral_normalize: + attention_mask = ( + np.array(attention_mask, dtype=np.int32) + if self._get_padding_strategies(padding, max_length=max_length) is not PaddingStrategy.DO_NOT_PAD + else None + ) + padded_inputs["input_features"] = self.normalize( + padded_inputs["input_features"], attention_mask=attention_mask + ) + + if return_tensors is not None: + padded_inputs = padded_inputs.convert_to_tensors(return_tensors) + + return padded_inputs + + +__all__ = ["Speech2TextFeatureExtractor"] diff --git a/third_party/transformers/src/transformers/models/speech_to_text/modeling_speech_to_text.py b/third_party/transformers/src/transformers/models/speech_to_text/modeling_speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..d66501a47048811fd512e5cd126c717c41208aa2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/modeling_speech_to_text.py @@ -0,0 +1,957 @@ +# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Speech2Text model.""" + +import math +from collections.abc import Callable + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, + logging, +) +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_speech_to_text import Speech2TextConfig + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.bart.modeling_bart.shift_tokens_right +def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): + """ + Shift input ids one token to the right. + """ + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +class Conv1dSubsampler(nn.Module): + """ + Convolutional subsampler: a stack of 1D convolution (along temporal dimension) followed by non-linear activation + via gated linear units (https://huggingface.co/papers/1911.08460) + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.num_layers = config.num_conv_layers + self.in_channels = config.input_feat_per_channel * config.input_channels + self.mid_channels = config.conv_channels + self.out_channels = config.d_model + self.kernel_sizes = config.conv_kernel_sizes + + self.conv_layers = nn.ModuleList( + nn.Conv1d( + self.in_channels if i == 0 else self.mid_channels // 2, + self.mid_channels if i < self.num_layers - 1 else self.out_channels * 2, + kernel_size=k, + stride=2, + padding=k // 2, + ) + for i, k in enumerate(self.kernel_sizes) + ) + + def forward(self, input_features): + hidden_states = input_features.transpose(1, 2).contiguous() # -> B x (C x D) x T + for conv in self.conv_layers: + hidden_states = conv(hidden_states) + hidden_states = nn.functional.glu(hidden_states, dim=1) + hidden_states = hidden_states.transpose(1, 2).contiguous() # -> T x B x (C x D) + return hidden_states + + +class Speech2TextSinusoidalPositionalEmbedding(nn.Module): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None): + super().__init__() + self.offset = 2 + self.num_positions = num_positions + self.embedding_dim = embedding_dim + self.padding_idx = padding_idx + self.make_weights(num_positions + self.offset, embedding_dim, padding_idx) + + def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: int | None = None): + emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx) + if hasattr(self, "weights"): + # in forward put the weights on the correct dtype and device of the param + emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device) + + self.register_buffer("weights", emb_weights, persistent=False) + + @staticmethod + def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: int | None = None): + """ + Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the + description in Section 3.5 of "Attention Is All You Need". + """ + half_dim = embedding_dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb) + emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1) + if embedding_dim % 2 == 1: + # zero pad + emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1) + if padding_idx is not None: + emb[padding_idx, :] = 0 + return emb.to(torch.get_default_dtype()) + + @torch.no_grad() + def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0): + bsz, seq_len = input_ids.size() + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length).to( + input_ids.device + ) + + # expand embeddings if needed + max_pos = self.padding_idx + 1 + seq_len + if max_pos > self.weights.size(0): + self.make_weights(max_pos + self.offset, self.embedding_dim, self.padding_idx) + + return self.weights.index_select(0, position_ids.view(-1)).view(bsz, seq_len, -1).detach() + + def create_position_ids_from_input_ids( + self, input_ids: torch.Tensor, padding_idx: int, past_key_values_length: int | None = 0 + ): + """ + Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding + symbols are ignored. This is modified from fairseq's `utils.make_positions`. + + Args: + x: torch.Tensor x: + Returns: torch.Tensor + """ + # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA. + mask = input_ids.ne(padding_idx).int() + incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask + return incremental_indices.long() + padding_idx + + +# Copied from transformers.models.bert.modeling_bert.eager_attention_forward +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float | None = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenAttention with Musicgen->Speech2Text +class Speech2TextAttention(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, + is_causal: bool | None = False, + config: Speech2TextConfig | None = None, + layer_idx: int | None = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + self.layer_idx = layer_idx + + 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 | None = False, + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2) + + is_updated = False + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_layer from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + + current_states = key_value_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = curr_past_key_values.layers[self.layer_idx].keys + value_states = curr_past_key_values.layers[self.layer_idx].values + else: + key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v_proj(current_states).view(*kv_input_shape).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 + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + output_attentions=output_attentions, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Speech2Text, MBART->SPEECH_TO_TEXT +class Speech2TextEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Speech2TextConfig): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = Speech2TextAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + config=config, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + if hidden_states.dtype == torch.float16: + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + return hidden_states + + +# copied from transformers.models.mbart.modeling_mbart.MBartDecoderLayer with MBart->Speech2Text, MBART->SPEECH_TO_TEXT +# TODO: change copy when applying cache class +class Speech2TextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Speech2TextConfig, layer_idx=None): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = Speech2TextAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + is_causal=True, + config=config, + layer_idx=layer_idx, + ) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.encoder_attn = Speech2TextAttention( + self.embed_dim, + config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + config=config, + layer_idx=layer_idx, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim) + self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + # Copied from transformers.models.musicgen.modeling_musicgen.MusicgenDecoderLayer.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = True, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(batch, seq_len, embed_dim)` + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_values (`Cache`): cached past key and value projection states + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + hidden_states, _ = self.encoder_attn( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring +class Speech2TextPreTrainedModel(PreTrainedModel): + config: Speech2TextConfig + base_model_prefix = "model" + main_input_name = "input_features" + input_modalities = ("text", "audio") + supports_gradient_checkpointing = True + # TODO: tests would need a rewrite to check for correct implementation + # Current tests always assume certain inputs to be passed + _supports_flash_attn = False + _supports_sdpa = False + _supports_flex_attn = False + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Speech2TextSinusoidalPositionalEmbedding): + emb_weights = module.get_embedding( + module.num_positions + module.offset, module.embedding_dim, module.padding_idx + ) + init.copy_(module.weights, emb_weights) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers + """ + for i in range(self.config.num_conv_layers): + input_lengths = (input_lengths - 1) // 2 + 1 + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask): + # generate creates 3D attention mask, because of the shape of input_features + # convert it to 2D if that's the case + if len(attention_mask.shape) > 2: + attention_mask = attention_mask[:, :, -1] + + subsampled_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)) + bsz = attention_mask.size()[0] + attention_mask = torch.zeros( + (bsz, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device + ) + + # these two operations makes sure that all values + # before the output lengths indices are attended to + attention_mask[(torch.arange(bsz, device=attention_mask.device), subsampled_lengths - 1)] = 1 + attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).long() + return attention_mask + + +class Speech2TextEncoder(Speech2TextPreTrainedModel): + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`Speech2TextEncoderLayer`]. + + Args: + config: Speech2TextConfig + embed_tokens (nn.Embedding): output embedding + """ + + _no_split_modules = ["Speech2TextEncoderLayer"] + _can_record_outputs = { + "hidden_states": Speech2TextEncoderLayer, + "attentions": Speech2TextAttention, + } + + def __init__(self, config: Speech2TextConfig): + 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_source_positions + self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 + + self.conv = Conv1dSubsampler(config) + + self.embed_positions = Speech2TextSinusoidalPositionalEmbedding( + self.max_source_positions, + embed_dim, + self.padding_idx, + ) + self.layers = nn.ModuleList([Speech2TextEncoderLayer(config) for _ in range(config.encoder_layers)]) + self.layer_norm = nn.LayerNorm(config.d_model) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_features, + attention_mask=None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + inputs_embeds = self.conv(input_features) + inputs_embeds = self.embed_scale * inputs_embeds + + # subsample attention mask if necessary + if attention_mask is not None: + attention_mask = self._get_feature_vector_attention_mask(inputs_embeds.shape[1], attention_mask) + padding_mask = attention_mask.ne(1).long() + else: + padding_mask = torch.zeros(inputs_embeds.shape[:2], dtype=torch.long, device=inputs_embeds.device) + + embed_pos = self.embed_positions(padding_mask) + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + for idx, encoder_layer in enumerate(self.layers): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if not to_drop: + hidden_states = encoder_layer( + hidden_states, + attention_mask, + **kwargs, + ) + + hidden_states = self.layer_norm(hidden_states) + + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +class Speech2TextDecoder(Speech2TextPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`Speech2TextDecoderLayer`] + + Args: + config: Speech2TextConfig + embed_tokens (nn.Embedding): output embedding + """ + + _can_record_outputs = { + "hidden_states": Speech2TextDecoderLayer, + "attentions": OutputRecorder(Speech2TextAttention, index=1, layer_name="self_attn"), + "cross_attentions": OutputRecorder(Speech2TextAttention, index=1, layer_name="encoder_attn"), + } + + def __init__(self, config: Speech2TextConfig): + 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_target_positions + self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + + self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx) + + self.embed_positions = Speech2TextSinusoidalPositionalEmbedding( + self.max_target_positions, + config.d_model, + self.padding_idx, + ) + + self.layers = nn.ModuleList( + [Speech2TextDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)] + ) + + self.layer_norm = nn.LayerNorm(config.d_model) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPastAndCrossAttentions: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale + + 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 + + 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_ids, past_key_values_length=past_key_values_length) + + hidden_states = inputs_embeds + positions + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + hidden_states = decoder_layer( + hidden_states, + attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.layer_norm(hidden_states) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Speech2TextModel(Speech2TextPreTrainedModel): + def __init__(self, config: Speech2TextConfig): + super().__init__(config) + + self.encoder = Speech2TextEncoder(config) + self.decoder = Speech2TextDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.decoder.embed_tokens + + def set_input_embeddings(self, value): + self.decoder.embed_tokens = value + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_features: 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, + past_key_values: Cache | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`SpeechToTextTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + SpeechToText 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_speech_to_text._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the + paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy. + + Example: + + ```python + >>> import torch + >>> from transformers import Speech2TextModel, AutoFeatureExtractor + >>> from datasets import load_dataset + + >>> model = Speech2TextModel.from_pretrained("facebook/s2t-small-librispeech-asr") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/s2t-small-librispeech-asr") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor( + ... ds[0]["audio"]["array"], sampling_rate=ds[0]["audio"]["sampling_rate"], return_tensors="pt" + ... ) + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 256] + ```""" + + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_features, + attention_mask=attention_mask, + **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, + ) + + # downsample encoder attention mask + if attention_mask is not None: + encoder_attention_mask = self._get_feature_vector_attention_mask( + encoder_outputs[0].shape[1], attention_mask + ) + else: + encoder_attention_mask = 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=encoder_attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The Speech2Text Model with a language modeling head. Can be used for summarization. + """ +) +class Speech2TextForConditionalGeneration(Speech2TextPreTrainedModel, GenerationMixin): + input_modalities = ("audio", "text") + base_model_prefix = "model" + _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"} + + def __init__(self, config: Speech2TextConfig): + super().__init__(config) + self.model = Speech2TextModel(config) + self.lm_head = nn.Linear(config.d_model, self.config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_features: 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, + past_key_values: Cache | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`SpeechToTextTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + SpeechToText 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_speech_to_text._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the + paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the 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 + >>> import torch + >>> from transformers import Speech2TextProcessor, Speech2TextForConditionalGeneration + >>> from datasets import load_dataset + + >>> model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-small-librispeech-asr") + >>> processor = Speech2TextProcessor.from_pretrained("facebook/s2t-small-librispeech-asr") + + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> inputs = processor( + ... ds[0]["audio"]["array"], sampling_rate=ds[0]["audio"]["sampling_rate"], return_tensors="pt" + ... ) + >>> input_features = inputs.input_features + + >>> generated_ids = model.generate(inputs=input_features) + + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> transcription + 'mister quilter is the apostle of the middle classes and we are glad to welcome his gospel' + ```""" + if labels is not None: + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + + outputs: Seq2SeqModelOutput = self.model( + input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + lm_logits = self.lm_head(outputs.last_hidden_state) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1)) + + return Seq2SeqLMOutput( + loss=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, + ) + + +__all__ = ["Speech2TextForConditionalGeneration", "Speech2TextModel", "Speech2TextPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/speech_to_text/processing_speech_to_text.py b/third_party/transformers/src/transformers/models/speech_to_text/processing_speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..eb381c942338d0e772db78afd49444c809ff647f --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/processing_speech_to_text.py @@ -0,0 +1,53 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Speech processor class for Speech2Text +""" + +from ...processing_utils import ProcessorMixin +from ...utils import auto_docstring + + +@auto_docstring +class Speech2TextProcessor(ProcessorMixin): + def __init__(self, feature_extractor, tokenizer): + super().__init__(feature_extractor, tokenizer) + + @auto_docstring + def __call__(self, *args, **kwargs): + audio = kwargs.pop("audio", None) + sampling_rate = kwargs.pop("sampling_rate", None) + text = kwargs.pop("text", None) + if len(args) > 0: + audio = args[0] + args = args[1:] + + if audio is None and text is None: + raise ValueError("You need to specify either an `audio` or `text` input to process.") + + if audio is not None: + inputs = self.feature_extractor(audio, *args, sampling_rate=sampling_rate, **kwargs) + if text is not None: + encodings = self.tokenizer(text, **kwargs) + + if text is None: + return inputs + elif audio is None: + return encodings + else: + inputs["labels"] = encodings["input_ids"] + return inputs + + +__all__ = ["Speech2TextProcessor"] diff --git a/third_party/transformers/src/transformers/models/speech_to_text/tokenization_speech_to_text.py b/third_party/transformers/src/transformers/models/speech_to_text/tokenization_speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..deb241f70864f894a9e875e8c763e680ce292575 --- /dev/null +++ b/third_party/transformers/src/transformers/models/speech_to_text/tokenization_speech_to_text.py @@ -0,0 +1,294 @@ +# Copyright 2021 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for Speech2Text.""" + +import json +import os +from pathlib import Path +from shutil import copyfile +from typing import Any + +import sentencepiece + +from ...tokenization_python import PreTrainedTokenizer +from ...utils import logging +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + +SPIECE_UNDERLINE = "▁" + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "spm_file": "sentencepiece.bpe.model", +} + + +MAX_MODEL_INPUT_SIZES = { + "facebook/s2t-small-librispeech-asr": 1024, +} + +MUSTC_LANGS = ["pt", "fr", "ru", "nl", "ro", "it", "es", "de"] + +LANGUAGES = {"mustc": MUSTC_LANGS} + + +@requires(backends=("sentencepiece",)) +class Speech2TextTokenizer(PreTrainedTokenizer): + """ + Construct an Speech2Text tokenizer. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains some of the main methods. Users should refer to + the superclass for more information regarding such methods. + + Args: + vocab_file (`str`): + File containing the vocabulary. + spm_file (`str`): + Path to the [SentencePiece](https://github.com/google/sentencepiece) model file + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sentence token. + eos_token (`str`, *optional*, defaults to `""`): + The end of sentence token. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + do_upper_case (`bool`, *optional*, defaults to `False`): + Whether or not to uppercase the output when decoding. + do_lower_case (`bool`, *optional*, defaults to `False`): + Whether or not to lowercase the input when tokenizing. + tgt_lang (`str`, *optional*): + A string representing the target language. + sp_model_kwargs (`dict`, *optional*): + Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for + SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things, + to set: + + - `enable_sampling`: Enable subword regularization. + - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout. + + - `nbest_size = {0,1}`: No sampling is performed. + - `nbest_size > 1`: samples from the nbest_size results. + - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice) + using forward-filtering-and-backward-sampling algorithm. + + - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for + BPE-dropout. + + **kwargs + Additional keyword arguments passed along to [`PreTrainedTokenizer`] + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + prefix_tokens: list[int] = [] + + def __init__( + self, + vocab_file, + spm_file, + bos_token="", + eos_token="", + pad_token="", + unk_token="", + do_upper_case=False, + do_lower_case=False, + tgt_lang=None, + lang_codes=None, + additional_special_tokens=None, + sp_model_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs + + self.do_upper_case = do_upper_case + self.do_lower_case = do_lower_case + + self.encoder = load_json(vocab_file) + self.decoder = {v: k for k, v in self.encoder.items()} + self.spm_file = spm_file + self.sp_model = load_spm(spm_file, self.sp_model_kwargs) + + if lang_codes is not None: + self.lang_codes = lang_codes + self.langs = LANGUAGES[lang_codes] + self.lang_tokens = [f"" for lang in self.langs] + self.lang_code_to_id = {lang: self.sp_model.PieceToId(f"") for lang in self.langs} + if additional_special_tokens is not None: + additional_special_tokens = self.lang_tokens + additional_special_tokens + else: + additional_special_tokens = self.lang_tokens + self._tgt_lang = tgt_lang if tgt_lang is not None else self.langs[0] + + self.set_tgt_lang_special_tokens(self._tgt_lang) + else: + self.lang_code_to_id = {} + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + do_upper_case=do_upper_case, + do_lower_case=do_lower_case, + tgt_lang=tgt_lang, + lang_codes=lang_codes, + sp_model_kwargs=self.sp_model_kwargs, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + @property + def vocab_size(self) -> int: + return len(self.encoder) + + def get_vocab(self) -> dict: + vocab = self.encoder.copy() + vocab.update(self.added_tokens_encoder) + return vocab + + @property + def tgt_lang(self) -> str: + return self._tgt_lang + + @tgt_lang.setter + def tgt_lang(self, new_tgt_lang) -> None: + self._tgt_lang = new_tgt_lang + self.set_tgt_lang_special_tokens(new_tgt_lang) + + def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None: + """Reset the special tokens to the target language setting. prefix=[eos, tgt_lang_code] and suffix=[eos].""" + lang_code_id = self.lang_code_to_id[tgt_lang] + self.prefix_tokens = [lang_code_id] + + def _tokenize(self, text: str) -> list[str]: + return self.sp_model.encode(text, out_type=str) + + def _convert_token_to_id(self, token): + return self.encoder.get(token, self.encoder[self.unk_token]) + + def _convert_id_to_token(self, index: int) -> str: + """Converts an index (integer) in a token (str) using the decoder.""" + return self.decoder.get(index, self.unk_token) + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + """Converts a sequence of tokens (strings for sub-words) in a single string.""" + current_sub_tokens = [] + out_string = "" + for token in tokens: + # make sure that special tokens are not decoded using sentencepiece model + if token in self.all_special_tokens: + decoded = self.sp_model.decode(current_sub_tokens) + out_string += (decoded.upper() if self.do_upper_case else decoded) + token + " " + current_sub_tokens = [] + else: + current_sub_tokens.append(token) + decoded = self.sp_model.decode(current_sub_tokens) + out_string += decoded.upper() if self.do_upper_case else decoded + return out_string.strip() + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None) -> list[int]: + """Build model inputs from a sequence by appending eos_token_id.""" + if token_ids_1 is None: + return self.prefix_tokens + token_ids_0 + [self.eos_token_id] + # We don't expect to process pairs, but leave the pair logic for API consistency + return self.prefix_tokens + token_ids_0 + token_ids_1 + [self.eos_token_id] + + 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 + ) + + prefix_ones = [1] * len(self.prefix_tokens) + suffix_ones = [1] + if token_ids_1 is None: + return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones + return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state["sp_model"] = None + return state + + def __setstate__(self, d: dict) -> None: + self.__dict__ = d + + # for backward compatibility + if not hasattr(self, "sp_model_kwargs"): + self.sp_model_kwargs = {} + + self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs) + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + save_dir = Path(save_directory) + assert save_dir.is_dir(), f"{save_directory} should be a directory" + vocab_save_path = save_dir / ( + (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"] + ) + spm_save_path = save_dir / ( + (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"] + ) + + save_json(self.encoder, vocab_save_path) + + if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file): + copyfile(self.spm_file, spm_save_path) + elif not os.path.isfile(self.spm_file): + with open(spm_save_path, "wb") as fi: + content_spiece_model = self.sp_model.serialized_model_proto() + fi.write(content_spiece_model) + + return (str(vocab_save_path), str(spm_save_path)) + + +def load_spm(path: str, sp_model_kwargs: dict[str, Any]) -> sentencepiece.SentencePieceProcessor: + spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs) + spm.Load(str(path)) + return spm + + +def load_json(path: str) -> dict | list: + with open(path, "r") as f: + return json.load(f) + + +def save_json(data, path: str) -> None: + with open(path, "w") as f: + json.dump(data, f, indent=2) + + +__all__ = ["Speech2TextTokenizer"] diff --git a/third_party/transformers/src/transformers/models/splinter/__init__.py b/third_party/transformers/src/transformers/models/splinter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..258d79a9c347991b6b10d66608c5cc2b3a9cecae --- /dev/null +++ b/third_party/transformers/src/transformers/models/splinter/__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_splinter import * + from .modeling_splinter import * + from .tokenization_splinter import * + from .tokenization_splinter_fast import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/splinter/configuration_splinter.py b/third_party/transformers/src/transformers/models/splinter/configuration_splinter.py new file mode 100644 index 0000000000000000000000000000000000000000..2fcb3b6a22c78af5df0990a4eced8098419d8d0b --- /dev/null +++ b/third_party/transformers/src/transformers/models/splinter/configuration_splinter.py @@ -0,0 +1,64 @@ +# Copyright 2021 Tel AViv University, AllenAI 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. +"""Splinter model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="tau/splinter-base") +@strict +class SplinterConfig(PreTrainedConfig): + r""" + question_token_id (`int`, *optional*, defaults to 104): + The id of the `[QUESTION]` token. + + Example: + + ```python + >>> from transformers import SplinterModel, SplinterConfig + + >>> # Initializing a Splinter tau/splinter-base style configuration + >>> configuration = SplinterConfig() + + >>> # Initializing a model from the tau/splinter-base style configuration + >>> model = SplinterModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "splinter" + + 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 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + question_token_id: int = 104 + + +__all__ = ["SplinterConfig"] diff --git a/third_party/transformers/src/transformers/models/splinter/modeling_splinter.py b/third_party/transformers/src/transformers/models/splinter/modeling_splinter.py new file mode 100644 index 0000000000000000000000000000000000000000..48f6d18fbda669a20c955d5537aee49906318fab --- /dev/null +++ b/third_party/transformers/src/transformers/models/splinter/modeling_splinter.py @@ -0,0 +1,742 @@ +# Copyright 2021 Tel AViv University, AllenAI 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 Splinter model.""" + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, ModelOutput, QuestionAnsweringModelOutput +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, torch_compilable_check +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_splinter import SplinterConfig + + +logger = logging.get_logger(__name__) + + +class SplinterEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> tuple: + 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] + + 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 = 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 + + +# 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->Splinter +class SplinterSelfAttention(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->Splinter +class SplinterSelfOutput(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->Splinter +class SplinterAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = SplinterSelfAttention(config) + self.output = SplinterSelfOutput(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 with Bert->Splinter +class SplinterIntermediate(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->Splinter +class SplinterOutput(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->Splinter +class SplinterLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = SplinterAttention(config) + self.intermediate = SplinterIntermediate(config) + self.output = SplinterOutput(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->Splinter +class SplinterEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([SplinterLayer(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, + ) + + +@auto_docstring +class SplinterPreTrainedModel(PreTrainedModel): + config: SplinterConfig + base_model_prefix = "splinter" + supports_gradient_checkpointing = True + _can_record_outputs = { + "hidden_states": SplinterLayer, + "attentions": SplinterSelfAttention, + } + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, SplinterEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@auto_docstring +class SplinterModel(SplinterPreTrainedModel): + """ + The model is an encoder (with only self-attention) 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): + super().__init__(config) + self.config = config + + self.embeddings = SplinterEmbeddings(config) + self.encoder = SplinterEncoder(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.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + r""" + token_type_ids (`torch.LongTensor` of shape `batch_size, 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, 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) + """ + 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 = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + **kwargs, + ) + sequence_output = encoder_outputs[0] + + return BaseModelOutput( + last_hidden_state=sequence_output, + ) + + +class SplinterFullyConnectedLayer(nn.Module): + def __init__(self, input_dim, output_dim, hidden_act="gelu"): + super().__init__() + + self.input_dim = input_dim + self.output_dim = output_dim + + self.dense = nn.Linear(self.input_dim, self.output_dim) + self.act_fn = ACT2FN[hidden_act] + self.LayerNorm = nn.LayerNorm(self.output_dim) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(inputs) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +class QuestionAwareSpanSelectionHead(nn.Module): + """ + Implementation of Question-Aware Span Selection (QASS) head, described in Splinter's paper: + + """ + + def __init__(self, config): + super().__init__() + + self.query_start_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size) + self.query_end_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size) + self.start_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size) + self.end_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size) + + self.start_classifier = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.end_classifier = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + + def forward(self, inputs, positions): + _, _, dim = inputs.size() + index = positions.unsqueeze(-1).repeat(1, 1, dim) # [batch_size, num_positions, dim] + gathered_reps = torch.gather(inputs, dim=1, index=index) # [batch_size, num_positions, dim] + + query_start_reps = self.query_start_transform(gathered_reps) # [batch_size, num_positions, dim] + query_end_reps = self.query_end_transform(gathered_reps) # [batch_size, num_positions, dim] + start_reps = self.start_transform(inputs) # [batch_size, seq_length, dim] + end_reps = self.end_transform(inputs) # [batch_size, seq_length, dim] + + hidden_states = self.start_classifier(query_start_reps) # [batch_size, num_positions, dim] + start_reps = start_reps.permute(0, 2, 1) # [batch_size, dim, seq_length] + start_logits = torch.matmul(hidden_states, start_reps) + + hidden_states = self.end_classifier(query_end_reps) + end_reps = end_reps.permute(0, 2, 1) + end_logits = torch.matmul(hidden_states, end_reps) + + return start_logits, end_logits + + +@auto_docstring +class SplinterForQuestionAnswering(SplinterPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.splinter = SplinterModel(config) + self.splinter_qass = QuestionAwareSpanSelectionHead(config) + self.question_token_id = config.question_token_id + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + question_positions: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | QuestionAnsweringModelOutput: + r""" + token_type_ids (`torch.LongTensor` of shape `batch_size, 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, 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) + question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*): + The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size, + num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be + the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size, + sequence_length)`. + """ + question_positions_were_none = False + if question_positions is None: + if input_ids is not None: + question_position_for_each_example = torch.argmax( + (torch.eq(input_ids, self.question_token_id)).int(), dim=-1 + ) + else: + question_position_for_each_example = torch.zeros( + inputs_embeds.size(0), dtype=torch.long, layout=inputs_embeds.layout, device=inputs_embeds.device + ) + question_positions = question_position_for_each_example.unsqueeze(-1) + question_positions_were_none = True + + outputs = self.splinter( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + sequence_output = outputs[0] + start_logits, end_logits = self.splinter_qass(sequence_output, question_positions) + + if question_positions_were_none: + start_logits, end_logits = start_logits.squeeze(1), end_logits.squeeze(1) + + if attention_mask is not None: + start_logits = start_logits + (1 - attention_mask) * torch.finfo(start_logits.dtype).min + end_logits = end_logits + (1 - attention_mask) * torch.finfo(end_logits.dtype).min + + 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.clamp_(0, ignored_index) + 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, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of Splinter as a span selection model. + """ +) +class SplinterForPreTrainingOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when start and end positions are provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`): + Span-end scores (before SoftMax). + """ + + loss: torch.FloatTensor | None = None + start_logits: torch.FloatTensor | None = None + end_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@auto_docstring( + custom_intro=""" + Splinter Model for the recurring span selection task as done during the pretraining. The difference to the QA task + is that we do not have a question, but multiple question tokens that replace the occurrences of recurring spans + instead. + """ +) +class SplinterForPreTraining(SplinterPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.splinter = SplinterModel(config) + self.splinter_qass = QuestionAwareSpanSelectionHead(config) + self.question_token_id = config.question_token_id + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + question_positions: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SplinterForPreTrainingOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, num_questions, 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_questions, 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_questions, 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_questions, 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. + start_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*): + The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size, + num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be + the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size, + sequence_length)`. + """ + if question_positions is None and start_positions is not None and end_positions is not None: + raise TypeError("question_positions must be specified in order to calculate the loss") + + elif question_positions is None and input_ids is None: + raise TypeError("question_positions must be specified when inputs_embeds is used") + + elif question_positions is None: + question_positions = self._prepare_question_positions(input_ids) + + outputs = self.splinter( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + sequence_output = outputs[0] + batch_size, sequence_length, dim = sequence_output.size() + # [batch_size, num_questions, sequence_length] + start_logits, end_logits = self.splinter_qass(sequence_output, question_positions) + + num_questions = question_positions.size(1) + if attention_mask is not None: + attention_mask_for_each_question = attention_mask.unsqueeze(1).expand( + batch_size, num_questions, sequence_length + ) + start_logits = start_logits + (1 - attention_mask_for_each_question) * torch.finfo(start_logits.dtype).min + end_logits = end_logits + (1 - attention_mask_for_each_question) * torch.finfo(end_logits.dtype).min + + total_loss = None + # [batch_size, num_questions, sequence_length] + if start_positions is not None and end_positions is not None: + # sometimes the start/end positions are outside our model inputs, we ignore these terms + start_positions.clamp_(0, max(0, sequence_length - 1)) + end_positions.clamp_(0, max(0, sequence_length - 1)) + + # Ignore zero positions in the loss. Splinter never predicts zero + # during pretraining and zero is used for padding question + # tokens as well as for start and end positions of padded + # question tokens. + loss_fct = CrossEntropyLoss(ignore_index=self.config.pad_token_id) + start_loss = loss_fct( + start_logits.view(batch_size * num_questions, sequence_length), + start_positions.view(batch_size * num_questions), + ) + end_loss = loss_fct( + end_logits.view(batch_size * num_questions, sequence_length), + end_positions.view(batch_size * num_questions), + ) + total_loss = (start_loss + end_loss) / 2 + + return SplinterForPreTrainingOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def _prepare_question_positions(self, input_ids: torch.Tensor) -> torch.Tensor: + rows, flat_positions = torch.where(input_ids == self.config.question_token_id) + num_questions = torch.bincount(rows) + positions = torch.full( + (input_ids.size(0), num_questions.max()), + self.config.pad_token_id, + dtype=torch.long, + device=input_ids.device, + ) + torch_compilable_check( + num_questions.size(0) == input_ids.size(0), + "All samples in the batch must have at least one question token.", + ) + cols = torch.cat([torch.arange(n) for n in num_questions]) + positions[rows, cols] = flat_positions + return positions + + +__all__ = [ + "SplinterForQuestionAnswering", + "SplinterForPreTraining", + "SplinterLayer", + "SplinterModel", + "SplinterPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/splinter/tokenization_splinter.py b/third_party/transformers/src/transformers/models/splinter/tokenization_splinter.py new file mode 100644 index 0000000000000000000000000000000000000000..43dab356265bd0684771f5760cc134d7a59d8cb6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/splinter/tokenization_splinter.py @@ -0,0 +1,175 @@ +# Copyright 2021 Tel AViv University, AllenAI and The HuggingFace Inc. team. All rights reserved. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for Splinter.""" + +import collections + +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import WordPiece + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"} + + +def load_vocab(vocab_file): + vocab = collections.OrderedDict() + with open(vocab_file, "r", encoding="utf-8") as reader: + tokens = reader.readlines() + for index, token in enumerate(tokens): + token = token.rstrip("\n") + vocab[token] = index + return vocab + + +class SplinterTokenizer(TokenizersBackend): + r""" + Construct a Splinter tokenizer (backed by HuggingFace's tokenizers library). Based on WordPiece. + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + Path to a vocabulary file. + tokenizer_file (`str`, *optional*): + Path to a tokenizers JSON file containing the serialization of a tokenizer. + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + unk_token (`str`, *optional*, defaults to `"[UNK]"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `"[SEP]"`): + The separator token, which is used when building a sequence from multiple sequences. + pad_token (`str`, *optional*, defaults to `"[PAD]"`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `"[CLS]"`): + The classifier token which is used when doing sequence classification. + mask_token (`str`, *optional*, defaults to `"[MASK]"`): + The token used for masking values. + question_token (`str`, *optional*, defaults to `"[QUESTION]"`): + The token used for constructing question representations. + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase`. + vocab (`str`, `dict` or `list`, *optional*): + Custom vocabulary dictionary. If not provided, a minimal vocabulary is created. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = WordPiece + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + do_lower_case: bool = True, + unk_token: str = "[UNK]", + sep_token: str = "[SEP]", + pad_token: str = "[PAD]", + cls_token: str = "[CLS]", + mask_token: str = "[MASK]", + question_token: str = "[QUESTION]", + tokenize_chinese_chars: bool = True, + strip_accents: bool | None = None, + **kwargs, + ): + self._vocab = ( + vocab + if vocab is not None + else { + str(pad_token): 0, + str(unk_token): 1, + str(cls_token): 2, + str(sep_token): 3, + str(mask_token): 4, + str(question_token): 5, + ".": 6, + } + ) + + self._tokenizer = Tokenizer(WordPiece(self._vocab, unk_token=str(unk_token))) + + self._tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + self._tokenizer.decoder = decoders.WordPiece(prefix="##") + + super().__init__( + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + question_token=question_token, + do_lower_case=do_lower_case, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + **kwargs, + ) + + self.do_lower_case = do_lower_case + self.tokenize_chinese_chars = tokenize_chinese_chars + self.strip_accents = strip_accents + self.question_token = question_token + if self.question_token not in self.all_special_tokens: + self.add_tokens([self.question_token], special_tokens=True) + self.update_post_processor() + + @property + def question_token_id(self): + return self.convert_tokens_to_ids(self.question_token) + + def update_post_processor(self): + cls = self.cls_token + sep = self.sep_token + question = self.question_token + dot = "." + cls_token_id = self.cls_token_id + sep_token_id = self.sep_token_id + question_token_id = self.question_token_id + dot_token_id = self.convert_tokens_to_ids(".") + + if cls is None or sep is None: + return + + if self.padding_side == "right": + pair = f"{cls}:0 $A:0 {question} {dot} {sep}:0 $B:1 {sep}:1" + else: + pair = f"{cls}:0 $A:0 {sep}:0 $B:1 {question} {dot} {sep}:1" + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=pair, + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + (question, question_token_id), + (dot, dot_token_id), + ], + ) + + +__all__ = ["SplinterTokenizer"] diff --git a/third_party/transformers/src/transformers/models/timm_backbone/__init__.py b/third_party/transformers/src/transformers/models/timm_backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a58f096b1e48227518b04a11250d66635c65701 --- /dev/null +++ b/third_party/transformers/src/transformers/models/timm_backbone/__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_timm_backbone import * + from .modeling_timm_backbone 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/timm_backbone/configuration_timm_backbone.py b/third_party/transformers/src/transformers/models/timm_backbone/configuration_timm_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..15c0a9cd8d280e835ea53a2b724831d2c1fe807c --- /dev/null +++ b/third_party/transformers/src/transformers/models/timm_backbone/configuration_timm_backbone.py @@ -0,0 +1,90 @@ +# 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. + +"""Configuration for Backbone models""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import BackboneConfigMixin +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="") +@strict +class TimmBackboneConfig(BackboneConfigMixin, PreTrainedConfig): + r""" + backbone (`str`, *optional*): + The timm checkpoint to load. + features_only (`bool`, *optional*, defaults to `True`): + Whether to output only the features or also the logits. + freeze_batch_norm_2d (`bool`, *optional*, defaults to `False`): + Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. + + Example: + ```python + >>> from transformers import TimmBackboneConfig, TimmBackbone + + >>> # Initializing a timm backbone + >>> configuration = TimmBackboneConfig("resnet50") + + >>> # Initializing a model from the configuration + >>> model = TimmBackbone(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "timm_backbone" + + backbone: str | None = None + num_channels: int = 3 + features_only: bool = True + _out_indices: list[int] | None = None + freeze_batch_norm_2d: bool = False + output_stride: int | None = None + + def __post_init__(self, **kwargs): + self.out_indices = self.out_indices if self.out_indices is not None else [-1] + super().__post_init__(**kwargs) + + @property + def out_indices(self): + return self._out_indices + + @out_indices.setter + def out_indices(self, out_indices: tuple[int, ...] | list[int]): + """ + Set the out_indices attribute. This will also update the out_features attribute to match the new out_indices. + """ + self._out_indices = list(out_indices) if out_indices is not None else out_indices + if getattr(self, "stage_names", None) is not None: + self.set_output_features_output_indices(out_features=None, out_indices=out_indices) + + @property + def out_features(self): + return self._out_features + + @out_features.setter + def out_features(self, out_features: list[str]): + """ + Set the out_features attribute. This will also update the out_indices attribute to match the new out_features. + """ + self._out_features = out_features + if getattr(self, "stage_names", None) is not None: + self.set_output_features_output_indices(out_features=out_features, out_indices=None) + + +__all__ = ["TimmBackboneConfig"] diff --git a/third_party/transformers/src/transformers/models/timm_backbone/modeling_timm_backbone.py b/third_party/transformers/src/transformers/models/timm_backbone/modeling_timm_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..c60606a7657ec9e9648ff808b5461337ae67a784 --- /dev/null +++ b/third_party/transformers/src/transformers/models/timm_backbone/modeling_timm_backbone.py @@ -0,0 +1,161 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...backbone_utils import BackboneMixin, filter_output_hidden_states +from ...modeling_outputs import BackboneOutput +from ...modeling_utils import PreTrainedModel +from ...utils import is_timm_available, requires_backends +from ...utils.generic import can_return_tuple +from .configuration_timm_backbone import TimmBackboneConfig + + +if is_timm_available(): + import timm + + +class TimmBackbone(BackboneMixin, PreTrainedModel): + """ + Wrapper class for timm models to be used as backbones. This enables using the timm models interchangeably with the + other models in the library keeping the same API. + """ + + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = False + config: TimmBackboneConfig + + def __init__(self, config, **kwargs): + requires_backends(self, "timm") + + if config.backbone is None: + raise ValueError("backbone is not set in the config. Please set it to a timm model name.") + + # We just take the final layer by default. This matches the default for the transformers models. + out_indices = config.out_indices if getattr(config, "out_indices", None) is not None else (-1,) + pretrained = kwargs.pop("pretrained", False) + in_chans = kwargs.pop("in_chans", config.num_channels) + + backbone = timm.create_model( + config.backbone, + pretrained=pretrained, + # This is currently not possible for transformer architectures. + features_only=config.features_only, + in_chans=in_chans, + out_indices=out_indices, + output_stride=config.output_stride, + **kwargs, + ) + + # Needs to be called after creating timm model, because `super()` will try to infer + # `stage_names` from model architecture + super().__init__(config, timm_backbone=backbone) + self._backbone = backbone + + # Converts all `BatchNorm2d` and `SyncBatchNorm` or `BatchNormAct2d` and `SyncBatchNormAct2d` layers of + # provided module into `FrozenBatchNorm2d` or `FrozenBatchNormAct2d` respectively + if getattr(config, "freeze_batch_norm_2d", False): + self.freeze_batch_norm_2d() + + # These are used to control the output of the model when called. If output_hidden_states is True, then + # return_layers is modified to include all layers. + self._return_layers = { + layer["module"]: str(layer["index"]) for layer in self._backbone.feature_info.get_dicts() + } + self._all_layers = {layer["module"]: str(i) for i, layer in enumerate(self._backbone.feature_info.info)} + + self.post_init() + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + requires_backends(cls, ["vision", "timm"]) + + config = kwargs.pop("config", TimmBackboneConfig()) + num_channels = kwargs.pop("num_channels", config.num_channels) + features_only = kwargs.pop("features_only", config.features_only) + out_indices = kwargs.pop("out_indices", config.out_indices) + config = TimmBackboneConfig( + backbone=pretrained_model_name_or_path, + num_channels=num_channels, + features_only=features_only, + out_indices=out_indices, + ) + return super()._from_config(config, pretrained=True, **kwargs) + + def freeze_batch_norm_2d(self): + timm.utils.model.freeze_batch_norm_2d(self._backbone) + + def unfreeze_batch_norm_2d(self): + timm.utils.model.unfreeze_batch_norm_2d(self._backbone) + + @torch.no_grad() + def _init_weights(self, module): + """We need to at least re-init the non-persistent buffers if the model was initialized on meta device (we + assume weights and persistent buffers will be part of checkpoint as we have no way to control timm inits)""" + if hasattr(module, "init_non_persistent_buffers"): + module.init_non_persistent_buffers() + elif isinstance(module, nn.BatchNorm2d): + # For non-pretrained models, always initialize buffers (handles both meta device and to_empty() cases) + running_mean = getattr(module, "running_mean", None) + if running_mean is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + + @can_return_tuple + @filter_output_hidden_states + def forward( + self, + pixel_values: torch.FloatTensor, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BackboneOutput | tuple[Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + + if output_attentions: + raise ValueError("Cannot output attentions for timm backbones at the moment") + + if output_hidden_states: + # We modify the return layers to include all the stages of the backbone + self._backbone.return_layers = self._all_layers + hidden_states = self._backbone(pixel_values, **kwargs) + self._backbone.return_layers = self._return_layers + feature_maps = tuple(hidden_states[i] for i in self.out_indices) + else: + feature_maps = self._backbone(pixel_values, **kwargs) + hidden_states = None + + feature_maps = tuple(feature_maps) + hidden_states = tuple(hidden_states) if hidden_states is not None else None + + if not return_dict: + output = (feature_maps,) + if output_hidden_states: + output = output + (hidden_states,) + return output + + return BackboneOutput(feature_maps=feature_maps, hidden_states=hidden_states, attentions=None) + + +__all__ = ["TimmBackbone"] diff --git a/third_party/transformers/src/transformers/models/unispeech_sat/__init__.py b/third_party/transformers/src/transformers/models/unispeech_sat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cc3144ed17d14cd7db2b93b05d21d0a0e3ffe3d0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/__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_unispeech_sat import * + from .modeling_unispeech_sat 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/unispeech_sat/configuration_unispeech_sat.py b/third_party/transformers/src/transformers/models/unispeech_sat/configuration_unispeech_sat.py new file mode 100644 index 0000000000000000000000000000000000000000..ba0caffc3297224cd18564a8f51a28809c9d8d3c --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/configuration_unispeech_sat.py @@ -0,0 +1,221 @@ +# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""UniSpeechSat model configuration""" + +import functools +import operator + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/unispeech-sat-base-100h-libri-ft") +@strict +class UniSpeechSatConfig(PreTrainedConfig): + r""" + feat_proj_dropout (`float`, *optional*, defaults to 0.0): + The dropout probability for output of the feature encoder. + feat_quantizer_dropout (`float`, *optional*, defaults to 0.0): + The dropout probability for the output of the feature encoder that's used by the quantizer. + final_dropout (`float`, *optional*, defaults to 0.1): + The dropout probability for the final projection layer of [`UniSpeechSatForCTC`]. + feat_extract_norm (`str`, *optional*, defaults to `"group"`): + The norm to be applied to 1D convolutional layers in feature encoder. One of `"group"` for group + normalization of only the first 1D convolutional layer or `"layer"` for layer normalization of all 1D + convolutional layers. + feat_extract_activation (`str, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the 1D convolutional layers of the feature + extractor. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` are supported. + conv_dim (`tuple[int]` or `list[int]`, *optional*, defaults to `(512, 512, 512, 512, 512, 512, 512)`): + A tuple of integers defining the number of input and output channels of each 1D convolutional layer in the + feature encoder. The length of *conv_dim* defines the number of 1D convolutional layers. + conv_stride (`tuple[int]` or `list[int]`, *optional*, defaults to `(5, 2, 2, 2, 2, 2, 2)`): + A tuple of integers defining the stride of each 1D convolutional layer in the feature encoder. The length + of *conv_stride* defines the number of convolutional layers and has to match the length of *conv_dim*. + conv_kernel (`tuple[int]` or `list[int]`, *optional*, defaults to `(10, 3, 3, 3, 3, 2, 2)`): + A tuple of integers defining the kernel size of each 1D convolutional layer in the feature encoder. The + length of *conv_kernel* defines the number of convolutional layers and has to match the length of + *conv_dim*. + conv_bias (`bool`, *optional*, defaults to `False`): + Whether the 1D convolutional layers have a bias. + num_conv_pos_embeddings (`int`, *optional*, defaults to 128): + Number of convolutional positional embeddings. Defines the kernel size of 1D convolutional positional + embeddings layer. + num_conv_pos_embedding_groups (`int`, *optional*, defaults to 16): + Number of groups of 1D convolutional positional embeddings layer. + do_stable_layer_norm (`bool`, *optional*, defaults to `False`): + Whether to apply *stable* layer norm architecture of the Transformer encoder. `do_stable_layer_norm is + True` corresponds to applying layer norm before the attention layer, whereas `do_stable_layer_norm is + False` corresponds to applying layer norm after the attention layer. + apply_spec_augment (`bool`, *optional*, defaults to `True`): + Whether to apply *SpecAugment* data augmentation to the outputs of the feature encoder. For reference see + [SpecAugment: A Simple Data Augmentation Method for Automatic Speech + Recognition](https://huggingface.co/papers/1904.08779). + mask_time_prob (`float`, *optional*, defaults to 0.05): + Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking + procedure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If + reasoning from the probability of each feature vector to be chosen as the start of the vector span to be + masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the + actual percentage of masked vectors. This is only relevant if `apply_spec_augment is True`. + mask_time_length (`int`, *optional*, defaults to 10): + Length of vector span along the time axis. + mask_time_min_masks (`int`, *optional*, defaults to 2): + The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step, + irrespectively of `mask_feature_prob`. Only relevant if ''mask_time_prob*len(time_axis)/mask_time_length < + mask_time_min_masks'' + mask_feature_prob (`float`, *optional*, defaults to 0.0): + Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The + masking procedure generates ''mask_feature_prob*len(feature_axis)/mask_time_length'' independent masks over + the axis. If reasoning from the probability of each feature vector to be chosen as the start of the vector + span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap + may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is + True`. + mask_feature_length (`int`, *optional*, defaults to 10): + Length of vector span along the feature axis. + mask_feature_min_masks (`int`, *optional*, defaults to 0): + The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time + step, irrespectively of `mask_feature_prob`. Only relevant if + ''mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks'' + num_codevectors_per_group (`int`, *optional*, defaults to 320): + Number of entries in each quantization codebook (group). + num_codevector_groups (`int`, *optional*, defaults to 2): + Number of codevector groups for product codevector quantization. + contrastive_logits_temperature (`float`, *optional*, defaults to 0.1): + The temperature *kappa* in the contrastive loss. + num_negatives (`int`, *optional*, defaults to 100): + Number of negative samples for the contrastive loss. + codevector_dim (`int`, *optional*, defaults to 256): + Dimensionality of the quantized feature vectors. + proj_codevector_dim (`int`, *optional*, defaults to 256): + Dimensionality of the final projection of both the quantized and the transformer features. + diversity_loss_weight (`int`, *optional*, defaults to 0.1): + The weight of the codebook diversity loss component. + ctc_zero_infinity (`bool`, *optional*, defaults to `False`): + Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly + occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance + of [`UniSpeechSatForCTC`]. + use_weighted_layer_sum (`bool`, *optional*, defaults to `False`): + Whether to use a weighted average of layer outputs with learned weights. Only relevant when using an + instance of [`UniSpeechSatForSequenceClassification`]. + classifier_proj_size (`int`, *optional*, defaults to 256): + Dimensionality of the projection before token mean-pooling for classification. + tdnn_dim (`tuple[int]` or `list[int]`, *optional*, defaults to `(512, 512, 512, 512, 1500)`): + A tuple of integers defining the number of output channels of each 1D convolutional layer in the *TDNN* + module of the *XVector* model. The length of *tdnn_dim* defines the number of *TDNN* layers. + tdnn_kernel (`tuple[int]` or `list[int]`, *optional*, defaults to `(5, 3, 3, 1, 1)`): + A tuple of integers defining the kernel size of each 1D convolutional layer in the *TDNN* module of the + *XVector* model. The length of *tdnn_kernel* has to match the length of *tdnn_dim*. + tdnn_dilation (`tuple[int]` or `list[int]`, *optional*, defaults to `(1, 2, 3, 1, 1)`): + A tuple of integers defining the dilation factor of each 1D convolutional layer in *TDNN* module of the + *XVector* model. The length of *tdnn_dilation* has to match the length of *tdnn_dim*. + xvector_output_dim (`int`, *optional*, defaults to 512): + Dimensionality of the *XVector* embedding vectors. + num_clusters (`int`, *optional*, defaults to 504): + Number of clusters for weak labeling. Only relevant when using an instance of + [`UniSpeechSatForPreTraining`]. + + Example: + + ```python + >>> from transformers import UniSpeechSatModel, UniSpeechSatConfig + + >>> # Initializing a UniSpeechSat microsoft/unispeech-sat-base-100h-libri-ft style configuration + >>> configuration = UniSpeechSatConfig() + + >>> # Initializing a model from the microsoft/unispeech-sat-base-100h-libri-ft style configuration + >>> model = UniSpeechSatModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "unispeech-sat" + + vocab_size: int = 32 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout: float | int = 0.1 + activation_dropout: float | int = 0.1 + attention_dropout: float | int = 0.1 + feat_proj_dropout: float | int = 0.0 + feat_quantizer_dropout: float | int = 0.0 + final_dropout: float | int = 0.1 + layerdrop: float | int = 0.1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + feat_extract_norm: str = "group" + feat_extract_activation: str = "gelu" + conv_dim: list[int] | tuple[int, ...] = (512, 512, 512, 512, 512, 512, 512) + conv_stride: list[int] | tuple[int, ...] = (5, 2, 2, 2, 2, 2, 2) + conv_kernel: list[int] | tuple[int, ...] = (10, 3, 3, 3, 3, 2, 2) + conv_bias: bool = False + num_conv_pos_embeddings: int = 128 + num_conv_pos_embedding_groups: int = 16 + do_stable_layer_norm: bool = False + apply_spec_augment: bool = True + mask_time_prob: float | int = 0.05 + mask_time_length: int = 10 + mask_time_min_masks: int = 2 + mask_feature_prob: float | int = 0.0 + mask_feature_length: int = 10 + mask_feature_min_masks: int = 0 + num_codevectors_per_group: int = 320 + num_codevector_groups: int = 2 + contrastive_logits_temperature: float = 0.1 + num_negatives: int = 100 + codevector_dim: int = 256 + proj_codevector_dim: int = 256 + diversity_loss_weight: float = 0.1 + ctc_loss_reduction: str = "mean" + ctc_zero_infinity: bool = False + use_weighted_layer_sum: bool = False + classifier_proj_size: int = 256 + tdnn_dim: list[int] | tuple[int, ...] = (512, 512, 512, 512, 1500) + tdnn_kernel: list[int] | tuple[int, ...] = (5, 3, 3, 1, 1) + tdnn_dilation: list[int] | tuple[int, ...] = (1, 2, 3, 1, 1) + xvector_output_dim: int = 512 + pad_token_id: int = 0 + bos_token_id: int = 1 + eos_token_id: int = 2 + num_clusters: int = 504 + + def __post_init__(self, **kwargs): + self.num_feat_extract_layers = len(self.conv_dim) + return super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if ( + (len(self.conv_stride) != self.num_feat_extract_layers) + or (len(self.conv_kernel) != self.num_feat_extract_layers) + or (len(self.conv_dim) != self.num_feat_extract_layers) + ): + raise ValueError( + "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` ==" + " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) =" + f" {len(self.conv_dim)}`, `len(config.conv_stride) = {len(self.conv_stride)}`," + f" `len(config.conv_kernel) = {len(self.conv_kernel)}`." + ) + + @property + def inputs_to_logits_ratio(self): + return functools.reduce(operator.mul, self.conv_stride, 1) + + +__all__ = ["UniSpeechSatConfig"] diff --git a/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_original_s3prl_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_original_s3prl_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b981fde16cb29c5ab58db496668a909b2dafdc --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_original_s3prl_checkpoint_to_pytorch.py @@ -0,0 +1,108 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert Hubert checkpoint.""" + +import argparse + +import torch + +from transformers import ( + UniSpeechSatConfig, + UniSpeechSatForAudioFrameClassification, + UniSpeechSatForSequenceClassification, + UniSpeechSatForXVector, + Wav2Vec2FeatureExtractor, + logging, +) + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def convert_classification(base_model_name, hf_config, downstream_dict): + model = UniSpeechSatForSequenceClassification.from_pretrained(base_model_name, config=hf_config) + model.projector.weight.data = downstream_dict["projector.weight"] + model.projector.bias.data = downstream_dict["projector.bias"] + model.classifier.weight.data = downstream_dict["model.post_net.linear.weight"] + model.classifier.bias.data = downstream_dict["model.post_net.linear.bias"] + return model + + +def convert_diarization(base_model_name, hf_config, downstream_dict): + model = UniSpeechSatForAudioFrameClassification.from_pretrained(base_model_name, config=hf_config) + model.classifier.weight.data = downstream_dict["model.linear.weight"] + model.classifier.bias.data = downstream_dict["model.linear.bias"] + return model + + +def convert_xvector(base_model_name, hf_config, downstream_dict): + model = UniSpeechSatForXVector.from_pretrained(base_model_name, config=hf_config) + model.projector.weight.data = downstream_dict["connector.weight"] + model.projector.bias.data = downstream_dict["connector.bias"] + for i, kernel_size in enumerate(hf_config.tdnn_kernel): + model.tdnn[i].kernel.weight.data = downstream_dict[ + f"model.framelevel_feature_extractor.module.{i}.kernel.weight" + ] + model.tdnn[i].kernel.bias.data = downstream_dict[f"model.framelevel_feature_extractor.module.{i}.kernel.bias"] + + model.feature_extractor.weight.data = downstream_dict["model.utterancelevel_feature_extractor.linear1.weight"] + model.feature_extractor.bias.data = downstream_dict["model.utterancelevel_feature_extractor.linear1.bias"] + model.classifier.weight.data = downstream_dict["model.utterancelevel_feature_extractor.linear2.weight"] + model.classifier.bias.data = downstream_dict["model.utterancelevel_feature_extractor.linear2.bias"] + model.objective.weight.data = downstream_dict["objective.W"] + return model + + +@torch.no_grad() +def convert_s3prl_checkpoint(base_model_name, config_path, checkpoint_path, model_dump_path): + """ + Copy/paste/tweak model's weights to transformers design. + """ + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + downstream_dict = checkpoint["Downstream"] + + hf_config = UniSpeechSatConfig.from_pretrained(config_path) + hf_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( + base_model_name, return_attention_mask=True, do_normalize=False + ) + + arch = hf_config.architectures[0] + if arch.endswith("ForSequenceClassification"): + hf_model = convert_classification(base_model_name, hf_config, downstream_dict) + elif arch.endswith("ForAudioFrameClassification"): + hf_model = convert_diarization(base_model_name, hf_config, downstream_dict) + elif arch.endswith("ForXVector"): + hf_model = convert_xvector(base_model_name, hf_config, downstream_dict) + else: + raise NotImplementedError(f"S3PRL weights conversion is not supported for {arch}") + + if hf_config.use_weighted_layer_sum: + hf_model.layer_weights.data = checkpoint["Featurizer"]["weights"] + + hf_feature_extractor.save_pretrained(model_dump_path) + hf_model.save_pretrained(model_dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--base_model_name", default=None, type=str, help="Name of the huggingface pretrained base model." + ) + parser.add_argument("--config_path", default=None, type=str, help="Path to the huggingface classifier config.") + parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to the s3prl checkpoint.") + parser.add_argument("--model_dump_path", default=None, type=str, help="Path to the final converted model.") + args = parser.parse_args() + convert_s3prl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path) diff --git a/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_sat_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_sat_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..674d2c20587f5bd746aecccf6d96e88932856271 --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/convert_unispeech_sat_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,223 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert UniSpeechSat checkpoint.""" + +import argparse + +import fairseq +import torch + +from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + +MAPPING = { + "post_extract_proj": "feature_projection.projection", + "encoder.pos_conv.0": "encoder.pos_conv_embed.conv", + "self_attn.k_proj": "encoder.layers.*.attention.k_proj", + "self_attn.v_proj": "encoder.layers.*.attention.v_proj", + "self_attn.q_proj": "encoder.layers.*.attention.q_proj", + "self_attn.out_proj": "encoder.layers.*.attention.out_proj", + "self_attn_layer_norm": "encoder.layers.*.layer_norm", + "fc1": "encoder.layers.*.feed_forward.intermediate_dense", + "fc2": "encoder.layers.*.feed_forward.output_dense", + "final_layer_norm": "encoder.layers.*.final_layer_norm", + "encoder.layer_norm": "encoder.layer_norm", + "encoder.layer_norm_for_extract": "layer_norm_for_extract", + "w2v_model.layer_norm": "feature_projection.layer_norm", + "quantizer.weight_proj": "quantizer.weight_proj", + "quantizer.vars": "quantizer.codevectors", + "project_q": "project_q", + "final_proj": "project_hid", + "w2v_encoder.proj": "lm_head", + "label_embs_concat": "label_embeddings_concat", + "mask_emb": "masked_spec_embed", + "spk_proj": "speaker_proj", +} +TOP_LEVEL_KEYS = [ + "lm_head", + "quantizer.weight_proj", + "quantizer.codevectors", + "project_q", + "project_hid", + "label_embeddings_concat", + "speaker_proj", + "layer_norm_for_extract", +] + + +def set_recursively(hf_pointer, key, value, full_name, weight_type): + for attribute in key.split("."): + hf_pointer = getattr(hf_pointer, attribute) + + if weight_type is not None: + hf_shape = getattr(hf_pointer, weight_type).shape + else: + hf_shape = hf_pointer.shape + + if hf_shape != value.shape: + raise ValueError( + f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be" + f" {value.shape} for {full_name}" + ) + + if weight_type == "weight": + hf_pointer.weight.data = value + elif weight_type == "weight_g": + hf_pointer.weight_g.data = value + elif weight_type == "weight_v": + hf_pointer.weight_v.data = value + elif weight_type == "bias": + hf_pointer.bias.data = value + else: + hf_pointer.data = value + + logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.") + + +def recursively_load_weights(fairseq_model, hf_model): + unused_weights = [] + fairseq_dict = fairseq_model.state_dict() + + feature_extractor = hf_model.unispeech_sat.feature_extractor + + for name, value in fairseq_dict.items(): + is_used = False + if "conv_layers" in name: + load_conv_layer( + name, + value, + feature_extractor, + unused_weights, + hf_model.config.feat_extract_norm == "group", + ) + is_used = True + else: + for key, mapped_key in MAPPING.items(): + mapped_key = "unispeech_sat." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key + if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]: + if "layer_norm_for_extract" in name and (".".join(name.split(".")[:-1]) != key): + # special case since naming is very similar + continue + is_used = True + if "*" in mapped_key: + layer_index = name.split(key)[0].split(".")[-2] + mapped_key = mapped_key.replace("*", layer_index) + if "weight_g" in name: + weight_type = "weight_g" + elif "weight_v" in name: + weight_type = "weight_v" + elif "bias" in name: + weight_type = "bias" + elif "weight" in name: + # TODO: don't match quantizer.weight_proj + weight_type = "weight" + else: + weight_type = None + set_recursively(hf_model, mapped_key, value, name, weight_type) + continue + if not is_used: + unused_weights.append(name) + + logger.warning(f"Unused weights: {unused_weights}") + + +def load_conv_layer(full_name, value, feature_extractor, unused_weights, use_group_norm): + name = full_name.split("conv_layers.")[-1] + items = name.split(".") + layer_id = int(items[0]) + type_id = int(items[1]) + + if type_id == 0: + if "bias" in name: + if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape: + raise ValueError( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].conv.bias.data = value + logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.") + elif "weight" in name: + if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape: + raise ValueError( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].conv.weight.data = value + logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.") + elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm): + if "bias" in name: + if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape: + raise ValueError( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor[layer_id].layer_norm.bias.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].layer_norm.bias.data = value + logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.") + elif "weight" in name: + if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape: + raise ValueError( + f"{full_name} has size {value.shape}, but" + f" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found." + ) + feature_extractor.conv_layers[layer_id].layer_norm.weight.data = value + logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.") + else: + unused_weights.append(full_name) + + +@torch.no_grad() +def convert_unispeech_sat_checkpoint( + checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True +): + """ + Copy/paste/tweak model's weights to transformers design. + """ + if config_path is not None: + config = UniSpeechSatConfig.from_pretrained(config_path) + else: + config = UniSpeechSatConfig() + + dict_path = "" + + if is_finetuned: + hf_wav2vec = UniSpeechSatForCTC(config) + else: + hf_wav2vec = UniSpeechSatForPreTraining(config) + + model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task( + [checkpoint_path], arg_overrides={"data": "/".join(dict_path.split("/")[:-1])} + ) + model = model[0].eval() + + recursively_load_weights(model, hf_wav2vec) + + hf_wav2vec.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") + parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint") + parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model") + parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") + parser.add_argument( + "--not_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not" + ) + args = parser.parse_args() + convert_unispeech_sat_checkpoint( + args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned + ) diff --git a/third_party/transformers/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py b/third_party/transformers/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py new file mode 100644 index 0000000000000000000000000000000000000000..bb97b659fae795b5fbe90272a3ea50e546b95ec5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py @@ -0,0 +1,1727 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/unispeech_sat/modular_unispeech_sat.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_unispeech_sat.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2021 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import warnings +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations.deepspeed import is_deepspeed_zero3_enabled +from ...integrations.fsdp import is_fsdp_managed_module +from ...masking_utils import create_bidirectional_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + CausalLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, + Wav2Vec2BaseModelOutput, + XVectorOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel, get_torch_context_manager_or_global_device +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, is_peft_available, logging +from .configuration_unispeech_sat import UniSpeechSatConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`UniSpeechSatForPreTrainingOutput`], with potential hidden states and attentions. + """ +) +class UniSpeechSatForPreTrainingOutput(ModelOutput): + r""" + loss (*optional*, returned when model is in train mode, `torch.FloatTensor` of shape `(1,)`): + Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official + paper](https://huggingface.co/papers/2006.11477). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`, *optional*): + Prediction scores of the contrastive loss model, i.e. the output of the model before the final softmax. + projected_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): + Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked + projected quantized states. + projected_quantized_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): + Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive + target vectors for contrastive loss. + codevector_perplexity (`torch.FloatTensor` of shape `(1,)`): + The perplexity of the codevector distribution, used to measure the diversity of the codebook. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + projected_states: torch.FloatTensor | None = None + projected_quantized_states: torch.FloatTensor | None = None + codevector_perplexity: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +class UniSpeechSatSamePadLayer(nn.Module): + def __init__(self, num_conv_pos_embeddings): + super().__init__() + self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 0 + + def forward(self, hidden_states): + if self.num_pad_remove > 0: + hidden_states = hidden_states[:, :, : -self.num_pad_remove] + return hidden_states + + +class UniSpeechSatPositionalConvEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + self.conv = nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size=config.num_conv_pos_embeddings, + padding=config.num_conv_pos_embeddings // 2, + groups=config.num_conv_pos_embedding_groups, + ) + + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0): + self.conv = weight_norm(self.conv, name="weight", dim=2) + if hasattr(self.conv, "parametrizations"): + weight_g = self.conv.parametrizations.weight.original0 + weight_v = self.conv.parametrizations.weight.original1 + else: + weight_g = self.conv.weight_g + weight_v = self.conv.weight_v + deepspeed.zero.register_external_parameter(self, weight_v) + deepspeed.zero.register_external_parameter(self, weight_g) + else: + self.conv = weight_norm(self.conv, name="weight", dim=2) + + self.padding = UniSpeechSatSamePadLayer(config.num_conv_pos_embeddings) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = hidden_states.transpose(1, 2) + + hidden_states = self.conv(hidden_states) + hidden_states = self.padding(hidden_states) + hidden_states = self.activation(hidden_states) + + hidden_states = hidden_states.transpose(1, 2) + return hidden_states + + +class UniSpeechSatNoLayerNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.activation(hidden_states) + return hidden_states + + +class UniSpeechSatLayerNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True) + self.activation = ACT2FN[config.feat_extract_activation] + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + + hidden_states = hidden_states.transpose(-2, -1) + hidden_states = self.layer_norm(hidden_states) + hidden_states = hidden_states.transpose(-2, -1) + + hidden_states = self.activation(hidden_states) + return hidden_states + + +class UniSpeechSatGroupNormConvLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1 + self.out_conv_dim = config.conv_dim[layer_id] + + self.conv = nn.Conv1d( + self.in_conv_dim, + self.out_conv_dim, + kernel_size=config.conv_kernel[layer_id], + stride=config.conv_stride[layer_id], + bias=config.conv_bias, + ) + self.activation = ACT2FN[config.feat_extract_activation] + + self.layer_norm = nn.GroupNorm(num_groups=self.out_conv_dim, num_channels=self.out_conv_dim, affine=True) + + def forward(self, hidden_states): + hidden_states = self.conv(hidden_states) + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.activation(hidden_states) + return hidden_states + + +class UniSpeechSatFeatureEncoder(nn.Module): + """Construct the features from raw audio waveform""" + + def __init__(self, config): + super().__init__() + + if config.feat_extract_norm == "group": + conv_layers = [UniSpeechSatGroupNormConvLayer(config, layer_id=0)] + [ + UniSpeechSatNoLayerNormConvLayer(config, layer_id=i + 1) + for i in range(config.num_feat_extract_layers - 1) + ] + elif config.feat_extract_norm == "layer": + conv_layers = [ + UniSpeechSatLayerNormConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers) + ] + else: + raise ValueError( + f"`config.feat_extract_norm` is {config.feat_extract_norm}, but has to be one of ['group', 'layer']" + ) + self.conv_layers = nn.ModuleList(conv_layers) + self.gradient_checkpointing = False + self._requires_grad = True + + def _freeze_parameters(self): + for param in self.parameters(): + param.requires_grad = False + self._requires_grad = False + + def forward(self, input_values): + hidden_states = input_values[:, None] + + # make sure hidden_states require grad for gradient_checkpointing + if self._requires_grad and self.training: + hidden_states.requires_grad = True + + for conv_layer in self.conv_layers: + hidden_states = conv_layer(hidden_states) + + return hidden_states + + +class UniSpeechSatFeatureProjection(nn.Module): + def __init__(self, config): + super().__init__() + self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps) + self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size) + self.dropout = nn.Dropout(config.feat_proj_dropout) + + def forward(self, hidden_states): + # non-projected hidden states are needed for quantization + norm_hidden_states = self.layer_norm(hidden_states) + hidden_states = self.projection(norm_hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states, norm_hidden_states + + +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 UniSpeechSatAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + is_causal: bool = False, + config: UniSpeechSatConfig | None = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + + self.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, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2) + + current_states = key_value_states if is_cross_attention else hidden_states + key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v_proj(current_states).view(*kv_input_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.dropout, + scaling=self.scaling, + output_attentions=output_attentions, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights, None + + +class UniSpeechSatFeedForward(nn.Module): + def __init__(self, config): + super().__init__() + self.intermediate_dropout = nn.Dropout(config.activation_dropout) + + self.intermediate_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 + + self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.output_dropout = nn.Dropout(config.hidden_dropout) + + def forward(self, hidden_states): + hidden_states = self.intermediate_dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + hidden_states = self.intermediate_dropout(hidden_states) + + hidden_states = self.output_dense(hidden_states) + hidden_states = self.output_dropout(hidden_states) + return hidden_states + + +class UniSpeechSatEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.attention = UniSpeechSatAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + dropout=config.attention_dropout, + is_decoder=False, + config=config, + ) + + self.dropout = nn.Dropout(config.hidden_dropout) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.feed_forward = UniSpeechSatFeedForward(config) + self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + attn_residual = hidden_states + hidden_states, attn_weights, _ = self.attention( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = self.dropout(hidden_states) + hidden_states = attn_residual + hidden_states + + hidden_states = self.layer_norm(hidden_states) + hidden_states = hidden_states + self.feed_forward(hidden_states) + hidden_states = self.final_layer_norm(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class UniSpeechSatEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pos_conv_embed = UniSpeechSatPositionalConvEmbedding(config) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout) + self.layers = nn.ModuleList([UniSpeechSatEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ): + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + if attention_mask is not None: + # make sure padded tokens output 0 + expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + hidden_states[~expand_attention_mask] = 0 + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + position_embeddings = self.pos_conv_embed(hidden_states) + hidden_states = hidden_states + position_embeddings.to(hidden_states.device) + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) + + for layer in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + dropout_probability = torch.rand([]) + + skip_the_layer = self.training and dropout_probability < self.config.layerdrop + if not skip_the_layer or synced_gpus: + # under fsdp or deepspeed zero3 all gpus must run in sync + layer_outputs = layer( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = layer_outputs[0] + + if skip_the_layer: + layer_outputs = (None, None) + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class UniSpeechSatAttnAdapterLayer(nn.Module): + def __init__(self, config): + """ + Implements adapter modules directly with 3D tensor weight as parameters and without using ModuleList to speed + up training throughput. + """ + super().__init__() + self.input_dim = config.adapter_attn_dim + self.hidden_dim = config.hidden_size + + self.norm = nn.LayerNorm(self.hidden_dim) + self.linear_1 = nn.Linear(self.hidden_dim, self.input_dim) + self.act_fn = nn.ReLU() + self.linear_2 = nn.Linear(self.input_dim, self.hidden_dim) + + def forward(self, hidden_states: torch.FloatTensor): + hidden_states = self.norm(hidden_states) + + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + + return hidden_states + + +class UniSpeechSatEncoderLayerStableLayerNorm(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.attention = UniSpeechSatAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + dropout=config.attention_dropout, + is_decoder=False, + config=config, + ) + self.dropout = nn.Dropout(config.hidden_dropout) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.feed_forward = UniSpeechSatFeedForward(config) + self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + if getattr(config, "adapter_attn_dim", None) is not None: + self.adapter_layer = UniSpeechSatAttnAdapterLayer(config) + else: + self.adapter_layer = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool = False, + ): + attn_residual = hidden_states + hidden_states = self.layer_norm(hidden_states) + hidden_states, attn_weights, _ = self.attention( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = self.dropout(hidden_states) + hidden_states = attn_residual + hidden_states + hidden_states = hidden_states + self.feed_forward(self.final_layer_norm(hidden_states)) + + if self.adapter_layer is not None: + hidden_states = hidden_states + self.adapter_layer(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class UniSpeechSatEncoderStableLayerNorm(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pos_conv_embed = UniSpeechSatPositionalConvEmbedding(config) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout) + self.layers = nn.ModuleList( + [UniSpeechSatEncoderLayerStableLayerNorm(config) for _ in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + 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 + + if attention_mask is not None: + # make sure padded tokens output 0 + expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + hidden_states[~expand_attention_mask] = 0 + + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + position_embeddings = self.pos_conv_embed(hidden_states) + hidden_states = hidden_states + position_embeddings + hidden_states = self.dropout(hidden_states) + + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) + + for layer in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + dropout_probability = torch.rand([]) + + skip_the_layer = self.training and dropout_probability < self.config.layerdrop + if not skip_the_layer or synced_gpus: + # under fsdp or deepspeed zero3 all gpus must run in sync + # XXX: could optimize this like synced_gpus in generate_utils but not sure if it's worth the code complication + layer_outputs = layer( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) + hidden_states = layer_outputs[0] + + if skip_the_layer: + layer_outputs = (None, None) + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + hidden_states = self.layer_norm(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class UniSpeechSatGumbelVectorQuantizer(nn.Module): + """ + Vector quantization using gumbel softmax. See `[CATEGORICAL REPARAMETERIZATION WITH + GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) for more information. + """ + + def __init__(self, config): + super().__init__() + self.num_groups = config.num_codevector_groups + self.num_vars = config.num_codevectors_per_group + + if config.codevector_dim % self.num_groups != 0: + raise ValueError( + f"`config.codevector_dim {config.codevector_dim} must be divisible " + f"by `config.num_codevector_groups` {self.num_groups} for concatenation" + ) + + # storage for codebook variables (codewords) + self.codevectors = nn.Parameter( + torch.FloatTensor(1, self.num_groups * self.num_vars, config.codevector_dim // self.num_groups) + ) + self.weight_proj = nn.Linear(config.hidden_size, self.num_groups * self.num_vars) + + # can be decayed for training + self.temperature = 2 + + @staticmethod + def _compute_perplexity(probs, mask=None): + marginal_probs = probs.mean(dim=0) + perplexity = torch.exp(-torch.sum(torch.xlogy(marginal_probs, marginal_probs), dim=-1)).sum() + return perplexity + + def forward(self, hidden_states): + batch_size, sequence_length, hidden_size = hidden_states.shape + + # project to codevector dim + hidden_states = self.weight_proj(hidden_states) + hidden_states = hidden_states.view(batch_size * sequence_length * self.num_groups, -1) + + if self.training: + # sample code vector probs via gumbel in differentiateable way + codevector_probs = nn.functional.gumbel_softmax( + hidden_states.float(), tau=self.temperature, hard=True + ).type_as(hidden_states) + + # compute perplexity + codevector_soft_dist = torch.softmax( + hidden_states.view(batch_size * sequence_length, self.num_groups, -1).float(), dim=-1 + ) + perplexity = self._compute_perplexity(codevector_soft_dist) + else: + # take argmax in non-differentiable way + # comptute hard codevector distribution (one hot) + codevector_idx = hidden_states.argmax(dim=-1) + codevector_probs = hidden_states.new_zeros(*hidden_states.shape).scatter_( + -1, codevector_idx.view(-1, 1), 1.0 + ) + codevector_probs = codevector_probs.view(batch_size * sequence_length, self.num_groups, -1) + + perplexity = self._compute_perplexity(codevector_probs) + + codevector_probs = codevector_probs.view(batch_size * sequence_length, -1) + # use probs to retrieve codevectors + codevectors_per_group = codevector_probs.unsqueeze(-1) * self.codevectors + codevectors = codevectors_per_group.view(batch_size * sequence_length, self.num_groups, self.num_vars, -1) + codevectors = codevectors.sum(-2).view(batch_size, sequence_length, -1) + + return codevectors, perplexity + + +@auto_docstring +class UniSpeechSatPreTrainedModel(PreTrainedModel): + config: UniSpeechSatConfig + base_model_prefix = "unispeech_sat" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + # gumbel softmax requires special init + if isinstance(module, UniSpeechSatGumbelVectorQuantizer): + init.normal_(module.weight_proj.weight, mean=0.0, std=1) + init.zeros_(module.weight_proj.bias) + init.uniform_(module.codevectors) + elif isinstance(module, UniSpeechSatPositionalConvEmbedding): + init.normal_( + module.conv.weight, + mean=0, + std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)), + ) + init.constant_(module.conv.bias, 0) + elif isinstance(module, UniSpeechSatFeatureProjection): + k = math.sqrt(1 / module.projection.in_features) + init.uniform_(module.projection.weight, a=-k, b=k) + init.uniform_(module.projection.bias, a=-k, b=k) + elif isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv1d): + init.kaiming_normal_(module.weight) + + if module.bias is not None: + k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) + init.uniform_(module.bias, a=-k, b=k) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor | int): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + # 1D convolutional layer output length formula taken + # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html + return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1 + + for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): + input_lengths = _conv_out_length(input_lengths, kernel_size, stride) + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor): + # Effectively attention_mask.sum(-1), but not inplace to be able to run + # on inference mode. + non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1] + output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths).to(torch.long) + batch_size = attention_mask.shape[0] + + attention_mask = torch.zeros( + (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device + ) + # these two operations makes sure that all values before the output lengths idxs are attended to + attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1 + attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool() + return attention_mask + + +def _compute_mask_indices( + shape: tuple[int, int], + mask_prob: float, + mask_length: int, + attention_mask: torch.LongTensor | None = None, + min_masks: int = 0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for + ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on + CPU as part of the preprocessing during training. + + Args: + shape: The shape for which to compute masks. This should be of a tuple of size 2 where + the first element is the batch size and the second element is the length of the axis to span. + mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of + independently generated mask spans of length `mask_length` is computed by + `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the + actual percentage will be smaller. + mask_length: size of the mask + min_masks: minimum number of masked spans + attention_mask: A (right-padded) attention mask which independently shortens the feature axis of + each batch dimension. + """ + batch_size, sequence_length = shape + + if mask_length < 1: + raise ValueError("`mask_length` has to be bigger than 0.") + + if mask_length > sequence_length: + raise ValueError( + f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}" + f" and `sequence_length`: {sequence_length}`" + ) + + # epsilon is used for probabilistic rounding + epsilon = np.random.rand(1).item() + + def compute_num_masked_span(input_length): + """Given input length, compute how many spans should be masked""" + num_masked_span = int(mask_prob * input_length / mask_length + epsilon) + num_masked_span = max(num_masked_span, min_masks) + + # make sure num masked span <= sequence_length + if num_masked_span * mask_length > sequence_length: + num_masked_span = sequence_length // mask_length + + # make sure num_masked span is also <= input_length - (mask_length - 1) + if input_length - (mask_length - 1) < num_masked_span: + num_masked_span = max(input_length - (mask_length - 1), 0) + + return num_masked_span + + # compute number of masked spans in batch + input_lengths = ( + attention_mask.detach().sum(-1).tolist() + if attention_mask is not None + else [sequence_length for _ in range(batch_size)] + ) + + # SpecAugment mask to fill + spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool) + spec_aug_mask_idxs = [] + + max_num_masked_span = compute_num_masked_span(sequence_length) + + if max_num_masked_span == 0: + return spec_aug_mask + + for input_length in input_lengths: + # compute num of masked spans for this input + num_masked_span = compute_num_masked_span(input_length) + + # get random indices to mask + spec_aug_mask_idx = np.random.choice( + np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False + ) + + # pick first sampled index that will serve as a dummy index to pad vector + # to ensure same dimension for all batches due to probabilistic rounding + # Picking first sample just pads those vectors twice. + if len(spec_aug_mask_idx) == 0: + # this case can only happen if `input_length` is strictly smaller then + # `sequence_length` in which case the last token has to be a padding + # token which we can use as a dummy mask id + dummy_mask_idx = sequence_length - 1 + else: + dummy_mask_idx = spec_aug_mask_idx[0] + + spec_aug_mask_idx = np.concatenate( + [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx] + ) + spec_aug_mask_idxs.append(spec_aug_mask_idx) + + spec_aug_mask_idxs = np.array(spec_aug_mask_idxs) + + # expand masked indices to masked spans + spec_aug_mask_idxs = np.broadcast_to( + spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length) + ) + spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length) + + # add offset to the starting indexes so that indexes now create a span + offsets = np.arange(mask_length)[None, None, :] + offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape( + batch_size, max_num_masked_span * mask_length + ) + spec_aug_mask_idxs = spec_aug_mask_idxs + offsets + + # ensure that we cannot have indices larger than sequence_length + if spec_aug_mask_idxs.max() > sequence_length - 1: + spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1 + + # scatter indices to mask + np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1) + + return spec_aug_mask + + +UniSpeechSatBaseModelOutput = Wav2Vec2BaseModelOutput + + +@auto_docstring +class UniSpeechSatModel(UniSpeechSatPreTrainedModel): + def __init__(self, config: UniSpeechSatConfig): + super().__init__(config) + self.config = config + self.feature_extractor = UniSpeechSatFeatureEncoder(config) + self.feature_projection = UniSpeechSatFeatureProjection(config) + + self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_()) + + if config.do_stable_layer_norm: + self.encoder = UniSpeechSatEncoderStableLayerNorm(config) + else: + self.encoder = UniSpeechSatEncoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def _mask_hidden_states( + self, + hidden_states: torch.FloatTensor, + mask_time_indices: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + ): + """ + Masks extracted features along time axis and/or along feature axis according to + [SpecAugment](https://huggingface.co/papers/1904.08779). + """ + + # `config.apply_spec_augment` can set masking to False + if not getattr(self.config, "apply_spec_augment", True): + return hidden_states + + # generate indices & apply SpecAugment along time axis + batch_size, sequence_length, hidden_size = hidden_states.size() + + if mask_time_indices is not None: + # apply SpecAugment along time axis with given mask_time_indices + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + elif self.config.mask_time_prob > 0 and self.training: + mask_time_indices = _compute_mask_indices( + (batch_size, sequence_length), + mask_prob=self.config.mask_time_prob, + mask_length=self.config.mask_time_length, + attention_mask=attention_mask, + min_masks=self.config.mask_time_min_masks, + ) + mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool) + hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype) + + if self.config.mask_feature_prob > 0 and self.training: + # generate indices & apply SpecAugment along feature axis + mask_feature_indices = _compute_mask_indices( + (batch_size, hidden_size), + mask_prob=self.config.mask_feature_prob, + mask_length=self.config.mask_feature_length, + min_masks=self.config.mask_feature_min_masks, + ) + mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool) + mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1) + hidden_states[mask_feature_indices] = 0 + + return hidden_states + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + mask_time_indices: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | UniSpeechSatBaseModelOutput: + r""" + mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict + masked extracted features in *config.proj_codevector_dim* space. + """ + 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 + + extract_features = self.feature_extractor(input_values) + extract_features = extract_features.transpose(1, 2) + + if attention_mask is not None: + # compute reduced attention_mask corresponding to feature vectors + attention_mask = self._get_feature_vector_attention_mask(extract_features.shape[1], attention_mask) + + hidden_states, extract_features = self.feature_projection(extract_features) + hidden_states = self._mask_hidden_states( + hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask + ) + + encoder_outputs = self.encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + + if not return_dict: + return (hidden_states, extract_features) + encoder_outputs[1:] + + return UniSpeechSatBaseModelOutput( + last_hidden_state=hidden_states, + extract_features=extract_features, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + UniSpeechSat Model with a vector-quantization module and ctc loss for pre-training. + """ +) +class UniSpeechSatForPreTraining(UniSpeechSatPreTrainedModel): + def __init__(self, config: UniSpeechSatConfig): + super().__init__(config) + self.unispeech_sat = UniSpeechSatModel(config) + self.dropout_features = nn.Dropout(config.feat_quantizer_dropout) + + self.quantizer = UniSpeechSatGumbelVectorQuantizer(config) + self.project_q = nn.Linear(config.codevector_dim, config.proj_codevector_dim) + self.project_hid = nn.Linear(config.hidden_size, config.proj_codevector_dim) + + self.dropout = nn.Dropout(config.final_dropout) + + self.speaker_proj = nn.Linear(config.hidden_size, config.codevector_dim) + self.label_embeddings_concat = nn.Parameter(torch.FloatTensor(config.num_clusters, config.codevector_dim)) + self.label_embeddings_concat.data.zero_() + + self.layer_norm_for_extract = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + if self.config.do_stable_layer_norm: + self.layer_norm_for_extract.requires_grad = False + + # Initialize weights and apply final processing + self.post_init() + + def set_gumbel_temperature(self, temperature: int): + """ + Set the Gumbel softmax temperature to a given value. Only necessary for training + """ + self.quantizer.temperature = temperature + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + @staticmethod + def compute_contrastive_logits( + target_features: torch.FloatTensor, + negative_features: torch.FloatTensor, + predicted_features: torch.FloatTensor, + temperature: int = 1, + ): + """ + Compute logits for contrastive loss based using cosine similarity as the distance measure between + `[positive_feature, negative_features]` and `[predicted_features]`. Additionally, temperature can be applied. + """ + target_features = torch.cat([target_features, negative_features], dim=0) + + logits = torch.cosine_similarity(predicted_features.float(), target_features.float(), dim=-1) + logits = logits.type_as(target_features) + + # apply temperature + logits = logits / temperature + return logits + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | UniSpeechSatForPreTrainingOutput: + r""" + Example: + + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, UniSpeechSatForPreTraining + >>> from transformers.models.unispeech_sat.modeling_unispeech_sat import _compute_mask_indices + + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/unispeech-sat-base") + >>> model = UniSpeechSatForPreTraining.from_pretrained("microsoft/unispeech-sat-base") + >>> # TODO: Add full pretraining example + ```""" + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + transformer_features = outputs[0] + + # quantize all (unmasked) extracted features and project to final vq dim + extract_features = self.dropout_features(outputs[1]) + + # TODO(PVP) - add pretraining logic and add to tests + logits = extract_features + loss = quantized_features = codevector_perplexity = None + + if not return_dict: + if loss is not None: + return (loss, logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:] + return (logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:] + + return UniSpeechSatForPreTrainingOutput( + loss=loss, + logits=logits, + projected_states=transformer_features, + projected_quantized_states=quantized_features, + codevector_perplexity=codevector_perplexity, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +_HIDDEN_STATES_START_POSITION = 2 + + +@auto_docstring( + custom_intro=""" + UniSpeechSat Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). + """ +) +class UniSpeechSatForCTC(UniSpeechSatPreTrainedModel): + def __init__(self, config, target_lang: str | None = None): + r""" + target_lang (`str`, *optional*): + Language id of adapter weights. Adapter weights are stored in the format adapter..safetensors or + adapter..bin. Only relevant when using an instance of [`UniSpeechSatForCTC`] with adapters. Uses 'eng' by + default. + """ + super().__init__(config) + + self.unispeech_sat = UniSpeechSatModel(config) + self.dropout = nn.Dropout(config.final_dropout) + + self.target_lang = target_lang + + if config.vocab_size is None: + raise ValueError( + f"You are trying to instantiate {self.__class__} with a configuration that " + "does not define the vocabulary size of the language model head. Please " + "instantiate the model as follows: `UniSpeechSatForCTC.from_pretrained(..., vocab_size=vocab_size)`. " + "or define `vocab_size` of your model's configuration." + ) + output_hidden_size = ( + config.output_hidden_size if hasattr(config, "add_adapter") and config.add_adapter else config.hidden_size + ) + self.lm_head = nn.Linear(output_hidden_size, config.vocab_size) + + # Initialize weights and apply final processing + self.post_init() + + def tie_weights(self, **kwargs): + """ + This method overwrites [`~PreTrainedModel.tie_weights`] so that adapter weights can be correctly loaded when + passing `target_lang=...` to `from_pretrained(...)`. + + This method is **not** supposed to be called by the user and is prone to be changed in the future. + """ + + if get_torch_context_manager_or_global_device() == torch.device("meta"): + return + + # Note that `tie_weights` is usually used to tie input and output embedding weights. The method is re-purposed to + # correctly load adapter layers for UniSpeechSat so that we do not have to introduce a new API to + # [`PreTrainedModel`]. While slightly hacky, UniSpeechSat never has to tie input and output embeddings, so that it is + # ok to repurpose this function here. + target_lang = self.target_lang + + if target_lang is not None and getattr(self.config, "adapter_attn_dim", None) is None: + raise ValueError(f"Cannot pass `target_lang`: {target_lang} if `config.adapter_attn_dim` is not defined.") + elif target_lang is None and getattr(self.config, "adapter_attn_dim", None) is not None: + logger.info("By default `target_lang` is set to 'eng'.") + elif target_lang is not None: + self.load_adapter(target_lang, force_load=True) + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.unispeech_sat.parameters(): + param.requires_grad = False + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + **kwargs, + ) -> tuple | CausalLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*): + Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to + the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. + All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., + config.vocab_size - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if labels is not None and labels.max() >= self.config.vocab_size: + raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}") + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.dropout(hidden_states) + + logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # retrieve loss input_lengths from attention_mask + attention_mask = ( + attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long) + ) + input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long) + + # assuming that padded tokens are filled with -100 + # when not being attended to + labels_mask = labels >= 0 + target_lengths = labels_mask.sum(-1) + flattened_targets = labels.masked_select(labels_mask) + + # ctc_loss doesn't support fp16 + log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1) + + with torch.backends.cudnn.flags(enabled=False): + loss = nn.functional.ctc_loss( + log_probs, + flattened_targets, + input_lengths, + target_lengths, + blank=self.config.pad_token_id, + reduction=self.config.ctc_loss_reduction, + zero_infinity=self.config.ctc_zero_infinity, + ) + + if not return_dict: + output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutput( + loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions + ) + + +@auto_docstring( + custom_intro=""" + UniSpeechSat Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like + SUPERB Keyword Spotting. + """ +) +class UniSpeechSatForSequenceClassification(UniSpeechSatPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + if hasattr(config, "add_adapter") and config.add_adapter: + raise ValueError( + "Sequence classification does not support the use of UniSpeechSat adapters (config.add_adapter=True)" + ) + self.unispeech_sat = UniSpeechSatModel(config) + num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings + if config.use_weighted_layer_sum: + self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers) + self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size) + self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.unispeech_sat.parameters(): + param.requires_grad = False + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutput: + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file + into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library + (`pip install torchcodec`) or the soundfile library (`pip install soundfile`). + To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion + into a tensor of type `torch.FloatTensor`. See [`UniSpeechSatProcessor.__call__`] for details. + 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 + output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.config.use_weighted_layer_sum: + hidden_states = outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = outputs[0] + + hidden_states = self.projector(hidden_states) + if attention_mask is None: + pooled_output = hidden_states.mean(dim=1) + else: + padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask) + expand_padding_mask = padding_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) + hidden_states[~expand_padding_mask] = 0.0 + pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1) + + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:] + 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 UniSpeechSatForAudioFrameClassification(UniSpeechSatPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + if hasattr(config, "add_adapter") and config.add_adapter: + raise ValueError( + "Audio frame classification does not support the use of UniSpeechSat adapters (config.add_adapter=True)" + ) + self.unispeech_sat = UniSpeechSatModel(config) + num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings + if config.use_weighted_layer_sum: + self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + self.num_labels = config.num_labels + + self.post_init() + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.unispeech_sat.parameters(): + param.requires_grad = False + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | TokenClassifierOutput: + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file + into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library + (`pip install torchcodec`) or the soundfile library (`pip install soundfile`). + To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion + into a tensor of type `torch.FloatTensor`. See [`UniSpeechSatProcessor.__call__`] for details. + 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 + output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.config.use_weighted_layer_sum: + hidden_states = outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = outputs[0] + + logits = self.classifier(hidden_states) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), torch.argmax(labels.view(-1, self.num_labels), axis=1)) + + if not return_dict: + output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:] + return output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class AMSoftmaxLoss(nn.Module): + def __init__(self, input_dim, num_labels, scale=30.0, margin=0.4): + super().__init__() + self.scale = scale + self.margin = margin + self.num_labels = num_labels + self.weight = nn.Parameter(torch.randn(input_dim, num_labels), requires_grad=True) + self.loss = nn.CrossEntropyLoss() + + def forward(self, hidden_states, labels): + labels = labels.flatten() + weight = nn.functional.normalize(self.weight, dim=0) + hidden_states = nn.functional.normalize(hidden_states, dim=1) + cos_theta = torch.mm(hidden_states, weight) + psi = cos_theta - self.margin + + onehot = nn.functional.one_hot(labels, self.num_labels) + logits = self.scale * torch.where(onehot.bool(), psi, cos_theta) + loss = self.loss(logits, labels) + + return loss + + +class TDNNLayer(nn.Module): + def __init__(self, config, layer_id=0): + super().__init__() + self.in_conv_dim = config.tdnn_dim[layer_id - 1] if layer_id > 0 else config.tdnn_dim[layer_id] + self.out_conv_dim = config.tdnn_dim[layer_id] + self.kernel_size = config.tdnn_kernel[layer_id] + self.dilation = config.tdnn_dilation[layer_id] + + self.kernel = nn.Linear(self.in_conv_dim * self.kernel_size, self.out_conv_dim) + self.activation = nn.ReLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if is_peft_available(): + from peft.tuners.lora import LoraLayer + + if is_peft_available(): + if isinstance(self.kernel, LoraLayer): + warnings.warn( + "Detected LoRA on TDNNLayer. LoRA weights won't be applied due to optimization. " + "You should exclude TDNNLayer from LoRA's target modules.", + ) + + # for backward compatibility, we keep nn.Linear but call F.conv1d for speed up + hidden_states = hidden_states.transpose(1, 2) + weight = self.kernel.weight.view(self.out_conv_dim, self.kernel_size, self.in_conv_dim).transpose(1, 2) + hidden_states = nn.functional.conv1d(hidden_states, weight, self.kernel.bias, dilation=self.dilation) + hidden_states = hidden_states.transpose(1, 2) + + hidden_states = self.activation(hidden_states) + return hidden_states + + +@auto_docstring( + custom_intro=""" + UniSpeechSat Model with an XVector feature extraction head on top for tasks like Speaker Verification. + """ +) +class UniSpeechSatForXVector(UniSpeechSatPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.unispeech_sat = UniSpeechSatModel(config) + num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings + if config.use_weighted_layer_sum: + self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers) + self.projector = nn.Linear(config.hidden_size, config.tdnn_dim[0]) + + tdnn_layers = [TDNNLayer(config, i) for i in range(len(config.tdnn_dim))] + self.tdnn = nn.ModuleList(tdnn_layers) + + self.feature_extractor = nn.Linear(config.tdnn_dim[-1] * 2, config.xvector_output_dim) + self.classifier = nn.Linear(config.xvector_output_dim, config.xvector_output_dim) + + self.objective = AMSoftmaxLoss(config.xvector_output_dim, config.num_labels) + + self.post_init() + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + def freeze_base_model(self): + """ + Calling this function will disable the gradient computation for the base model so that its parameters will not + be updated during training. Only the classification head will be updated. + """ + for param in self.unispeech_sat.parameters(): + param.requires_grad = False + + def _get_tdnn_output_lengths(self, input_lengths: torch.LongTensor | int): + """ + Computes the output length of the TDNN layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + # 1D convolutional layer output length formula taken + # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html + return (input_length - kernel_size) // stride + 1 + + for kernel_size in self.config.tdnn_kernel: + input_lengths = _conv_out_length(input_lengths, kernel_size, 1) + + return input_lengths + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + **kwargs, + ) -> tuple | XVectorOutput: + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file + into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library + (`pip install torchcodec`) or the soundfile library (`pip install soundfile`). + To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion + into a tensor of type `torch.FloatTensor`. See [`UniSpeechSatProcessor.__call__`] for details. + 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 + output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.config.use_weighted_layer_sum: + hidden_states = outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = outputs[0] + + hidden_states = self.projector(hidden_states) + + for tdnn_layer in self.tdnn: + hidden_states = tdnn_layer(hidden_states) + + # Statistic Pooling + if attention_mask is None: + mean_features = hidden_states.mean(dim=1) + std_features = hidden_states.std(dim=1) + else: + feat_extract_output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(dim=1)) + tdnn_output_lengths = self._get_tdnn_output_lengths(feat_extract_output_lengths) + mean_features = [] + std_features = [] + for i, length in enumerate(tdnn_output_lengths): + mean_features.append(hidden_states[i, :length].mean(dim=0)) + std_features.append(hidden_states[i, :length].std(dim=0)) + mean_features = torch.stack(mean_features) + std_features = torch.stack(std_features) + statistic_pooling = torch.cat([mean_features, std_features], dim=-1) + + output_embeddings = self.feature_extractor(statistic_pooling) + logits = self.classifier(output_embeddings) + + loss = None + if labels is not None: + loss = self.objective(logits, labels) + + if not return_dict: + output = (logits, output_embeddings) + outputs[_HIDDEN_STATES_START_POSITION:] + return ((loss,) + output) if loss is not None else output + + return XVectorOutput( + loss=loss, + logits=logits, + embeddings=output_embeddings, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "UniSpeechSatForAudioFrameClassification", + "UniSpeechSatForCTC", + "UniSpeechSatForPreTraining", + "UniSpeechSatForSequenceClassification", + "UniSpeechSatForXVector", + "UniSpeechSatModel", + "UniSpeechSatPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/unispeech_sat/modular_unispeech_sat.py b/third_party/transformers/src/transformers/models/unispeech_sat/modular_unispeech_sat.py new file mode 100644 index 0000000000000000000000000000000000000000..c445c42b91395dff160738fd44e4fc750dbc1e48 --- /dev/null +++ b/third_party/transformers/src/transformers/models/unispeech_sat/modular_unispeech_sat.py @@ -0,0 +1,444 @@ +# Copyright 2021 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch UniSpeechSat model.""" + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...modeling_outputs import ModelOutput, Wav2Vec2BaseModelOutput +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from ..wav2vec2.modeling_wav2vec2 import ( + Wav2Vec2Encoder, + Wav2Vec2EncoderStableLayerNorm, + Wav2Vec2FeatureEncoder, + Wav2Vec2FeatureProjection, + Wav2Vec2ForAudioFrameClassification, + Wav2Vec2ForCTC, + Wav2Vec2ForSequenceClassification, + Wav2Vec2ForXVector, + Wav2Vec2GumbelVectorQuantizer, + Wav2Vec2Model, + Wav2Vec2PositionalConvEmbedding, +) +from .configuration_unispeech_sat import UniSpeechSatConfig + + +logger = logging.get_logger(__name__) + + +_HIDDEN_STATES_START_POSITION = 2 + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`UniSpeechSatForPreTrainingOutput`], with potential hidden states and attentions. + """ +) +class UniSpeechSatForPreTrainingOutput(ModelOutput): + r""" + loss (*optional*, returned when model is in train mode, `torch.FloatTensor` of shape `(1,)`): + Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official + paper](https://huggingface.co/papers/2006.11477). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`, *optional*): + Prediction scores of the contrastive loss model, i.e. the output of the model before the final softmax. + projected_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): + Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked + projected quantized states. + projected_quantized_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): + Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive + target vectors for contrastive loss. + codevector_perplexity (`torch.FloatTensor` of shape `(1,)`): + The perplexity of the codevector distribution, used to measure the diversity of the codebook. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + projected_states: torch.FloatTensor | None = None + projected_quantized_states: torch.FloatTensor | None = None + codevector_perplexity: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +class UniSpeechSatPositionalConvEmbedding(Wav2Vec2PositionalConvEmbedding): + pass + + +class UniSpeechSatFeatureEncoder(Wav2Vec2FeatureEncoder): + pass + + +class UniSpeechSatFeatureProjection(Wav2Vec2FeatureProjection): + pass + + +class UniSpeechSatEncoder(Wav2Vec2Encoder): + pass + + +class UniSpeechSatEncoderStableLayerNorm(Wav2Vec2EncoderStableLayerNorm): + pass + + +class UniSpeechSatGumbelVectorQuantizer(Wav2Vec2GumbelVectorQuantizer): + def __init__(self, config): + super().__init__(config) + self.weight_proj = nn.Linear(config.hidden_size, self.num_groups * self.num_vars) + + @staticmethod + def _compute_perplexity(probs, mask=None): + marginal_probs = probs.mean(dim=0) + perplexity = torch.exp(-torch.sum(torch.xlogy(marginal_probs, marginal_probs), dim=-1)).sum() + return perplexity + + def forward(self, hidden_states): + batch_size, sequence_length, hidden_size = hidden_states.shape + + # project to codevector dim + hidden_states = self.weight_proj(hidden_states) + hidden_states = hidden_states.view(batch_size * sequence_length * self.num_groups, -1) + + if self.training: + # sample code vector probs via gumbel in differentiateable way + codevector_probs = nn.functional.gumbel_softmax( + hidden_states.float(), tau=self.temperature, hard=True + ).type_as(hidden_states) + + # compute perplexity + codevector_soft_dist = torch.softmax( + hidden_states.view(batch_size * sequence_length, self.num_groups, -1).float(), dim=-1 + ) + perplexity = self._compute_perplexity(codevector_soft_dist) + else: + # take argmax in non-differentiable way + # comptute hard codevector distribution (one hot) + codevector_idx = hidden_states.argmax(dim=-1) + codevector_probs = hidden_states.new_zeros(*hidden_states.shape).scatter_( + -1, codevector_idx.view(-1, 1), 1.0 + ) + codevector_probs = codevector_probs.view(batch_size * sequence_length, self.num_groups, -1) + + perplexity = self._compute_perplexity(codevector_probs) + + codevector_probs = codevector_probs.view(batch_size * sequence_length, -1) + # use probs to retrieve codevectors + codevectors_per_group = codevector_probs.unsqueeze(-1) * self.codevectors + codevectors = codevectors_per_group.view(batch_size * sequence_length, self.num_groups, self.num_vars, -1) + codevectors = codevectors.sum(-2).view(batch_size, sequence_length, -1) + + return codevectors, perplexity + + +@auto_docstring +class UniSpeechSatPreTrainedModel(PreTrainedModel): + config: UniSpeechSatConfig + base_model_prefix = "unispeech_sat" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + # gumbel softmax requires special init + if isinstance(module, UniSpeechSatGumbelVectorQuantizer): + init.normal_(module.weight_proj.weight, mean=0.0, std=1) + init.zeros_(module.weight_proj.bias) + init.uniform_(module.codevectors) + elif isinstance(module, UniSpeechSatPositionalConvEmbedding): + init.normal_( + module.conv.weight, + mean=0, + std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)), + ) + init.constant_(module.conv.bias, 0) + elif isinstance(module, UniSpeechSatFeatureProjection): + k = math.sqrt(1 / module.projection.in_features) + init.uniform_(module.projection.weight, a=-k, b=k) + init.uniform_(module.projection.bias, a=-k, b=k) + elif isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv1d): + init.kaiming_normal_(module.weight) + + if module.bias is not None: + k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) + init.uniform_(module.bias, a=-k, b=k) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor | int): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + # 1D convolutional layer output length formula taken + # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html + return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1 + + for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): + input_lengths = _conv_out_length(input_lengths, kernel_size, stride) + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor): + # Effectively attention_mask.sum(-1), but not inplace to be able to run + # on inference mode. + non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1] + output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths).to(torch.long) + batch_size = attention_mask.shape[0] + + attention_mask = torch.zeros( + (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device + ) + # these two operations makes sure that all values before the output lengths idxs are attended to + attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1 + attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool() + return attention_mask + + +UniSpeechSatBaseModelOutput = Wav2Vec2BaseModelOutput + + +class UniSpeechSatModel(UniSpeechSatPreTrainedModel, Wav2Vec2Model): + def __init__(self, config: UniSpeechSatConfig): + UniSpeechSatPreTrainedModel.__init__(self, config) + self.config = config + self.feature_extractor = UniSpeechSatFeatureEncoder(config) + self.feature_projection = UniSpeechSatFeatureProjection(config) + + self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_()) + + if config.do_stable_layer_norm: + self.encoder = UniSpeechSatEncoderStableLayerNorm(config) + else: + self.encoder = UniSpeechSatEncoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def freeze_feature_encoder(self): + raise AttributeError("Not needed for UniSpeechSat") + + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + mask_time_indices: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | UniSpeechSatBaseModelOutput: + r""" + mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict + masked extracted features in *config.proj_codevector_dim* space. + """ + 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 + + extract_features = self.feature_extractor(input_values) + extract_features = extract_features.transpose(1, 2) + + if attention_mask is not None: + # compute reduced attention_mask corresponding to feature vectors + attention_mask = self._get_feature_vector_attention_mask(extract_features.shape[1], attention_mask) + + hidden_states, extract_features = self.feature_projection(extract_features) + hidden_states = self._mask_hidden_states( + hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask + ) + + encoder_outputs = self.encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + + if not return_dict: + return (hidden_states, extract_features) + encoder_outputs[1:] + + return UniSpeechSatBaseModelOutput( + last_hidden_state=hidden_states, + extract_features=extract_features, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + UniSpeechSat Model with a vector-quantization module and ctc loss for pre-training. + """ +) +class UniSpeechSatForPreTraining(UniSpeechSatPreTrainedModel): + def __init__(self, config: UniSpeechSatConfig): + super().__init__(config) + self.unispeech_sat = UniSpeechSatModel(config) + self.dropout_features = nn.Dropout(config.feat_quantizer_dropout) + + self.quantizer = UniSpeechSatGumbelVectorQuantizer(config) + self.project_q = nn.Linear(config.codevector_dim, config.proj_codevector_dim) + self.project_hid = nn.Linear(config.hidden_size, config.proj_codevector_dim) + + self.dropout = nn.Dropout(config.final_dropout) + + self.speaker_proj = nn.Linear(config.hidden_size, config.codevector_dim) + self.label_embeddings_concat = nn.Parameter(torch.FloatTensor(config.num_clusters, config.codevector_dim)) + self.label_embeddings_concat.data.zero_() + + self.layer_norm_for_extract = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + if self.config.do_stable_layer_norm: + self.layer_norm_for_extract.requires_grad = False + + # Initialize weights and apply final processing + self.post_init() + + def set_gumbel_temperature(self, temperature: int): + """ + Set the Gumbel softmax temperature to a given value. Only necessary for training + """ + self.quantizer.temperature = temperature + + def freeze_feature_encoder(self): + """ + Calling this function will disable the gradient computation for the feature encoder so that its parameter will + not be updated during training. + """ + self.unispeech_sat.feature_extractor._freeze_parameters() + + @staticmethod + def compute_contrastive_logits( + target_features: torch.FloatTensor, + negative_features: torch.FloatTensor, + predicted_features: torch.FloatTensor, + temperature: int = 1, + ): + """ + Compute logits for contrastive loss based using cosine similarity as the distance measure between + `[positive_feature, negative_features]` and `[predicted_features]`. Additionally, temperature can be applied. + """ + target_features = torch.cat([target_features, negative_features], dim=0) + + logits = torch.cosine_similarity(predicted_features.float(), target_features.float(), dim=-1) + logits = logits.type_as(target_features) + + # apply temperature + logits = logits / temperature + return logits + + @auto_docstring + def forward( + self, + input_values: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | UniSpeechSatForPreTrainingOutput: + r""" + Example: + + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, UniSpeechSatForPreTraining + >>> from transformers.models.unispeech_sat.modeling_unispeech_sat import _compute_mask_indices + + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/unispeech-sat-base") + >>> model = UniSpeechSatForPreTraining.from_pretrained("microsoft/unispeech-sat-base") + >>> # TODO: Add full pretraining example + ```""" + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.unispeech_sat( + input_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + transformer_features = outputs[0] + + # quantize all (unmasked) extracted features and project to final vq dim + extract_features = self.dropout_features(outputs[1]) + + # TODO(PVP) - add pretraining logic and add to tests + logits = extract_features + loss = quantized_features = codevector_perplexity = None + + if not return_dict: + if loss is not None: + return (loss, logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:] + return (logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:] + + return UniSpeechSatForPreTrainingOutput( + loss=loss, + logits=logits, + projected_states=transformer_features, + projected_quantized_states=quantized_features, + codevector_perplexity=codevector_perplexity, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class UniSpeechSatForCTC(Wav2Vec2ForCTC): + pass + + +class UniSpeechSatForSequenceClassification(Wav2Vec2ForSequenceClassification): + pass + + +class UniSpeechSatForAudioFrameClassification(Wav2Vec2ForAudioFrameClassification): + pass + + +class UniSpeechSatForXVector(Wav2Vec2ForXVector): + pass + + +__all__ = [ + "UniSpeechSatForAudioFrameClassification", + "UniSpeechSatForCTC", + "UniSpeechSatForPreTraining", + "UniSpeechSatForSequenceClassification", + "UniSpeechSatForXVector", + "UniSpeechSatModel", + "UniSpeechSatPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/__init__.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c79b85ae96cff49e5c38362b248925b802d8e95 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_vibevoice_acoustic_tokenizer import * + from .feature_extraction_vibevoice_acoustic_tokenizer import * + from .modeling_vibevoice_acoustic_tokenizer 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/vibevoice_acoustic_tokenizer/configuration_vibevoice_acoustic_tokenizer.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/configuration_vibevoice_acoustic_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..5408c5aeeb4549519ffaf4a05da3a6386ca7483b --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/configuration_vibevoice_acoustic_tokenizer.py @@ -0,0 +1,162 @@ +# Copyright 2026 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 huggingface_hub.dataclasses import strict + +from ...configuration_utils import PretrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/VibeVoice-1.5B") +@strict +class VibeVoiceAcousticTokenizerConfig(PretrainedConfig): + r""" + channels (`int`, *optional*, defaults to 1): + Number of input channels. + hidden_size (`int`, *optional*, defaults to 64): + Dimensionality of latent representations. + kernel_size (`int`, *optional*, defaults to 7): + Kernel size for convolutional layers. + num_filters (`int`, *optional*, defaults to 32): + Number of filters in initial convolutional layer, and doubles after each downsampling. + downsampling_ratios (`List[int]`, *optional*, defaults to `[2, 2, 4, 5, 5, 8]`): + Downsampling ratios for each layer. + depths (`List[int]`, *optional*, defaults to `[3, 3, 3, 3, 3, 3, 8]`): + Number of ConvNeXt blocks at each stage. + ffn_expansion (`int`, *optional*, defaults to 4): + Expansion factor for feed-forward networks. + vae_std (`float`, *optional*, defaults to 0.625): + Standard deviation used for VAE sampling after encoder. + + Example: + + ```python + >>> from transformers import VibeVoiceAcousticTokenizerModel, VibeVoiceAcousticTokenizerConfig + + >>> # Initializing a VibeVoice Acoustic Tokenizer configuration + >>> configuration = VibeVoiceAcousticTokenizerConfig() + + >>> # Initializing a model (with random weights) + >>> model = VibeVoiceAcousticTokenizerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "vibevoice_acoustic_tokenizer" + + channels: int = 1 + hidden_size: int = 64 + kernel_size: int = 7 + rms_norm_eps: float = 1e-5 + layer_scale_init_value: float = 1e-6 + initializer_range: float = 1e-2 + num_filters: int = 32 + downsampling_ratios: list[int] | tuple[int, ...] = (2, 2, 4, 5, 5, 8) + depths: list[int] | tuple[int, ...] = (3, 3, 3, 3, 3, 3, 8) + hidden_act: str = "gelu" + ffn_expansion: int = 4 + vae_std: float = 0.625 + + @property + def hop_length(self): + return int(math.prod(self.downsampling_ratios)) + + @property + def encoder_config(self): + return VibeVoiceAcousticTokenizerEncoderConfig(**self.to_dict()) + + @property + def decoder_config(self): + config_dict = self.to_dict() + config_dict["depths"] = list(reversed(config_dict["depths"])) + return VibeVoiceAcousticTokenizerDecoderConfig(**config_dict) + + +@auto_docstring(checkpoint="microsoft/VibeVoice-1.5B") +@strict +class VibeVoiceAcousticTokenizerEncoderConfig(VibeVoiceAcousticTokenizerConfig): + r""" + channels (`int`, *optional*, defaults to 1): + Number of input channels. + hidden_size (`int`, *optional*, defaults to 64): + Dimensionality of latent representations. + kernel_size (`int`, *optional*, defaults to 7): + Kernel size for convolutional layers. + num_filters (`int`, *optional*, defaults to 32): + Number of filters in initial convolutional layer, and doubles after each downsampling. + downsampling_ratios (`List[int]`, *optional*, defaults to `[2, 2, 4, 5, 5, 8]`): + Downsampling ratios for each layer. + depths (`List[int]`, *optional*, defaults to `[3, 3, 3, 3, 3, 3, 8]`): + Number of ConvNeXt blocks at each stage. + ffn_expansion (`int`, *optional*, defaults to 4): + Expansion factor for feed-forward networks. + vae_std (`float`, *optional*, defaults to 0.625): + Standard deviation used for VAE sampling after encoder. + """ + + model_type = "vibevoice_acoustic_tokenizer_encoder" + base_config_key = "encoder_config" + + @property + def encoder_config(self): + return None + + +@auto_docstring(checkpoint="microsoft/VibeVoice-1.5B") +@strict +class VibeVoiceAcousticTokenizerDecoderConfig(VibeVoiceAcousticTokenizerConfig): + r""" + channels (`int`, *optional*, defaults to 1): + Number of input channels. + hidden_size (`int`, *optional*, defaults to 64): + Dimensionality of latent representations. + kernel_size (`int`, *optional*, defaults to 7): + Kernel size for convolutional layers. + num_filters (`int`, *optional*, defaults to 32): + Number of filters in initial convolutional layer, and doubles after each downsampling. + downsampling_ratios (`List[int]`, *optional*, defaults to `[2, 2, 4, 5, 5, 8]`): + Downsampling ratios for each layer. + depths (`List[int]`, *optional*, defaults to `[3, 3, 3, 3, 3, 3, 8]`): + Number of ConvNeXt blocks at each stage. + ffn_expansion (`int`, *optional*, defaults to 4): + Expansion factor for feed-forward networks. + vae_std (`float`, *optional*, defaults to 0.625): + Standard deviation used for VAE sampling after encoder. + """ + + model_type = "vibevoice_acoustic_tokenizer_decoder" + base_config_key = "decoder_config" + + @property + def encoder_config(self): + return None + + @property + def decoder_config(self): + return None + + @property + def upsampling_ratios(self): + return list(reversed(self.downsampling_ratios)) + + +__all__ = [ + "VibeVoiceAcousticTokenizerConfig", + "VibeVoiceAcousticTokenizerEncoderConfig", + "VibeVoiceAcousticTokenizerDecoderConfig", +] diff --git a/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/convert_vibevoice_acoustic_tokenizer_to_hf.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/convert_vibevoice_acoustic_tokenizer_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..f4b8ae7d9d2898900f29592077023994e1564733 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/convert_vibevoice_acoustic_tokenizer_to_hf.py @@ -0,0 +1,249 @@ +# Copyright 2026 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 gc +import json +import logging +import re +from typing import Any + +import torch +from safetensors.torch import load_file + +from transformers import ( + AutoFeatureExtractor, + AutoModel, + VibeVoiceAcousticTokenizerConfig, + VibeVoiceAcousticTokenizerFeatureExtractor, + VibeVoiceAcousticTokenizerModel, +) + + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + +# fmt: off +STATE_DICT_MAPPING = { + # Encoder + r"^model\.acoustic_tokenizer\.encoder\.downsample_layers\.0\.0\.conv\.": r"encoder.stem.conv.conv.", + r"^model\.acoustic_tokenizer\.encoder\.stages\.0\.": r"encoder.stem.stage.", + r"^model\.acoustic_tokenizer\.encoder\.downsample_layers\.(\d+)\.0\.conv\.": (r"encoder.conv_layers.\1.conv.conv.", -1), + r"^model\.acoustic_tokenizer\.encoder\.stages\.(\d+)\.": (r"encoder.conv_layers.\1.stage.", -1), + r"^model\.acoustic_tokenizer\.encoder\.head\.conv\.": r"encoder.head.", + + # Decoder + r"^model\.acoustic_tokenizer\.decoder\.upsample_layers\.0\.0\.conv\.conv\.": r"decoder.stem.conv.conv.", + r"^model\.acoustic_tokenizer\.decoder\.stages\.0\.": r"decoder.stem.stage.", + r"^model\.acoustic_tokenizer\.decoder\.upsample_layers\.(\d+)\.0\.convtr\.convtr\.": (r"decoder.conv_layers.\1.convtr.convtr.", -1), + r"^model\.acoustic_tokenizer\.decoder\.stages\.(\d+)\.": (r"decoder.conv_layers.\1.stage.", -1), + r"^model\.acoustic_tokenizer\.decoder\.head\.conv\.": r"decoder.head.", + + # Common patterns (apply after specific patterns) + r"mixer\.conv\.conv\.conv\.": r"mixer.conv.", + r"\.conv\.conv\.conv\.": r".conv.conv.", +} +# fmt: on + + +def map_old_key_to_new(old_key: str) -> str: + new_key = old_key + + # Apply all regex patterns + for pattern, replacement in STATE_DICT_MAPPING.items(): + # Check if replacement needs index shifting + if isinstance(replacement, tuple): + replacement_pattern, index_shift = replacement + + # Use callback to handle index shifting + def shift_index(match): + result = replacement_pattern + for i, group in enumerate(match.groups(), 1): + if group and group.isdigit(): + shifted_idx = int(group) + index_shift + result = result.replace(f"\\{i}", str(shifted_idx)) + else: + result = result.replace(f"\\{i}", group) + return result + + new_key, n = re.subn(pattern, shift_index, new_key) + else: + new_key, n = re.subn(pattern, replacement, new_key) + + return new_key + + +def convert_state_dict(original_state_dict: dict[str, Any]) -> dict[str, Any]: + new_state_dict = {} + + for old_key, tensor in original_state_dict.items(): + new_key = map_old_key_to_new(old_key) + new_state_dict[new_key] = tensor + if old_key != new_key: + logger.debug(f"Converted: {old_key} -> {new_key}") + + return new_state_dict + + +def convert_checkpoint(checkpoint, config_path, push_to_hub, bfloat16, processor_config=None): + if bfloat16: + dtype = torch.bfloat16 + else: + dtype = torch.float32 + + # 1) Load state dict from safetensors checkpoint + logger.info(f"Loading checkpoint from {checkpoint}") + original_state_dict = load_file(checkpoint) + + # 2) Prepare feature extractor + audio_config = {} + if processor_config is not None: + with open(processor_config, "r") as f: + processor_config = json.load(f) + audio_config = processor_config.get("audio_processor", {}) + if "sampling_rate" not in audio_config: + audio_config["sampling_rate"] = 24000 + if "normalize_audio" not in audio_config: + audio_config["normalize_audio"] = True + if "target_dB_FS" not in audio_config: + audio_config["target_dB_FS"] = -25 + if "eps" not in audio_config: + audio_config["eps"] = 1e-6 + feature_extractor = VibeVoiceAcousticTokenizerFeatureExtractor(**audio_config) + + # 3) Prepare model configuration + with open(config_path, "r") as f: + model_config = json.load(f) + + # Clean up acoustic tokenizer config + acoustic_config_dict = model_config["acoustic_tokenizer_config"].copy() + if "encoder_depths" in acoustic_config_dict and isinstance(acoustic_config_dict["encoder_depths"], str): + acoustic_config_dict["encoder_depths"] = list(map(int, acoustic_config_dict["encoder_depths"].split("-"))) + if "layernorm_eps" in acoustic_config_dict: + acoustic_config_dict["rms_norm_eps"] = acoustic_config_dict.pop("layernorm_eps") + if "encoder_ratios" in acoustic_config_dict: + acoustic_config_dict["downsampling_ratios"] = list(reversed(acoustic_config_dict.pop("encoder_ratios"))) + if "encoder_n_filters" in acoustic_config_dict: + acoustic_config_dict["num_filters"] = acoustic_config_dict.pop("encoder_n_filters") + if "encoder_depths" in acoustic_config_dict: + acoustic_config_dict["depths"] = acoustic_config_dict.pop("encoder_depths") + if "vae_dim" in acoustic_config_dict: + acoustic_config_dict["hidden_size"] = acoustic_config_dict.pop("vae_dim") + if "fix_std" in acoustic_config_dict: + # Original hardcodes a scaling factor for vae_std + acoustic_config_dict["vae_std"] = acoustic_config_dict.pop("fix_std") / 0.8 + + # Remove unused/constant parameters + for key in [ + "decoder_depths", + "decoder_n_filters", + "decoder_ratios", + "std_dist_type", + "pad_mode", + "conv_bias", + "causal", + "mixer_layer", + "layernorm", + "disable_last_norm", + "conv_norm", + "corpus_normalize", + "layernorm_elementwise_affine", + ]: + acoustic_config_dict.pop(key, None) + + # 4) Convert state dict to match HF model structure + logger.info("Converting state dict") + converted_state_dict = convert_state_dict(original_state_dict) + + # 5) Filter for acoustic tokenizer weights + acoustic_state_dict = { + k: v for k, v in converted_state_dict.items() if k.startswith("encoder.") or k.startswith("decoder.") + } + + # 6) Create and save acoustic tokenizer + logger.info("Creating acoustic tokenizer model") + acoustic_config = VibeVoiceAcousticTokenizerConfig(**acoustic_config_dict) + acoustic_model = VibeVoiceAcousticTokenizerModel(acoustic_config).to(dtype) + + # Load weights into HF model + logger.info("Loading weights into model") + missing, unexpected = acoustic_model.load_state_dict(acoustic_state_dict, strict=False) + if len(unexpected) != 0: + raise ValueError(f"Unexpected keys: {unexpected}") + if len(missing) != 0: + raise ValueError(f"Missing keys: {missing}") + + if push_to_hub: + logger.info(f"Pushing to hub as {push_to_hub}") + feature_extractor.push_to_hub(push_to_hub) + acoustic_model.push_to_hub(push_to_hub) + + gc.collect() + logger.info("Verifying conversion by reloading model") + AutoFeatureExtractor.from_pretrained(push_to_hub) + AutoModel.from_pretrained(push_to_hub, dtype=torch.bfloat16, device_map="auto") + logger.info("Model reloaded successfully!") + logger.info("Conversion complete!") + + +""" +Conversion script to extract the acoustic tokenizer from the original VibeVoice model checkpoint and push a checkpoint +for an `VibeVoiceAcousticTokenizerModel` object. + + +1) download 1.5B model. +```bash +# -- download checkpoint and configs +# -- download script here: https://gist.github.com/ebezzam/507dfd544e0a0f12402966503cbc73e6#file-download_vibevoice_checkpoint-py +python src/transformers/models/vibevoice/download_vibevoice_checkpoint.py +wget https://huggingface.co/microsoft/VibeVoice-1.5B/resolve/main/config.json -P /raid/eric/vibevoice +wget https://huggingface.co/microsoft/VibeVoice-1.5B/resolve/main/preprocessor_config.json -P /raid/eric/vibevoice +``` + +2) run conversion with: +``` +python src/transformers/models/vibevoice_acoustic_tokenizer/convert_vibevoice_acoustic_tokenizer_to_hf.py \ + --checkpoint /raid/eric/vibevoice/VibeVoice-1.5B-combined.safetensors \ + --config_path /raid/eric/vibevoice/config.json \ + --processor_config /raid/eric/vibevoice/preprocessor_config.json \ + --push_to_hub bezzam/VibeVoice-AcousticTokenizer +``` + +A checkpoint will be pushed to `bezzam/VibeVoice-AcousticTokenizer` on the HF Hub. +""" +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint", required=True, default=None, type=str, help="Original VibeVoice model checkpoint." + ) + parser.add_argument("--config_path", default=None, type=str, help="Path to config.json of model to convert") + parser.add_argument( + "--processor_config", default=None, type=str, help="Path to preprocessor_config.json of model to convert" + ) + parser.add_argument( + "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub." + ) + parser.add_argument( + "--float32", action="store_true", help="Whether to use float32 precision. Default is bfloat16." + ) + + args = parser.parse_args() + convert_checkpoint( + args.checkpoint, + args.config_path, + args.push_to_hub, + bfloat16=not args.float32, + processor_config=args.processor_config, + ) diff --git a/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/feature_extraction_vibevoice_acoustic_tokenizer.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/feature_extraction_vibevoice_acoustic_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..9990852d83cf810c8db4c389b20395b1a06f0ba2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/feature_extraction_vibevoice_acoustic_tokenizer.py @@ -0,0 +1,150 @@ +# Copyright 2026 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 ...audio_utils import AudioInput, make_list_of_audio +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import PaddingStrategy, logging +from ...utils.import_utils import is_torch_available, requires + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +@requires(backends=("torch",)) +class VibeVoiceAcousticTokenizerFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs a VibeVoiceAcousticTokenizer feature extractor. + + Args: + feature_size (`int`, *optional*, defaults to 1): + The number of channels. + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized, expressed in hertz (Hz). + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used for padding. + normalize_audio (`bool`, *optional*, defaults to `True`): + Whether to normalize audio to a target dB FS. + target_dB_FS (`float`, *optional*, defaults to -25): + Target dB FS for normalization. + eps (`float`, *optional*, defaults to 1e-06): + A small value to avoid division by zero when normalizing. + + """ + + model_input_names = ["input_values", "padding_mask"] + + def __init__( + self, + feature_size=1, + sampling_rate=24000, + padding_value=0.0, + normalize_audio=True, + target_dB_FS=-25, + eps=1e-6, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + + self.normalize_audio = normalize_audio + self.target_dB_FS = target_dB_FS + self.eps = eps + + def __call__( + self, + audio: AudioInput, + sampling_rate: int | None = None, + padding: bool | str | PaddingStrategy | None = True, + pad_to_multiple_of: int | None = None, + max_length: int | None = None, + return_attention_mask: bool | None = True, + return_tensors: str | None = "pt", + ) -> BatchFeature: + """ + Args: + audio (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`: + The sequence or batch of sequences to be processed. Each sequence can be a numpy array, a torch tensor, + a list of numpy arrays or a list of torch tensors. + sampling_rate (`int`, *optional*): + The sampling rate at which the `audio` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of" + f" {self.sampling_rate}. Please make sure that the provided audio input was sampled with" + f" {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + if return_tensors != "pt": + raise ValueError(f"{self.__class__.__name__} only supports `return_tensors='pt'`.") + + # Ensure batch of mono tensors + audio = make_list_of_audio(audio) + for idx, example in enumerate(audio): + example = torch.tensor(example, dtype=torch.float32) + if example.ndim != 1: + raise ValueError(f"Audio should be mono, got shape: {example.shape}") + audio[idx] = example + + if self.normalize_audio: + for idx, example in enumerate(audio): + rms = torch.sqrt(torch.mean(example**2)) + example *= 10 ** (self.target_dB_FS / 20) / (rms + self.eps) + max_val = torch.max(torch.abs(example)) + if max_val > 1.0: + example = example / (max_val + self.eps) + audio[idx] = example + + output_values = BatchFeature({"input_values": audio}) + if padding or pad_to_multiple_of: + output_values = self.pad( + output_values, + padding=padding, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + max_length=max_length, + ) + if return_attention_mask: + output_values["padding_mask"] = output_values.pop("attention_mask") + + # add channel dimension + output_values["input_values"] = output_values["input_values"][:, None, :] + + return output_values + + +__all__ = ["VibeVoiceAcousticTokenizerFeatureExtractor"] diff --git a/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modeling_vibevoice_acoustic_tokenizer.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modeling_vibevoice_acoustic_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..bda302f0c94fb3972c5f7179a58b01d36e1dc706 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modeling_vibevoice_acoustic_tokenizer.py @@ -0,0 +1,577 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/vibevoice_acoustic_tokenizer/modular_vibevoice_acoustic_tokenizer.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_vibevoice_acoustic_tokenizer.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 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 +from typing import Optional + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations import use_kernel_forward_from_hub +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, can_return_tuple, is_torchdynamo_compiling +from ..auto.modeling_auto import AutoModel +from .configuration_vibevoice_acoustic_tokenizer import ( + VibeVoiceAcousticTokenizerConfig, + VibeVoiceAcousticTokenizerDecoderConfig, + VibeVoiceAcousticTokenizerEncoderConfig, +) + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Decoded audio. + latents (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Projected latents (continuous representations for acoustic tokens) at the output of the encoder. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each decoder + layer that can be passed to subsequent forward calls. + """ + + audio: torch.FloatTensor | None = None + latents: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerEncoderOutput(ModelOutput): + r""" + latents (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Projected latents (continuous representations for acoustic tokens) at the output of the encoder. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each encoder + layer that can be passed to subsequent forward calls. + """ + + latents: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerDecoderOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Decoded audio. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each decoder + layer that can be passed to subsequent forward calls. + """ + + audio: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +@use_kernel_forward_from_hub("RMSNorm") +class VibeVoiceAcousticTokenizerRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + VibeVoiceAcousticTokenizerRMSNorm 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 VibeVoiceAcousticTokenizerFeedForward(nn.Module): + def __init__(self, config, hidden_size): + super().__init__() + self.linear1 = nn.Linear(hidden_size, config.ffn_expansion * hidden_size) + self.activation = ACT2FN[config.hidden_act] + self.linear2 = nn.Linear(config.ffn_expansion * hidden_size, hidden_size) + + def forward(self, hidden_states): + return self.linear2(self.activation(self.linear1(hidden_states))) + + +class VibeVoiceAcousticTokenizerConv1dCacheLayer: + def __init__(self): + self.cache: torch.Tensor | None = None + self.is_initialized: bool = False + + def lazy_initialization(self, hidden_states, conv_module): + self.left_pad = conv_module.left_pad + self.in_channels = conv_module.in_channels + self.cache = torch.zeros( + hidden_states.shape[0], + self.in_channels, + self.left_pad, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cache) + + self.is_initialized = True + + def update(self, hidden_states, conv_module=None): + if not self.is_initialized and conv_module is not None: + self.lazy_initialization(hidden_states, conv_module) + elif not self.is_initialized: + raise ValueError( + "VibeVoiceAcousticTokenizerConv1dCacheLayer is not initialized. Make sure to provide conv_module to the update method." + ) + + # get the padding states + if self.left_pad > 0: + shortfall = max(0, self.left_pad - hidden_states.shape[-1]) + if shortfall > 0: + padding_states = torch.cat([self.cache[:, :, -shortfall:], hidden_states], dim=-1) + else: + padding_states = hidden_states[:, :, -self.left_pad :] + else: + padding_states = torch.empty( + hidden_states.shape[0], self.in_channels, 0, dtype=hidden_states.dtype, device=hidden_states.device + ) + + current_cache = self.cache.clone() + self.cache.copy_(padding_states) + + return current_cache + + +class VibeVoiceAcousticTokenizerConv1dPaddingCache: + def __init__(self): + self.layers = {} + + def update(self, hidden_states, cache_key, conv_module): + if cache_key not in self.layers: + self.layers[cache_key] = VibeVoiceAcousticTokenizerConv1dCacheLayer() + + padding_states = self.layers[cache_key].update(hidden_states, conv_module) + padded_hidden_states = torch.cat([padding_states, hidden_states], dim=-1) + return padded_hidden_states + + +# TODO: @eustlb, @ebezzam this should be latter factorized with other causalconv1d (e.g. VoxtralRealtimeCausalConv1d) +class VibeVoiceAcousticTokenizerCausalConv1d(nn.Module): + """Conv1d with built-in causal padding and optional streaming support through a cache.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + cache_key: str, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + ): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation, groups=groups) + self.causal_padding = (kernel_size - 1) * dilation - (stride - 1) + if self.causal_padding < 0: + raise ValueError( + f"Invalid causal padding {self.causal_padding} for kernel_size={kernel_size}, " + f"dilation={dilation}, stride={stride}." + ) + self.cache_key = cache_key + self.in_channels = in_channels + self.left_pad = self.causal_padding + + def forward( + self, + hidden_states: torch.Tensor, + padding_cache: VibeVoiceAcousticTokenizerConv1dPaddingCache | None = None, + ) -> torch.Tensor: + if padding_cache is not None: + hidden_states = padding_cache.update(hidden_states, self.cache_key, self) + else: + hidden_states = nn.functional.pad(hidden_states, (self.left_pad, 0)) + + return self.conv(hidden_states) + + +class VibeVoiceAcousticTokenizerCausalConvTranspose1d(nn.Module): + """ConvTranspose1d with built-in causal padding and optional streaming support through a cache.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + cache_key: str, + stride: int = 1, + ): + super().__init__() + self.convtr = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride) + + self.stride = stride + self.cache_key = cache_key + self.in_channels = in_channels + self.padding_total = kernel_size - stride + self.causal_padding = kernel_size - 1 + self.left_pad = self.causal_padding + + def forward( + self, + hidden_states: torch.Tensor, + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None, + ) -> torch.Tensor: + time_dim = hidden_states.shape[-1] + + if padding_cache is not None: + hidden_states = padding_cache.update(hidden_states, self.cache_key, self) + hidden_states = self.convtr(hidden_states) + + # Remove extra padding at the right side + if self.padding_total > 0: + hidden_states = hidden_states[..., : -self.padding_total] + + if padding_cache is not None: + # For first chunk return full output, for subsequent chunks return only new output + expected_new_output = time_dim * self.stride + if hidden_states.shape[2] >= expected_new_output: + hidden_states = hidden_states[:, :, -expected_new_output:] + return hidden_states + + +class VibeVoiceAcousticTokenizerConvNext1dLayer(nn.Module): + """ConvNeXt-like block adapted for 1D convolutions.""" + + def __init__(self, config, hidden_size, dilation=1, stride=1, layer_idx=None): + super().__init__() + + self.norm = VibeVoiceAcousticTokenizerRMSNorm(hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = VibeVoiceAcousticTokenizerRMSNorm(hidden_size, eps=config.rms_norm_eps) + self.ffn = VibeVoiceAcousticTokenizerFeedForward(config, hidden_size) + self.gamma = nn.Parameter(config.layer_scale_init_value * torch.ones(hidden_size), requires_grad=True) + self.ffn_gamma = nn.Parameter(config.layer_scale_init_value * torch.ones(hidden_size), requires_grad=True) + self.mixer = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=hidden_size, + out_channels=hidden_size, + kernel_size=config.kernel_size, + cache_key=f"convnext_layer_{layer_idx}", + groups=hidden_size, + dilation=dilation, + stride=stride, + ) + + def forward(self, hidden_states, padding_cache=None): + # mixer + residual = hidden_states + hidden_states = self.norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = self.mixer(hidden_states, padding_cache=padding_cache) + hidden_states = hidden_states * self.gamma.unsqueeze(-1) + hidden_states = residual + hidden_states + + # ffn + residual = hidden_states + hidden_states = self.ffn_norm(hidden_states.transpose(1, 2)) + hidden_states = self.ffn(hidden_states).transpose(1, 2) + hidden_states = hidden_states * self.ffn_gamma.unsqueeze(-1) + return residual + hidden_states + + +class VibeVoiceAcousticTokenizerEncoderStem(nn.Module): + def __init__(self, config): + super().__init__() + + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.channels, + out_channels=config.num_filters, + kernel_size=config.kernel_size, + cache_key="encoder_stem", + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, + hidden_size=config.num_filters, + layer_idx=layer_idx, + ) + for layer_idx in range(1, config.depths[0] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerEncoderLayer(nn.Module): + def __init__(self, config, stage_idx): + super().__init__() + + depth_idx = stage_idx + 1 # first depth is for stem layer + layer_idx = sum(depth + 1 for depth in config.depths[:depth_idx]) + intermediate_channels = int(config.num_filters * (2 ** (depth_idx))) + + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=int(config.num_filters * (2**stage_idx)), + out_channels=intermediate_channels, + kernel_size=int(config.downsampling_ratios[stage_idx] * 2), + cache_key=f"encoder_layer_{stage_idx}", + stride=config.downsampling_ratios[stage_idx], + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, hidden_size=intermediate_channels, layer_idx=layer_idx + offset + ) + for offset in range(1, config.depths[depth_idx] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +@auto_docstring +class VibeVoiceAcousticTokenizerPreTrainedModel(PreTrainedModel): + config: VibeVoiceAcousticTokenizerConfig + base_model_prefix = "vibevoice_acoustic_tokenizer" + main_input_name = "input_values" + _no_split_modules = ["VibeVoiceAcousticTokenizerEncoderModel", "VibeVoiceAcousticTokenizerDecoderModel"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, VibeVoiceAcousticTokenizerConvNext1dLayer): + init.constant_(module.gamma, self.config.layer_scale_init_value) + init.constant_(module.ffn_gamma, self.config.layer_scale_init_value) + + +class VibeVoiceAcousticTokenizerEncoderModel(VibeVoiceAcousticTokenizerPreTrainedModel): + config: VibeVoiceAcousticTokenizerEncoderConfig + + def __init__(self, config): + super().__init__(config) + + self.stem = VibeVoiceAcousticTokenizerEncoderStem(config) + self.conv_layers = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerEncoderLayer(config, stage_idx) + for stage_idx in range(len(config.downsampling_ratios)) + ] + ) + self.head = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=int(config.num_filters * (2 ** len(config.downsampling_ratios))), + out_channels=config.hidden_size, + kernel_size=config.kernel_size, + cache_key="encoder_head", + ) + self.post_init() + + def forward(self, hidden_states, padding_cache=None, use_cache=False, **kwargs): + if use_cache and padding_cache is None: + padding_cache = VibeVoiceAcousticTokenizerConv1dPaddingCache() + + hidden_states = self.stem(hidden_states, padding_cache=padding_cache) + for layer in self.conv_layers: + hidden_states = layer(hidden_states, padding_cache=padding_cache) + hidden_states = self.head(hidden_states, padding_cache=padding_cache) + latents = hidden_states.permute(0, 2, 1) + return VibeVoiceAcousticTokenizerEncoderOutput(latents=latents, padding_cache=padding_cache) + + +class VibeVoiceAcousticTokenizerDecoderStem(nn.Module): + def __init__(self, config): + super().__init__() + + intermediate_channels = int(config.num_filters * 2 ** (len(config.depths) - 1)) + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.hidden_size, + out_channels=intermediate_channels, + kernel_size=config.kernel_size, + cache_key="decoder_stem", + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, + hidden_size=intermediate_channels, + layer_idx=layer_idx, + ) + for layer_idx in range(1, config.depths[0] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerDecoderLayer(nn.Module): + def __init__(self, config, stage_idx): + super().__init__() + + depth_idx = stage_idx + 1 # first depth is for stem layer + layer_idx = sum(depth + 1 for depth in config.depths[:depth_idx]) + intermediate_channels = int(config.num_filters * (2 ** (len(config.depths) - 2 - stage_idx))) + + self.convtr = VibeVoiceAcousticTokenizerCausalConvTranspose1d( + in_channels=int(config.num_filters * (2 ** (len(config.depths) - 1 - stage_idx))), + out_channels=intermediate_channels, + kernel_size=int(config.upsampling_ratios[stage_idx] * 2), + cache_key=f"decoder_layer_{stage_idx}", + stride=config.upsampling_ratios[stage_idx], + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, hidden_size=intermediate_channels, layer_idx=layer_idx + offset + ) + for offset in range(1, config.depths[depth_idx] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.convtr(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerDecoderModel(VibeVoiceAcousticTokenizerPreTrainedModel): + config: VibeVoiceAcousticTokenizerDecoderConfig + + def __init__(self, config): + super().__init__(config) + + self.stem = VibeVoiceAcousticTokenizerDecoderStem(config) + self.conv_layers = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerDecoderLayer(config, stage_idx) + for stage_idx in range(len(config.upsampling_ratios)) + ] + ) + self.head = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.num_filters, + out_channels=config.channels, + kernel_size=config.kernel_size, + cache_key="decoder_head", + ) + self.post_init() + + def forward(self, hidden_states, padding_cache=None, use_cache=False, **kwargs): + if use_cache and padding_cache is None: + padding_cache = VibeVoiceAcousticTokenizerConv1dPaddingCache() + + hidden_states = self.stem(hidden_states, padding_cache=padding_cache) + for layer in self.conv_layers: + hidden_states = layer(hidden_states, padding_cache=padding_cache) + hidden_states = self.head(hidden_states, padding_cache=padding_cache) + return VibeVoiceAcousticTokenizerDecoderOutput(audio=hidden_states, padding_cache=padding_cache) + + +@auto_docstring( + custom_intro=""" + VibeVoice acoustic tokenizer with an encoder and decoder for continuous acoustic tokens. + """ +) +class VibeVoiceAcousticTokenizerModel(VibeVoiceAcousticTokenizerPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.encoder = AutoModel.from_config(config.encoder_config) + self.decoder = AutoModel.from_config(config.decoder_config) + self.post_init() + + @can_return_tuple + @auto_docstring + def encode(self, input_values, padding_cache=None, use_cache=None, sample=True): + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input audio waveform to be encoded into latent representation. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + sample (`bool`, *optional*): + Whether to sample from the VAE. If False, no noise is added. + """ + encoder_output = self.encoder(input_values, padding_cache=padding_cache, use_cache=use_cache) + + if sample: + noise_std = self.config.vae_std * torch.randn( + encoder_output.latents.shape[0], + device=encoder_output.latents.device, + dtype=encoder_output.latents.dtype, + ) + encoder_output.latents = encoder_output.latents + noise_std[:, None, None] * torch.randn_like( + encoder_output.latents + ) + return encoder_output + + @can_return_tuple + @auto_docstring + def decode(self, latents, padding_cache=None, use_cache=False): + r""" + latents (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input latent representation to be decoded back into audio. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + """ + latents = latents.permute(0, 2, 1) + return self.decoder(latents, padding_cache=padding_cache, use_cache=use_cache) + + @can_return_tuple + @auto_docstring + def forward(self, input_values, padding_cache=None, use_cache=False, sample=True, **kwargs): + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input audio waveform to be encoded into latent representation. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. Note only used by decoder. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + sample (`bool`, *optional*): + Whether to sample from the VAE latent distribution. If False, no noise is added to the latents. + """ + encoder_output = self.encode(input_values, use_cache=use_cache, sample=sample) + decoder_output = self.decode(encoder_output.latents, padding_cache=padding_cache, use_cache=use_cache) + return VibeVoiceAcousticTokenizerOutput( + audio=decoder_output.audio, + latents=encoder_output.latents, + padding_cache=decoder_output.padding_cache, + ) + + +__all__ = [ + "VibeVoiceAcousticTokenizerModel", + "VibeVoiceAcousticTokenizerEncoderModel", + "VibeVoiceAcousticTokenizerDecoderModel", + "VibeVoiceAcousticTokenizerPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modular_vibevoice_acoustic_tokenizer.py b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modular_vibevoice_acoustic_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..041409493062c0a1eed879f6978a9d56d29ee6f1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vibevoice_acoustic_tokenizer/modular_vibevoice_acoustic_tokenizer.py @@ -0,0 +1,499 @@ +# Copyright 2026 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 +from typing import Optional + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, can_return_tuple +from ..auto.modeling_auto import AutoModel +from ..llama.modeling_llama import LlamaRMSNorm +from ..voxtral_realtime.modeling_voxtral_realtime import VoxtralRealtimeConv1dPaddingCache +from .configuration_vibevoice_acoustic_tokenizer import ( + VibeVoiceAcousticTokenizerConfig, + VibeVoiceAcousticTokenizerDecoderConfig, + VibeVoiceAcousticTokenizerEncoderConfig, +) + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Decoded audio. + latents (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Projected latents (continuous representations for acoustic tokens) at the output of the encoder. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each decoder + layer that can be passed to subsequent forward calls. + """ + + audio: torch.FloatTensor | None = None + latents: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerEncoderOutput(ModelOutput): + r""" + latents (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Projected latents (continuous representations for acoustic tokens) at the output of the encoder. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each encoder + layer that can be passed to subsequent forward calls. + """ + + latents: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +@dataclass +@auto_docstring +class VibeVoiceAcousticTokenizerDecoderOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Decoded audio. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*, returned when `use_cache=True` is passed): + A [`VibeVoiceAcousticTokenizerConv1dPaddingCache`] instance containing cached convolution states for each decoder + layer that can be passed to subsequent forward calls. + """ + + audio: torch.FloatTensor | None = None + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None + + +class VibeVoiceAcousticTokenizerRMSNorm(LlamaRMSNorm): + pass + + +class VibeVoiceAcousticTokenizerFeedForward(nn.Module): + def __init__(self, config, hidden_size): + super().__init__() + self.linear1 = nn.Linear(hidden_size, config.ffn_expansion * hidden_size) + self.activation = ACT2FN[config.hidden_act] + self.linear2 = nn.Linear(config.ffn_expansion * hidden_size, hidden_size) + + def forward(self, hidden_states): + return self.linear2(self.activation(self.linear1(hidden_states))) + + +class VibeVoiceAcousticTokenizerConv1dPaddingCache(VoxtralRealtimeConv1dPaddingCache): + pass + + +# TODO: @eustlb, @ebezzam this should be latter factorized with other causalconv1d (e.g. VoxtralRealtimeCausalConv1d) +class VibeVoiceAcousticTokenizerCausalConv1d(nn.Module): + """Conv1d with built-in causal padding and optional streaming support through a cache.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + cache_key: str, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + ): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation, groups=groups) + self.causal_padding = (kernel_size - 1) * dilation - (stride - 1) + if self.causal_padding < 0: + raise ValueError( + f"Invalid causal padding {self.causal_padding} for kernel_size={kernel_size}, " + f"dilation={dilation}, stride={stride}." + ) + self.cache_key = cache_key + self.in_channels = in_channels + self.left_pad = self.causal_padding + + def forward( + self, + hidden_states: torch.Tensor, + padding_cache: VibeVoiceAcousticTokenizerConv1dPaddingCache | None = None, + ) -> torch.Tensor: + if padding_cache is not None: + hidden_states = padding_cache.update(hidden_states, self.cache_key, self) + else: + hidden_states = nn.functional.pad(hidden_states, (self.left_pad, 0)) + + return self.conv(hidden_states) + + +class VibeVoiceAcousticTokenizerCausalConvTranspose1d(nn.Module): + """ConvTranspose1d with built-in causal padding and optional streaming support through a cache.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + cache_key: str, + stride: int = 1, + ): + super().__init__() + self.convtr = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride) + + self.stride = stride + self.cache_key = cache_key + self.in_channels = in_channels + self.padding_total = kernel_size - stride + self.causal_padding = kernel_size - 1 + self.left_pad = self.causal_padding + + def forward( + self, + hidden_states: torch.Tensor, + padding_cache: Optional["VibeVoiceAcousticTokenizerConv1dPaddingCache"] = None, + ) -> torch.Tensor: + time_dim = hidden_states.shape[-1] + + if padding_cache is not None: + hidden_states = padding_cache.update(hidden_states, self.cache_key, self) + hidden_states = self.convtr(hidden_states) + + # Remove extra padding at the right side + if self.padding_total > 0: + hidden_states = hidden_states[..., : -self.padding_total] + + if padding_cache is not None: + # For first chunk return full output, for subsequent chunks return only new output + expected_new_output = time_dim * self.stride + if hidden_states.shape[2] >= expected_new_output: + hidden_states = hidden_states[:, :, -expected_new_output:] + return hidden_states + + +class VibeVoiceAcousticTokenizerConvNext1dLayer(nn.Module): + """ConvNeXt-like block adapted for 1D convolutions.""" + + def __init__(self, config, hidden_size, dilation=1, stride=1, layer_idx=None): + super().__init__() + + self.norm = VibeVoiceAcousticTokenizerRMSNorm(hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = VibeVoiceAcousticTokenizerRMSNorm(hidden_size, eps=config.rms_norm_eps) + self.ffn = VibeVoiceAcousticTokenizerFeedForward(config, hidden_size) + self.gamma = nn.Parameter(config.layer_scale_init_value * torch.ones(hidden_size), requires_grad=True) + self.ffn_gamma = nn.Parameter(config.layer_scale_init_value * torch.ones(hidden_size), requires_grad=True) + self.mixer = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=hidden_size, + out_channels=hidden_size, + kernel_size=config.kernel_size, + cache_key=f"convnext_layer_{layer_idx}", + groups=hidden_size, + dilation=dilation, + stride=stride, + ) + + def forward(self, hidden_states, padding_cache=None): + # mixer + residual = hidden_states + hidden_states = self.norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = self.mixer(hidden_states, padding_cache=padding_cache) + hidden_states = hidden_states * self.gamma.unsqueeze(-1) + hidden_states = residual + hidden_states + + # ffn + residual = hidden_states + hidden_states = self.ffn_norm(hidden_states.transpose(1, 2)) + hidden_states = self.ffn(hidden_states).transpose(1, 2) + hidden_states = hidden_states * self.ffn_gamma.unsqueeze(-1) + return residual + hidden_states + + +class VibeVoiceAcousticTokenizerEncoderStem(nn.Module): + def __init__(self, config): + super().__init__() + + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.channels, + out_channels=config.num_filters, + kernel_size=config.kernel_size, + cache_key="encoder_stem", + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, + hidden_size=config.num_filters, + layer_idx=layer_idx, + ) + for layer_idx in range(1, config.depths[0] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerEncoderLayer(nn.Module): + def __init__(self, config, stage_idx): + super().__init__() + + depth_idx = stage_idx + 1 # first depth is for stem layer + layer_idx = sum(depth + 1 for depth in config.depths[:depth_idx]) + intermediate_channels = int(config.num_filters * (2 ** (depth_idx))) + + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=int(config.num_filters * (2**stage_idx)), + out_channels=intermediate_channels, + kernel_size=int(config.downsampling_ratios[stage_idx] * 2), + cache_key=f"encoder_layer_{stage_idx}", + stride=config.downsampling_ratios[stage_idx], + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, hidden_size=intermediate_channels, layer_idx=layer_idx + offset + ) + for offset in range(1, config.depths[depth_idx] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +@auto_docstring +class VibeVoiceAcousticTokenizerPreTrainedModel(PreTrainedModel): + config: VibeVoiceAcousticTokenizerConfig + base_model_prefix = "vibevoice_acoustic_tokenizer" + main_input_name = "input_values" + _no_split_modules = ["VibeVoiceAcousticTokenizerEncoderModel", "VibeVoiceAcousticTokenizerDecoderModel"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, VibeVoiceAcousticTokenizerConvNext1dLayer): + init.constant_(module.gamma, self.config.layer_scale_init_value) + init.constant_(module.ffn_gamma, self.config.layer_scale_init_value) + + +class VibeVoiceAcousticTokenizerEncoderModel(VibeVoiceAcousticTokenizerPreTrainedModel): + config: VibeVoiceAcousticTokenizerEncoderConfig + + def __init__(self, config): + super().__init__(config) + + self.stem = VibeVoiceAcousticTokenizerEncoderStem(config) + self.conv_layers = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerEncoderLayer(config, stage_idx) + for stage_idx in range(len(config.downsampling_ratios)) + ] + ) + self.head = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=int(config.num_filters * (2 ** len(config.downsampling_ratios))), + out_channels=config.hidden_size, + kernel_size=config.kernel_size, + cache_key="encoder_head", + ) + self.post_init() + + def forward(self, hidden_states, padding_cache=None, use_cache=False, **kwargs): + if use_cache and padding_cache is None: + padding_cache = VibeVoiceAcousticTokenizerConv1dPaddingCache() + + hidden_states = self.stem(hidden_states, padding_cache=padding_cache) + for layer in self.conv_layers: + hidden_states = layer(hidden_states, padding_cache=padding_cache) + hidden_states = self.head(hidden_states, padding_cache=padding_cache) + latents = hidden_states.permute(0, 2, 1) + return VibeVoiceAcousticTokenizerEncoderOutput(latents=latents, padding_cache=padding_cache) + + +class VibeVoiceAcousticTokenizerDecoderStem(nn.Module): + def __init__(self, config): + super().__init__() + + intermediate_channels = int(config.num_filters * 2 ** (len(config.depths) - 1)) + self.conv = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.hidden_size, + out_channels=intermediate_channels, + kernel_size=config.kernel_size, + cache_key="decoder_stem", + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, + hidden_size=intermediate_channels, + layer_idx=layer_idx, + ) + for layer_idx in range(1, config.depths[0] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.conv(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerDecoderLayer(nn.Module): + def __init__(self, config, stage_idx): + super().__init__() + + depth_idx = stage_idx + 1 # first depth is for stem layer + layer_idx = sum(depth + 1 for depth in config.depths[:depth_idx]) + intermediate_channels = int(config.num_filters * (2 ** (len(config.depths) - 2 - stage_idx))) + + self.convtr = VibeVoiceAcousticTokenizerCausalConvTranspose1d( + in_channels=int(config.num_filters * (2 ** (len(config.depths) - 1 - stage_idx))), + out_channels=intermediate_channels, + kernel_size=int(config.upsampling_ratios[stage_idx] * 2), + cache_key=f"decoder_layer_{stage_idx}", + stride=config.upsampling_ratios[stage_idx], + ) + self.stage = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerConvNext1dLayer( + config, hidden_size=intermediate_channels, layer_idx=layer_idx + offset + ) + for offset in range(1, config.depths[depth_idx] + 1) + ] + ) + + def forward(self, hidden_states, padding_cache=None): + hidden_states = self.convtr(hidden_states, padding_cache=padding_cache) + for block in self.stage: + hidden_states = block(hidden_states, padding_cache=padding_cache) + return hidden_states + + +class VibeVoiceAcousticTokenizerDecoderModel(VibeVoiceAcousticTokenizerPreTrainedModel): + config: VibeVoiceAcousticTokenizerDecoderConfig + + def __init__(self, config): + super().__init__(config) + + self.stem = VibeVoiceAcousticTokenizerDecoderStem(config) + self.conv_layers = nn.ModuleList( + [ + VibeVoiceAcousticTokenizerDecoderLayer(config, stage_idx) + for stage_idx in range(len(config.upsampling_ratios)) + ] + ) + self.head = VibeVoiceAcousticTokenizerCausalConv1d( + in_channels=config.num_filters, + out_channels=config.channels, + kernel_size=config.kernel_size, + cache_key="decoder_head", + ) + self.post_init() + + def forward(self, hidden_states, padding_cache=None, use_cache=False, **kwargs): + if use_cache and padding_cache is None: + padding_cache = VibeVoiceAcousticTokenizerConv1dPaddingCache() + + hidden_states = self.stem(hidden_states, padding_cache=padding_cache) + for layer in self.conv_layers: + hidden_states = layer(hidden_states, padding_cache=padding_cache) + hidden_states = self.head(hidden_states, padding_cache=padding_cache) + return VibeVoiceAcousticTokenizerDecoderOutput(audio=hidden_states, padding_cache=padding_cache) + + +@auto_docstring( + custom_intro=""" + VibeVoice acoustic tokenizer with an encoder and decoder for continuous acoustic tokens. + """ +) +class VibeVoiceAcousticTokenizerModel(VibeVoiceAcousticTokenizerPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.encoder = AutoModel.from_config(config.encoder_config) + self.decoder = AutoModel.from_config(config.decoder_config) + self.post_init() + + @can_return_tuple + @auto_docstring + def encode(self, input_values, padding_cache=None, use_cache=None, sample=True): + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input audio waveform to be encoded into latent representation. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + sample (`bool`, *optional*): + Whether to sample from the VAE. If False, no noise is added. + """ + encoder_output = self.encoder(input_values, padding_cache=padding_cache, use_cache=use_cache) + + if sample: + noise_std = self.config.vae_std * torch.randn( + encoder_output.latents.shape[0], + device=encoder_output.latents.device, + dtype=encoder_output.latents.dtype, + ) + encoder_output.latents = encoder_output.latents + noise_std[:, None, None] * torch.randn_like( + encoder_output.latents + ) + return encoder_output + + @can_return_tuple + @auto_docstring + def decode(self, latents, padding_cache=None, use_cache=False): + r""" + latents (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input latent representation to be decoded back into audio. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + """ + latents = latents.permute(0, 2, 1) + return self.decoder(latents, padding_cache=padding_cache, use_cache=use_cache) + + @can_return_tuple + @auto_docstring + def forward(self, input_values, padding_cache=None, use_cache=False, sample=True, **kwargs): + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`): + Input audio waveform to be encoded into latent representation. + padding_cache (`VibeVoiceAcousticTokenizerConv1dPaddingCache`, *optional*): + Cache object for streaming mode to maintain convolution states across layers. Note only used by decoder. + use_cache (`bool`, *optional*): + Whether to use caching for convolution states. + sample (`bool`, *optional*): + Whether to sample from the VAE latent distribution. If False, no noise is added to the latents. + """ + encoder_output = self.encode(input_values, use_cache=use_cache, sample=sample) + decoder_output = self.decode(encoder_output.latents, padding_cache=padding_cache, use_cache=use_cache) + return VibeVoiceAcousticTokenizerOutput( + audio=decoder_output.audio, + latents=encoder_output.latents, + padding_cache=decoder_output.padding_cache, + ) + + +__all__ = [ + "VibeVoiceAcousticTokenizerModel", + "VibeVoiceAcousticTokenizerEncoderModel", + "VibeVoiceAcousticTokenizerDecoderModel", + "VibeVoiceAcousticTokenizerPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/vilt/__init__.py b/third_party/transformers/src/transformers/models/vilt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..962b2785e8cfd5cec4efad46a4994bd794fe956c --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_vilt import * + from .feature_extraction_vilt import * + from .image_processing_pil_vilt import * + from .image_processing_vilt import * + from .modeling_vilt import * + from .processing_vilt 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/vilt/configuration_vilt.py b/third_party/transformers/src/transformers/models/vilt/configuration_vilt.py new file mode 100644 index 0000000000000000000000000000000000000000..4227dfd02f3464c8d2472295e61aa2a97a9c18dd --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/configuration_vilt.py @@ -0,0 +1,82 @@ +# 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. +"""VilT model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="dandelin/vilt-b32-mlm") +@strict +class ViltConfig(PreTrainedConfig): + r""" + modality_type_vocab_size (`int`, *optional*, defaults to 2): + The vocabulary size of the modalities passed when calling [`ViltModel`]. This is used after concatenating the + embeddings of the text and image modalities. + max_image_length (`int`, *optional*, defaults to -1): + The maximum number of patches to take as input for the Transformer encoder. If set to a positive integer, + the encoder will sample `max_image_length` patches at maximum. If set to -1, will not be taken into + account. + num_images (`int`, *optional*, defaults to -1): + The number of images to use for natural language visual reasoning. If set to a positive integer, will be + used by [`ViltForImagesAndTextClassification`] for defining the classifier head. + + Example: + + ```python + >>> from transformers import ViLTModel, ViLTConfig + + >>> # Initializing a ViLT dandelin/vilt-b32-mlm style configuration + >>> configuration = ViLTConfig() + + >>> # Initializing a model from the dandelin/vilt-b32-mlm style configuration + >>> model = ViLTModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "vilt" + + vocab_size: int = 30522 + type_vocab_size: int = 2 + modality_type_vocab_size: int = 2 + max_position_embeddings: int = 40 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.0 + attention_probs_dropout_prob: float | int = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + image_size: int | list[int] | tuple[int, int] = 384 + patch_size: int | list[int] | tuple[int, int] = 32 + num_channels: int = 3 + qkv_bias: bool = True + max_image_length: int = -1 + tie_word_embeddings: bool = True + num_images: int = -1 + pad_token_id: int | None = None + + def __post_init__(self, **kwargs): + kwargs.pop("tie_word_embeddings", None) + self.tie_word_embeddings = True # force it + super().__post_init__(**kwargs) + + +__all__ = ["ViltConfig"] diff --git a/third_party/transformers/src/transformers/models/vilt/convert_vilt_original_to_pytorch.py b/third_party/transformers/src/transformers/models/vilt/convert_vilt_original_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..3102572b79bb3553245c0bda9fa9f24504c0fbf1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/convert_vilt_original_to_pytorch.py @@ -0,0 +1,303 @@ +# 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 ViLT checkpoints from the original Github repository.""" + +import argparse +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import ( + BertTokenizer, + ViltConfig, + ViltForImageAndTextRetrieval, + ViltForImagesAndTextClassification, + ViltForMaskedLM, + ViltForQuestionAnswering, + ViltImageProcessor, + ViltProcessor, +) +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def create_rename_keys(config, vqa_model=False, nlvr_model=False, irtr_model=False): + rename_keys = [] + for i in range(config.num_hidden_layers): + # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms + rename_keys.append((f"transformer.blocks.{i}.norm1.weight", f"vilt.encoder.layer.{i}.layernorm_before.weight")) + rename_keys.append((f"transformer.blocks.{i}.norm1.bias", f"vilt.encoder.layer.{i}.layernorm_before.bias")) + rename_keys.append( + (f"transformer.blocks.{i}.attn.proj.weight", f"vilt.encoder.layer.{i}.attention.output.dense.weight") + ) + rename_keys.append( + (f"transformer.blocks.{i}.attn.proj.bias", f"vilt.encoder.layer.{i}.attention.output.dense.bias") + ) + rename_keys.append((f"transformer.blocks.{i}.norm2.weight", f"vilt.encoder.layer.{i}.layernorm_after.weight")) + rename_keys.append((f"transformer.blocks.{i}.norm2.bias", f"vilt.encoder.layer.{i}.layernorm_after.bias")) + rename_keys.append( + (f"transformer.blocks.{i}.mlp.fc1.weight", f"vilt.encoder.layer.{i}.intermediate.dense.weight") + ) + rename_keys.append((f"transformer.blocks.{i}.mlp.fc1.bias", f"vilt.encoder.layer.{i}.intermediate.dense.bias")) + rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.weight", f"vilt.encoder.layer.{i}.output.dense.weight")) + rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.bias", f"vilt.encoder.layer.{i}.output.dense.bias")) + + # embeddings + rename_keys.extend( + [ + # text embeddings + ("text_embeddings.word_embeddings.weight", "vilt.embeddings.text_embeddings.word_embeddings.weight"), + ( + "text_embeddings.position_embeddings.weight", + "vilt.embeddings.text_embeddings.position_embeddings.weight", + ), + ("text_embeddings.position_ids", "vilt.embeddings.text_embeddings.position_ids"), + ( + "text_embeddings.token_type_embeddings.weight", + "vilt.embeddings.text_embeddings.token_type_embeddings.weight", + ), + ("text_embeddings.LayerNorm.weight", "vilt.embeddings.text_embeddings.LayerNorm.weight"), + ("text_embeddings.LayerNorm.bias", "vilt.embeddings.text_embeddings.LayerNorm.bias"), + # patch embeddings + ("transformer.cls_token", "vilt.embeddings.cls_token"), + ("transformer.patch_embed.proj.weight", "vilt.embeddings.patch_embeddings.projection.weight"), + ("transformer.patch_embed.proj.bias", "vilt.embeddings.patch_embeddings.projection.bias"), + ("transformer.pos_embed", "vilt.embeddings.position_embeddings"), + # token type embeddings + ("token_type_embeddings.weight", "vilt.embeddings.token_type_embeddings.weight"), + ] + ) + + # final layernorm + pooler + rename_keys.extend( + [ + ("transformer.norm.weight", "vilt.layernorm.weight"), + ("transformer.norm.bias", "vilt.layernorm.bias"), + ("pooler.dense.weight", "vilt.pooler.dense.weight"), + ("pooler.dense.bias", "vilt.pooler.dense.bias"), + ] + ) + + # classifier head(s) + if vqa_model: + # classification head + rename_keys.extend( + [ + ("vqa_classifier.0.weight", "classifier.0.weight"), + ("vqa_classifier.0.bias", "classifier.0.bias"), + ("vqa_classifier.1.weight", "classifier.1.weight"), + ("vqa_classifier.1.bias", "classifier.1.bias"), + ("vqa_classifier.3.weight", "classifier.3.weight"), + ("vqa_classifier.3.bias", "classifier.3.bias"), + ] + ) + elif nlvr_model: + # classification head + rename_keys.extend( + [ + ("nlvr2_classifier.0.weight", "classifier.0.weight"), + ("nlvr2_classifier.0.bias", "classifier.0.bias"), + ("nlvr2_classifier.1.weight", "classifier.1.weight"), + ("nlvr2_classifier.1.bias", "classifier.1.bias"), + ("nlvr2_classifier.3.weight", "classifier.3.weight"), + ("nlvr2_classifier.3.bias", "classifier.3.bias"), + ] + ) + else: + pass + + return rename_keys + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config): + for i in range(config.num_hidden_layers): + prefix = "vilt." + # read in weights + bias of input projection layer (in timm, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"transformer.blocks.{i}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"transformer.blocks.{i}.attn.qkv.bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[ + : config.hidden_size, : + ] + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size] + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + config.hidden_size : config.hidden_size * 2, : + ] + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[ + config.hidden_size : config.hidden_size * 2 + ] + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[ + -config.hidden_size :, : + ] + state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :] + + +def remove_classification_head_(state_dict): + ignore_keys = ["head.weight", "head.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +@torch.no_grad() +def convert_vilt_checkpoint(checkpoint_url, pytorch_dump_folder_path): + """ + Copy/paste/tweak model's weights to our ViLT structure. + """ + + # define configuration and initialize HuggingFace model + config = ViltConfig(image_size=384, patch_size=32, tie_word_embeddings=False) + mlm_model = False + vqa_model = False + nlvr_model = False + irtr_model = False + if "vqa" in checkpoint_url: + vqa_model = True + config.num_labels = 3129 + repo_id = "huggingface/label-files" + filename = "vqa2-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()} + model = ViltForQuestionAnswering(config) + elif "nlvr" in checkpoint_url: + nlvr_model = True + config.num_labels = 2 + config.id2label = {0: "False", 1: "True"} + config.label2id = {v: k for k, v in config.id2label.items()} + config.modality_type_vocab_size = 3 + model = ViltForImagesAndTextClassification(config) + elif "irtr" in checkpoint_url: + irtr_model = True + model = ViltForImageAndTextRetrieval(config) + elif "mlm_itm" in checkpoint_url: + mlm_model = True + model = ViltForMaskedLM(config) + else: + raise ValueError("Unknown model type") + + # load state_dict of original model, remove and rename some keys + state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["state_dict"] + rename_keys = create_rename_keys(config, vqa_model, nlvr_model, irtr_model) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + read_in_q_k_v(state_dict, config) + if mlm_model or irtr_model: + ignore_keys = ["itm_score.fc.weight", "itm_score.fc.bias"] + for k in ignore_keys: + state_dict.pop(k, None) + + # load state dict into HuggingFace model + model.eval() + if mlm_model: + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + assert missing_keys == ["mlm_score.decoder.bias"] + else: + model.load_state_dict(state_dict) + + # Define processor + image_processor = ViltImageProcessor(size=384) + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + processor = ViltProcessor(image_processor, tokenizer) + + # Forward pass on example inputs (image + text) + if nlvr_model: + url = "https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg" + with httpx.stream("GET", url) as response: + image1 = Image.open(BytesIO(response.read())) + image2 = Image.open(BytesIO(response.read())) + text = ( + "The left image contains twice the number of dogs as the right image, and at least two dogs in total are" + " standing." + ) + encoding_1 = processor(image1, text, return_tensors="pt") + encoding_2 = processor(image2, text, return_tensors="pt") + outputs = model( + input_ids=encoding_1.input_ids, + pixel_values=encoding_1.pixel_values, + pixel_values_2=encoding_2.pixel_values, + ) + else: + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + if mlm_model: + text = "a bunch of [MASK] laying on a [MASK]." + else: + text = "How many cats are there?" + encoding = processor(image, text, return_tensors="pt") + outputs = model(**encoding) + + # Verify outputs + if mlm_model: + expected_shape = torch.Size([1, 11, 30522]) + expected_slice = torch.tensor([-12.5061, -12.5123, -12.5174]) + assert outputs.logits.shape == expected_shape + assert torch.allclose(outputs.logits[0, 0, :3], expected_slice, atol=1e-4) + + # verify masked token prediction equals "cats" + predicted_id = outputs.logits[0, 4, :].argmax(-1).item() + assert tokenizer.decode([predicted_id]) == "cats" + elif vqa_model: + expected_shape = torch.Size([1, 3129]) + expected_slice = torch.tensor([-15.9495, -18.1472, -10.3041]) + assert torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4) + assert outputs.logits.shape == expected_shape + assert torch.allclose(outputs.logits[0, 0, :3], expected_slice, atol=1e-4) + + # verify vqa prediction equals "2" + predicted_idx = outputs.logits.argmax(-1).item() + assert model.config.id2label[predicted_idx] == "2" + elif nlvr_model: + expected_shape = torch.Size([1, 2]) + expected_slice = torch.tensor([-2.8721, 2.1291]) + assert torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4) + assert outputs.logits.shape == expected_shape + + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + 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 __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--checkpoint_url", + default="https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt", + type=str, + help="URL of the checkpoint you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory." + ) + + args = parser.parse_args() + convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path) diff --git a/third_party/transformers/src/transformers/models/vilt/image_processing_pil_vilt.py b/third_party/transformers/src/transformers/models/vilt/image_processing_pil_vilt.py new file mode 100644 index 0000000000000000000000000000000000000000..88948e0ef538bb5ebbb659ae9bf87bb195d15a09 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/image_processing_pil_vilt.py @@ -0,0 +1,245 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for Vilt.""" + +from collections.abc import Iterable +from typing import Any + +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, + ChannelDimension, + PILImageResampling, + SizeDict, + get_image_size, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs +from ...utils import ( + TensorType, + auto_docstring, +) + + +# Set maximum size based on the typical aspect ratio of the COCO dataset +MAX_LONGER_EDGE = 1333 +MAX_SHORTER_EDGE = 800 + + +def max_across_indices(values: Iterable[Any]) -> list[Any]: + """ + Return the maximum value across all indices of an iterable of values. + """ + return [max(values_i) for values_i in zip(*values)] + + +def make_pixel_mask( + image: np.ndarray, output_size: tuple[int, int], input_data_format: str | ChannelDimension | None = None +) -> np.ndarray: + """ + Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + + Args: + image (`np.ndarray`): + Image to make the pixel mask for. + output_size (`tuple[int, int]`): + Output size of the mask. + """ + input_height, input_width = get_image_size(image, channel_dim=input_data_format) + mask = np.zeros(output_size, dtype=np.int64) + mask[:input_height, :input_width] = 1 + return mask + + +def get_resize_output_image_size( + input_image: np.ndarray, + shorter: int = 800, + longer: int = 1333, + size_divisor: int = 32, + input_data_format: str | ChannelDimension | None = None, +) -> tuple[int, int]: + input_height, input_width = get_image_size(input_image, input_data_format) + min_size, max_size = shorter, longer + + scale = min_size / min(input_height, input_width) + + if input_height < input_width: + new_height = min_size + new_width = scale * input_width + else: + new_height = scale * input_height + new_width = min_size + + if max(new_height, new_width) > max_size: + scale = max_size / max(new_height, new_width) + new_height = scale * new_height + new_width = scale * new_width + + new_height, new_width = int(new_height + 0.5), int(new_width + 0.5) + new_height = new_height // size_divisor * size_divisor + new_width = new_width // size_divisor * size_divisor + + return new_height, new_width + + +# Adapted from transformers.models.vilt.image_processing_vilt.ViltImageProcessorKwargs +class ViltImageProcessorKwargs(ImagesKwargs, total=False): + r""" + size_divisor (`int`, *optional*, defaults to `self.size_divisor`): + The size by which to make sure both the height and width can be divided. Only has an effect if `do_resize` + is set to `True`. + """ + + size_divisor: int + + +@auto_docstring +class ViltImageProcessorPil(PilBackend): + valid_kwargs = ViltImageProcessorKwargs + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"shortest_edge": 384} + do_resize = True + do_rescale = True + do_normalize = True + size_divisor = 32 + do_pad = True + default_to_square = False + model_input_names = ["pixel_values", "pixel_mask"] + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: "PILImageResampling | None" = None, + size_divisor: int | None = None, + ) -> np.ndarray: + """ + Resize an image to specified size. + + Args: + image (`np.ndarray`): Image to resize. + size (`SizeDict`): Size dictionary with shortest_edge key. + resample (`PILImageResampling | int`, *optional*): Interpolation method to use. + size_divisor (`int`, *optional*): Value to ensure height/width are divisible by. + + Returns: + `np.ndarray`: Resized image. + """ + if not hasattr(size, "shortest_edge") or size.shortest_edge is None: + raise ValueError(f"The `size` dictionary must contain the key `shortest_edge`. Got {size}") + shorter = size.shortest_edge + longer = int(MAX_LONGER_EDGE / MAX_SHORTER_EDGE * shorter) + output_size = get_resize_output_image_size( + image, + shorter=shorter, + longer=longer, + size_divisor=size_divisor or self.size_divisor, + input_data_format=ChannelDimension.FIRST, + ) + + return super().resize(image, SizeDict(height=output_size[0], width=output_size[1]), resample=resample) + + def _pad_batch( + self, + images: list[np.ndarray], + return_tensors: str | TensorType | None, + ) -> tuple: + """ + Pad a batch of images to the same size based on the maximum dimensions. + + Args: + images (`list[np.ndarray]`): List of images to pad. + return_tensors (`str` or `TensorType`, *optional*): The type of tensors to return. + + Returns: + `tuple`: Tuple containing padded images and pixel masks. + """ + # Calculate global maximum dimensions across all images + max_size = get_max_height_width(images, input_data_format=ChannelDimension.FIRST) + + padded_images = [] + pixel_masks = [] + + for image in images: + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + needs_padding = input_height != max_size[0] or input_width != max_size[1] + + if needs_padding: + pad_bottom = max_size[0] - input_height + pad_right = max_size[1] - input_width + padding = ((0, pad_bottom), (0, pad_right)) + + padded_image = pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=0, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + pixel_mask = make_pixel_mask(image, max_size, input_data_format=ChannelDimension.FIRST) + else: + padded_image = image + pixel_mask = np.ones(max_size, dtype=np.int64) + + padded_images.append(padded_image) + pixel_masks.append(pixel_mask) + + return padded_images, pixel_masks + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + return_tensors: str | TensorType | None, + size_divisor: int | None = None, + **kwargs, + ) -> BatchFeature: + processed_images = [] + for image in images: + if do_resize: + image = self.resize(image, size, resample, size_divisor) + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + + # Handle padding if required + data = {} + if do_pad: + pixel_values, pixel_mask = self._pad_batch(processed_images, return_tensors) + data = {"pixel_values": pixel_values, "pixel_mask": pixel_mask} + else: + data = {"pixel_values": processed_images} + + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["ViltImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/vilt/image_processing_vilt.py b/third_party/transformers/src/transformers/models/vilt/image_processing_vilt.py new file mode 100644 index 0000000000000000000000000000000000000000..488691aef0f28154630ec9b3ac36d75253592361 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/image_processing_vilt.py @@ -0,0 +1,234 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for Vilt.""" + +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, + PILImageResampling, + SizeDict, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs +from ...utils import ( + TensorType, + auto_docstring, +) + + +# Set maximum size based on the typical aspect ratio of the COCO dataset +MAX_LONGER_EDGE = 1333 +MAX_SHORTER_EDGE = 800 + + +class ViltImageProcessorKwargs(ImagesKwargs, total=False): + r""" + size_divisor (`int`, *optional*, defaults to `self.size_divisor`): + The size by which to make sure both the height and width can be divided. Only has an effect if `do_resize` + is set to `True`. + """ + + size_divisor: int + + +@auto_docstring +class ViltImageProcessor(TorchvisionBackend): + valid_kwargs = ViltImageProcessorKwargs + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"shortest_edge": 384} + do_resize = True + do_rescale = True + do_normalize = True + size_divisor = 32 + do_pad = True + default_to_square = False + model_input_names = ["pixel_values", "pixel_mask"] + + def resize( + self, + images: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + size_divisor: int | None = None, + ) -> "torch.Tensor": + """ + Resize an image or batch of images to specified size. + + Args: + images (`torch.Tensor`): Image or batch of images to resize. + size (`SizeDict`): Size dictionary with shortest_edge key. + resample (`PILImageResampling | tvF.InterpolationMode | int`, *optional*): Interpolation method to use. + size_divisor (`int`, *optional*): Value to ensure height/width are divisible by. + + Returns: + `torch.Tensor`: Resized image or batch of images. + """ + + # Resize with aspect ratio preservation + shorter = size.shortest_edge + longer = int(MAX_LONGER_EDGE / MAX_SHORTER_EDGE * shorter) + + heights = images.shape[-2] + widths = images.shape[-1] + + # Determine the new dimensions + if heights < widths: + new_heights = shorter + new_widths = widths * (shorter / heights) + else: + new_heights = heights * (shorter / widths) + new_widths = shorter + + # Check if the longer side exceeds max size + if max(new_heights, new_widths) > longer: + scale = longer / max(new_heights, new_widths) + new_heights = new_heights * scale + new_widths = new_widths * scale + + new_heights = int(new_heights + 0.5) + new_widths = int(new_widths + 0.5) + + # Make dimensions divisible by size_divisor + if size_divisor is not None: + new_heights = new_heights // size_divisor * size_divisor + new_widths = new_widths // size_divisor * size_divisor + + # Resize the image + return super().resize(images, SizeDict(height=new_heights, width=new_widths), resample=resample) + + def _pad_batch( + self, + images: list["torch.Tensor"], + return_tensors: str | TensorType | None, + disable_grouping: bool | None, + ) -> tuple: + """ + Pad a batch of images to the same size based on the maximum dimensions. + + Args: + images (`list[torch.Tensor]`): List of images to pad. + return_tensors (`str` or `TensorType`, *optional*): The type of tensors to return. + + Returns: + `tuple`: Tuple containing padded images and pixel masks. + """ + # Calculate global maximum dimensions across all images + max_size = get_max_height_width(images) + + # Group images by shape before padding + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + processed_images = {} + processed_masks = {} + + for shape, stacked_images in grouped_images.items(): + # Create mask template for efficient masking + if return_tensors == "pt" and len(stacked_images) > 0: + device = stacked_images.device + mask_template = torch.zeros(max_size, dtype=torch.int64, device=device) + + original_size = stacked_images.shape[-2:] + needs_padding = original_size[0] != max_size[0] or original_size[1] != max_size[1] + + if needs_padding: + padding_bottom = max_size[0] - original_size[0] + padding_right = max_size[1] - original_size[1] + padding = [0, 0, padding_right, padding_bottom] + + padded_images = tvF.pad(stacked_images, padding, fill=0) + pixel_mask = mask_template.clone() + pixel_mask[: original_size[0], : original_size[1]].fill_(1) + pixel_masks = pixel_mask.unsqueeze(0).repeat(stacked_images.shape[0], 1, 1) + else: + padded_images = stacked_images + pixel_masks = torch.ones( + (stacked_images.shape[0], max_size[0], max_size[1]), + dtype=torch.int64, + device=stacked_images.device, + ) + + # Store processed group + processed_images[shape] = padded_images + processed_masks[shape] = pixel_masks + + # Reorder images back to original order + padded_images = reorder_images(processed_images, grouped_images_index) + pixel_masks = reorder_images(processed_masks, grouped_images_index) + + return padded_images, pixel_masks + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + size_divisor: int | None = 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 = {} + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size, resample, size_divisor) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + # 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) + + # Handle padding if required + data = {} + if do_pad: + pixel_values, pixel_mask = self._pad_batch( + processed_images, return_tensors, disable_grouping=disable_grouping + ) + data = {"pixel_values": pixel_values, "pixel_mask": pixel_mask} + else: + # If no padding, just return the processed images + if return_tensors == "pt": + processed_images = torch.stack(processed_images) + data = {"pixel_values": processed_images} + + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["ViltImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/vilt/modeling_vilt.py b/third_party/transformers/src/transformers/models/vilt/modeling_vilt.py new file mode 100644 index 0000000000000000000000000000000000000000..02b9a9f3f2ebd0f7e5db13d7a0e4bfd4d0bb708c --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/modeling_vilt.py @@ -0,0 +1,1297 @@ +# Copyright 2022 NAVER AI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch ViLT model.""" + +import collections.abc +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 ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_vilt import ViltConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`ViltForImagesAndTextClassification`]. + """ +) +class ViltForImagesAndTextClassificationOutput(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). + hidden_states (`list[tuple(torch.FloatTensor)]`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + List of tuples of `torch.FloatTensor` (one for each image-text pair, each tuple containing the output of + the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: list[tuple[torch.FloatTensor]] | None = None + attentions: list[tuple[torch.FloatTensor]] | None = None + + +class ViltEmbeddings(nn.Module): + """ + Construct the text and patch embeddings. + + Text embeddings are equivalent to BERT embeddings. + + Patch embeddings are equivalent to ViT embeddings. + """ + + def __init__(self, config): + super().__init__() + + # text embeddings + self.text_embeddings = TextEmbeddings(config) + # patch embeddings + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.patch_embeddings = ViltPatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size)) + # modality type (text/patch) embeddings + self.token_type_embeddings = nn.Embedding(config.modality_type_vocab_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def visual_embed(self, pixel_values, pixel_mask, max_image_length=200): + _, _, ph, pw = self.patch_embeddings.projection.weight.shape + + x = self.patch_embeddings(pixel_values) + x_mask = pixel_mask[:, None, :, :].float() + x_mask = nn.functional.interpolate(x_mask, size=(x.shape[2], x.shape[3])).long() + x_h = x_mask[:, 0].sum(dim=1)[:, 0] + x_w = x_mask[:, 0].sum(dim=2)[:, 0] + + batch_size, num_channels, height, width = x.shape + patch_dim = self.config.image_size // self.config.patch_size + spatial_pos = self.position_embeddings[:, 1:, :].transpose(1, 2).view(1, num_channels, patch_dim, patch_dim) + pos_embed = torch.cat( + [ + nn.functional.pad( + nn.functional.interpolate( + spatial_pos, + size=(h, w), + mode="bilinear", + align_corners=True, + ), + (0, width - w, 0, height - h), + ) + for h, w in zip(x_h, x_w) + ], + dim=0, + ) + + pos_embed = pos_embed.flatten(2).transpose(1, 2) + x = x.flatten(2).transpose(1, 2) + # Set `device` here, otherwise `patch_index` will always be on `CPU` and will fail near the end for torch>=1.13 + patch_index = torch.stack( + torch.meshgrid(torch.arange(x_mask.shape[-2]), torch.arange(x_mask.shape[-1]), indexing="ij"), dim=-1 + ).to(device=x_mask.device) + patch_index = patch_index[None, None, :, :, :] + patch_index = patch_index.expand(x_mask.shape[0], x_mask.shape[1], -1, -1, -1) + patch_index = patch_index.flatten(1, 3) + x_mask = x_mask.flatten(1) + + if max_image_length < 0 or max_image_length is None or not isinstance(max_image_length, int): + # suppose aug is 800 x 1333, then, maximum effective res is 800 x 1333 (if one side gets bigger, the other will be constrained and be shrunk) + # (800 // self.patch_size) * (1333 // self.patch_size) is the maximum number of patches that single image can get. + # if self.patch_size = 32, 25 * 41 = 1025 + # if res is 384 x 640, 12 * 20 = 240 + effective_resolution = x_h * x_w + max_image_length = effective_resolution.max() + else: + effective_resolution = x_h * x_w + max_image_length = min(effective_resolution.max(), max_image_length) + + valid_idx = x_mask.nonzero(as_tuple=False) + non_valid_idx = (1 - x_mask).nonzero(as_tuple=False) + unique_rows = valid_idx[:, 0].unique() + valid_row_idx = [valid_idx[valid_idx[:, 0] == u] for u in unique_rows] + non_valid_row_idx = [non_valid_idx[non_valid_idx[:, 0] == u] for u in unique_rows] + + valid_nums = [v.size(0) for v in valid_row_idx] + non_valid_nums = [v.size(0) for v in non_valid_row_idx] + pad_nums = [max_image_length - v for v in valid_nums] + + select = [] + for i, (v, nv, p) in enumerate(zip(valid_nums, non_valid_nums, pad_nums)): + if p <= 0: + valid_choice = torch.multinomial(torch.ones(v).float(), max_image_length) + select.append(valid_row_idx[i][valid_choice]) + else: + pad_choice = torch.multinomial(torch.ones(nv).float(), p, replacement=True) + select.append(torch.cat([valid_row_idx[i], non_valid_row_idx[i][pad_choice]], dim=0)) + + select = torch.cat(select, dim=0) + x = x[select[:, 0], select[:, 1]].view(batch_size, -1, num_channels) + x_mask = x_mask[select[:, 0], select[:, 1]].view(batch_size, -1) + # `patch_index` should be on the same device as `select`, which is ensured at definition time. + patch_index = patch_index[select[:, 0], select[:, 1]].view(batch_size, -1, 2) + pos_embed = pos_embed[select[:, 0], select[:, 1]].view(batch_size, -1, num_channels) + + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + pos_embed = torch.cat( + (self.position_embeddings[:, 0, :][:, None, :].expand(batch_size, -1, -1), pos_embed), dim=1 + ) + x = x + pos_embed + x = self.dropout(x) + + x_mask = torch.cat([torch.ones(x_mask.shape[0], 1).to(x_mask), x_mask], dim=1) + + return x, x_mask, (patch_index, (height, width)) + + def forward( + self, + input_ids, + attention_mask, + token_type_ids, + pixel_values, + pixel_mask, + inputs_embeds, + image_embeds, + image_token_type_idx=1, + ): + # PART 1: text embeddings + text_embeds = self.text_embeddings( + input_ids=input_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds + ) + + # PART 2: patch embeddings (with interpolated position encodings) + if image_embeds is None: + image_embeds, image_masks, patch_index = self.visual_embed( + pixel_values, pixel_mask, max_image_length=self.config.max_image_length + ) + else: + image_masks = pixel_mask.flatten(1) + + # PART 3: add modality type embeddings + # 0 indicates text, 1 indicates image, 2 is optionally used when a second image is provided (NLVR2) + if image_token_type_idx is None: + image_token_type_idx = 1 + text_embeds = text_embeds + self.token_type_embeddings( + torch.zeros_like(attention_mask, dtype=torch.long, device=text_embeds.device) + ) + image_embeds = image_embeds + self.token_type_embeddings( + torch.full_like(image_masks, image_token_type_idx, dtype=torch.long, device=text_embeds.device) + ) + + # PART 4: concatenate + embeddings = torch.cat([text_embeds, image_embeds], dim=1) + masks = torch.cat([attention_mask, image_masks], dim=1) + + return embeddings, masks + + +class TextEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + self.register_buffer( + "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False + ) + + def forward(self, input_ids=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 ViltPatchEmbeddings(nn.Module): + """ + Image to Patch Embedding. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values): + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + target_dtype = self.projection.weight.dtype + x = self.projection(pixel_values.to(dtype=target_dtype)) + return x + + +class ViltSelfAttention(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.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + batch_size, seq_length, _ = hidden_states.shape + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in BertModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + attention_probs = nn.Softmax(dim=-1)(attention_scores) + + # 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 + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Vilt +class ViltSelfOutput(nn.Module): + """ + The residual connection is defined in ViltLayer instead of here (as is the case with other models), due to the + layernorm applied before each block. + """ + + def __init__(self, config: ViltConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class ViltAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.attention = ViltSelfAttention(config) + self.output = ViltSelfOutput(config) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + self_outputs = self.attention(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.vit.modeling_vit.ViTIntermediate with ViT->Vilt +class ViltIntermediate(nn.Module): + def __init__(self, config: ViltConfig): + 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.vit.modeling_vit.ViTOutput with ViT->Vilt +class ViltOutput(nn.Module): + def __init__(self, config: ViltConfig): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +class ViltLayer(GradientCheckpointingLayer): + """This corresponds to the Block class in the timm implementation.""" + + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = ViltAttention(config) + self.intermediate = ViltIntermediate(config) + self.output = ViltOutput(config) + self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + self_attention_outputs = self.attention( + self.layernorm_before(hidden_states), # in ViLT, layernorm is applied before self-attention + attention_mask, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + # first residual connection + hidden_states = attention_output + hidden_states.to(attention_output.device) + + # in ViLT, layernorm is also applied after self-attention + layer_output = self.layernorm_after(hidden_states) + layer_output = self.intermediate(layer_output) + + # second residual connection is done here + layer_output = self.output(layer_output, hidden_states) + + outputs = (layer_output,) + outputs + + return outputs + + +class ViltEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([ViltLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + 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 + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_outputs = layer_module(hidden_states, attention_mask, output_attentions) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring +class ViltPreTrainedModel(PreTrainedModel): + config: ViltConfig + base_model_prefix = "vilt" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["ViltEmbeddings", "ViltSelfAttention"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, TextEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + init.zeros_(module.token_type_ids) + + +@auto_docstring +class ViltModel(ViltPreTrainedModel): + def __init__(self, config, add_pooling_layer=True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + self.embeddings = ViltEmbeddings(config) + self.encoder = ViltEncoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.pooler = ViltPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.text_embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.text_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, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + image_token_type_idx: int | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BaseModelOutputWithPooling | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + image_token_type_idx (`int`, *optional*): + - The token type ids for images. + + Examples: + + ```python + >>> from transformers import ViltProcessor, ViltModel + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # prepare image and text + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> text = "hello world" + + >>> processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-mlm") + >>> model = ViltModel.from_pretrained("dandelin/vilt-b32-mlm") + + >>> inputs = processor(image, text, return_tensors="pt") + >>> outputs = model(**inputs) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + 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") + + text_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(((text_batch_size, seq_length)), device=device) + + if pixel_values is not None and image_embeds is not None: + raise ValueError("You cannot specify both pixel_values and image_embeds at the same time") + elif pixel_values is None and image_embeds is None: + raise ValueError("You have to specify either pixel_values or image_embeds") + + image_batch_size = pixel_values.shape[0] if pixel_values is not None else image_embeds.shape[0] + if image_batch_size != text_batch_size: + raise ValueError("The text inputs and image inputs need to have the same batch size") + if pixel_mask is None: + pixel_mask = torch.ones((image_batch_size, self.config.image_size, self.config.image_size), device=device) + + embedding_output, attention_mask = self.embeddings( + input_ids, + attention_mask, + token_type_ids, + pixel_values, + pixel_mask, + inputs_embeds, + image_embeds, + image_token_type_idx=image_token_type_idx, + ) + + # 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) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class ViltPooler(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): + # 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( + custom_intro=""" + ViLT Model with a language modeling head on top as done during pretraining. + """ +) +class ViltForMaskedLM(ViltPreTrainedModel): + _tied_weights_keys = { + "mlm_score.decoder.weight": "vilt.embeddings.text_embeddings.word_embeddings.weight", + } + + def __init__(self, config): + super().__init__(config) + + self.vilt = ViltModel(config) + self.mlm_score = ViltMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.mlm_score.decoder + + def set_output_embeddings(self, new_embeddings): + self.mlm_score.decoder = new_embeddings + self.mlm_score.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, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskedLMOutput | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + labels (*torch.LongTensor* of shape *(batch_size, sequence_length)*, *optional*): + Labels for computing the masked language modeling loss. Indices should be in *[-100, 0, ..., + config.vocab_size]* (see *input_ids* docstring) Tokens with indices set to *-100* are ignored (masked), the + loss is only computed for the tokens with labels in *[0, ..., config.vocab_size]* + + Examples: + + ```python + >>> from transformers import ViltProcessor, ViltForMaskedLM + >>> import httpx + >>> from io import BytesIO + >>> from PIL import Image + >>> import re + >>> import torch + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> text = "a bunch of [MASK] laying on a [MASK]." + + >>> processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-mlm") + >>> model = ViltForMaskedLM.from_pretrained("dandelin/vilt-b32-mlm") + + >>> # prepare inputs + >>> encoding = processor(image, text, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**encoding) + + >>> tl = len(re.findall("\[MASK\]", text)) + >>> inferred_token = [text] + + >>> # gradually fill in the MASK tokens, one by one + >>> with torch.no_grad(): + ... for i in range(tl): + ... encoded = processor.tokenizer(inferred_token) + ... input_ids = torch.tensor(encoded.input_ids) + ... encoded = encoded["input_ids"][0][1:-1] + ... outputs = model(input_ids=input_ids, pixel_values=encoding.pixel_values) + ... mlm_logits = outputs.logits[0] # shape (seq_len, vocab_size) + ... # only take into account text features (minus CLS and SEP token) + ... mlm_logits = mlm_logits[1 : input_ids.shape[1] - 1, :] + ... mlm_values, mlm_ids = mlm_logits.softmax(dim=-1).max(dim=-1) + ... # only take into account text + ... mlm_values[torch.tensor(encoded) != 103] = 0 + ... select = mlm_values.argmax().item() + ... encoded[select] = mlm_ids[select].item() + ... inferred_token = [processor.decode(encoded)] + + >>> selected_token = "" + >>> encoded = processor.tokenizer(inferred_token) + >>> output = processor.decode(encoded.input_ids[0], skip_special_tokens=True) + >>> print(output) + a bunch of cats laying on a couch. + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_mask=pixel_mask, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output, pooled_output = outputs[:2] + # split up final hidden states into text and image features + text_seq_len = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + text_features, _ = (sequence_output[:, :text_seq_len], sequence_output[:, text_seq_len:]) + + mlm_logits = self.mlm_score(text_features) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + # move labels to correct device to enable PP + labels = labels.to(mlm_logits.device) + masked_lm_loss = loss_fct(mlm_logits.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (mlm_logits,) + outputs[2:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=mlm_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class ViltPredictionHeadTransform(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + if isinstance(config.hidden_act, str): + self.transform_act_fn = ACT2FN[config.hidden_act] + else: + self.transform_act_fn = config.hidden_act + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states): + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +class ViltMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.transform = ViltPredictionHeadTransform(config) + self.decoder = nn.Linear(config.hidden_size, config.vocab_size) + + def forward(self, x): + x = self.transform(x) + x = self.decoder(x) + return x + + +@auto_docstring( + custom_intro=""" + Vilt Model transformer with a classifier head on top (a linear layer on top of the final hidden state of the [CLS] + token) for visual question answering, e.g. for VQAv2. + """ +) +class ViltForQuestionAnswering(ViltPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + self.vilt = ViltModel(config) + + # Classifier head + self.classifier = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size * 2), + nn.LayerNorm(config.hidden_size * 2), + nn.GELU(), + nn.Linear(config.hidden_size * 2, 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, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> SequenceClassifierOutput | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + labels (`torch.FloatTensor` of shape `(batch_size, num_labels)`, *optional*): + Labels for computing the visual question answering loss. This tensor must be either a one-hot encoding of + all answers that are applicable for a given example in the batch, or a soft encoding indicating which + answers are applicable, where 1.0 is the highest score. + + Examples: + + ```python + >>> from transformers import ViltProcessor, ViltForQuestionAnswering + >>> import httpx + >>> from io import BytesIO + >>> from PIL import Image + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> text = "How many cats are there?" + + >>> processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa") + >>> model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa") + + >>> # prepare inputs + >>> encoding = processor(image, text, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**encoding) + >>> logits = outputs.logits + >>> idx = logits.argmax(-1).item() + >>> print("Predicted answer:", model.config.id2label[idx]) + Predicted answer: 2 + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_mask=pixel_mask, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooler_output = outputs.pooler_output if return_dict else outputs[1] + + logits = self.classifier(pooler_output) + + loss = None + if labels is not None: + # move labels to correct device to enable PP + labels = labels.to(logits.device) + loss = nn.functional.binary_cross_entropy_with_logits(logits, labels) * labels.shape[1] + # see https://github.com/jnhwkim/ban-vqa/blob/master/train.py#L19 + + 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( + custom_intro=""" + Vilt Model transformer with a classifier head on top (a linear layer on top of the final hidden state of the [CLS] + token) for image-to-text or text-to-image retrieval, e.g. MSCOCO and F30K. + """ +) +class ViltForImageAndTextRetrieval(ViltPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.vilt = ViltModel(config) + + # Classifier head + self.rank_output = 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, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> SequenceClassifierOutput | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels are currently not supported. + + Examples: + + ```python + >>> from transformers import ViltProcessor, ViltForImageAndTextRetrieval + >>> import httpx + >>> from io import BytesIO + >>> from PIL import Image + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> texts = ["An image of two cats chilling on a couch", "A football player scoring a goal"] + + >>> processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-coco") + >>> model = ViltForImageAndTextRetrieval.from_pretrained("dandelin/vilt-b32-finetuned-coco") + + >>> # forward pass + >>> scores = dict() + >>> for text in texts: + ... # prepare inputs + ... encoding = processor(image, text, return_tensors="pt") + ... outputs = model(**encoding) + ... scores[text] = outputs.logits[0, :].item() + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + loss = None + if labels is not None: + raise NotImplementedError("Training is not yet supported.") + + outputs = self.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_mask=pixel_mask, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooler_output = outputs.pooler_output if return_dict else outputs[1] + + logits = self.rank_output(pooler_output) + + 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( + custom_intro=""" + Vilt Model transformer with a classifier head on top for natural language visual reasoning, e.g. NLVR2. + """ +) +class ViltForImagesAndTextClassification(ViltPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + self.vilt = ViltModel(config) + + # Classifier head + num_images = config.num_images + self.classifier = nn.Sequential( + nn.Linear(config.hidden_size * num_images, config.hidden_size * num_images), + nn.LayerNorm(config.hidden_size * num_images), + nn.GELU(), + nn.Linear(config.hidden_size * num_images, 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, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> ViltForImagesAndTextClassificationOutput | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Binary classification labels. + + Examples: + + ```python + >>> from transformers import ViltProcessor, ViltForImagesAndTextClassification + >>> import httpx + >>> from io import BytesIO + >>> from PIL import Image + + >>> url_1 = "https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg" + >>> with httpx.stream("GET", url_1) as response: + ... image_1 = Image.open(BytesIO(response.read())) + + >>> url_2 = "https://lil.nlp.cornell.edu/nlvr/exs/ex0_1.jpg" + >>> with httpx.stream("GET", url_2) as response: + ... image_2 = Image.open(BytesIO(response.read())) + + >>> text = "The left image contains twice the number of dogs as the right image." + + >>> processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-nlvr2") + >>> model = ViltForImagesAndTextClassification.from_pretrained("dandelin/vilt-b32-finetuned-nlvr2") + + >>> # prepare inputs + >>> encoding = processor([image_1, image_2], text, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(input_ids=encoding.input_ids, pixel_values=encoding.pixel_values.unsqueeze(0)) + >>> logits = outputs.logits + >>> idx = logits.argmax(-1).item() + >>> print("Predicted answer:", model.config.id2label[idx]) + Predicted answer: True + ```""" + 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 not None and pixel_values.ndim == 4: + # add dummy num_images dimension + pixel_values = pixel_values.unsqueeze(1) + + if image_embeds is not None and image_embeds.ndim == 3: + # add dummy num_images dimension + image_embeds = image_embeds.unsqueeze(1) + + num_images = pixel_values.shape[1] if pixel_values is not None else None + if num_images is None: + num_images = image_embeds.shape[1] if image_embeds is not None else None + if num_images != self.config.num_images: + raise ValueError( + "Make sure to match the number of images in the model with the number of images in the input." + ) + pooler_outputs = [] + hidden_states = [] if output_hidden_states else None + attentions = [] if output_attentions else None + for i in range(num_images): + # forward every image through the model + outputs = self.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values[:, i, :, :, :] if pixel_values is not None else None, + pixel_mask=pixel_mask[:, i, :, :] if pixel_mask is not None else None, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds[:, i, :, :] if image_embeds is not None else None, + image_token_type_idx=i + 1, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + pooler_output = outputs.pooler_output if return_dict else outputs[1] + pooler_outputs.append(pooler_output) + if output_hidden_states: + hidden_states.append(outputs.hidden_states) + if output_attentions: + attentions.append(outputs.attentions) + + pooled_output = torch.cat(pooler_outputs, dim=-1) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(logits.device) + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits, hidden_states, attentions) + return ((loss,) + output) if loss is not None else output + + return ViltForImagesAndTextClassificationOutput( + loss=loss, + logits=logits, + hidden_states=hidden_states, + attentions=attentions, + ) + + +@auto_docstring +class ViltForTokenClassification(ViltPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + self.vilt = ViltModel(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, + pixel_values: torch.FloatTensor | None = None, + pixel_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + image_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> TokenClassifierOutput | tuple[torch.FloatTensor]: + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*): + Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `pixel_values` into patch embeddings. + labels (`torch.LongTensor` of shape `(batch_size, text_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.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_mask=pixel_mask, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + text_input_size = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output[:, :text_input_size]) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(logits.device) + 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, + ) + + +__all__ = [ + "ViltForImageAndTextRetrieval", + "ViltForImagesAndTextClassification", + "ViltForTokenClassification", + "ViltForMaskedLM", + "ViltForQuestionAnswering", + "ViltLayer", + "ViltModel", + "ViltPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/vilt/processing_vilt.py b/third_party/transformers/src/transformers/models/vilt/processing_vilt.py new file mode 100644 index 0000000000000000000000000000000000000000..be47b2e6ee754c9b00660040bb7407f69b32483e --- /dev/null +++ b/third_party/transformers/src/transformers/models/vilt/processing_vilt.py @@ -0,0 +1,45 @@ +# 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 ViLT. +""" + +from ...processing_utils import ProcessingKwargs, ProcessorMixin +from ...utils import auto_docstring + + +class ViltProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "add_special_tokens": True, + "padding": False, + "stride": 0, + "return_overflowing_tokens": False, + "return_special_tokens_mask": False, + "return_offsets_mapping": False, + "return_length": False, + "verbose": True, + }, + } + + +@auto_docstring +class ViltProcessor(ProcessorMixin): + valid_processor_kwargs = ViltProcessorKwargs + + def __init__(self, image_processor=None, tokenizer=None, **kwargs): + super().__init__(image_processor, tokenizer) + + +__all__ = ["ViltProcessor"] diff --git a/third_party/transformers/src/transformers/models/wav2vec2_bert/__init__.py b/third_party/transformers/src/transformers/models/wav2vec2_bert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7520263c51bcd3b63747ee0a55e1baff24a777f9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/wav2vec2_bert/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_wav2vec2_bert import * + from .modeling_wav2vec2_bert import * + from .processing_wav2vec2_bert import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/wav2vec2_bert/configuration_wav2vec2_bert.py b/third_party/transformers/src/transformers/models/wav2vec2_bert/configuration_wav2vec2_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..1961e4506f6d791344d841731235cc3454d3e8c8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/wav2vec2_bert/configuration_wav2vec2_bert.py @@ -0,0 +1,204 @@ +# Copyright 2024 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Wav2Vec2Bert model configuration""" + +from typing import Literal + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="facebook/wav2vec2-bert-rel-pos-large") +@strict +class Wav2Vec2BertConfig(PreTrainedConfig): + r""" + feature_projection_input_dim (`int`, *optional*, defaults to 160): + Input dimension of this model, i.e the dimension after processing input audios with [`SeamlessM4TFeatureExtractor`] or [`Wav2Vec2BertProcessor`]. + feat_proj_dropout (`float`, *optional*, defaults to 0.0): + The dropout probability for the feature projection. + final_dropout (`float`, *optional*, defaults to 0.1): + The dropout probability for the final projection layer of [`Wav2Vec2BertForCTC`]. + apply_spec_augment (`bool`, *optional*, defaults to `True`): + Whether to apply *SpecAugment* data augmentation to the outputs of the feature encoder. For reference see + [SpecAugment: A Simple Data Augmentation Method for Automatic Speech + Recognition](https://huggingface.co/papers/1904.08779). + mask_time_prob (`float`, *optional*, defaults to 0.05): + Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking + procedure generates `mask_time_prob*len(time_axis)/mask_time_length ``independent masks over the axis. If + reasoning from the probability of each feature vector to be chosen as the start of the vector span to be + masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the + actual percentage of masked vectors. This is only relevant if `apply_spec_augment is True`. + mask_time_length (`int`, *optional*, defaults to 10): + Length of vector span along the time axis. + mask_time_min_masks (`int`, *optional*, defaults to 2): + The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step, + irrespectively of `mask_feature_prob`. Only relevant if `mask_time_prob*len(time_axis)/mask_time_length < + mask_time_min_masks`. + mask_feature_prob (`float`, *optional*, defaults to 0.0): + Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The + masking procedure generates `mask_feature_prob*len(feature_axis)/mask_time_length` independent masks over + the axis. If reasoning from the probability of each feature vector to be chosen as the start of the vector + span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap + may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is + True`. + mask_feature_length (`int`, *optional*, defaults to 10): + Length of vector span along the feature axis. + mask_feature_min_masks (`int`, *optional*, defaults to 0): + The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time + step, irrespectively of `mask_feature_prob`. Only relevant if + `mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks`. + ctc_zero_infinity (`bool`, *optional*, defaults to `False`): + Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly + occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance + of [`Wav2Vec2BertForCTC`]. + use_weighted_layer_sum (`bool`, *optional*, defaults to `False`): + Whether to use a weighted average of layer outputs with learned weights. Only relevant when using an + instance of [`Wav2Vec2BertForSequenceClassification`]. + classifier_proj_size (`int`, *optional*, defaults to 768): + Dimensionality of the projection before token mean-pooling for classification. + tdnn_dim (`tuple[int]` or `list[int]`, *optional*, defaults to `(512, 512, 512, 512, 1500)`): + A tuple of integers defining the number of output channels of each 1D convolutional layer in the *TDNN* + module of the *XVector* model. The length of *tdnn_dim* defines the number of *TDNN* layers. + tdnn_kernel (`tuple[int]` or `list[int]`, *optional*, defaults to `(5, 3, 3, 1, 1)`): + A tuple of integers defining the kernel size of each 1D convolutional layer in the *TDNN* module of the + *XVector* model. The length of *tdnn_kernel* has to match the length of *tdnn_dim*. + tdnn_dilation (`tuple[int]` or `list[int]`, *optional*, defaults to `(1, 2, 3, 1, 1)`): + A tuple of integers defining the dilation factor of each 1D convolutional layer in *TDNN* module of the + *XVector* model. The length of *tdnn_dilation* has to match the length of *tdnn_dim*. + xvector_output_dim (`int`, *optional*, defaults to 512): + Dimensionality of the *XVector* embedding vectors. + add_adapter (`bool`, *optional*, defaults to `False`): + Whether a convolutional attention network should be stacked on top of the Wav2Vec2Bert Encoder. Can be very + useful for warm-starting Wav2Vec2Bert for SpeechEncoderDecoder models. + adapter_kernel_size (`int`, *optional*, defaults to 3): + Kernel size of the convolutional layers in the adapter network. Only relevant if `add_adapter is True`. + adapter_stride (`int`, *optional*, defaults to 2): + Stride of the convolutional layers in the adapter network. Only relevant if `add_adapter is True`. + num_adapter_layers (`int`, *optional*, defaults to 1): + Number of convolutional layers that should be used in the adapter network. Only relevant if `add_adapter is + True`. + adapter_act (`str` or `function`, *optional*, defaults to `"relu"`): + The non-linear activation function (function or string) in the adapter layers. If string, `"gelu"`, + `"relu"`, `"selu"`, `"swish"` and `"gelu_new"` are supported. + use_intermediate_ffn_before_adapter (`bool`, *optional*, defaults to `False`): + Whether an intermediate feed-forward block should be stacked on top of the Wav2Vec2Bert Encoder and before the adapter network. + Only relevant if `add_adapter is True`. + output_hidden_size (`int`, *optional*): + Dimensionality of the encoder output layer. If not defined, this defaults to *hidden-size*. Only relevant + if `add_adapter is True`. + position_embeddings_type (`str`, *optional*, defaults to `"relative_key"`): + Can be specified to : + - `rotary`, for rotary position embeddings. + - `relative`, for relative position embeddings. + - `relative_key`, for relative position embeddings as defined by Shaw in [Self-Attention + with Relative Position Representations (Shaw et al.)](https://huggingface.co/papers/1803.02155). + If left to `None`, no relative position embeddings is applied. + rotary_embedding_base (`int`, *optional*, defaults to 10000): + If `"rotary"` position embeddings are used, defines the size of the embedding base. + max_source_positions (`int`, *optional*, defaults to 5000): + if `"relative"` position embeddings are used, defines the maximum source input positions. + left_max_position_embeddings (`int`, *optional*, defaults to 64): + If `"relative_key"` (aka Shaw) position embeddings are used, defines the left clipping value for relative positions. + right_max_position_embeddings (`int`, *optional*, defaults to 8): + If `"relative_key"` (aka Shaw) position embeddings are used, defines the right clipping value for relative positions. + conv_depthwise_kernel_size (`int`, *optional*, defaults to 31): + Kernel size of convolutional depthwise 1D layer in Conformer blocks. + conformer_conv_dropout (`float`, *optional*, defaults to 0.1): + The dropout probability for all convolutional layers in Conformer blocks. + + Example: + + ```python + >>> from transformers import Wav2Vec2BertConfig, Wav2Vec2BertModel + + >>> # Initializing a Wav2Vec2Bert facebook/wav2vec2-bert-rel-pos-large style configuration + >>> configuration = Wav2Vec2BertConfig() + + >>> # Initializing a model (with random weights) from the facebook/wav2vec2-bert-rel-pos-large style configuration + >>> model = Wav2Vec2BertModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "wav2vec2-bert" + + vocab_size: int | None = None + hidden_size: int = 1024 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + intermediate_size: int = 4096 + feature_projection_input_dim: int = 160 + hidden_act: str = "swish" + hidden_dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + attention_dropout: float | int = 0.0 + feat_proj_dropout: float | int = 0.0 + final_dropout: float | int = 0.1 + layerdrop: float | int = 0.1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + apply_spec_augment: bool = True + mask_time_prob: float | int = 0.05 + mask_time_length: int = 10 + mask_time_min_masks: int = 2 + mask_feature_prob: float | int = 0.0 + mask_feature_length: int = 10 + mask_feature_min_masks: int = 0 + ctc_loss_reduction: str = "sum" + ctc_zero_infinity: bool = False + use_weighted_layer_sum: bool = False + classifier_proj_size: int = 768 + tdnn_dim: list[int] | tuple[int, ...] = (512, 512, 512, 512, 1500) + tdnn_kernel: list[int] | tuple[int, ...] = (5, 3, 3, 1, 1) + tdnn_dilation: list[int] | tuple[int, ...] = (1, 2, 3, 1, 1) + xvector_output_dim: int = 512 + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + add_adapter: bool = False + adapter_kernel_size: int = 3 + adapter_stride: int = 2 + num_adapter_layers: int = 1 + adapter_act: str = "relu" + use_intermediate_ffn_before_adapter: bool = False + output_hidden_size: int | None = None + position_embeddings_type: Literal["rotary", "relative", "relative_key"] | None = "relative_key" + rotary_embedding_base: int = 10000 + max_source_positions: int = 5000 + left_max_position_embeddings: int = 64 + right_max_position_embeddings: int = 8 + conv_depthwise_kernel_size: int = 31 + conformer_conv_dropout: float | int = 0.1 + + def __post_init__(self, **kwargs): + self.output_hidden_size = self.output_hidden_size or self.hidden_size + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.use_intermediate_ffn_before_adapter and not self.add_adapter: + raise ValueError("`use_intermediate_ffn_before_adapter` is `True` but `add_adapter` is `False`.") + + @property + def inputs_to_logits_ratio(self): + ratio = self.feature_projection_input_dim * 2 + if self.add_adapter: + ratio = ratio * (self.adapter_stride**self.num_adapter_layers) + return ratio + + +__all__ = ["Wav2Vec2BertConfig"] diff --git a/third_party/transformers/src/transformers/models/wav2vec2_bert/convert_wav2vec2_seamless_checkpoint.py b/third_party/transformers/src/transformers/models/wav2vec2_bert/convert_wav2vec2_seamless_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..3355d1555f85e0c6a7d288f0e157d23e19bfc670 --- /dev/null +++ b/third_party/transformers/src/transformers/models/wav2vec2_bert/convert_wav2vec2_seamless_checkpoint.py @@ -0,0 +1,215 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert Wav2Vec2Bert BERT checkpoint.""" + +import argparse + +import torch +import torchaudio +from fairseq2.data import Collater +from fairseq2.data.audio import WaveformToFbankConverter +from fairseq2.nn.padding import get_seqs_and_padding_mask +from seamless_communication.models.conformer_shaw import load_conformer_shaw_model + +from transformers import ( + SeamlessM4TFeatureExtractor, + Wav2Vec2BertConfig, + Wav2Vec2BertModel, + logging, +) + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +wav2vec_convert_list = [ + ("encoder_frontend.model_dim_proj", "feature_projection.projection"), + ("encoder_frontend.post_extract_layer_norm", "feature_projection.layer_norm"), + ("encoder_frontend.pos_encoder.conv", "encoder.pos_conv_embed.conv"), + ("encoder.inner.layers", "encoder.layers"), + ("encoder.inner_layer_norm", "encoder.layer_norm"), + ("encoder.adaptor_layers", "adapter.layers"), + ("inner_proj", "intermediate_dense"), + ("self_attn.output_proj", "self_attn.linear_out"), + ("output_proj", "output_dense"), + ("self_attn.k_proj", "self_attn.linear_k"), + ("self_attn.v_proj", "self_attn.linear_v"), + ("self_attn.q_proj", "self_attn.linear_q"), + ("self_attn.sdpa.u_bias", "self_attn.pos_bias_u"), + ("self_attn.sdpa.v_bias", "self_attn.pos_bias_v"), + ("self_attn.sdpa.rel_k_embed", "self_attn.distance_embedding"), + ("self_attn.sdpa.r_proj", "self_attn.linear_pos"), + ("conv.pointwise_conv1", "conv_module.pointwise_conv1"), + ("conv.pointwise_conv2", "conv_module.pointwise_conv2"), + ("conv.depthwise_conv", "conv_module.depthwise_conv"), + ("conv.layer_norm", "conv_module.depthwise_layer_norm"), + ("conv_layer_norm", "conv_module.layer_norm"), + ("encoder.proj1", "intermediate_ffn.intermediate_dense"), + ("encoder.proj2", "intermediate_ffn.output_dense"), + ("encoder.layer_norm", "inner_layer_norm"), + ("masker.temporal_mask_embed", "masked_spec_embed"), +] + +keys_to_remove = { + "quantizer.entry_proj", + "final_proj", + "final_target_proj", + "quantizer.entries", + "quantizer.num_updates", +} + + +def param_count(model): + return sum(p[1].numel() for p in model.named_parameters() if "final_proj" not in p[0]) + + +def _convert_model( + original_model, + hf_model, + convert_list, +): + state_dict = original_model.state_dict() + + for k, v in list(state_dict.items()): + new_key = k + for old_layer_name, new_layer_name in convert_list: + if old_layer_name in new_key: + new_key = new_key.replace(old_layer_name, new_layer_name) + + # must do it by hand + if ".layer_norm" in new_key and new_key.split(".layer_norm")[0][-1].isnumeric(): + new_key = new_key.replace("layer_norm", "final_layer_norm") + + add_key = True + for key in keys_to_remove: + if key in new_key: + state_dict.pop(k) + add_key = False + break + + if add_key: + state_dict[new_key] = state_dict.pop(k) + + extra_keys = set(state_dict.keys()) - set(hf_model.state_dict().keys()) + extra_keys = set({k for k in extra_keys if "num_updates" not in k}) # filter unnecessary param + missing_keys = set(hf_model.state_dict().keys()) - set(state_dict.keys()) + if len(extra_keys) != 0: + raise ValueError(f"extra keys found: {extra_keys}") + if len(missing_keys) != 0: + raise ValueError(f"missing keys: {missing_keys}") + hf_model.load_state_dict(state_dict, strict=True) + n_params = param_count(hf_model) + + logger.info(f"model loaded: {round(n_params / 1e6, 1)}M params") + + hf_model.eval() + del state_dict + + return hf_model + + +@torch.no_grad() +def convert_wav2vec2_bert_checkpoint( + checkpoint_path, + pytorch_dump_folder_path, + config_path=None, + repo_id=None, +): + """ + Copy/paste/tweak model's weights to transformers design. + """ + if config_path is not None: + config = Wav2Vec2BertConfig.from_pretrained(config_path, hidden_act="swish") + else: + config = Wav2Vec2BertConfig(apply_spec_augment=False) + + hf_wav2vec = Wav2Vec2BertModel(config) + + model = load_conformer_shaw_model(checkpoint_path, dtype=torch.float32) + model.eval() + + hf_wav2vec = _convert_model(model, hf_wav2vec, wav2vec_convert_list) + + hf_wav2vec.save_pretrained(pytorch_dump_folder_path) + + if repo_id: + hf_wav2vec.push_to_hub(repo_id, create_pr=True) + + # save feature extractor + fe = SeamlessM4TFeatureExtractor(padding_value=1) + fe.save_pretrained(pytorch_dump_folder_path) + + if repo_id: + fe.push_to_hub(repo_id, create_pr=True) + + if args.audio_path: + waveform, sample_rate = torchaudio.load(args.audio_path) + waveform = torchaudio.functional.resample(waveform, sample_rate, fe.sampling_rate) + + fbank_converter = WaveformToFbankConverter( + num_mel_bins=80, + waveform_scale=2**15, + channel_last=True, + standardize=True, + dtype=torch.float32, + ) + collater = Collater(pad_value=1) + + decoded_audio = {"waveform": waveform.T, "sample_rate": fe.sampling_rate, "format": -1} + src = collater(fbank_converter(decoded_audio))["fbank"] + seqs, padding_mask = get_seqs_and_padding_mask(src) + + with torch.inference_mode(): + seqs, padding_mask = model.encoder_frontend(seqs, padding_mask) + original_output, padding_mask = model.encoder(seqs, padding_mask) + + hf_wav2vec.eval() + + inputs = fe(waveform, return_tensors="pt", padding=True) + with torch.no_grad(): + outputs = hf_wav2vec(**inputs) + + torch.testing.assert_close(original_output, outputs.last_hidden_state, rtol=5e-3, atol=5e-3) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + help="Path to the output PyTorch model.", + ) + parser.add_argument( + "--checkpoint_path", default="conformer_shaw", type=str, help="Path to seamless communication checkpoint" + ) + parser.add_argument( + "--config_path", + default=None, + type=str, + help="Path to hf config.json of model to convert", + ) + parser.add_argument("--repo_id", default=None, type=str, help="Push to this repo id if precised.") + parser.add_argument( + "--audio_path", + default=None, + type=str, + help="If specified, check that the original model and the converted model produce the same outputs.", + ) + + args = parser.parse_args() + convert_wav2vec2_bert_checkpoint( + args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.repo_id + ) diff --git a/third_party/transformers/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py b/third_party/transformers/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..6023d856798b1aed1d64e3c37759be073d292bfc --- /dev/null +++ b/third_party/transformers/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py @@ -0,0 +1,1529 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.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_wav2vec2_bert.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +import math +import warnings + +import numpy as np +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...integrations.deepspeed import is_deepspeed_zero3_enabled +from ...integrations.fsdp import is_fsdp_managed_module +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + CausalLMOutput, + SequenceClassifierOutput, + TokenClassifierOutput, + Wav2Vec2BaseModelOutput, + XVectorOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, is_peft_available +from .configuration_wav2vec2_bert import Wav2Vec2BertConfig + + +class Wav2Vec2BertRotaryPositionalEmbedding(nn.Module): + """Rotary positional embedding + Reference : https://blog.eleuther.ai/rotary-embeddings/ Paper: https://huggingface.co/papers/2104.09864 + """ + + def __init__(self, config): + super().__init__() + dim = config.hidden_size // config.num_attention_heads + base = config.rotary_embedding_base + + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) + # Ignore copy + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.cached_sequence_length = None + self.cached_rotary_positional_embedding = None + + def forward(self, hidden_states): + sequence_length = hidden_states.shape[1] + + if sequence_length == self.cached_sequence_length and self.cached_rotary_positional_embedding is not None: + return self.cached_rotary_positional_embedding + + self.cached_sequence_length = sequence_length + # Embeddings are computed in the dtype of the inv_freq constant + time_stamps = torch.arange(sequence_length).type_as(self.inv_freq) + freqs = torch.einsum("i,j->ij", time_stamps, self.inv_freq) + embeddings = torch.cat((freqs, freqs), dim=-1) + + cos_embeddings = embeddings.cos()[:, None, None, :] + sin_embeddings = embeddings.sin()[:, None, None, :] + # Computed embeddings are cast to the dtype of the hidden state inputs + self.cached_rotary_positional_embedding = torch.stack([cos_embeddings, sin_embeddings]).type_as(hidden_states) + return self.cached_rotary_positional_embedding + + +class Wav2Vec2BertRelPositionalEmbedding(nn.Module): + """Relative positional encoding module.""" + + def __init__(self, config): + super().__init__() + self.max_len = config.max_source_positions + self.d_model = config.hidden_size + self.register_buffer("pe", self.extend_pe(torch.tensor(0.0).expand(1, self.max_len)), persistent=False) + + def extend_pe(self, x, pe=None): + # Reset the positional encodings + if pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if pe.size(1) >= x.size(1) * 2 - 1: + if pe.dtype != x.dtype or pe.device != x.device: + pe = pe.to(dtype=x.dtype, device=x.device) + return pe + # Suppose `i` is the position of query vector and `j` is the + # position of key vector. We use positive relative positions when keys + # are to the left (i>j) and negative relative positions otherwise (i