diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..40f84ccb59be91ad440ba413366dd726e555ca26
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/__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_dac import *
+ from .feature_extraction_dac import *
+ from .modeling_dac import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/configuration_dac.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/configuration_dac.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d11a341f2625991174a105a3d460ad7a20206a7
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/configuration_dac.py
@@ -0,0 +1,78 @@
+# Copyright 2024 Descript 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.
+"""Dac model configuration"""
+
+import math
+
+import numpy as np
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="descript/dac_16khz")
+@strict
+class DacConfig(PreTrainedConfig):
+ r"""
+ downsampling_ratios (`list[int]`, *optional*, defaults to `[2, 4, 8, 8]`):
+ Ratios for downsampling in the encoder. These are used in reverse order for upsampling in the decoder.
+ quantizer_dropout (`bool`, *optional*, defaults to 0):
+ Whether to apply dropout to the quantizer.
+ commitment_loss_weight (float, *optional*, defaults to 0.25):
+ Weight of the commitment loss term in the VQVAE loss function.
+ codebook_loss_weight (float, *optional*, defaults to 1.0):
+ Weight of the codebook loss term in the VQVAE loss function.
+
+ Example:
+
+ ```python
+ >>> from transformers import DacModel, DacConfig
+
+ >>> # Initializing a "descript/dac_16khz" style configuration
+ >>> configuration = DacConfig()
+
+ >>> # Initializing a model (with random weights) from the "descript/dac_16khz" style configuration
+ >>> model = DacModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "dac"
+
+ encoder_hidden_size: int = 64
+ downsampling_ratios: list[int] | tuple[int, ...] = (2, 4, 8, 8)
+ decoder_hidden_size: int = 1536
+ n_codebooks: int = 9
+ codebook_size: int = 1024
+ codebook_dim: int = 8
+ quantizer_dropout: float | int = 0.0
+ commitment_loss_weight: float = 0.25
+ codebook_loss_weight: float = 1.0
+ sampling_rate: int = 16000
+
+ def __post_init__(self, **kwargs):
+ self.upsampling_ratios = self.downsampling_ratios[::-1]
+ self.hidden_size = self.encoder_hidden_size * (2 ** len(self.downsampling_ratios))
+ self.hop_length = int(np.prod(self.downsampling_ratios))
+ super().__post_init__(**kwargs)
+
+ @property
+ def frame_rate(self) -> int:
+ hop_length = np.prod(self.upsampling_ratios)
+ return math.ceil(self.sampling_rate / hop_length)
+
+
+__all__ = ["DacConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/feature_extraction_dac.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/feature_extraction_dac.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f910f57f09fa3ed1cbce0e1795aac622fe43bc9
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/feature_extraction_dac.py
@@ -0,0 +1,170 @@
+# Copyright 2024 Descript 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.
+"""Feature extractor class for DAC"""
+
+import numpy as np
+
+from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
+from ...feature_extraction_utils import BatchFeature
+from ...utils import PaddingStrategy, TensorType, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class DacFeatureExtractor(SequenceFeatureExtractor):
+ r"""
+ Constructs an Dac 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.
+
+ Args:
+ feature_size (`int`, *optional*, defaults to 1):
+ The feature dimension of the extracted features. Use 1 for mono, 2 for stereo.
+ sampling_rate (`int`, *optional*, defaults to 16000):
+ 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.
+ hop_length (`int`, *optional*, defaults to 512):
+ Overlap length between successive windows.
+ """
+
+ model_input_names = ["input_values", "n_quantizers"]
+
+ def __init__(
+ self,
+ feature_size: int = 1,
+ sampling_rate: int = 16000,
+ padding_value: float = 0.0,
+ hop_length: int = 512,
+ **kwargs,
+ ):
+ super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
+ self.hop_length = hop_length
+
+ def __call__(
+ self,
+ raw_audio: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
+ padding: bool | str | PaddingStrategy | None = None,
+ truncation: bool | None = False,
+ max_length: int | None = None,
+ return_tensors: str | TensorType | None = None,
+ sampling_rate: int | None = None,
+ ) -> BatchFeature:
+ """
+ Main method to featurize and prepare for the model one or several sequence(s).
+
+ Args:
+ raw_audio (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`):
+ The sequence or batch of sequences to be processed. 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. The numpy array must be of shape
+ `(num_samples,)` for mono audio (`feature_size = 1`), or `(2, num_samples)` for stereo audio
+ (`feature_size = 2`).
+ 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).
+ truncation (`bool`, *optional*, defaults to `False`):
+ Activates truncation to cut input sequences longer than `max_length` to `max_length`.
+ max_length (`int`, *optional*):
+ Maximum length of the returned list and optionally padding length (see above).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*, default to 'pt'):
+ 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 `audio` input was sampled. It is strongly recommended to pass
+ `sampling_rate` at the forward call to prevent silent errors.
+ """
+ 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 padding and truncation:
+ raise ValueError("Both padding and truncation were set. Make sure you only set one.")
+ elif padding is None:
+ # by default let's pad the inputs
+ padding = True
+
+ is_batched = bool(
+ isinstance(raw_audio, (list, tuple)) and (isinstance(raw_audio[0], (np.ndarray, tuple, list)))
+ )
+
+ if is_batched:
+ raw_audio = [np.asarray(audio, dtype=np.float32).T for audio in raw_audio]
+ elif not is_batched and not isinstance(raw_audio, np.ndarray):
+ raw_audio = np.asarray(raw_audio, dtype=np.float32)
+ elif isinstance(raw_audio, np.ndarray) and raw_audio.dtype is np.dtype(np.float64):
+ raw_audio = raw_audio.astype(np.float32)
+
+ # always return batch
+ if not is_batched:
+ raw_audio = [np.asarray(raw_audio).T]
+
+ # verify inputs are valid
+ for idx, example in enumerate(raw_audio):
+ if example.ndim > 2:
+ raise ValueError(f"Expected input shape (channels, length) but got shape {example.shape}")
+ if self.feature_size == 1 and example.ndim != 1:
+ raise ValueError(f"Expected mono audio but example has {example.shape[-1]} channels")
+ if self.feature_size == 2:
+ raise ValueError("Stereo audio isn't supported for now")
+
+ input_values = BatchFeature({"input_values": raw_audio})
+
+ # normal padding on batch
+ padded_inputs = self.pad(
+ input_values,
+ max_length=max_length,
+ truncation=truncation,
+ padding=padding,
+ return_attention_mask=padding,
+ pad_to_multiple_of=self.hop_length,
+ )
+ if padding:
+ padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask")
+ if padding:
+ padded_inputs.input_values = padded_inputs.input_values[:, np.newaxis, :]
+
+ input_values = []
+ for example in padded_inputs.pop("input_values"):
+ if self.feature_size == 1:
+ example = example[..., None]
+ input_values.append(example.T)
+
+ padded_inputs["input_values"] = input_values
+ if return_tensors is not None:
+ padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
+
+ return padded_inputs
+
+
+__all__ = ["DacFeatureExtractor"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/modeling_dac.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/modeling_dac.py
new file mode 100644
index 0000000000000000000000000000000000000000..92cab160944d55203bbf6618245b1dc4aeba37b4
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/dac/modeling_dac.py
@@ -0,0 +1,689 @@
+# Copyright 2024 Descript 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.
+"""Transformers DAC model."""
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...modeling_utils import PreTrainedAudioTokenizerBase
+from ...utils import ModelOutput, auto_docstring
+from .configuration_dac import DacConfig
+
+
+@auto_docstring
+@dataclass
+class DacOutput(ModelOutput):
+ r"""
+ loss (`torch.Tensor`):
+ Loss from the encoder model, comprising the weighted combination of the commitment and codebook losses.
+ audio_values (`torch.Tensor` of shape `(batch_size, input_length)`):
+ Reconstructed audio data.
+ quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`):
+ Quantized continuous representation of input.
+ audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, time_steps)`):
+ Codebook indices for each codebook (quantized discrete representation of input).
+ projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`):
+ Projected latents (continuous representation of input before quantization).
+ """
+
+ loss: torch.FloatTensor | None = None
+ audio_values: torch.FloatTensor | None = None
+ quantized_representation: torch.FloatTensor | None = None
+ audio_codes: torch.LongTensor | None = None
+ projected_latents: torch.FloatTensor | None = None
+
+
+@auto_docstring
+@dataclass
+class DacEncoderOutput(ModelOutput):
+ r"""
+ loss (`torch.Tensor`):
+ Loss from the encoder model, comprising the weighted combination of the commitment and codebook losses.
+ quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`, *optional*):
+ Quantized continuous representation of input.
+ audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`, *optional*):
+ Codebook indices for each codebook (quantized discrete representation of input).
+ projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`, *optional*):
+ Projected latents (continuous representation of input before quantization).
+ """
+
+ loss: torch.FloatTensor | None = None
+ quantized_representation: torch.FloatTensor | None = None
+ audio_codes: torch.FloatTensor | None = None
+ projected_latents: torch.FloatTensor | None = None
+
+
+@auto_docstring
+@dataclass
+# Copied from transformers.models.encodec.modeling_encodec.EncodecDecoderOutput with Encodec->Dac, segment_length->input_length
+class DacDecoderOutput(ModelOutput):
+ r"""
+ audio_values (`torch.FloatTensor` of shape `(batch_size, input_length)`, *optional*):
+ Decoded audio values, obtained using the decoder part of Dac.
+ """
+
+ audio_values: torch.FloatTensor | None = None
+
+
+class Snake1d(nn.Module):
+ """
+ A 1-dimensional Snake activation function module.
+ """
+
+ def __init__(self, hidden_dim):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.ones(1, hidden_dim, 1))
+
+ def forward(self, hidden_states):
+ shape = hidden_states.shape
+ hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
+ hidden_states = hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2)
+ hidden_states = hidden_states.reshape(shape)
+ return hidden_states
+
+
+class DacVectorQuantize(nn.Module):
+ """
+ Implementation of VQ similar to Karpathy's repo (https://github.com/karpathy/deep-vector-quantization)
+
+ Additionally uses following tricks from improved VQGAN
+ (https://huggingface.co/papers/2110.04627):
+ 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+ for improved codebook usage
+ 2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+ improves training stability
+ """
+
+ def __init__(self, config: DacConfig):
+ super().__init__()
+
+ self.codebook_dim = config.codebook_dim
+ self.in_proj = nn.Conv1d(config.hidden_size, config.codebook_dim, kernel_size=1)
+ self.out_proj = nn.Conv1d(config.codebook_dim, config.hidden_size, kernel_size=1)
+ self.codebook = nn.Embedding(config.codebook_size, config.codebook_dim)
+
+ def forward(self, hidden_state):
+ """
+ Quantizes the input tensor using a fixed codebook and returns the corresponding codebook vectors.
+
+ Args:
+ hidden_state (`torch.FloatTensor` of shape `(batch_size, dimension, time_steps)`):
+ Input tensor.
+
+ Returns:
+ quantized_representation (`torch.Tensor`of shape `(batch_size, dimension, time_steps)`):
+ Quantized continuous representation of input.
+ commitment_loss (`torch.FloatTensor`of shape `(1)`):
+ Commitment loss to train encoder to predict vectors closer to codebook entries.
+ codebook_loss (`torch.FloatTensor`of shape `(1)`):
+ Codebook loss to update the codebook.
+ audio_codes (`torch.LongTensor` of shape `(batch_size, time_steps)`):
+ Codebook indices for each codebook, quantized discrete representation of input.
+ projected_latents (torch.FloatTensor of shape `(batch_size, num_codebooks * dimension, time_steps)`):
+ Projected latents (continuous representation of input before quantization).
+ """
+
+ projected_latents = self.in_proj(hidden_state)
+ quantized_representation, audio_codes = self.decode_latents(projected_latents)
+
+ commitment_loss = F.mse_loss(projected_latents, quantized_representation.detach(), reduction="mean")
+ codebook_loss = F.mse_loss(quantized_representation, projected_latents.detach(), reduction="mean")
+ # noop in forward pass, straight-through gradient estimator in backward pass
+ quantized_representation = projected_latents + (quantized_representation - projected_latents).detach()
+ quantized_representation = self.out_proj(quantized_representation)
+
+ return quantized_representation, commitment_loss, codebook_loss, audio_codes, projected_latents
+
+ def decode_latents(self, hidden_states):
+ batch_size, hidden_dim, sequence_length = hidden_states.shape
+ encodings = hidden_states.permute(0, 2, 1).reshape(batch_size * sequence_length, hidden_dim)
+ codebook = self.codebook.weight # codebook: (N x D)
+
+ # L2 normalize encodings and codebook (ViT-VQGAN)
+ encodings = F.normalize(encodings)
+ codebook = F.normalize(codebook)
+
+ # Compute euclidean distance with codebook
+ l2_norm = encodings.pow(2).sum(1, keepdim=True)
+ dist = -(l2_norm - 2 * encodings @ codebook.t()) + codebook.pow(2).sum(1, keepdim=True).t()
+
+ indices = dist.max(1)[1]
+ indices = indices.reshape(hidden_states.size(0), -1)
+ quantized_representation = self.codebook(indices).transpose(1, 2)
+ return quantized_representation, indices
+
+
+class DacResidualUnit(nn.Module):
+ """
+ A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations.
+ """
+
+ def __init__(self, dimension: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+
+ self.snake1 = Snake1d(dimension)
+ self.conv1 = nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad)
+ self.snake2 = Snake1d(dimension)
+ self.conv2 = nn.Conv1d(dimension, dimension, kernel_size=1)
+
+ def forward(self, hidden_state):
+ """
+ Forward pass through the residual unit.
+
+ Args:
+ hidden_state (`torch.Tensor` of shape `(batch_size, channels, time_steps)`):
+ Input tensor .
+
+ Returns:
+ output_tensor (`torch.Tensor` of shape `(batch_size, channels, time_steps)`):
+ Input tensor after passing through the residual unit.
+ """
+ output_tensor = hidden_state
+ output_tensor = self.conv1(self.snake1(output_tensor))
+ output_tensor = self.conv2(self.snake2(output_tensor))
+
+ padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2
+ if padding > 0:
+ hidden_state = hidden_state[..., padding:-padding]
+ output_tensor = hidden_state + output_tensor
+ return output_tensor
+
+
+class DacEncoderBlock(nn.Module):
+ """Encoder block used in DAC encoder."""
+
+ def __init__(self, config: DacConfig, stride: int = 1, stride_index: int = 1):
+ super().__init__()
+
+ dimension = config.encoder_hidden_size * 2**stride_index
+ self.res_unit1 = DacResidualUnit(dimension // 2, dilation=1)
+ self.res_unit2 = DacResidualUnit(dimension // 2, dilation=3)
+ self.res_unit3 = DacResidualUnit(dimension // 2, dilation=9)
+ self.snake1 = Snake1d(dimension // 2)
+ self.conv1 = nn.Conv1d(
+ dimension // 2, dimension, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)
+ )
+
+ def forward(self, hidden_state):
+ hidden_state = self.res_unit1(hidden_state)
+ hidden_state = self.res_unit2(hidden_state)
+ hidden_state = self.snake1(self.res_unit3(hidden_state))
+ hidden_state = self.conv1(hidden_state)
+
+ return hidden_state
+
+
+class DacDecoderBlock(nn.Module):
+ """Decoder block used in DAC decoder."""
+
+ def __init__(self, config: DacConfig, stride: int = 1, stride_index: int = 1):
+ super().__init__()
+
+ input_dim = config.decoder_hidden_size // 2**stride_index
+ output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1)
+ self.snake1 = Snake1d(input_dim)
+ self.conv_t1 = nn.ConvTranspose1d(
+ input_dim,
+ output_dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ )
+
+ self.res_unit1 = DacResidualUnit(output_dim, dilation=1)
+ self.res_unit2 = DacResidualUnit(output_dim, dilation=3)
+ self.res_unit3 = DacResidualUnit(output_dim, dilation=9)
+
+ def forward(self, hidden_state):
+ hidden_state = self.snake1(hidden_state)
+ hidden_state = self.conv_t1(hidden_state)
+ hidden_state = self.res_unit1(hidden_state)
+ hidden_state = self.res_unit2(hidden_state)
+ hidden_state = self.res_unit3(hidden_state)
+
+ return hidden_state
+
+
+class DacResidualVectorQuantizer(nn.Module):
+ """
+ ResidualVectorQuantize block - Introduced in SoundStream: An end2end neural audio codec (https://huggingface.co/papers/2107.03312)
+ """
+
+ def __init__(self, config: DacConfig):
+ super().__init__()
+
+ n_codebooks = config.n_codebooks
+ quantizer_dropout = config.quantizer_dropout
+
+ self.n_codebooks = n_codebooks
+
+ self.quantizers = nn.ModuleList([DacVectorQuantize(config) for i in range(config.n_codebooks)])
+ self.quantizer_dropout = quantizer_dropout
+
+ def forward(self, hidden_state, n_quantizers: int | None = None):
+ """
+ Quantizes the input tensor using a fixed set of codebooks and returns corresponding codebook vectors.
+ Args:
+ hidden_state (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`):
+ Input tensor to be quantized.
+ n_quantizers (`int`, *optional*):
+ Number of quantizers to use. If specified and `self.quantizer_dropout` is True,
+ this argument is ignored during training, and a random number of quantizers is used.
+
+ Returns:
+ quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`):
+ Quantized continuous representation of input.
+ audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`):
+ Codebook indices for each codebook (quantized discrete representation of input).
+ projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`):
+ Projected latents (continuous representation of input before quantization).
+ commitment_loss (`torch.Tensor` of shape `(1)`):
+ Commitment loss to train the encoder to predict vectors closer to codebook entries.
+ codebook_loss (`torch.Tensor` of shape `(1)`):
+ Codebook loss to update the codebook.
+ """
+
+ quantized_representation = 0
+ residual = hidden_state
+ commitment_loss = 0
+ codebook_loss = 0
+
+ audio_codes = []
+ projected_latents = []
+
+ n_quantizers = n_quantizers if n_quantizers is not None else self.n_codebooks
+ if self.training:
+ n_quantizers = torch.ones((hidden_state.shape[0],)) * self.n_codebooks + 1
+ dropout = torch.randint(1, self.n_codebooks + 1, (hidden_state.shape[0],))
+ n_dropout = int(hidden_state.shape[0] * self.quantizer_dropout)
+ n_quantizers[:n_dropout] = dropout[:n_dropout]
+ n_quantizers = n_quantizers.to(hidden_state.device)
+
+ for i, quantizer in enumerate(self.quantizers):
+ if self.training is False and i >= n_quantizers:
+ break
+
+ quantized_representation_i, commitment_loss_i, codebook_loss_i, indices_i, projected_latents_i = quantizer(
+ residual
+ )
+
+ # Create mask to apply quantizer dropout
+ mask = torch.full((hidden_state.shape[0],), i, device=hidden_state.device, dtype=torch.long) < n_quantizers
+ quantized_representation = quantized_representation + quantized_representation_i * mask[:, None, None]
+ residual = residual - quantized_representation_i
+
+ # Sum losses
+ commitment_loss += commitment_loss_i * mask
+ codebook_loss += codebook_loss_i * mask
+
+ audio_codes.append(indices_i)
+ projected_latents.append(projected_latents_i)
+
+ audio_codes = torch.stack(audio_codes, dim=1)
+ projected_latents = torch.cat(projected_latents, dim=1)
+
+ return quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss
+
+ def from_codes(self, audio_codes: torch.Tensor):
+ """
+ Reconstructs the continuous representation from quantized codes.
+
+ Args:
+ audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`):
+ Quantized discrete representation of input.
+
+ Returns:
+ quantized_representation (`torch.Tensor`):
+ Quantized continuous representation of input.
+ projected_latents (`torch.Tensor`):
+ List of projected latents (continuous representations of input before quantization)
+ for each codebook.
+ audio_codes (`torch.Tensor`):
+ Codebook indices for each codebook.
+ """
+ quantized_representation = 0.0
+ projected_latents = []
+ n_codebooks = audio_codes.shape[1]
+ for i in range(n_codebooks):
+ projected_latents_i = self.quantizers[i].codebook(audio_codes[:, i, :]).transpose(1, 2)
+ projected_latents.append(projected_latents_i)
+ quantized_representation += self.quantizers[i].out_proj(projected_latents_i)
+ return quantized_representation, torch.cat(projected_latents, dim=1), audio_codes
+
+ def from_latents(self, latents: torch.Tensor):
+ """Reconstructs the quantized representation from unquantized latents.
+
+ Args:
+ latents (`torch.Tensor` of shape `(batch_size, total_latent_dimension, time_steps)`):
+ Continuous representation of input after projection.
+
+ Returns:
+ quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`):
+ Quantized representation of the full-projected space.
+ quantized_latents (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`):
+ Quantized representation of the latent space (continuous representation before quantization).
+ """
+ quantized_representation = 0
+ quantized_latents = []
+ codes = []
+ codebook_dims_tensor = torch.tensor([0] + [q.codebook_dim for q in self.quantizers])
+ dims = torch.cumsum(codebook_dims_tensor, dim=0)
+
+ n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]
+ for i in range(n_codebooks):
+ hidden_dim_j, hidden_dim_k = dims[i], dims[i + 1]
+ latent_chunk = latents[:, hidden_dim_j:hidden_dim_k, :]
+ quantized_latents_i, codes_i = self.quantizers[i].decode_latents(latent_chunk)
+ quantized_latents.append(quantized_latents_i)
+ codes.append(codes_i)
+
+ quantized_with_ste = latent_chunk + (quantized_latents_i - latent_chunk)
+ quantized_representation_i = self.quantizers[i].out_proj(quantized_with_ste)
+ quantized_representation = quantized_representation + quantized_representation_i
+
+ return quantized_representation, torch.cat(quantized_latents, dim=1)
+
+
+class DacDecoder(nn.Module):
+ """DAC Decoder"""
+
+ def __init__(self, config: DacConfig):
+ super().__init__()
+
+ input_channel = config.hidden_size
+ channels = config.decoder_hidden_size
+ strides = config.upsampling_ratios
+
+ # Add first conv layer
+ self.conv1 = nn.Conv1d(input_channel, channels, kernel_size=7, padding=3)
+
+ # Add upsampling + MRF blocks
+ block = []
+ for stride_index, stride in enumerate(strides):
+ block += [DacDecoderBlock(config, stride, stride_index)]
+
+ self.block = nn.ModuleList(block)
+ output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1)
+ self.snake1 = Snake1d(output_dim)
+ self.conv2 = nn.Conv1d(output_dim, 1, kernel_size=7, padding=3)
+ self.tanh = nn.Tanh()
+
+ def forward(self, hidden_state):
+ hidden_state = self.conv1(hidden_state)
+
+ for layer in self.block:
+ hidden_state = layer(hidden_state)
+
+ hidden_state = self.snake1(hidden_state)
+ hidden_state = self.conv2(hidden_state)
+ hidden_state = self.tanh(hidden_state)
+
+ return hidden_state
+
+
+class DacEncoder(nn.Module):
+ """DAC Encoder"""
+
+ def __init__(self, config: DacConfig):
+ super().__init__()
+
+ strides = config.downsampling_ratios
+ # Create first convolution
+ self.conv1 = nn.Conv1d(1, config.encoder_hidden_size, kernel_size=7, padding=3)
+
+ self.block = []
+ # Create EncoderBlocks that double channels as they downsample by `stride`
+ for stride_index, stride in enumerate(strides):
+ stride_index = stride_index + 1
+ self.block += [DacEncoderBlock(config, stride=stride, stride_index=stride_index)]
+
+ self.block = nn.ModuleList(self.block)
+ d_model = config.encoder_hidden_size * 2**stride_index
+ self.snake1 = Snake1d(d_model)
+ self.conv2 = nn.Conv1d(d_model, config.hidden_size, kernel_size=3, padding=1)
+
+ def forward(self, hidden_state):
+ hidden_state = self.conv1(hidden_state)
+
+ for module in self.block:
+ hidden_state = module(hidden_state)
+
+ hidden_state = self.snake1(hidden_state)
+ hidden_state = self.conv2(hidden_state)
+
+ return hidden_state
+
+
+@auto_docstring
+class DacPreTrainedModel(PreTrainedAudioTokenizerBase):
+ config: DacConfig
+ base_model_prefix = "dac"
+ main_input_name = "input_values"
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ if isinstance(module, nn.Conv1d):
+ init.trunc_normal_(module.weight, std=0.02)
+ init.constant_(module.bias, 0)
+ elif isinstance(module, Snake1d):
+ init.ones_(module.alpha)
+ elif isinstance(module, nn.ConvTranspose1d):
+ module.reset_parameters()
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=0.02)
+
+ def apply_weight_norm(self):
+ weight_norm = nn.utils.weight_norm
+ if hasattr(nn.utils.parametrizations, "weight_norm"):
+ weight_norm = nn.utils.parametrizations.weight_norm
+
+ for layer in self.quantizer.quantizers:
+ weight_norm(layer.in_proj)
+ weight_norm(layer.out_proj)
+
+ weight_norm(self.encoder.conv1)
+ weight_norm(self.encoder.conv2)
+
+ for layer in self.encoder.block:
+ weight_norm(layer.conv1)
+ weight_norm(layer.res_unit1.conv1)
+ weight_norm(layer.res_unit1.conv2)
+ weight_norm(layer.res_unit2.conv1)
+ weight_norm(layer.res_unit2.conv2)
+ weight_norm(layer.res_unit3.conv1)
+ weight_norm(layer.res_unit3.conv2)
+
+ weight_norm(self.decoder.conv1)
+ weight_norm(self.decoder.conv2)
+
+ for layer in self.decoder.block:
+ weight_norm(layer.conv_t1)
+ weight_norm(layer.res_unit1.conv1)
+ weight_norm(layer.res_unit1.conv2)
+ weight_norm(layer.res_unit2.conv1)
+ weight_norm(layer.res_unit2.conv2)
+ weight_norm(layer.res_unit3.conv1)
+ weight_norm(layer.res_unit3.conv2)
+
+ def remove_weight_norm(self):
+ for layer in self.quantizer.quantizers:
+ nn.utils.remove_weight_norm(layer.in_proj)
+ nn.utils.remove_weight_norm(layer.out_proj)
+
+ nn.utils.remove_weight_norm(self.encoder.conv1)
+ nn.utils.remove_weight_norm(self.encoder.conv2)
+
+ for layer in self.encoder.block:
+ nn.utils.remove_weight_norm(layer.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit1.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit1.conv2)
+ nn.utils.remove_weight_norm(layer.res_unit2.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit2.conv2)
+ nn.utils.remove_weight_norm(layer.res_unit3.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit3.conv2)
+
+ nn.utils.remove_weight_norm(self.decoder.conv1)
+ nn.utils.remove_weight_norm(self.decoder.conv2)
+
+ for layer in self.decoder.block:
+ nn.utils.remove_weight_norm(layer.conv_t1)
+ nn.utils.remove_weight_norm(layer.res_unit1.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit1.conv2)
+ nn.utils.remove_weight_norm(layer.res_unit2.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit2.conv2)
+ nn.utils.remove_weight_norm(layer.res_unit3.conv1)
+ nn.utils.remove_weight_norm(layer.res_unit3.conv2)
+
+
+@auto_docstring(
+ custom_intro="""
+ The DAC (Descript Audio Codec) model.
+ """
+)
+class DacModel(DacPreTrainedModel):
+ input_modalities = "audio"
+
+ def __init__(self, config: DacConfig):
+ super().__init__(config)
+ self.config = config
+
+ self.encoder = DacEncoder(config)
+ self.decoder = DacDecoder(config)
+
+ self.quantizer = DacResidualVectorQuantizer(config)
+
+ self.bits_per_codebook = int(math.log2(self.config.codebook_size))
+ if 2**self.bits_per_codebook != self.config.codebook_size:
+ raise ValueError("The codebook_size must be a power of 2.")
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def encode(
+ self,
+ input_values: torch.Tensor,
+ n_quantizers: int | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | DacEncoderOutput:
+ r"""
+ input_values (`torch.Tensor of shape `(batch_size, 1, time_steps)`):
+ Input audio data to encode,
+ n_quantizers (int, *optional*):
+ Number of quantizers to use. If None, all quantizers are used. Default is None.
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ quantized_representation = self.encoder(input_values)
+ quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss = self.quantizer(
+ quantized_representation, n_quantizers
+ )
+
+ loss = self.config.commitment_loss_weight * commitment_loss + self.config.codebook_loss_weight * codebook_loss
+
+ if not return_dict:
+ return (loss, quantized_representation, audio_codes, projected_latents)
+
+ return DacEncoderOutput(loss, quantized_representation, audio_codes, projected_latents)
+
+ @auto_docstring
+ def decode(
+ self,
+ quantized_representation: torch.Tensor | None = None,
+ audio_codes: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | DacDecoderOutput:
+ r"""
+ quantized_representation (torch.Tensor of shape `(batch_size, dimension, time_steps)`, *optional*):
+ Quantized continuous representation of input.
+ audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`, *optional*):
+ The codebook indices for each codebook, representing the quantized discrete
+ representation of the input. This parameter should be provided if you want
+ to decode directly from the audio codes (it will overwrite quantized_representation).
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`DacDecoderOutput`] instead of a plain tuple.
+ """
+
+ if quantized_representation is None and audio_codes is None:
+ raise ValueError("Either `quantized_representation` or `audio_codes` must be provided.")
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ if audio_codes is not None:
+ quantized_representation = self.quantizer.from_codes(audio_codes)[0]
+
+ audio_values = self.decoder(quantized_representation).squeeze(1)
+
+ if not return_dict:
+ return (audio_values,)
+
+ return DacDecoderOutput(audio_values)
+
+ @auto_docstring
+ def forward(
+ self,
+ input_values: torch.Tensor,
+ n_quantizers: int | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | DacOutput:
+ r"""
+ input_values (`torch.Tensor` of shape `(batch_size, 1, time_steps)`):
+ Audio data to encode.
+ n_quantizers (`int`, *optional*):
+ Number of quantizers to use. If `None`, all quantizers are used. Default is `None`.
+
+ Examples:
+
+ ```python
+ >>> from datasets import load_dataset, Audio
+ >>> from transformers import DacModel, AutoProcessor
+ >>> librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+ >>> model = DacModel.from_pretrained("descript/dac_16khz")
+ >>> processor = AutoProcessor.from_pretrained("descript/dac_16khz")
+ >>> librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
+ >>> audio_sample = librispeech_dummy[-1]["audio"]["array"]
+ >>> inputs = processor(raw_audio=audio_sample, sampling_rate=processor.sampling_rate, return_tensors="pt")
+
+ >>> encoder_outputs = model.encode(inputs["input_values"])
+ >>> # Get the intermediate audio codes
+ >>> audio_codes = encoder_outputs.audio_codes
+ >>> # Reconstruct the audio from its quantized representation
+ >>> audio_values = model.decode(encoder_outputs.quantized_representation)
+ >>> # or the equivalent with a forward pass
+ >>> audio_values = model(inputs["input_values"]).audio_values
+ ```"""
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ length = input_values.shape[-1]
+
+ loss, quantized_representation, audio_codes, projected_latents = self.encode(
+ input_values, n_quantizers, return_dict=False
+ )
+ audio_values = self.decode(quantized_representation, return_dict=False)[0][..., :length]
+
+ if not return_dict:
+ return (loss, audio_values, quantized_representation, audio_codes, projected_latents)
+
+ return DacOutput(loss, audio_values, quantized_representation, audio_codes, projected_latents)
+
+
+__all__ = ["DacModel", "DacPreTrainedModel"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ddca6bff1eb2098b2d32e6822887a96b51da1cc
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/__init__.py
@@ -0,0 +1,31 @@
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_paddleocr_vl import *
+ from .image_processing_paddleocr_vl import *
+ from .image_processing_pil_paddleocr_vl import *
+ from .modeling_paddleocr_vl import *
+ from .processing_paddleocr_vl import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ca58c6502876dbb8ff5e2ee752c1de52be6a28b
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py
@@ -0,0 +1,1695 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/paddleocr_vl/modular_paddleocr_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_paddleocr_vl.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The PaddlePaddle 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
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN, GELUActivation
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+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, 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, torch_int
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import (
+ accepts_precomputed_kwargs,
+ is_flash_attention_requested,
+ maybe_autocast,
+ merge_with_config_defaults,
+)
+from ...utils.output_capturing import capture_outputs
+from ...vision_utils import get_vision_cu_seqlens, get_vision_position_ids
+from .configuration_paddleocr_vl import PaddleOCRTextConfig, PaddleOCRVisionConfig, PaddleOCRVLConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class PaddleOCRProjector(nn.Module):
+ def __init__(self, config: PaddleOCRVLConfig):
+ super().__init__()
+ self.merge_kernel_size = (config.vision_config.spatial_merge_size, config.vision_config.spatial_merge_size)
+
+ hidden_size = config.vision_config.hidden_size * self.merge_kernel_size[0] * self.merge_kernel_size[1]
+
+ self.pre_norm = torch.nn.LayerNorm(config.vision_config.hidden_size, eps=1e-05)
+ self.linear_1 = nn.Linear(hidden_size, hidden_size, bias=True)
+ self.act = GELUActivation()
+ self.linear_2 = nn.Linear(hidden_size, config.text_config.hidden_size, bias=True)
+
+ def forward(self, image_features: torch.Tensor, image_grid_thw: torch.Tensor) -> torch.Tensor:
+ image_features_chunks = image_features.split(image_grid_thw.prod(dim=1).tolist(), dim=0)
+ m1, m2 = self.merge_kernel_size
+
+ processed_features = []
+ for image_feature, image_grid in zip(image_features_chunks, image_grid_thw):
+ image_feature = self.pre_norm(image_feature)
+ t, h, w = image_grid
+ d = image_feature.shape[-1]
+ h_block = h // m1
+ w_block = w // m2
+
+ image_feature = image_feature.reshape(t, h_block, m1, w_block, m2, d)
+ image_feature = image_feature.transpose(2, 3)
+ image_feature = image_feature.reshape(t * h_block * w_block, m1 * m2 * d)
+
+ hidden_states = self.linear_1(image_feature)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ processed_features.append(hidden_states)
+
+ return torch.cat(processed_features, dim=0)
+
+
+class PaddleOCRVisionRotaryEmbedding(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, position_ids: torch.Tensor) -> torch.Tensor:
+ return (position_ids.unsqueeze(-1) * self.inv_freq).flatten(1)
+
+
+class PaddleOCRRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: PaddleOCRVLConfig, 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: PaddleOCRVLConfig | 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, PaddleOCR 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 PaddleOCRMLP(nn.Module):
+ def __init__(self, config: PaddleOCRTextConfig):
+ 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.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
+
+
+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
+
+
+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_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 PaddleOCRAttention(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: PaddleOCRVLConfig, 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 = getattr(config, "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 = 0.0
+ self.rope_parameters = config.rope_parameters
+ self.scaling = self.head_dim**-0.5
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.use_bias)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_bias)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_bias)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias)
+ 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
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class PaddleOCRRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ PaddleOCRRMSNorm 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 PaddleOCRDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: PaddleOCRTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = PaddleOCRAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = PaddleOCRMLP(config)
+ self.input_layernorm = PaddleOCRRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = PaddleOCRRMSNorm(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 PaddleOCRVLPreTrainedModel(PreTrainedModel):
+ config: PaddleOCRVLConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["PaddleOCRDecoderLayer"]
+ _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": PaddleOCRDecoderLayer,
+ "attentions": PaddleOCRAttention,
+ }
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, PaddleOCRVisionEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, PaddleOCRVisionRotaryEmbedding):
+ inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+
+
+@auto_docstring
+class PaddleOCRTextModel(PaddleOCRVLPreTrainedModel):
+ def __init__(self, config: PaddleOCRTextConfig):
+ 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(
+ [PaddleOCRDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = PaddleOCRRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = PaddleOCRRotaryEmbedding(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.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)
+
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ text_position_ids = None
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=text_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=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,
+ )
+
+
+class PaddleOCRVisionEmbeddings(nn.Module):
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ padding="valid",
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.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 and no class embeddings.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+ num_positions = self.position_embedding.weight.shape[0]
+
+ patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
+
+ dim = embeddings.shape[-1]
+
+ 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=(height, width),
+ mode="bilinear",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+ return patch_pos_embed
+
+ @deprecate_kwarg("image_grid_thw", new_name="grid_thw", version="5.11.0")
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ grid_thw: torch.LongTensor | None = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`):
+ The tensors corresponding to the input images.
+ grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ batch_size, squence_len, channel, height, width = pixel_values.shape
+ target_dtype = self.patch_embedding.weight.dtype
+ pixel_values = pixel_values.reshape(batch_size * squence_len, channel, height, width)
+ patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
+ embeddings = patch_embeds.flatten(-2).squeeze(-1)
+ embeddings = embeddings.reshape(batch_size, squence_len, -1)
+
+ start = 0
+ embeddings = embeddings.squeeze(0)
+ tmp_embeddings = []
+ for t, h, w in grid_thw:
+ end = start + t * h * w
+ image_embeddings = embeddings[start:end, :]
+ position_embedding = self.interpolate_pos_encoding(image_embeddings, h, w).squeeze(0).repeat(t, 1)
+ image_embeddings = image_embeddings + position_embedding
+ tmp_embeddings.append(image_embeddings)
+ start = end
+ embeddings = torch.concat(tmp_embeddings, dim=0)
+
+ return embeddings
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+class PaddleOCRVisionAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.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)
+ self.num_key_value_groups = 1
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """
+ Args:
+ hidden_states (`torch.Tensor`):
+ Input to the layer of shape `(seq_len, embed_dim)`.
+ cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`):
+ The cumulative sequence lengths of each image or video feature.
+ position_embeddings (`tuple(torch.Tensor, torch.Tensor)` of shape `(num_patches, head_dim // 2)`):
+ The cosine and sine position embeddings for vision attention.
+ """
+ seq_length = hidden_states.shape[0]
+ query_states = self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)
+ key_states = self.k_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)
+ value_states = self.v_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)
+
+ 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 2: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, attn_weights = 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, attn_weights = [], []
+ for q, k, v in zip(*splits):
+ attn_output, attn_weight = 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,
+ )
+ attn_outputs.append(attn_output)
+ attn_weights.append(attn_weight)
+
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class PaddleOCRVisionMLP(nn.Module):
+ def __init__(self, config: PaddleOCRVisionConfig):
+ 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 PaddleOCRVisionEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.self_attn = PaddleOCRVisionAttention(config=config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = PaddleOCRVisionMLP(config=config)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`):
+ The cumulative sequence lengths of each image or video feature.
+ position_embeddings (`tuple(torch.Tensor, torch.Tensor)` of shape `(num_patches, head_dim // 2)`):
+ The cosine and sine position embeddings for vision attention.
+ """
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class PaddleOCRVisionEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`PaddleOCRVisionEncoderLayer`].
+
+ Args:
+ config: PaddleOCRVisionConfig
+ """
+
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([PaddleOCRVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+ embed_dim = config.hidden_size
+ num_heads = config.num_attention_heads
+ head_dim = embed_dim // num_heads
+ self.rotary_pos_emb = PaddleOCRVisionRotaryEmbedding(head_dim // 2)
+
+ # Ignore copy
+ @can_return_tuple
+ @auto_docstring
+ @deprecate_kwarg("image_grid_thw", new_name="grid_thw", version="5.11.0")
+ def forward(
+ self,
+ inputs_embeds: torch.FloatTensor,
+ attention_mask: torch.Tensor | None = None,
+ grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ r"""
+ inputs_embeds (`torch.FloatTensor` of shape `(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.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ The attention_mask used in forward function shape [batch_size X sequence_length] if not None.
+ grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ # Use merge_size=1: PaddleOCR merges patches in the projector (after the encoder),
+ # unlike Qwen which merges inside the encoder, so rotary positions here are simple (row, col).
+ position_ids = get_vision_position_ids(grid_thw, 1, kwargs=kwargs)
+ cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs)
+
+ hidden_states = inputs_embeds
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+ rotary_embeddings = self.rotary_pos_emb(position_ids)
+ rotary_embeddings = rotary_embeddings.repeat(1, 2)
+ position_embeddings = (rotary_embeddings.cos(), rotary_embeddings.sin())
+
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class PaddleOCRVisionTransformer(PaddleOCRVLPreTrainedModel):
+ config: PaddleOCRVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = "image"
+ _can_record_outputs = {
+ "hidden_states": PaddleOCRVisionEncoderLayer,
+ "attentions": PaddleOCRVisionAttention,
+ }
+
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__(config)
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = PaddleOCRVisionEmbeddings(config)
+ self.encoder = PaddleOCRVisionEncoder(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)
+ @deprecate_kwarg("image_grid_thw", new_name="grid_thw", version="5.11.0")
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ attention_mask: torch.Tensor | None = None,
+ grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPooling:
+ """
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, patch_size * patch_size * image_channels)`):
+ The tensors corresponding to the input images.
+ attention_mask (`torch.Tensor`, *optional*):
+ The attention_mask used in forward function shape [batch_size X sequence_length] if not None.
+ grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ hidden_states = self.embeddings(pixel_values, grid_thw=grid_thw)
+ encoder_outputs: BaseModelOutput = self.encoder(
+ inputs_embeds=hidden_states,
+ grid_thw=grid_thw,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs.last_hidden_state
+ last_hidden_state = self.post_layernorm(last_hidden_state)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=None,
+ )
+
+
+class PaddleOCRVisionModel(PaddleOCRVLPreTrainedModel):
+ config: PaddleOCRVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = "image"
+
+ def __init__(self, config: PaddleOCRVisionConfig):
+ super().__init__(config)
+
+ self.vision_model = PaddleOCRVisionTransformer(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @deprecate_kwarg("image_grid_thw", new_name="grid_thw", version="5.11.0")
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ """
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`):
+ The tensors corresponding to the input images.
+ 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.vision_model(pixel_values=pixel_values, grid_thw=grid_thw, **kwargs)
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Llava outputs, with hidden states and attentions.
+ """
+)
+@dataclass
+class PaddleOCRVLModelOutputWithPast(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
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for PaddleOCRVL causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class PaddleOCRVLCausalLMOutputWithPast(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
+
+
+@auto_docstring
+class PaddleOCRVLModel(PaddleOCRVLPreTrainedModel):
+ base_model_prefix = "model"
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ _keys_to_ignore_on_load_unexpected = ["packing_position_embedding", "vision_model.head"]
+
+ def __init__(self, config: PaddleOCRVLConfig):
+ super().__init__(config)
+ self.visual = PaddleOCRVisionModel._from_config(config.vision_config)
+ self.language_model = PaddleOCRTextModel._from_config(config.text_config)
+ self.rope_deltas = None
+ self.projector = PaddleOCRProjector(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ 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,
+ )
+
+ # Add `start_position` after arange for compile
+ position_temporal = torch.arange(llm_grid_t, device=device) * time_interval
+ position_width = torch.arange(llm_grid_w, device=device) + start_position
+ position_height = torch.arange(llm_grid_h, device=device) + start_position
+
+ # Repeat the positions per each grid and per video frame. Repeat patterns are important
+ # do not modify without checking values!
+ position_width = position_width.repeat(llm_grid_h * llm_grid_t)
+ position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t)
+ # Important: add `start_positions` after applying `time_interval`, order matters
+ position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position
+ vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
+
+ return vision_position_ids
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor,
+ mm_token_type_ids: torch.IntTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text`
+ sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred
+ position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width)
+ while text tokens use standard 1D RoPE.
+
+ Example:
+ Temporal patches: 3; Height patches: 2; Width patches: 2
+ Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total.
+
+ Temporal position IDs are spaced by:
+ `interval = tokens_per_second * temporal_patch_size / fps`
+
+ If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch:
+ `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]`
+
+ Height IDs repeat per row: `[0, 0, 1, 1, ...]`
+ Width IDs alternate per column: `[0, 1, 0, 1, ...]`
+ Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1`
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
+ Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2).
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+
+ mrope_position_deltas = []
+ position_ids = torch.zeros(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ grid_iters = {
+ 1: iter(image_grid_thw) if image_grid_thw is not None else None,
+ 2: iter(video_grid_thw) if video_grid_thw is not None else None,
+ }
+
+ for batch_idx, current_input_ids in enumerate(input_ids):
+ input_token_type = mm_token_type_ids[batch_idx]
+ if attention_mask is not None:
+ current_input_ids = current_input_ids[attention_mask[batch_idx].bool()]
+ input_token_type = input_token_type[attention_mask[batch_idx].bool()]
+
+ input_type_group = []
+ for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]):
+ group = list(group)
+ start_index = group[0][0]
+ end_index = group[-1][0] + 1
+ input_type_group.append((key, start_index, end_index))
+
+ current_pos = 0
+ 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])
+ vision_position_ids = self.get_vision_position_ids(
+ current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device
+ )
+ llm_pos_ids_list.append(vision_position_ids)
+ current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ if attention_mask is not None:
+ position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device)
+ else:
+ position_ids[:, batch_idx] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+
+ @accepts_precomputed_kwargs(modality="image")
+ @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).unsqueeze(0)
+ vision_outputs = self.visual(pixel_values=pixel_values, grid_thw=image_grid_thw, **kwargs)
+ image_embeds = vision_outputs.last_hidden_state
+ image_embeds = self.projector(image_embeds, image_grid_thw)
+ 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
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ 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
+
+ def compute_3d_position_ids(
+ self,
+ input_ids: torch.Tensor | None,
+ inputs_embeds: torch.Tensor | None,
+ image_grid_thw: torch.Tensor | None = None,
+ video_grid_thw: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: torch.Tensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ ) -> torch.Tensor | None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ has_multimodal = image_grid_thw is not None or video_grid_thw is not None
+ if has_multimodal and mm_token_type_ids is None and input_ids is not None:
+ raise ValueError(
+ "Multimodal data was passed (via `image_grid_thw` or `video_grid_thw`) but `mm_token_type_ids` is "
+ "missing. Please pass `mm_token_type_ids` to the model so that multimodal RoPE (M-RoPE) can be "
+ "computed correctly. `mm_token_type_ids` is returned by the processor alongside `input_ids`."
+ )
+ can_compute_mrope = input_ids is not None and mm_token_type_ids is not None and has_multimodal
+
+ if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ attention_mask=attention_mask,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+ self.rope_deltas = rope_deltas
+ # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+ # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+ # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+ # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+ elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+ position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device)
+ else:
+ position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length)
+ position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device)
+ delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0)
+ position_ids = position_ids + delta.to(device=inputs_embeds.device)
+ else:
+ # Can't build correct 3D positions. Let the model infer it
+ position_ids = None
+ return position_ids
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: list[torch.FloatTensor] | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple | PaddleOCRVLModelOutputWithPast:
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+ if inputs_embeds is None:
+ inputs_embeds = self.language_model.embed_tokens(input_ids)
+
+ if pixel_values is not None:
+ image_embeds = self.get_image_features(
+ pixel_values, image_grid_thw, return_dict=True, **kwargs
+ ).pooler_output
+ image_embeds = image_embeds.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 position_ids is None:
+ position_ids = self.compute_3d_position_ids(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ output = PaddleOCRVLModelOutputWithPast(
+ 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,
+ )
+
+ return output
+
+
+class PaddleOCRVLForConditionalGeneration(PaddleOCRVLPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+ _keys_to_ignore_on_load_unexpected = ["packing_position_embedding", "vision_model.head"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = PaddleOCRVLModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+ @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,
+ image_grid_thw: torch.LongTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | PaddleOCRVLCausalLMOutputWithPast:
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, PaddleOCRVLForConditionalGeneration
+
+ >>> model = PaddleOCRVLForConditionalGeneration.from_pretrained("PaddlePaddle/PaddleOCR-VL", dtype="bfloat16")
+ >>> processor = AutoProcessor.from_pretrained("PaddlePaddle/PaddleOCR-VL")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "image": "https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/ocr_demo.jpg",
+ },
+ {"type": "text", "text": "OCR:"},
+ ],
+ }
+ ]
+
+ >>> inputs = processor.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ return_dict=True,
+ return_tensors="pt"
+ ).to(model.device)
+
+ >>> # 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: PaddleOCRVLModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ image_grid_thw=image_grid_thw,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ pixel_values=pixel_values,
+ rope_deltas=rope_deltas,
+ mm_token_type_ids=mm_token_type_ids,
+ **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=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return PaddleOCRVLCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=outputs.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs):
+ # Overwritten -- requires 3D position ids
+
+ text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)
+
+ # Early exit in case we are continuing generation from past kv
+ past_length = 0
+ if (cache := model_kwargs.get("past_key_values")) is not None:
+ past_length = cache.get_seq_length()
+ if past_length != 0 and self.model.rope_deltas is not None:
+ position_ids = text_positions[None, ...] + self.model.rope_deltas
+ return position_ids
+
+ # Otherwise compute 3d position ids for vision tokens and concat with text position ids
+ if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+ inputs_tensor = model_kwargs["input_ids"]
+
+ is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+ if (
+ is_input_ids
+ and model_kwargs.get("mm_token_type_ids") is not None
+ and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None)
+ ):
+ model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"}
+ vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs)
+ self.model.rope_deltas = rope_deltas
+ else:
+ vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1)
+ self.model.rope_deltas = torch.zeros(
+ inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device
+ )
+
+ # Concatenate "text + vision" positions into [4, bs, seq-len]
+ text_positions = text_positions[None, ...]
+ position_ids = torch.cat([text_positions, vision_positions], dim=0)
+
+ return position_ids
+
+ def _get_image_nums_and_video_nums(
+ self,
+ input_ids: torch.LongTensor | None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+ These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Returns:
+ image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+ video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+ """
+ 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__ = [
+ "PaddleOCRVLForConditionalGeneration",
+ "PaddleOCRVLModel",
+ "PaddleOCRVLPreTrainedModel",
+ "PaddleOCRVisionTransformer",
+ "PaddleOCRTextModel",
+ "PaddleOCRVisionModel",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/configuration_perceiver.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/configuration_perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..02d7547af9ad4bd68a67d9d48ba1409cda5fba53
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/configuration_perceiver.py
@@ -0,0 +1,114 @@
+# Copyright Deepmind 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.
+"""Perceiver model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="deepmind/language-perceiver")
+@strict
+class PerceiverConfig(PreTrainedConfig):
+ r"""
+ num_latents (`int`, *optional*, defaults to 256):
+ The number of latents.
+ d_latents (`int`, *optional*, defaults to 1280):
+ Dimension of the latent embeddings.
+ num_blocks (`int`, *optional*, defaults to 1):
+ Number of blocks in the Transformer encoder.
+ num_self_attends_per_block (`int`, *optional*, defaults to 26):
+ The number of self-attention layers per block.
+ num_self_attention_heads (`int`, *optional*, defaults to 8):
+ Number of attention heads for each self-attention layer in the Transformer encoder.
+ num_cross_attention_heads (`int`, *optional*, defaults to 8):
+ Number of attention heads for each cross-attention layer in the Transformer encoder.
+ qk_channels (`int`, *optional*):
+ Dimension to project the queries + keys before applying attention in the cross-attention and self-attention
+ layers of the encoder. Will default to preserving the dimension of the queries if not specified.
+ v_channels (`int`, *optional*):
+ Dimension to project the values before applying attention in the cross-attention and self-attention layers
+ of the encoder. Will default to preserving the dimension of the queries if not specified.
+ cross_attention_shape_for_attention (`str`, *optional*, defaults to `"kv"`):
+ Dimension to use when downsampling the queries and keys in the cross-attention layer of the encoder.
+ self_attention_widening_factor (`int`, *optional*, defaults to 1):
+ Dimension of the feed-forward layer in the cross-attention layer of the Transformer encoder.
+ cross_attention_widening_factor (`int`, *optional*, defaults to 1):
+ Dimension of the feed-forward layer in the self-attention layers of the Transformer encoder.
+ use_query_residual (`float`, *optional*, defaults to `True`):
+ Whether to add a query residual in the cross-attention layer of the encoder.
+ image_size (`int`, *optional*, defaults to 56):
+ Size of the images after preprocessing, for [`PerceiverForImageClassificationLearned`].
+ train_size (`list[int]`, *optional*, defaults to `[368, 496]`):
+ Training size of the images for the optical flow model.
+ num_frames (`int`, *optional*, defaults to 16):
+ Number of video frames used for the multimodal autoencoding model.
+ audio_samples_per_frame (`int`, *optional*, defaults to 1920):
+ Number of audio samples per frame for the multimodal autoencoding model.
+ samples_per_patch (`int`, *optional*, defaults to 16):
+ Number of audio samples per patch when preprocessing the audio for the multimodal autoencoding model.
+ output_shape (`list[int]`, *optional*, defaults to `[1, 16, 224, 224]`):
+ Shape of the output (batch_size, num_frames, height, width) for the video decoder queries of the multimodal
+ autoencoding model. This excludes the channel dimension.
+ output_num_channels (`int`, *optional*, defaults to 512):
+ Number of output channels for each modalitiy decoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import PerceiverModel, PerceiverConfig
+
+ >>> # Initializing a Perceiver deepmind/language-perceiver style configuration
+ >>> configuration = PerceiverConfig()
+
+ >>> # Initializing a model from the deepmind/language-perceiver style configuration
+ >>> model = PerceiverModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "perceiver"
+
+ num_latents: int = 256
+ d_latents: int = 1280
+ d_model: int = 768
+ num_blocks: int = 1
+ num_self_attends_per_block: int = 26
+ num_self_attention_heads: int = 8
+ num_cross_attention_heads: int = 8
+ qk_channels: int | None = None
+ v_channels: int | None = None
+ cross_attention_shape_for_attention: str = "kv"
+ self_attention_widening_factor: int = 1
+ cross_attention_widening_factor: int = 1
+ hidden_act: str = "gelu"
+ attention_probs_dropout_prob: float | int = 0.1
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ use_query_residual: bool = True
+ vocab_size: int = 262
+ max_position_embeddings: int = 2048
+ image_size: int | list[int] | tuple[int, int] = 56
+ train_size: list[int] | tuple[int, ...] = (368, 496)
+ num_frames: int = 16
+ audio_samples_per_frame: int = 1920
+ samples_per_patch: int = 16
+ output_shape: list[int] | tuple[int, ...] = (1, 16, 224, 224)
+ output_num_channels: int = 512
+ _label_trainable_num_channels: int = 1024
+
+
+__all__ = ["PerceiverConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_perceiver.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6b8e5896d341a7b30b01bd17e938381236686e2
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_perceiver.py
@@ -0,0 +1,124 @@
+# 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 Perceiver."""
+
+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_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+@auto_docstring
+class PerceiverImageProcessor(TorchvisionBackend):
+ """Torchvision backend for Perceiver with custom center crop."""
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = IMAGENET_DEFAULT_MEAN
+ image_std = IMAGENET_DEFAULT_STD
+ size = {"height": 224, "width": 224}
+ crop_size = {"height": 256, "width": 256}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+
+ def __init__(self, **kwargs: Unpack[ImagesKwargs]):
+ super().__init__(**kwargs)
+
+ def center_crop(
+ self,
+ image: "torch.Tensor",
+ size: SizeDict,
+ crop_size: SizeDict,
+ **kwargs,
+ ) -> "torch.Tensor":
+ """
+ Center crop an image to ((size.height / crop_size.height) * min_dim, (size.width / crop_size.width) * min_dim),
+ where min_dim is the minimum of the image height and width.
+ If the requested crop size exceeds the image dimensions along any edge, the image is padded with zeros before
+ center cropping.
+ """
+ if size.height is None or size.width is None:
+ raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}")
+ if crop_size.height is None or crop_size.width is None:
+ raise ValueError(f"The crop_size dictionary must have keys 'height' and 'width'. Got {crop_size.keys()}")
+ height, width = image.shape[-2:]
+ min_dim = min(height, width)
+ cropped_height = int((size.height / crop_size.height) * min_dim)
+ cropped_width = int((size.width / crop_size.width) * min_dim)
+ return super().center_crop(
+ image,
+ SizeDict(height=cropped_height, width=cropped_width),
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for Perceiver: center_crop -> resize -> rescale and normalize."""
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ cropped_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_center_crop:
+ stacked_images = self.center_crop(stacked_images, size=size, crop_size=crop_size)
+ cropped_images_grouped[shape] = stacked_images
+ cropped_images = reorder_images(cropped_images_grouped, grouped_images_index)
+
+ grouped_images, grouped_images_index = group_images_by_shape(cropped_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)
+
+ 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
+ )
+ 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__ = ["PerceiverImageProcessor"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_pil_perceiver.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_pil_perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1ab681929ddfdf67e824cca5505b6c3d16ce7ed
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/image_processing_pil_perceiver.py
@@ -0,0 +1,107 @@
+# 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 Perceiver."""
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_utils import (
+ IMAGENET_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+@auto_docstring
+class PerceiverImageProcessorPil(PilBackend):
+ """PIL backend for Perceiver with custom center crop."""
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = IMAGENET_DEFAULT_MEAN
+ image_std = IMAGENET_DEFAULT_STD
+ size = {"height": 224, "width": 224}
+ crop_size = {"height": 256, "width": 256}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+
+ def __init__(self, **kwargs: Unpack[ImagesKwargs]):
+ super().__init__(**kwargs)
+
+ def center_crop(
+ self,
+ image: np.ndarray,
+ size: SizeDict,
+ crop_size: SizeDict,
+ **kwargs,
+ ) -> np.ndarray:
+ """
+ Center crop an image to ((size.height / crop_size.height) * min_dim, (size.width / crop_size.width) * min_dim),
+ where min_dim is the minimum of the image height and width.
+ If the requested crop size exceeds the image dimensions along any edge, the image is padded with zeros before
+ center cropping.
+ """
+ if size.height is None or size.width is None:
+ raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}")
+ if crop_size.height is None or crop_size.width is None:
+ raise ValueError(f"The crop_size dictionary must have keys 'height' and 'width'. Got {crop_size.keys()}")
+ height, width = image.shape[-2:]
+ min_dim = min(height, width)
+ cropped_height = int((size.height / crop_size.height) * min_dim)
+ cropped_width = int((size.width / crop_size.width) * min_dim)
+ return super().center_crop(
+ image,
+ SizeDict(height=cropped_height, width=cropped_width),
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for Perceiver: center_crop -> resize -> rescale and normalize."""
+ processed_images = []
+ for image in images:
+ if do_center_crop:
+ image = self.center_crop(image, size=size, crop_size=crop_size)
+ if do_resize:
+ image = self.resize(image=image, size=size, resample=resample)
+ 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__ = ["PerceiverImageProcessorPil"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/modeling_perceiver.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/modeling_perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..db0088b96d67ac6544f43f481595750db2184995
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/modeling_perceiver.py
@@ -0,0 +1,3307 @@
+# Copyright 2021 Deepmind 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 Perceiver model."""
+
+import abc
+import math
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import reduce
+from operator import __add__
+from typing import Any, Optional
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...masking_utils import create_bidirectional_mask
+from ...modeling_outputs import BaseModelOutputWithCrossAttentions
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import ModelOutput, auto_docstring, logging, torch_int
+from .configuration_perceiver import PerceiverConfig
+
+
+ModalitySizeType = Mapping[str, int]
+PreprocessorOutputType = tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]
+PreprocessorType = Callable[..., PreprocessorOutputType]
+PostprocessorType = Callable[..., Any]
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Perceiver base model's outputs, with potential hidden states, attentions and cross-attentions.
+ """
+)
+@dataclass
+class PerceiverModelOutput(ModelOutput):
+ r"""
+ logits (`torch.FloatTensor` of shape `(batch_size, num_labels)`):
+ Classification (or regression if config.num_labels==1) scores (before SoftMax).
+ """
+
+ logits: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Perceiver decoder outputs, with potential cross-attentions.
+ """
+)
+@dataclass
+class PerceiverDecoderOutput(ModelOutput):
+ r"""
+ logits (`torch.FloatTensor` of shape `(batch_size, num_labels)`):
+ Output of the basic decoder.
+ """
+
+ logits: torch.FloatTensor | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Perceiver's masked language model outputs.
+ """
+)
+@dataclass
+class PerceiverMaskedLMOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Masked language modeling (MLM) 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).
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Perceiver's outputs of sequence/image classification models, optical flow and multimodal
+ autoencoding.
+ """
+)
+@dataclass
+class PerceiverClassifierOutput(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).
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: 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 PerceiverEmbeddings(nn.Module):
+ """Construct the latent embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.latents = nn.Parameter(torch.randn(config.num_latents, config.d_latents))
+
+ def forward(self, batch_size: int):
+ return self.latents.expand(batch_size, -1, -1) # Thanks, Phil Wang
+
+
+class PerceiverSelfAttention(nn.Module):
+ """Multi-headed {cross, self}-attention. Can be used both in the encoder as well as in the decoder."""
+
+ def __init__(
+ self,
+ config,
+ is_cross_attention=False,
+ qk_channels=None,
+ v_channels=None,
+ num_heads=1,
+ q_dim=None,
+ kv_dim=None,
+ ):
+ super().__init__()
+ self.num_heads = num_heads
+ # Q and K must have the same number of channels.
+ # Default to preserving Q's input's shape.
+ if qk_channels is None:
+ qk_channels = q_dim
+ # V's num_channels determines the shape of the output of QKV-attention.
+ # Default to the same number of channels used in the key-query operation.
+ if v_channels is None:
+ v_channels = qk_channels
+ if qk_channels % num_heads != 0:
+ raise ValueError(f"qk_channels ({qk_channels}) must be divisible by num_heads ({num_heads}).")
+ if v_channels % num_heads != 0:
+ raise ValueError(f"v_channels ({v_channels}) must be divisible by num_heads ({num_heads}).")
+
+ self.qk_channels = qk_channels
+ self.v_channels = v_channels
+ self.qk_channels_per_head = self.qk_channels // num_heads
+ self.v_channels_per_head = self.v_channels // num_heads
+
+ # Layer normalization
+ self.layernorm1 = nn.LayerNorm(q_dim)
+ self.layernorm2 = nn.LayerNorm(kv_dim) if is_cross_attention else nn.Identity()
+
+ # Projection matrices
+ self.query = nn.Linear(q_dim, qk_channels)
+ self.key = nn.Linear(kv_dim, qk_channels)
+ self.value = nn.Linear(kv_dim, v_channels)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def transpose_for_scores(self, x, channels_per_head):
+ new_x_shape = x.size()[:-1] + (self.num_heads, channels_per_head)
+ x = x.view(*new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs: torch.FloatTensor | None = None,
+ inputs_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ hidden_states = self.layernorm1(hidden_states)
+ inputs = self.layernorm2(inputs)
+
+ # Project queries, keys and values to a common feature dimension. If this is instantiated as a cross-attention module,
+ # the keys and values come from the inputs; the attention mask needs to be such that the inputs's non-relevant tokens are not attended to.
+ is_cross_attention = inputs is not None
+ queries = self.query(hidden_states)
+
+ if is_cross_attention:
+ keys = self.key(inputs)
+ values = self.value(inputs)
+ attention_mask = inputs_mask
+ else:
+ keys = self.key(hidden_states)
+ values = self.value(hidden_states)
+
+ # Reshape channels for multi-head attention.
+ # We reshape from (batch_size, time, channels) to (batch_size, num_heads, time, channels per head)
+ queries = self.transpose_for_scores(queries, self.qk_channels_per_head)
+ keys = self.transpose_for_scores(keys, self.qk_channels_per_head)
+ values = self.transpose_for_scores(values, self.v_channels_per_head)
+
+ # Take the dot product between the queries and keys to get the raw attention scores.
+ attention_scores = torch.matmul(queries, keys.transpose(-1, -2))
+
+ batch_size, num_heads, seq_len, q_head_dim = queries.shape
+ _, _, _, v_head_dim = values.shape
+ hiddens = self.num_heads * v_head_dim
+
+ attention_scores = attention_scores / math.sqrt(q_head_dim)
+
+ if attention_mask is not None:
+ # Apply the attention mask (precomputed for all layers in PerceiverModel 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, values)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (hiddens,)
+ context_layer = context_layer.view(*new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+ return outputs
+
+
+class PerceiverSelfOutput(nn.Module):
+ def __init__(self, config, input_channels, output_channels):
+ super().__init__()
+ self.dense = nn.Linear(input_channels, output_channels)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ return hidden_states
+
+
+class PerceiverAttention(nn.Module):
+ """Attention module, including a dense block."""
+
+ def __init__(
+ self,
+ config,
+ is_cross_attention=False,
+ qk_channels=None,
+ v_channels=None,
+ num_heads=1,
+ q_dim=None,
+ kv_dim=None,
+ use_query_residual=True,
+ ):
+ super().__init__()
+ # MultiHead attention
+ if is_cross_attention and qk_channels is None:
+ if config.cross_attention_shape_for_attention == "q":
+ qk_channels = q_dim
+ elif config.cross_attention_shape_for_attention == "kv":
+ qk_channels = kv_dim
+ else:
+ raise ValueError(
+ f"Unknown value {config.cross_attention_shape_for_attention} for "
+ "cross_attention_shape_for_attention."
+ )
+ else:
+ if qk_channels is None:
+ qk_channels = q_dim
+ if v_channels is None:
+ v_channels = qk_channels
+ self.self = PerceiverSelfAttention(
+ config,
+ is_cross_attention=is_cross_attention,
+ qk_channels=qk_channels,
+ v_channels=v_channels,
+ num_heads=num_heads,
+ q_dim=q_dim,
+ kv_dim=kv_dim,
+ )
+ # dense block
+ output_channels = None
+ if is_cross_attention:
+ output_channels = q_dim
+ else:
+ if output_channels is None:
+ output_channels = v_channels
+ self.output = PerceiverSelfOutput(config, input_channels=self.self.v_channels, output_channels=output_channels)
+ self.use_query_residual = use_query_residual
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs: torch.FloatTensor | None = None,
+ inputs_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ self_outputs = self.self(
+ hidden_states,
+ attention_mask,
+ inputs,
+ inputs_mask,
+ output_attentions,
+ )
+
+ # Output projection
+ attention_output = self.output(self_outputs[0])
+
+ # Optionally include a residual to the original queries.
+ # Consider omitting the residual if the semantics of query and output
+ # are different, e.g. if queries are positions and outputs are pixels.
+ if self.use_query_residual:
+ attention_output = attention_output + hidden_states
+
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class PerceiverMLP(nn.Module):
+ """A Transformer-style dense module to follow attention."""
+
+ def __init__(self, config, input_size, widening_factor):
+ super().__init__()
+ self.dense1 = nn.Linear(input_size, widening_factor * input_size)
+ 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(widening_factor * input_size, input_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense1(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ hidden_states = self.dense2(hidden_states)
+ return hidden_states
+
+
+class PerceiverLayer(nn.Module):
+ def __init__(
+ self,
+ config,
+ is_cross_attention=False,
+ qk_channels=None,
+ v_channels=None,
+ num_heads=1,
+ q_dim=None,
+ kv_dim=None,
+ widening_factor=4,
+ use_query_residual=True,
+ ):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = PerceiverAttention(
+ config,
+ is_cross_attention=is_cross_attention,
+ qk_channels=qk_channels,
+ v_channels=v_channels,
+ num_heads=num_heads,
+ q_dim=q_dim,
+ kv_dim=kv_dim,
+ use_query_residual=use_query_residual,
+ )
+ self.layernorm = nn.LayerNorm(q_dim)
+ self.mlp = PerceiverMLP(config, input_size=q_dim, widening_factor=widening_factor)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs: torch.FloatTensor | None = None,
+ inputs_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ attention_outputs = self.attention(
+ hidden_states,
+ attention_mask,
+ inputs,
+ inputs_mask,
+ output_attentions,
+ )
+ attention_output = attention_outputs[0]
+
+ outputs = attention_outputs[1:] # add attentions if we output attention weights
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+
+ layer_output = layer_output + attention_output # residual connection
+
+ outputs = (layer_output,) + outputs
+
+ return outputs
+
+ def feed_forward_chunk(self, attention_output):
+ layer_output = self.layernorm(attention_output)
+ layer_output = self.mlp(layer_output)
+ return layer_output
+
+
+class PerceiverEncoder(nn.Module):
+ """The Perceiver Encoder: a scalable, fully attentional encoder."""
+
+ def __init__(self, config, kv_dim=None):
+ super().__init__()
+ self.config = config
+
+ # Check that we can use multihead-attention with these shapes.
+ if config.d_latents % config.num_self_attention_heads != 0:
+ raise ValueError(
+ f"num_z_channels ({config.d_latents}) must be divisible by"
+ f" num_self_attend_heads ({config.num_self_attention_heads})."
+ )
+ if config.d_latents % config.num_cross_attention_heads != 0:
+ raise ValueError(
+ f"num_z_channels ({config.d_latents}) must be divisible by"
+ f" num_cross_attend_heads ({config.num_cross_attention_heads})."
+ )
+
+ # Construct the cross attention layer.
+ self.cross_attention = PerceiverLayer(
+ config,
+ is_cross_attention=True,
+ qk_channels=config.qk_channels,
+ v_channels=config.v_channels,
+ num_heads=config.num_cross_attention_heads,
+ q_dim=config.d_latents,
+ kv_dim=kv_dim,
+ widening_factor=config.cross_attention_widening_factor,
+ use_query_residual=config.use_query_residual,
+ )
+
+ # Construct a single block of self-attention layers.
+ # We get deeper architectures by applying this block more than once.
+ self_attention_layers = []
+ for _ in range(config.num_self_attends_per_block):
+ layer = PerceiverLayer(
+ config,
+ is_cross_attention=False,
+ qk_channels=config.qk_channels,
+ v_channels=config.v_channels,
+ num_heads=config.num_self_attention_heads,
+ q_dim=config.d_latents,
+ kv_dim=config.d_latents,
+ widening_factor=config.self_attention_widening_factor,
+ )
+ self_attention_layers.append(layer)
+
+ self.self_attends = nn.ModuleList(self_attention_layers)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs: torch.FloatTensor | None = None,
+ inputs_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ output_hidden_states: bool | None = False,
+ return_dict: bool | None = True,
+ ) -> tuple | BaseModelOutputWithCrossAttentions:
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+ all_cross_attentions = () if output_attentions else None
+
+ # Apply the cross-attention between the latents (hidden_states) and inputs:
+ layer_outputs = self.cross_attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ inputs=inputs,
+ inputs_mask=inputs_mask,
+ output_attentions=output_attentions,
+ )
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ all_cross_attentions = all_cross_attentions + (layer_outputs[1],)
+
+ # Apply the block of self-attention layers more than once:
+ for _ in range(self.config.num_blocks):
+ for i, layer_module in enumerate(self.self_attends):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_outputs = layer_module(
+ hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, all_hidden_states, all_self_attentions, all_cross_attentions]
+ if v is not None
+ )
+ return BaseModelOutputWithCrossAttentions(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+@auto_docstring
+class PerceiverPreTrainedModel(PreTrainedModel):
+ config: PerceiverConfig
+ base_model_prefix = "perceiver"
+ main_input_name = "inputs"
+ input_modalities = ("image",) # techinically can be anything but HF impl has only image processor
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif hasattr(module, "latents"):
+ init.normal_(module.latents, mean=0.0, std=self.config.initializer_range)
+ elif hasattr(module, "position_embeddings") and isinstance(module, PerceiverTrainablePositionEncoding):
+ init.normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, nn.ParameterDict):
+ for modality in module:
+ init.normal_(module[modality], mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, (nn.LayerNorm, nn.BatchNorm2d)):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+
+
+@auto_docstring(
+ custom_intro="""
+ The Perceiver: a scalable, fully attentional architecture.
+
+
+
+ Note that it's possible to fine-tune Perceiver on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """
+)
+class PerceiverModel(PerceiverPreTrainedModel):
+ def __init__(
+ self,
+ config,
+ decoder: Optional["PerceiverAbstractDecoder"] = None,
+ input_preprocessor: PreprocessorType = None,
+ output_postprocessor: PostprocessorType = None,
+ ):
+ r"""
+ decoder (`PerceiverDecoder`, *optional*):
+ Decoder module that transforms latent representations into task predictions.
+ input_preprocessor (`PreprocessorType`, *optional*):
+ Preprocessor that encodes raw inputs into tensors for the model.
+ output_postprocessor (`PostprocessorType`, *optional*):
+ Postprocessor that transforms model outputs into final predictions.
+ """
+ super().__init__(config)
+ self.config = config
+
+ self.input_preprocessor = input_preprocessor
+ self.output_postprocessor = output_postprocessor
+ self.embeddings = PerceiverEmbeddings(config)
+ self.encoder = PerceiverEncoder(
+ config, kv_dim=input_preprocessor.num_channels if input_preprocessor is not None else config.d_model
+ )
+ self.decoder = decoder
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.latents
+
+ def set_input_embeddings(self, value):
+ self.embeddings.latents = value
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.FloatTensor,
+ attention_mask: torch.FloatTensor | None = None,
+ subsampled_output_points: dict[str, torch.Tensor] | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverModelOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ subsampled_output_points (`dict[str, torch.Tensor]`, *optional*):
+ Dictionary of tensors used as queries for the decoder. The decoder maps these queries to the latent
+ representation of the model. Used for subsampled decoding, e.g. when only decoding certain image patches.
+
+ Examples:
+
+ ```python
+ >>> from transformers import PerceiverConfig, PerceiverTokenizer, PerceiverImageProcessor, PerceiverModel
+ >>> from transformers.models.perceiver.modeling_perceiver import (
+ ... PerceiverTextPreprocessor,
+ ... PerceiverImagePreprocessor,
+ ... PerceiverClassificationDecoder,
+ ... )
+ >>> import torch
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> # EXAMPLE 1: using the Perceiver to classify texts
+ >>> # - we define a TextPreprocessor, which can be used to embed tokens
+ >>> # - we define a ClassificationDecoder, which can be used to decode the
+ >>> # final hidden states of the latents to classification logits
+ >>> # using trainable position embeddings
+ >>> config = PerceiverConfig()
+ >>> preprocessor = PerceiverTextPreprocessor(config)
+ >>> decoder = PerceiverClassificationDecoder(
+ ... config,
+ ... num_channels=config.d_latents,
+ ... trainable_position_encoding_kwargs=dict(num_channels=config.d_latents, index_dims=1),
+ ... use_query_residual=True,
+ ... )
+ >>> model = PerceiverModel(config, input_preprocessor=preprocessor, decoder=decoder)
+
+ >>> # you can then do a forward pass as follows:
+ >>> tokenizer = PerceiverTokenizer()
+ >>> text = "hello world"
+ >>> inputs = tokenizer(text, return_tensors="pt").input_ids
+
+ >>> with torch.no_grad():
+ ... outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 2]
+
+ >>> # to train, one can train the model using standard cross-entropy:
+ >>> criterion = torch.nn.CrossEntropyLoss()
+
+ >>> labels = torch.tensor([1])
+ >>> loss = criterion(logits, labels)
+
+ >>> # EXAMPLE 2: using the Perceiver to classify images
+ >>> # - we define an ImagePreprocessor, which can be used to embed images
+ >>> config = PerceiverConfig(image_size=224)
+ >>> preprocessor = PerceiverImagePreprocessor(
+ ... config,
+ ... prep_type="conv1x1",
+ ... spatial_downsample=1,
+ ... out_channels=256,
+ ... position_encoding_type="trainable",
+ ... concat_or_add_pos="concat",
+ ... project_pos_dim=256,
+ ... trainable_position_encoding_kwargs=dict(
+ ... num_channels=256,
+ ... index_dims=config.image_size**2,
+ ... ),
+ ... )
+
+ >>> model = PerceiverModel(
+ ... config,
+ ... input_preprocessor=preprocessor,
+ ... decoder=PerceiverClassificationDecoder(
+ ... config,
+ ... num_channels=config.d_latents,
+ ... trainable_position_encoding_kwargs=dict(num_channels=config.d_latents, index_dims=1),
+ ... use_query_residual=True,
+ ... ),
+ ... )
+
+ >>> # you can then do a forward pass as follows:
+ >>> image_processor = PerceiverImageProcessor()
+ >>> 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").pixel_values
+
+ >>> with torch.no_grad():
+ ... outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 2]
+
+ >>> # to train, one can train the model using standard cross-entropy:
+ >>> criterion = torch.nn.CrossEntropyLoss()
+
+ >>> labels = torch.tensor([1])
+ >>> loss = criterion(logits, labels)
+ ```"""
+ 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 self.input_preprocessor is not None:
+ inputs, modality_sizes, inputs_without_pos = self.input_preprocessor(
+ inputs, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+ else:
+ modality_sizes = None
+ inputs_without_pos = None
+ if inputs.size()[-1] != self.config.d_model:
+ raise ValueError(
+ f"Last dimension of the inputs: {inputs.size()[-1]} doesn't correspond to config.d_model:"
+ f" {self.config.d_model}. Make sure to set config.d_model appropriately."
+ )
+
+ batch_size, seq_length, _ = inputs.size()
+ device = inputs.device
+
+ # If no attention mask is provided, make them all ones
+ if attention_mask is None:
+ attention_mask = torch.ones((batch_size, seq_length), device=device)
+
+ embedding_output = self.embeddings(batch_size=batch_size)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=None,
+ inputs=inputs,
+ inputs_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+
+ logits = None
+ if self.decoder:
+ if subsampled_output_points is not None:
+ output_modality_sizes = {
+ "audio": subsampled_output_points["audio"].shape[0],
+ "image": subsampled_output_points["image"].shape[0],
+ "label": 1,
+ }
+ else:
+ output_modality_sizes = modality_sizes
+ decoder_query = self.decoder.decoder_query(
+ inputs, modality_sizes, inputs_without_pos, subsampled_points=subsampled_output_points
+ )
+ decoder_outputs = self.decoder(
+ decoder_query,
+ z=sequence_output,
+ query_mask=attention_mask,
+ output_attentions=output_attentions,
+ )
+ logits = decoder_outputs.logits
+
+ # add cross-attentions of decoder
+ if output_attentions and decoder_outputs.cross_attentions is not None:
+ if return_dict:
+ encoder_outputs.cross_attentions = (
+ encoder_outputs.cross_attentions + decoder_outputs.cross_attentions
+ )
+ else:
+ encoder_outputs = encoder_outputs + decoder_outputs.cross_attentions
+
+ if self.output_postprocessor:
+ logits = self.output_postprocessor(logits, modality_sizes=output_modality_sizes)
+
+ if not return_dict:
+ if logits is not None:
+ return (logits, sequence_output) + encoder_outputs[1:]
+ else:
+ return (sequence_output,) + encoder_outputs[1:]
+
+ return PerceiverModelOutput(
+ logits=logits,
+ last_hidden_state=sequence_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for masked language modeling.
+ """
+)
+class PerceiverForMaskedLM(PerceiverPreTrainedModel):
+ def __init__(self, config: PerceiverConfig):
+ super().__init__(config)
+
+ text_preprocessor = PerceiverTextPreprocessor(config)
+
+ trainable_position_encoding_kwargs_decoder = {
+ "num_channels": text_preprocessor.num_channels,
+ "index_dims": config.max_position_embeddings,
+ }
+
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=text_preprocessor,
+ decoder=PerceiverBasicDecoder(
+ config,
+ output_num_channels=config.d_latents,
+ output_index_dims=config.max_position_embeddings, # we need to define the seq_len of the inputs beforehand
+ num_channels=text_preprocessor.num_channels,
+ qk_channels=8 * 32,
+ v_channels=text_preprocessor.num_channels,
+ num_heads=8,
+ use_query_residual=False,
+ final_project=False,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_decoder,
+ ),
+ )
+ self.embedding_decoder = PerceiverEmbeddingDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ input_ids: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverMaskedLMOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
+ loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, PerceiverForMaskedLM
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("deepmind/language-perceiver")
+ >>> model = PerceiverForMaskedLM.from_pretrained("deepmind/language-perceiver")
+
+ >>> # training
+ >>> text = "This is an incomplete sentence where some words are missing."
+ >>> inputs = tokenizer(text, padding="max_length", return_tensors="pt")
+ >>> # mask " missing."
+ >>> inputs["input_ids"][0, 52:61] = tokenizer.mask_token_id
+ >>> labels = tokenizer(text, padding="max_length", return_tensors="pt").input_ids
+
+ >>> outputs = model(**inputs, labels=labels)
+ >>> loss = outputs.loss
+ >>> round(loss.item(), 2)
+ 19.87
+
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 2048, 262]
+
+ >>> # inference
+ >>> text = "This is an incomplete sentence where some words are missing."
+ >>> encoding = tokenizer(text, padding="max_length", return_tensors="pt")
+
+ >>> # mask bytes corresponding to " missing.". Note that the model performs much better if the masked span starts with a space.
+ >>> encoding["input_ids"][0, 52:61] = tokenizer.mask_token_id
+
+ >>> # forward pass
+ >>> with torch.no_grad():
+ ... outputs = model(**encoding)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 2048, 262]
+
+ >>> masked_tokens_predictions = logits[0, 52:61].argmax(dim=-1).tolist()
+ >>> tokenizer.decode(masked_tokens_predictions)
+ ' missing.'
+ ```"""
+ if inputs is not None and input_ids is not None:
+ raise ValueError("You cannot use both `inputs` and `input_ids`")
+ elif inputs is None and input_ids is not None:
+ inputs = input_ids
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ logits = self.embedding_decoder(
+ outputs.logits if return_dict else outputs[0], embedding_layer=self.perceiver.input_preprocessor.embeddings
+ )
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
+ masked_lm_loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+ return PerceiverMaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for text classification.
+ """
+)
+class PerceiverForSequenceClassification(PerceiverPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ trainable_position_encoding_kwargs_decoder = {"num_channels": config.d_latents, "index_dims": 1}
+
+ self.num_labels = config.num_labels
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=PerceiverTextPreprocessor(config),
+ decoder=PerceiverClassificationDecoder(
+ config,
+ num_channels=config.d_latents,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_decoder,
+ use_query_residual=True,
+ ),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ input_ids: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the classification/regression loss. Indices should be in `[0, ..., config.num_labels -
+ 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels >
+ 1` a classification loss is computed (Cross-Entropy).
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, PerceiverForSequenceClassification
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("deepmind/language-perceiver")
+ >>> model = PerceiverForSequenceClassification.from_pretrained("deepmind/language-perceiver")
+
+ >>> text = "hello world"
+ >>> inputs = tokenizer(text, return_tensors="pt").input_ids
+ >>> outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 2]
+ ```"""
+ if inputs is not None and input_ids is not None:
+ raise ValueError("You cannot use both `inputs` and `input_ids`")
+ elif inputs is None and input_ids is not None:
+ inputs = input_ids
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ logits = outputs.logits if return_dict else outputs[0]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for image classification, for tasks such as ImageNet.
+
+ This model uses learned position embeddings. In other words, this model is not given any privileged information about
+ the structure of images. As shown in the paper, this model can achieve a top-1 accuracy of 72.7 on ImageNet.
+
+ [`PerceiverForImageClassificationLearned`] uses [`~models.perceiver.modeling_perceiver.PerceiverImagePreprocessor`]
+ (with `prep_type="conv1x1"`) to preprocess the input images, and
+ [`~models.perceiver.modeling_perceiver.PerceiverClassificationDecoder`] to decode the latent representation of
+ [`PerceiverModel`] into classification logits.
+ """
+)
+class PerceiverForImageClassificationLearned(PerceiverPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ trainable_position_encoding_kwargs_preprocessor = {"num_channels": 256, "index_dims": config.image_size**2}
+ trainable_position_encoding_kwargs_decoder = {"num_channels": config.d_latents, "index_dims": 1}
+
+ self.num_labels = config.num_labels
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=PerceiverImagePreprocessor(
+ config,
+ prep_type="conv1x1",
+ spatial_downsample=1,
+ out_channels=256,
+ position_encoding_type="trainable",
+ concat_or_add_pos="concat",
+ project_pos_dim=256,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_preprocessor,
+ ),
+ decoder=PerceiverClassificationDecoder(
+ config,
+ num_channels=config.d_latents,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_decoder,
+ use_query_residual=True,
+ ),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ 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).
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, PerceiverForImageClassificationLearned
+ >>> 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("deepmind/vision-perceiver-learned")
+ >>> model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
+
+ >>> inputs = image_processor(images=image, return_tensors="pt").pixel_values
+ >>> outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 1000]
+
+ >>> # model predicts one of the 1000 ImageNet classes
+ >>> predicted_class_idx = logits.argmax(-1).item()
+ >>> print("Predicted class:", model.config.id2label[predicted_class_idx])
+ Predicted class: tabby, tabby cat
+ ```"""
+ if inputs is not None and pixel_values is not None:
+ raise ValueError("You cannot use both `inputs` and `pixel_values`")
+ elif inputs is None and pixel_values is not None:
+ inputs = pixel_values
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=return_dict,
+ )
+ logits = outputs.logits if return_dict else outputs[0]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for image classification, for tasks such as ImageNet.
+
+ This model uses fixed 2D Fourier position embeddings. As shown in the paper, this model can achieve a top-1 accuracy of
+ 79.0 on ImageNet, and 84.5 when pre-trained on a large-scale dataset (i.e. JFT).
+
+ [`PerceiverForImageClassificationLearned`] uses [`~models.perceiver.modeling_perceiver.PerceiverImagePreprocessor`]
+ (with `prep_type="pixels"`) to preprocess the input images, and
+ [`~models.perceiver.modeling_perceiver.PerceiverClassificationDecoder`] to decode the latent representation of
+ [`PerceiverModel`] into classification logits.
+ """
+)
+class PerceiverForImageClassificationFourier(PerceiverPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ fourier_position_encoding_kwargs_preprocessor = {
+ "concat_pos": True,
+ "max_resolution": (224, 224),
+ "num_bands": 64,
+ "sine_only": False,
+ }
+ trainable_position_encoding_kwargs_decoder = {"num_channels": config.d_latents, "index_dims": 1}
+
+ self.num_labels = config.num_labels
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=PerceiverImagePreprocessor(
+ config,
+ prep_type="pixels",
+ spatial_downsample=1,
+ fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_preprocessor,
+ ),
+ decoder=PerceiverClassificationDecoder(
+ config,
+ num_channels=config.d_latents,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_decoder,
+ use_query_residual=True,
+ ),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ 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).
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, PerceiverForImageClassificationFourier
+ >>> 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("deepmind/vision-perceiver-fourier")
+ >>> model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
+
+ >>> inputs = image_processor(images=image, return_tensors="pt").pixel_values
+ >>> outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 1000]
+
+ >>> # model predicts one of the 1000 ImageNet classes
+ >>> predicted_class_idx = logits.argmax(-1).item()
+ >>> print("Predicted class:", model.config.id2label[predicted_class_idx])
+ Predicted class: tabby, tabby cat
+ ```"""
+ if inputs is not None and pixel_values is not None:
+ raise ValueError("You cannot use both `inputs` and `pixel_values`")
+ elif inputs is None and pixel_values is not None:
+ inputs = pixel_values
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ logits = outputs.logits if return_dict else outputs[0]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for image classification, for tasks such as ImageNet.
+
+ This model uses a 2D conv+maxpool preprocessing network. As shown in the paper, this model can achieve a top-1 accuracy
+ of 82.1 on ImageNet.
+
+ [`PerceiverForImageClassificationLearned`] uses [`~models.perceiver.modeling_perceiver.PerceiverImagePreprocessor`]
+ (with `prep_type="conv"`) to preprocess the input images, and
+ [`~models.perceiver.modeling_perceiver.PerceiverClassificationDecoder`] to decode the latent representation of
+ [`PerceiverModel`] into classification logits.
+ """
+)
+class PerceiverForImageClassificationConvProcessing(PerceiverPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ fourier_position_encoding_kwargs_preprocessor = {
+ "concat_pos": True,
+ "max_resolution": (56, 56),
+ "num_bands": 64,
+ "sine_only": False,
+ }
+ trainable_position_encoding_kwargs_decoder = {"num_channels": config.d_latents, "index_dims": 1}
+
+ self.num_labels = config.num_labels
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=PerceiverImagePreprocessor(
+ config,
+ prep_type="conv",
+ spatial_downsample=1,
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_preprocessor,
+ ),
+ decoder=PerceiverClassificationDecoder(
+ config,
+ num_channels=config.d_latents,
+ trainable_position_encoding_kwargs=trainable_position_encoding_kwargs_decoder,
+ use_query_residual=True,
+ ),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ 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).
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, PerceiverForImageClassificationConvProcessing
+ >>> 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("deepmind/vision-perceiver-conv")
+ >>> model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
+
+ >>> inputs = image_processor(images=image, return_tensors="pt").pixel_values
+ >>> outputs = model(inputs=inputs)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 1000]
+
+ >>> # model predicts one of the 1000 ImageNet classes
+ >>> predicted_class_idx = logits.argmax(-1).item()
+ >>> print("Predicted class:", model.config.id2label[predicted_class_idx])
+ Predicted class: tabby, tabby cat
+ ```"""
+ if inputs is not None and pixel_values is not None:
+ raise ValueError("You cannot use both `inputs` and `pixel_values`")
+ elif inputs is None and pixel_values is not None:
+ inputs = pixel_values
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ logits = outputs.logits if return_dict else outputs[0]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for optical flow, for tasks such as Sintel and KITTI. [`PerceiverForOpticalFlow`] uses
+ [`~models.perceiver.modeling_perceiver.PerceiverImagePreprocessor`] (with *prep_type="patches"*) to preprocess the
+ input images, and [`~models.perceiver.modeling_perceiver.PerceiverOpticalFlowDecoder`] to decode the latent
+ representation of [`PerceiverModel`].
+
+ As input, one concatenates 2 subsequent frames along the channel dimension and extract a 3 x 3 patch around each pixel
+ (leading to 3 x 3 x 3 x 2 = 54 values for each pixel). Fixed Fourier position encodings are used to encode the position
+ of each pixel in the patch. Next, one applies the Perceiver encoder. To decode, one queries the latent representation
+ using the same encoding used for the input.
+ """
+)
+class PerceiverForOpticalFlow(PerceiverPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ fourier_position_encoding_kwargs_preprocessor = {
+ "num_bands": 64,
+ "max_resolution": config.train_size,
+ "sine_only": False,
+ "concat_pos": True,
+ }
+ fourier_position_encoding_kwargs_decoder = {
+ "concat_pos": True,
+ "max_resolution": config.train_size,
+ "num_bands": 64,
+ "sine_only": False,
+ }
+
+ image_preprocessor = PerceiverImagePreprocessor(
+ config,
+ prep_type="patches",
+ spatial_downsample=1,
+ conv_after_patching=True,
+ conv_after_patching_in_channels=54,
+ temporal_downsample=2,
+ position_encoding_type="fourier",
+ # position_encoding_kwargs
+ fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_preprocessor,
+ )
+
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=image_preprocessor,
+ decoder=PerceiverOpticalFlowDecoder(
+ config,
+ num_channels=image_preprocessor.num_channels,
+ output_image_shape=config.train_size,
+ rescale_factor=100.0,
+ # decoder kwargs
+ use_query_residual=False,
+ output_num_channels=2,
+ # We query the decoder using the first frame features
+ # rather than a standard decoder position encoding.
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_decoder,
+ ),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the optical flow loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+
+ Examples:
+
+ ```python
+ >>> from transformers import PerceiverForOpticalFlow
+ >>> import torch
+
+ >>> model = PerceiverForOpticalFlow.from_pretrained("deepmind/optical-flow-perceiver")
+
+ >>> # in the Perceiver IO paper, the authors extract a 3 x 3 patch around each pixel,
+ >>> # leading to 3 x 3 x 3 = 27 values for each pixel (as each pixel also has 3 color channels)
+ >>> # patches have shape (batch_size, num_frames, num_channels, height, width)
+ >>> # the authors train on resolutions of 368 x 496
+ >>> patches = torch.randn(1, 2, 27, 368, 496)
+ >>> outputs = model(inputs=patches)
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 368, 496, 2]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ loss = None
+ if labels is not None:
+ raise NotImplementedError("Optical flow training is not yet supported")
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ logits = outputs.logits if return_dict else outputs[0]
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Example use of Perceiver for multimodal (video) autoencoding, for tasks such as Kinetics-700.
+
+ [`PerceiverForMultimodalAutoencoding`] uses [`~models.perceiver.modeling_perceiver.PerceiverMultimodalPreprocessor`] to
+ preprocess the 3 modalities: images, audio and class labels. This preprocessor uses modality-specific preprocessors to
+ preprocess every modality separately, after which they are concatenated. Trainable position embeddings are used to pad
+ each modality to the same number of channels to make concatenation along the time dimension possible. Next, one applies
+ the Perceiver encoder.
+
+ [`~models.perceiver.modeling_perceiver.PerceiverMultimodalDecoder`] is used to decode the latent representation of
+ [`PerceiverModel`]. This decoder uses each modality-specific decoder to construct queries. The decoder queries are
+ created based on the inputs after preprocessing. However, autoencoding an entire video in a single forward pass is
+ computationally infeasible, hence one only uses parts of the decoder queries to do cross-attention with the latent
+ representation. This is determined by the subsampled indices for each modality, which can be provided as additional
+ input to the forward pass of [`PerceiverForMultimodalAutoencoding`].
+
+ [`~models.perceiver.modeling_perceiver.PerceiverMultimodalDecoder`] also pads the decoder queries of the different
+ modalities to the same number of channels, in order to concatenate them along the time dimension. Next, cross-attention
+ is performed with the latent representation of [`PerceiverModel`].
+
+ Finally, [`~models.perceiver.modeling_perceiver.PerceiverMultiModalPostprocessor`] is used to turn this tensor into an
+ actual video. It first splits up the output into the different modalities, and then applies the respective
+ postprocessor for each modality.
+
+ Note that, by masking the classification label during evaluation (i.e. simply providing a tensor of zeros for the
+ "label" modality), this auto-encoding model becomes a Kinetics 700 video classifier.
+ """
+)
+class PerceiverForMultimodalAutoencoding(PerceiverPreTrainedModel):
+ def __init__(self, config: PerceiverConfig):
+ super().__init__(config)
+
+ n_audio_samples = config.num_frames * config.audio_samples_per_frame
+
+ input_preprocessor = PerceiverMultimodalPreprocessor(
+ min_padding_size=4,
+ modalities={
+ "audio": PerceiverAudioPreprocessor(
+ config,
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs={
+ "num_bands": 192,
+ "max_resolution": (n_audio_samples,),
+ "sine_only": False,
+ "concat_pos": True,
+ },
+ prep_type="patches",
+ samples_per_patch=config.samples_per_patch,
+ ),
+ "image": PerceiverImagePreprocessor(
+ config,
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs={
+ "num_bands": 32,
+ "max_resolution": (config.num_frames, config.image_size, config.image_size),
+ "sine_only": False,
+ "concat_pos": True,
+ },
+ prep_type="patches",
+ spatial_downsample=4,
+ temporal_downsample=1,
+ ),
+ "label": PerceiverOneHotPreprocessor(config),
+ },
+ mask_probs={"image": 0.0, "audio": 0.0, "label": 1.0},
+ )
+
+ image_decoder = PerceiverBasicVideoAutoencodingDecoder(
+ config,
+ # Autoencoding, don't pass inputs to the queries.
+ concat_preprocessed_input=False,
+ output_shape=config.output_shape,
+ output_num_channels=config.output_num_channels,
+ use_query_residual=False,
+ position_encoding_only=True,
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs={
+ "num_bands": 32,
+ "max_resolution": (config.num_frames, config.image_size, config.image_size),
+ "sine_only": False,
+ "concat_pos": True,
+ },
+ )
+
+ decoder = PerceiverMultimodalDecoder(
+ config,
+ # Autoencoding, don't pass inputs to the queries.
+ concat_preprocessed_input=False,
+ # Modality specific decoders are used ONLY to generate queries.
+ # All modalties are decoded together using a unified decoder.
+ modalities={
+ "audio": PerceiverBasicDecoder(
+ config,
+ # Autoencoding, don't pass inputs to the queries.
+ concat_preprocessed_input=False,
+ output_index_dims=(n_audio_samples // config.samples_per_patch,),
+ output_num_channels=config.output_num_channels,
+ use_query_residual=False,
+ position_encoding_only=True,
+ position_encoding_type="fourier",
+ fourier_position_encoding_kwargs={
+ "num_bands": 192,
+ "max_resolution": (n_audio_samples,),
+ "sine_only": False,
+ "concat_pos": True,
+ },
+ ),
+ "image": image_decoder,
+ "label": PerceiverClassificationDecoder(
+ config,
+ # Autoencoding, don't pass inputs to the queries.
+ concat_preprocessed_input=False,
+ use_query_residual=False,
+ position_encoding_only=True,
+ position_encoding_type="trainable",
+ trainable_position_encoding_kwargs={
+ "num_channels": config._label_trainable_num_channels,
+ "index_dims": 1,
+ },
+ ),
+ },
+ num_outputs=None,
+ output_num_channels=config.output_num_channels,
+ use_query_residual=False,
+ )
+
+ output_postprocessor = PerceiverMultimodalPostprocessor(
+ modalities={
+ "audio": PerceiverAudioPostprocessor(config, in_channels=config.output_num_channels),
+ "image": PerceiverProjectionPostprocessor(in_channels=config.output_num_channels, out_channels=3),
+ "label": PerceiverClassificationPostprocessor(config, in_channels=config.output_num_channels),
+ }
+ )
+
+ self.perceiver = PerceiverModel(
+ config,
+ input_preprocessor=input_preprocessor,
+ decoder=decoder,
+ output_postprocessor=output_postprocessor,
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ subsampled_output_points: dict[str, torch.Tensor] | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | PerceiverClassifierOutput:
+ r"""
+ inputs (`torch.FloatTensor`):
+ Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
+ subsampled_output_points (`dict[str, torch.Tensor]`, *optional*):
+ Dictionary of tensors used as queries for the decoder. The decoder maps these queries to the latent
+ representation of the model. Used for subsampled decoding, e.g. when only decoding certain image patches.
+ 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).
+
+ Examples:
+
+ ```python
+ >>> from transformers import PerceiverForMultimodalAutoencoding
+ >>> import torch
+ >>> import numpy as np
+
+ >>> # create multimodal inputs
+ >>> images = torch.randn((1, 16, 3, 224, 224))
+ >>> audio = torch.randn((1, 30720, 1))
+ >>> inputs = dict(image=images, audio=audio, label=torch.zeros((images.shape[0], 700)))
+
+ >>> model = PerceiverForMultimodalAutoencoding.from_pretrained("deepmind/multimodal-perceiver")
+
+ >>> # in the Perceiver IO paper, videos are auto-encoded in chunks
+ >>> # each chunk subsamples different index dimensions of the image and audio modality decoder queries
+ >>> nchunks = 128
+ >>> image_chunk_size = np.prod((16, 224, 224)) // nchunks
+ >>> audio_chunk_size = audio.shape[1] // model.config.samples_per_patch // nchunks
+ >>> # process the first chunk
+ >>> chunk_idx = 0
+ >>> subsampling = {
+ ... "image": torch.arange(image_chunk_size * chunk_idx, image_chunk_size * (chunk_idx + 1)),
+ ... "audio": torch.arange(audio_chunk_size * chunk_idx, audio_chunk_size * (chunk_idx + 1)),
+ ... "label": None,
+ ... }
+
+ >>> outputs = model(inputs=inputs, subsampled_output_points=subsampling)
+ >>> logits = outputs.logits
+ >>> list(logits["audio"].shape)
+ [1, 240]
+
+ >>> list(logits["image"].shape)
+ [1, 6272, 3]
+
+ >>> list(logits["label"].shape)
+ [1, 700]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ loss = None
+ if labels is not None:
+ raise NotImplementedError("Multimodal autoencoding training is not yet supported")
+
+ outputs = self.perceiver(
+ inputs=inputs,
+ attention_mask=attention_mask,
+ subsampled_output_points=subsampled_output_points,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ logits = outputs.logits if return_dict else outputs[0]
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return PerceiverClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+# Below: position encodings
+
+
+def build_position_encoding(
+ position_encoding_type,
+ out_channels=None,
+ project_pos_dim=-1,
+ trainable_position_encoding_kwargs=None,
+ fourier_position_encoding_kwargs=None,
+):
+ """
+ Builds the position encoding.
+
+ Args:
+ - out_channels: refers to the number of channels of the position encodings.
+ - project_pos_dim: if specified, will project the position encodings to this dimension.
+
+ """
+
+ if position_encoding_type == "trainable":
+ if not trainable_position_encoding_kwargs:
+ raise ValueError("Make sure to pass trainable_position_encoding_kwargs")
+ output_pos_enc = PerceiverTrainablePositionEncoding(**trainable_position_encoding_kwargs)
+ elif position_encoding_type == "fourier":
+ # We don't use the index_dims argument, as this is only known during the forward pass
+ if not fourier_position_encoding_kwargs:
+ raise ValueError("Make sure to pass fourier_position_encoding_kwargs")
+ output_pos_enc = PerceiverFourierPositionEncoding(**fourier_position_encoding_kwargs)
+ else:
+ raise ValueError(f"Unknown position encoding type: {position_encoding_type}.")
+
+ # Optionally, project the position encoding to a target dimension:
+ positions_projection = nn.Linear(out_channels, project_pos_dim) if project_pos_dim > 0 else nn.Identity()
+
+ return output_pos_enc, positions_projection
+
+
+# Below: Perceiver decoders
+
+
+class PerceiverAbstractDecoder(nn.Module, metaclass=abc.ABCMeta):
+ """Perceiver abstract decoder."""
+
+ @abc.abstractmethod
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ raise NotImplementedError
+
+ @property
+ @abc.abstractmethod
+ def num_query_channels(self):
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def forward(self, query, z, query_mask=None):
+ raise NotImplementedError
+
+
+class PerceiverProjectionDecoder(PerceiverAbstractDecoder):
+ """
+ Baseline projection decoder (no cross-attention).
+
+ Args:
+ config ([`PerceiverConfig`]):
+ Model configuration.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.classifier = nn.Linear(config.d_latents, config.num_labels)
+
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ return None
+
+ def forward(
+ self, query: torch.Tensor, z: torch.FloatTensor, query_mask: torch.FloatTensor | None = None
+ ) -> torch.FloatTensor:
+ # (batch_size, num_latents, d_latents) -> (batch_size, d_latents)
+ z = torch.mean(z, dim=1)
+ # (batch_size, d_latents) -> (batch_size, config.num_labels)
+ logits = self.classifier(z)
+ return logits
+
+
+class PerceiverBasicDecoder(PerceiverAbstractDecoder):
+ """
+ Cross-attention-based decoder. This class can be used to decode the final hidden states of the latents using a
+ cross-attention operation, in which the latents produce keys and values.
+
+ The shape of the output of this class depends on how one defines the output queries (also called decoder queries).
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ output_num_channels (`int`, *optional*):
+ The number of channels in the output. Will only be used in case *final_project* is set to `True`.
+ position_encoding_type (`str`, *optional*, defaults to "trainable"):
+ The type of position encoding to use. Can be either "trainable", "fourier", or "none".
+ output_index_dims (`int`, *optional*):
+ The number of dimensions of the output queries. Ignored if 'position_encoding_type' == 'none'.
+ num_channels (`int`, *optional*, defaults to 128):
+ The number of channels of the decoder queries. Ignored if 'position_encoding_type' == 'none'.
+ qk_channels (`int`, *optional*):
+ The number of channels of the queries and keys in the cross-attention layer.
+ v_channels (`int`, *optional*):
+ The number of channels of the values in the cross-attention layer.
+ num_heads (`int`, *optional*, defaults to 1):
+ The number of attention heads in the cross-attention layer.
+ widening_factor (`int`, *optional*, defaults to 1):
+ The widening factor of the cross-attention layer.
+ use_query_residual (`bool`, *optional*, defaults to `False`):
+ Whether to use a residual connection between the query and the output of the cross-attention layer.
+ concat_preprocessed_input (`bool`, *optional*, defaults to `False`):
+ Whether to concatenate the preprocessed input to the query.
+ final_project (`bool`, *optional*, defaults to `True`):
+ Whether to project the output of the cross-attention layer to a target dimension.
+ position_encoding_only (`bool`, *optional*, defaults to `False`):
+ Whether to only use this class to define output queries.
+ """
+
+ def __init__(
+ self,
+ config: PerceiverConfig,
+ output_num_channels: int,
+ position_encoding_type: str | None = "trainable",
+ # The following 2 arguments are ignored if position_encoding_type == 'none':
+ output_index_dims: int | None = None,
+ num_channels: int | None = 128,
+ subsampled_index_dims: int | None = None,
+ qk_channels: int | None = None,
+ v_channels: int | None = None,
+ num_heads: int | None = 1,
+ widening_factor: int | None = 1,
+ use_query_residual: bool | None = False,
+ concat_preprocessed_input: bool | None = False,
+ final_project: bool | None = True,
+ position_encoding_only: bool | None = False,
+ **position_encoding_kwargs,
+ ) -> None:
+ super().__init__()
+
+ self.output_num_channels = output_num_channels
+ # If `none`, the decoder will not construct any position encodings.
+ # You should construct your own when querying the decoder.
+ self.output_position_encodings = None
+ self.position_encoding_type = position_encoding_type
+ self.position_encoding_kwargs = position_encoding_kwargs
+ if position_encoding_type != "none":
+ self.output_position_encodings, self.positions_projection = build_position_encoding(
+ position_encoding_type=position_encoding_type, **position_encoding_kwargs
+ )
+
+ self.output_index_dims = output_index_dims
+ self.num_channels = num_channels
+ if subsampled_index_dims is None:
+ subsampled_index_dims = output_index_dims
+ self.subsampled_index_dims = subsampled_index_dims
+ self.concat_preprocessed_input = concat_preprocessed_input
+ self.final_project = final_project
+ self.position_encoding_only = position_encoding_only
+
+ # for multimodal autoencoding, we don't need the decoder cross-attention and final layer
+ # so then we will set position_encoding_only to True
+ if not self.position_encoding_only:
+ self.decoding_cross_attention = PerceiverLayer(
+ config,
+ is_cross_attention=True,
+ qk_channels=qk_channels,
+ v_channels=v_channels,
+ num_heads=num_heads,
+ q_dim=num_channels,
+ kv_dim=config.d_latents,
+ widening_factor=widening_factor,
+ use_query_residual=use_query_residual,
+ )
+ self.final_layer = nn.Linear(num_channels, output_num_channels) if final_project else nn.Identity()
+
+ @property
+ def num_query_channels(self) -> int:
+ if self.position_encoding_type == "none": # Queries come from elsewhere
+ raise ValueError(
+ "You cannot calculate number of decoder query channels when position_encoding_type is set to none"
+ )
+ if self.position_encoding_only:
+ if "project_pos_dim" in self.position_encoding_kwargs:
+ return self.position_encoding_kwargs["project_pos_dim"]
+ return self.output_position_encodings.output_size()
+ if self.final_project:
+ return self.output_num_channels
+ return self.num_channels
+
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ if self.position_encoding_type == "none": # Queries come from elsewhere
+ raise ValueError("You cannot construct decoder queries when position_encoding_type is set to none")
+ if subsampled_points is not None:
+ # subsampled_points are the indices if the inputs would be flattened
+ # however, the inputs aren't flattened, that's why we use unravel_index
+ # to get the indices for the unflattened array
+ # unravel_index returns a tuple (x_idx, y_idx, ...)
+ # stack to get the [n, d] tensor of coordinates
+ indices = torch.unravel_index(subsampled_points, self.output_index_dims)
+ pos = torch.stack(indices, dim=1)
+ batch_size = inputs.shape[0]
+ # Map these coordinates to [-1, 1]
+ pos = -1 + 2 * pos / torch.tensor(self.output_index_dims)[None, :]
+ pos = torch.broadcast_to(pos[None], [batch_size, pos.shape[0], pos.shape[1]])
+ # Construct the position encoding.
+ if self.position_encoding_type == "trainable":
+ pos_emb = self.output_position_encodings(batch_size)
+ elif self.position_encoding_type == "fourier":
+ pos_emb = self.output_position_encodings(
+ self.output_index_dims, batch_size=batch_size, device=inputs.device, dtype=inputs.dtype, pos=pos
+ )
+
+ # Optionally project them to a target dimension.
+ pos_emb = self.positions_projection(pos_emb)
+ pos_emb = torch.reshape(pos_emb, [pos_emb.shape[0], -1, pos_emb.shape[-1]])
+ else:
+ batch_size = inputs.shape[0]
+ index_dims = inputs.shape[2:]
+
+ # Construct the position encoding.
+ if self.position_encoding_type == "trainable":
+ pos_emb = self.output_position_encodings(batch_size)
+ elif self.position_encoding_type == "fourier":
+ pos_emb = self.output_position_encodings(
+ index_dims, batch_size, device=inputs.device, dtype=inputs.dtype
+ )
+
+ # Optionally project them to a target dimension.
+ pos_emb = self.positions_projection(pos_emb)
+
+ if self.concat_preprocessed_input:
+ if inputs_without_pos is None:
+ raise ValueError("Value is required for inputs_without_pos if concat_preprocessed_input is True")
+ pos_emb = torch.cat([inputs_without_pos, pos_emb], dim=-1)
+
+ return pos_emb
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ z: torch.FloatTensor,
+ query_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> PerceiverDecoderOutput:
+ # Cross-attention decoding.
+ # key, value: B x N x K; query: B x M x K
+ # Attention maps -> B x N x M
+ # Output -> B x M x K
+ cross_attentions = () if output_attentions else None
+
+ layer_outputs = self.decoding_cross_attention(
+ query,
+ attention_mask=query_mask,
+ inputs=z,
+ inputs_mask=None,
+ output_attentions=output_attentions,
+ )
+ output = layer_outputs[0]
+
+ if output_attentions:
+ cross_attentions = cross_attentions + (layer_outputs[1],)
+
+ logits = self.final_layer(output)
+
+ return PerceiverDecoderOutput(logits=logits, cross_attentions=cross_attentions)
+
+
+class PerceiverClassificationDecoder(PerceiverAbstractDecoder):
+ """
+ Cross-attention based classification decoder. Light-weight wrapper of [`PerceiverBasicDecoder`] for logit output.
+ Will turn the output of the Perceiver encoder which is of shape (batch_size, num_latents, d_latents) to a tensor of
+ shape (batch_size, num_labels). The queries are of shape (batch_size, 1, num_labels).
+
+ Args:
+ config ([`PerceiverConfig`]):
+ Model configuration.
+ """
+
+ def __init__(self, config, **decoder_kwargs):
+ super().__init__()
+
+ self.num_labels = config.num_labels
+ self.decoder = PerceiverBasicDecoder(
+ config,
+ output_num_channels=self.num_labels,
+ output_index_dims=1, # Predict a single logit array.
+ **decoder_kwargs,
+ )
+
+ @property
+ def num_query_channels(self) -> int:
+ return self.decoder.num_query_channels
+
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ return self.decoder.decoder_query(
+ inputs, modality_sizes, inputs_without_pos, subsampled_points=subsampled_points
+ )
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ z: torch.FloatTensor,
+ query_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> PerceiverDecoderOutput:
+ decoder_outputs = self.decoder(query, z, output_attentions=output_attentions)
+
+ # B x 1 x num_classes -> B x num_classes
+ logits = decoder_outputs.logits[:, 0, :]
+
+ return PerceiverDecoderOutput(logits=logits, cross_attentions=decoder_outputs.cross_attentions)
+
+
+class PerceiverOpticalFlowDecoder(PerceiverAbstractDecoder):
+ """Cross-attention based optical flow decoder."""
+
+ def __init__(self, config, output_image_shape, output_num_channels=2, rescale_factor=100.0, **decoder_kwargs):
+ super().__init__()
+
+ self.output_image_shape = output_image_shape
+ self.output_num_channels = output_num_channels
+ self.rescale_factor = rescale_factor
+ self.decoder = PerceiverBasicDecoder(config, output_num_channels=output_num_channels, **decoder_kwargs)
+
+ @property
+ def num_query_channels(self) -> int:
+ return self.decoder.num_query_channels
+
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ if subsampled_points is not None:
+ raise ValueError("FlowDecoder doesn't support subsampling yet.")
+ return inputs
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ z: torch.FloatTensor,
+ query_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> PerceiverDecoderOutput:
+ decoder_outputs = self.decoder(query, z, output_attentions=output_attentions)
+ preds = decoder_outputs.logits
+ # Output flow and rescale.
+ preds /= self.rescale_factor
+ preds = preds.reshape([preds.shape[0]] + list(self.output_image_shape) + [preds.shape[-1]])
+ return PerceiverDecoderOutput(logits=preds, cross_attentions=decoder_outputs.cross_attentions)
+
+
+class PerceiverBasicVideoAutoencodingDecoder(PerceiverAbstractDecoder):
+ """
+ Cross-attention based video-autoencoding decoder. Light-weight wrapper of [*PerceiverBasicDecoder*] with video
+ reshaping logic.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ output_shape (`list[int]`):
+ Shape of the output as (batch_size, num_frames, height, width), excluding the channel dimension.
+ position_encoding_type (`str`):
+ The type of position encoding to use. Can be either "trainable", "fourier", or "none".
+ """
+
+ def __init__(
+ self, config: PerceiverConfig, output_shape: list[int], position_encoding_type: str, **decoder_kwargs
+ ) -> None:
+ super().__init__()
+ if len(output_shape) != 4: # B, T, H, W
+ raise ValueError(f"Expected rank 4 output_shape, got {output_shape}.")
+ # Build the decoder components:
+ self.output_shape = output_shape
+ self.output_num_channels = decoder_kwargs["output_num_channels"]
+
+ self.decoder = PerceiverBasicDecoder(
+ config,
+ output_index_dims=self.output_shape[1:4], # T*H*W
+ position_encoding_type=position_encoding_type,
+ **decoder_kwargs,
+ )
+
+ @property
+ def num_query_channels(self) -> int:
+ return self.decoder.num_query_channels
+
+ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, subsampled_points=None):
+ return self.decoder.decoder_query(
+ inputs,
+ modality_sizes=modality_sizes,
+ inputs_without_pos=inputs_without_pos,
+ subsampled_points=subsampled_points,
+ )
+
+ def forward(
+ self, query: torch.Tensor, z: torch.FloatTensor, query_mask: torch.FloatTensor | None = None
+ ) -> PerceiverDecoderOutput:
+ decoder_outputs = self.decoder(query, z)
+ logits = decoder_outputs.logits
+
+ logits = torch.reshape(logits, self.output_shape + [logits.shape[-1]])
+ return PerceiverDecoderOutput(logits=logits, cross_attentions=decoder_outputs.cross_attentions)
+
+
+def restructure(modality_sizes: ModalitySizeType, inputs: torch.Tensor) -> Mapping[str, torch.Tensor]:
+ """
+ Partitions a [B, N, C] tensor into tensors for each modality.
+
+ Args:
+ modality_sizes
+ dict specifying the size of the modality
+ inputs:
+ input tensor
+
+ Returns:
+ dict mapping name of modality to its associated tensor.
+ """
+ outputs = {}
+ index = 0
+ # Apply a predictable ordering to the modalities
+ for modality in sorted(modality_sizes.keys()):
+ size = modality_sizes[modality]
+ inp = inputs[:, index : index + size]
+ index += size
+ outputs[modality] = inp
+ return outputs
+
+
+class PerceiverMultimodalDecoder(PerceiverAbstractDecoder):
+ """
+ Multimodal decoding by composing uni-modal decoders. The *modalities* argument of the constructor is a dictionary
+ mapping modality name to the decoder of that modality. That decoder will be used to construct queries for that
+ modality. Modality-specific queries are padded with trainable modality-specific parameters, after which they are
+ concatenated along the time dimension.
+
+ Next, there is a shared cross attention operation across all modalities.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ modalities (`dict[str, PerceiverAbstractDecoder]`):
+ Dictionary mapping modality name to the decoder of that modality.
+ num_outputs (`int`):
+ The number of outputs of the decoder.
+ output_num_channels (`int`):
+ The number of channels in the output.
+ min_padding_size (`int`, *optional*, defaults to 2):
+ The minimum padding size for all modalities. The final output will have num_channels equal to the maximum
+ channels across all modalities plus min_padding_size.
+ subsampled_index_dims (`dict[str, PerceiverAbstractDecoder]`, *optional*):
+ Dictionary mapping modality name to the subsampled index dimensions to use for the decoder query of that
+ modality.
+ """
+
+ def __init__(
+ self,
+ config: PerceiverConfig,
+ modalities: dict[str, PerceiverAbstractDecoder],
+ num_outputs: int,
+ output_num_channels: int,
+ min_padding_size: int | None = 2,
+ subsampled_index_dims: dict[str, PerceiverAbstractDecoder] | None = None,
+ **decoder_kwargs,
+ ) -> None:
+ super().__init__()
+ self.modalities = nn.ModuleDict(modalities)
+ self.subsampled_index_dims = subsampled_index_dims
+ self.min_padding_size = min_padding_size
+ self.output_num_channels = output_num_channels
+ self.num_outputs = num_outputs
+ self.decoder = PerceiverBasicDecoder(
+ config,
+ output_index_dims=(num_outputs,),
+ output_num_channels=output_num_channels,
+ position_encoding_type="none",
+ num_channels=self.num_query_channels,
+ **decoder_kwargs,
+ )
+ self.padding = nn.ParameterDict(
+ {
+ modality: nn.Parameter(torch.randn(1, self.num_query_channels - decoder.num_query_channels))
+ for modality, decoder in modalities.items()
+ }
+ )
+
+ @property
+ def num_query_channels(self) -> int:
+ max_channel_size = max(decoder.num_query_channels for _, decoder in self.modalities.items())
+ common_channel_size = max_channel_size + self.min_padding_size
+ return common_channel_size
+
+ def decoder_query(self, inputs, modality_sizes, inputs_without_pos=None, subsampled_points=None):
+ # Partition the flat inputs among the different modalities
+ inputs = restructure(modality_sizes, inputs)
+
+ # Obtain modality-specific decoders' queries
+ subsampled_points = subsampled_points or {}
+
+ decoder_queries = {}
+ for modality, decoder in self.modalities.items():
+ # Get input_without_pos for this modality if it exists.
+ input_without_pos = None
+ if inputs_without_pos is not None:
+ input_without_pos = inputs_without_pos.get(modality, None)
+ query = decoder.decoder_query(
+ inputs=inputs[modality],
+ modality_sizes=None,
+ inputs_without_pos=input_without_pos,
+ subsampled_points=subsampled_points.get(modality, None),
+ )
+ decoder_queries[modality] = query
+
+ # Pad all queries with trainable position encodings to make them have the same channels
+
+ def embed(modality, x):
+ x = torch.reshape(x, [x.shape[0], np.prod(x.shape[1:-1]), x.shape[-1]])
+ pos = self.padding[modality]
+ pos = torch.broadcast_to(pos, [x.shape[0], x.shape[1], self.num_query_channels - x.shape[2]])
+ return torch.cat([x, pos], dim=2)
+
+ # Apply a predictable ordering to the modalities
+ return torch.cat(
+ [embed(modality, decoder_queries[modality]) for modality in sorted(self.modalities.keys())], dim=1
+ )
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ z: torch.FloatTensor,
+ query_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> torch.Tensor:
+ # B x 1 x num_classes -> B x num_classes
+ decoder_outputs = self.decoder(query, z, output_attentions=output_attentions)
+
+ return decoder_outputs
+
+
+# Below: IO pre- and post-processor classes for Perceiver.
+def space_to_depth(frames: torch.Tensor, temporal_block_size: int = 1, spatial_block_size: int = 1) -> torch.Tensor:
+ """
+ Space to depth transform. Rearranges blocks of spatial data, into depth.
+
+ This function assumes the channels to be first, but will place the channels last after transformation.
+ """
+ if len(frames.shape) == 4:
+ batch_size, num_channels, height, width = frames.shape
+ # split up dimensions (height by spatial_block_size, width by spatial_block_size)
+ frames = frames.view(
+ batch_size,
+ num_channels,
+ height // spatial_block_size,
+ spatial_block_size,
+ width // spatial_block_size,
+ spatial_block_size,
+ )
+ # move blocks to last dimension: (batch_size, H//bs, W//bs, bs, bs, C)
+ frames = frames.permute(0, 2, 4, 3, 5, 1).contiguous()
+ # concatenate blocks along channel dimension: (batch_size, H//bs, W//bs, bs*bs*C)
+ frames = frames.view(
+ batch_size,
+ height // spatial_block_size,
+ width // spatial_block_size,
+ (spatial_block_size**2) * num_channels,
+ )
+ return frames
+ elif len(frames.shape) == 5:
+ batch_size, time, num_channels, height, width = frames.shape
+ # split up dimensions (time by temporal_block_size, height by spatial_block_size, width by spatial_block_size)
+ frames = frames.view(
+ batch_size,
+ time // temporal_block_size,
+ temporal_block_size,
+ num_channels,
+ height // spatial_block_size,
+ spatial_block_size,
+ width // spatial_block_size,
+ spatial_block_size,
+ )
+ # move blocks to last dimension: (batch_size, T//ts, H//bs, W//bs, ts, bs, bs, C)
+ frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous()
+ # concatenate blocks along channel dimension: (batch_size, T//ts, H//bs, W//bs, ts*bs*bs*C)
+ frames = frames.view(
+ batch_size,
+ time // temporal_block_size,
+ height // spatial_block_size,
+ width // spatial_block_size,
+ temporal_block_size * (spatial_block_size**2) * num_channels,
+ )
+ return frames
+ else:
+ raise ValueError(
+ "Frames should be of rank 4 (batch, channels, height, width)"
+ " or rank 5 (batch, time, channels, height, width)"
+ )
+
+
+class Conv2dSamePadding(nn.Conv2d):
+ """
+ Conv2d layer with padding="same" support. Source:
+ https://gist.github.com/sumanmichael/4de9dee93f972d47c80c4ade8e149ea6
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.zero_pad_2d = nn.ZeroPad2d(
+ reduce(__add__, [(k // 2 + (k - 2 * (k // 2)) - 1, k // 2) for k in self.kernel_size[::-1]])
+ )
+
+ def forward(self, input):
+ return self._conv_forward(self.zero_pad_2d(input), self.weight, self.bias)
+
+
+class Conv2DDownsample(nn.Module):
+ """Downsamples 4x by applying a 2D convolution and doing max pooling."""
+
+ def __init__(
+ self,
+ num_layers: int = 1,
+ in_channels: int = 3,
+ out_channels: int = 64,
+ use_batchnorm: bool = True,
+ ):
+ """
+ Constructs a Conv2DDownsample model.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 64):
+ The number of conv output channels.
+ use_batchnorm (`bool`, *optional*, defaults to `True`):
+ Whether to use batchnorm.
+ """
+ super().__init__()
+
+ self.conv = Conv2dSamePadding(
+ in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=2, bias=False
+ )
+ self.batchnorm = nn.BatchNorm2d(num_features=out_channels) if use_batchnorm else nn.Identity()
+ self.relu = nn.ReLU()
+ self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2)
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ out = self.conv(inputs)
+ out = self.batchnorm(out)
+ out = self.relu(out)
+ out = self.max_pool(out)
+ return out
+
+
+def generate_fourier_features(pos, num_bands, max_resolution=(224, 224), concat_pos=True, sine_only=False):
+ """
+ Generate a Fourier frequency position encoding with linear spacing.
+
+ Args:
+ pos (`torch.LongTensor` of shape `(batch_size, sequence_length, dim)`):
+ The Tensor containing the position of n points in d dimensional space.
+ num_bands (`int`):
+ The number of frequency bands (K) to use.
+ max_resolution (`tuple[int]`, *optional*, defaults to (224, 224)):
+ The maximum resolution (i.e. the number of pixels per dim). A tuple representing resolution for each dimension.
+ concat_pos (`bool`, *optional*, defaults to `True`):
+ Whether to concatenate the input position encoding to the Fourier features.
+ sine_only (`bool`, *optional*, defaults to `False`):
+ Whether to use a single phase (sin) or two (sin/cos) for each frequency band.
+
+ Returns:
+ `torch.FloatTensor` of shape `(batch_size, sequence_length, n_channels)`: The Fourier position embeddings. If
+ `concat_pos` is `True` and `sine_only` is `False`, output dimensions are ordered as: [dim_1, dim_2, ..., dim_d,
+ sin(pi*f_1*dim_1), ..., sin(pi*f_K*dim_1), ..., sin(pi*f_1*dim_d), ..., sin(pi*f_K*dim_d), cos(pi*f_1*dim_1),
+ ..., cos(pi*f_K*dim_1), ..., cos(pi*f_1*dim_d), ..., cos(pi*f_K*dim_d)], where dim_i is pos[:, i] and f_k is the
+ kth frequency band.
+ """
+
+ batch_size = pos.shape[0]
+
+ min_freq = 1.0
+ # Nyquist frequency at the target resolution:
+ freq_bands = torch.stack(
+ [torch.linspace(start=min_freq, end=res / 2, steps=num_bands) for res in max_resolution], dim=0
+ )
+
+ # Get frequency bands for each spatial dimension.
+ # Output is size [n, d * num_bands]
+ per_pos_features = pos[0, :, :][:, :, None] * freq_bands[None, :, :]
+ per_pos_features = torch.reshape(per_pos_features, [-1, np.prod(per_pos_features.shape[1:])])
+
+ if sine_only:
+ # Output is size [n, d * num_bands]
+ per_pos_features = torch.sin(np.pi * (per_pos_features))
+ else:
+ # Output is size [n, 2 * d * num_bands]
+ per_pos_features = torch.cat(
+ [torch.sin(np.pi * per_pos_features), torch.cos(np.pi * per_pos_features)], dim=-1
+ )
+ # Concatenate the raw input positions.
+ if concat_pos:
+ # Adds d bands to the encoding.
+ per_pos_features = torch.cat([pos, per_pos_features.expand(batch_size, -1, -1)], dim=-1)
+ return per_pos_features
+
+
+def build_linear_positions(index_dims, output_range=(-1.0, 1.0)):
+ """
+ Generate an array of position indices for an N-D input array.
+
+ Args:
+ index_dims (`list[int]`):
+ The shape of the index dimensions of the input array.
+ output_range (`tuple[float]`, *optional*, defaults to `(-1.0, 1.0)`):
+ The min and max values taken by each input index dimension.
+
+ Returns:
+ `torch.FloatTensor` of shape `(index_dims[0], index_dims[1], .., index_dims[-1], N)`.
+ """
+
+ def _linspace(n_xels_per_dim):
+ return torch.linspace(start=output_range[0], end=output_range[1], steps=n_xels_per_dim, dtype=torch.float32)
+
+ dim_ranges = [_linspace(n_xels_per_dim) for n_xels_per_dim in index_dims]
+ array_index_grid = torch.meshgrid(*dim_ranges, indexing="ij")
+
+ return torch.stack(array_index_grid, dim=-1)
+
+
+class PerceiverAbstractPositionEncoding(nn.Module, metaclass=abc.ABCMeta):
+ """Perceiver abstract position encoding."""
+
+ @property
+ @abc.abstractmethod
+ def num_dimensions(self) -> int:
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def output_size(self, *args, **kwargs) -> int:
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def forward(self, batch_size, pos):
+ raise NotImplementedError
+
+
+class PerceiverTrainablePositionEncoding(PerceiverAbstractPositionEncoding):
+ """Trainable position encoding."""
+
+ def __init__(self, index_dims, num_channels=128):
+ super().__init__()
+ self._num_channels = num_channels
+ self._index_dims = index_dims
+ index_dim = np.prod(index_dims)
+ self.position_embeddings = nn.Parameter(torch.randn(index_dim, num_channels))
+
+ @property
+ def num_dimensions(self) -> int:
+ if isinstance(self._index_dims, int):
+ return 1
+ return len(self._index_dims)
+
+ def output_size(self, *args, **kwargs) -> int:
+ return self._num_channels
+
+ def interpolate_pos_encoding(self, position_embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ num_positions = position_embeddings.shape[0]
+ new_height = new_width = torch_int(num_positions**0.5)
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and height == new_height and width == new_width:
+ return position_embeddings
+
+ position_embeddings = position_embeddings.reshape(1, new_height, new_width, self._num_channels).permute(
+ 0, 3, 1, 2
+ )
+
+ position_embeddings = nn.functional.interpolate(
+ position_embeddings,
+ size=(height, width),
+ mode="bicubic",
+ align_corners=False,
+ )
+ position_embeddings = position_embeddings.reshape(1, self._num_channels, -1).permute(0, 2, 1).squeeze(0)
+ return position_embeddings
+
+ def forward(
+ self, batch_size: int, interpolate_pos_encoding: bool = False, input_size: torch.Size | None = None
+ ) -> torch.Tensor:
+ position_embeddings = self.position_embeddings
+
+ if interpolate_pos_encoding:
+ height, width = input_size
+ position_embeddings = self.interpolate_pos_encoding(position_embeddings, height, width)
+
+ if batch_size is not None:
+ position_embeddings = position_embeddings.expand(batch_size, -1, -1)
+ return position_embeddings
+
+
+def _check_or_build_spatial_positions(pos, index_dims, batch_size):
+ """
+ Checks or builds spatial position features (x, y, ...).
+
+ Args:
+ pos (`torch.FloatTensor`):
+ None, or an array of position features. If None, position features are built. Otherwise, their size is checked.
+ index_dims (`list[int]`):
+ An iterable giving the spatial/index size of the data to be featurized.
+ batch_size (`int`):
+ The batch size of the data to be featurized.
+
+ Returns:
+ `torch.FloatTensor` of shape `(batch_size, prod(index_dims))` an array of position features.
+ """
+ if pos is None:
+ pos = build_linear_positions(index_dims)
+ # equivalent to `torch.broadcast_to(pos[None], (batch_size,) + pos.shape)`
+ # but `torch.broadcast_to` cannot be converted to ONNX
+ pos = pos[None].expand((batch_size,) + pos.shape)
+ pos = torch.reshape(pos, [batch_size, np.prod(index_dims), -1])
+ else:
+ # Just a warning label: you probably don't want your spatial features to
+ # have a different spatial layout than your pos coordinate system.
+ # But feel free to override if you think it'll work!
+ if pos.shape[-1] != len(index_dims):
+ raise ValueError("Spatial features have the wrong number of dimensions.")
+ return pos
+
+
+class PerceiverFourierPositionEncoding(PerceiverAbstractPositionEncoding):
+ """Fourier (Sinusoidal) position encoding."""
+
+ def __init__(self, num_bands, max_resolution, concat_pos=True, sine_only=False):
+ super().__init__()
+ self.num_bands = num_bands
+ self.max_resolution = max_resolution
+ self.concat_pos = concat_pos
+ self.sine_only = sine_only
+
+ @property
+ def num_dimensions(self) -> int:
+ return len(self.max_resolution)
+
+ def output_size(self):
+ """Returns size of positional encodings last dimension."""
+ num_dims = len(self.max_resolution)
+ encoding_size = self.num_bands * num_dims
+ if not self.sine_only:
+ encoding_size *= 2
+ if self.concat_pos:
+ encoding_size += self.num_dimensions
+
+ return encoding_size
+
+ def forward(
+ self,
+ index_dims: list[int],
+ batch_size: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ pos: torch.FloatTensor | None = None,
+ ) -> torch.FloatTensor:
+ pos = _check_or_build_spatial_positions(pos, index_dims, batch_size)
+ fourier_pos_enc = generate_fourier_features(
+ pos,
+ num_bands=self.num_bands,
+ max_resolution=self.max_resolution,
+ concat_pos=self.concat_pos,
+ sine_only=self.sine_only,
+ ).to(device=device, dtype=dtype)
+ return fourier_pos_enc
+
+
+class AbstractPreprocessor(nn.Module):
+ @property
+ def num_channels(self) -> int:
+ """Returns size of preprocessor output."""
+ raise NotImplementedError()
+
+
+class PerceiverTextPreprocessor(AbstractPreprocessor):
+ """
+ Text preprocessing for Perceiver Encoder. Can be used to embed `inputs` and add positional encodings.
+
+ The dimensionality of the embeddings is determined by the `d_model` attribute of the configuration.
+
+ Args:
+ config ([`PerceiverConfig`]):
+ Model configuration.
+ """
+
+ def __init__(self, config: PerceiverConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.embeddings = nn.Embedding(num_embeddings=config.vocab_size, embedding_dim=config.d_model)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.d_model)
+
+ @property
+ def num_channels(self) -> int:
+ return self.config.d_model
+
+ def forward(
+ self,
+ inputs: torch.LongTensor,
+ pos: torch.Tensor | None = None,
+ network_input_is_1d: bool = True,
+ interpolate_pos_encoding: bool = False,
+ ):
+ embeddings_without_pos = self.embeddings(inputs)
+
+ seq_length = inputs.shape[1]
+ position_ids = torch.arange(0, seq_length, device=inputs.device)
+ embeddings = embeddings_without_pos + self.position_embeddings(position_ids)
+
+ return embeddings, None, embeddings_without_pos
+
+
+class PerceiverEmbeddingDecoder(nn.Module):
+ """
+ Module to decode embeddings (for masked language modeling).
+
+ Args:
+ config ([`PerceiverConfig`]):
+ Model configuration.
+ """
+
+ def __init__(self, config: PerceiverConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.vocab_size = config.vocab_size
+ self.bias = nn.Parameter(torch.zeros(self.vocab_size))
+
+ def forward(self, hidden_states: torch.Tensor, embedding_layer: torch.Tensor) -> torch.Tensor:
+ batch_size, seq_len, d_model = hidden_states.shape
+ # Flatten batch dim
+ output = torch.matmul(hidden_states.reshape([-1, d_model]), embedding_layer.weight.transpose(0, 1))
+ output = output + self.bias
+
+ return output.reshape([batch_size, seq_len, self.vocab_size])
+
+
+class PerceiverMultimodalPostprocessor(nn.Module):
+ """
+ Multimodal postprocessing for Perceiver. Can be used to combine modality-specific postprocessors into a single
+ postprocessor.
+
+ Args:
+ modalities (`Mapping[str, PostprocessorType]`):
+ Dictionary mapping modality name to postprocessor class for that modality.
+ input_is_dict (`bool`, *optional*, defaults to `False`):
+ If True, input is assumed to be dictionary structured, and outputs keep the same dictionary shape. If
+ False, input is a tensor which is sliced up during postprocessing by *modality_sizes*.
+ """
+
+ def __init__(self, modalities: Mapping[str, PostprocessorType], input_is_dict: bool = False):
+ super().__init__()
+ self.modalities = nn.ModuleDict(modalities)
+ self.input_is_dict = input_is_dict
+
+ def forward(
+ self, inputs: torch.Tensor, pos: torch.Tensor | None = None, modality_sizes=None
+ ) -> Mapping[str, torch.Tensor]:
+ if not self.input_is_dict:
+ # Slice up modalities by their sizes.
+ if modality_sizes is None:
+ raise ValueError("Modality sizes should be specified if input is not a dictionary.")
+ inputs = restructure(modality_sizes=modality_sizes, inputs=inputs)
+
+ outputs = {
+ modality: postprocessor(inputs[modality], pos=pos, modality_sizes=None)
+ for modality, postprocessor in self.modalities.items()
+ }
+ return outputs
+
+
+class PerceiverClassificationPostprocessor(nn.Module):
+ """
+ Classification postprocessing for Perceiver. Can be used to convert the decoder output to classification logits.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ in_channels (`int`):
+ Number of channels in the input.
+ """
+
+ def __init__(self, config: PerceiverConfig, in_channels: int) -> None:
+ super().__init__()
+ self.classifier = nn.Linear(in_channels, config.num_labels)
+
+ def forward(self, inputs, pos: torch.Tensor | None = None, modality_sizes=None) -> torch.Tensor:
+ logits = self.classifier(inputs)
+ return logits[:, 0, :]
+
+
+class PerceiverAudioPostprocessor(nn.Module):
+ """
+ Audio postprocessing for Perceiver. Can be used to convert the decoder output to audio features.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ in_channels (`int`):
+ Number of channels in the input.
+ postproc_type (`str`, *optional*, defaults to `"patches"`):
+ Postprocessor type to use. Currently, only "patches" is supported.
+ """
+
+ def __init__(self, config: PerceiverConfig, in_channels: int, postproc_type: str = "patches") -> None:
+ super().__init__()
+
+ if postproc_type != "patches": # to be supported: 'conv', 'patches', 'pixels'
+ raise ValueError("Invalid postproc_type!")
+
+ # Architecture parameters:
+ self.classifier = nn.Linear(in_channels, config.samples_per_patch)
+
+ def forward(self, inputs: torch.Tensor, pos: torch.Tensor | None = None, modality_sizes=None) -> torch.Tensor:
+ logits = self.classifier(inputs)
+ return torch.reshape(logits, [inputs.shape[0], -1])
+
+
+class PerceiverProjectionPostprocessor(nn.Module):
+ """
+ Projection postprocessing for Perceiver. Can be used to project the channels of the decoder output to a lower
+ dimension.
+
+ Args:
+ in_channels (`int`):
+ Number of channels in the input.
+ out_channels (`int`):
+ Number of channels in the output.
+ """
+
+ def __init__(self, in_channels: int, out_channels: int) -> None:
+ super().__init__()
+ self.classifier = nn.Linear(in_channels, out_channels)
+
+ def forward(self, inputs: torch.Tensor, pos: torch.Tensor | None = None, modality_sizes=None) -> torch.Tensor:
+ logits = self.classifier(inputs)
+ return logits
+
+
+class PerceiverImagePreprocessor(AbstractPreprocessor):
+ """
+ Image preprocessing for Perceiver Encoder.
+
+ Note: the *out_channels* argument refers to the output channels of a convolutional layer, if *prep_type* is set to
+ "conv1x1" or "conv". If one adds absolute position embeddings, one must make sure the *num_channels* of the
+ position encoding kwargs are set equal to the *out_channels*.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ prep_type (`str`, *optional*, defaults to `"conv"`):
+ Preprocessing type. Can be "conv1x1", "conv", "patches", "pixels".
+ spatial_downsample (`int`, *optional*, defaults to 4):
+ Spatial downsampling factor.
+ temporal_downsample (`int`, *optional*, defaults to 1):
+ Temporal downsampling factor (only relevant in case a time dimension is present).
+ position_encoding_type (`str`, *optional*, defaults to `"fourier"`):
+ Position encoding type. Can be "fourier" or "trainable".
+ in_channels (`int`, *optional*, defaults to 3):
+ Number of channels in the input.
+ out_channels (`int`, *optional*, defaults to 64):
+ Number of channels in the output.
+ conv_after_patching (`bool`, *optional*, defaults to `False`):
+ Whether to apply a convolutional layer after patching.
+ conv_after_patching_in_channels (`int`, *optional*, defaults to 54):
+ Number of channels in the input of the convolutional layer after patching.
+ conv2d_use_batchnorm (`bool`, *optional*, defaults to `True`):
+ Whether to use batch normalization in the convolutional layer.
+ concat_or_add_pos (`str`, *optional*, defaults to `"concat"`):
+ How to concatenate the position encoding to the input. Can be "concat" or "add".
+ project_pos_dim (`int`, *optional*, defaults to -1):
+ Dimension of the position encoding to project to. If -1, no projection is applied.
+ **position_encoding_kwargs (`Dict`, *optional*):
+ Keyword arguments for the position encoding.
+ """
+
+ def __init__(
+ self,
+ config,
+ prep_type="conv",
+ spatial_downsample: int = 4,
+ temporal_downsample: int = 1,
+ position_encoding_type: str = "fourier",
+ in_channels: int = 3,
+ out_channels: int = 64,
+ conv_after_patching: bool = False,
+ conv_after_patching_in_channels: int = 54, # only relevant when conv_after_patching = True
+ conv2d_use_batchnorm: bool = True,
+ concat_or_add_pos: str = "concat",
+ project_pos_dim: int = -1,
+ **position_encoding_kwargs,
+ ):
+ super().__init__()
+ self.config = config
+
+ if prep_type not in ("conv", "patches", "pixels", "conv1x1"):
+ raise ValueError(f"Prep_type {prep_type} is invalid")
+
+ if concat_or_add_pos not in ["concat", "add"]:
+ raise ValueError(f"Invalid value {concat_or_add_pos} for concat_or_add_pos.")
+
+ self.in_channels = in_channels
+ self.prep_type = prep_type
+ self.spatial_downsample = spatial_downsample
+ self.temporal_downsample = temporal_downsample
+ self.position_encoding_type = position_encoding_type
+ self.concat_or_add_pos = concat_or_add_pos
+ self.conv_after_patching = conv_after_patching
+ self.out_channels = out_channels
+
+ if self.prep_type == "conv":
+ # Downsampling with conv is currently restricted
+ convnet_num_layers = math.log(spatial_downsample, 4)
+ convnet_num_layers_is_int = convnet_num_layers == np.round(convnet_num_layers)
+ if not convnet_num_layers_is_int or temporal_downsample != 1:
+ raise ValueError(
+ "Only powers of 4 expected for spatial and 1 expected for temporal downsampling with conv."
+ )
+ self.convnet = Conv2DDownsample(
+ in_channels=in_channels,
+ num_layers=int(convnet_num_layers),
+ out_channels=out_channels,
+ use_batchnorm=conv2d_use_batchnorm,
+ )
+
+ elif self.prep_type == "conv1x1":
+ if temporal_downsample != 1:
+ raise ValueError("Conv1x1 does not downsample in time.")
+ self.convnet_1x1 = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(1, 1),
+ # spatial_downsample is unconstrained for 1x1 convolutions.
+ stride=(spatial_downsample, spatial_downsample),
+ )
+
+ # Position embeddings
+ self.project_pos_dim = project_pos_dim
+ self.position_embeddings, self.positions_projection = build_position_encoding(
+ position_encoding_type=position_encoding_type,
+ out_channels=out_channels,
+ project_pos_dim=project_pos_dim,
+ **position_encoding_kwargs,
+ )
+
+ # Optional convolutional layer after patches.
+ self.conv_after_patches = (
+ nn.Linear(conv_after_patching_in_channels, self.out_channels) if conv_after_patching else nn.Identity()
+ )
+
+ @property
+ def num_channels(self) -> int:
+ # Let's assume that the number of resolutions (in the context of image preprocessing)
+ # of the input data is 2 or 3 depending on whether we are processing image or video respectively.
+ # In this case, for convenience, we will declare is_temporal variable,
+ # which will show whether the data has a temporal dimension or not.
+ is_temporal = self.position_embeddings.num_dimensions > 2
+
+ # position embedding
+ if self.project_pos_dim > 0:
+ pos_dim = self.project_pos_dim
+ else:
+ pos_dim = self.position_embeddings.output_size()
+ if self.concat_or_add_pos == "add":
+ return pos_dim
+
+ # inputs
+ if self.conv_after_patching or self.prep_type in ("conv1x1", "conv"):
+ inp_dim = self.out_channels
+ elif self.prep_type == "pixels":
+ inp_dim = self.in_channels
+ if not is_temporal:
+ inp_dim = math.ceil(inp_dim / self.spatial_downsample)
+ elif self.prep_type == "patches":
+ if self.conv_after_patching:
+ inp_dim = self.out_channels
+ else:
+ inp_dim = self.in_channels * self.spatial_downsample**2
+ if is_temporal:
+ inp_dim *= self.temporal_downsample
+
+ return inp_dim + pos_dim
+
+ def _build_network_inputs(
+ self, inputs: torch.Tensor, network_input_is_1d: bool = True, interpolate_pos_encoding: bool = False
+ ):
+ """
+ Construct the final input, including position encoding.
+
+ This method expects the inputs to always have channels as last dimension.
+
+ """
+ batch_size = inputs.shape[0]
+ input_size = inputs.shape[1:3]
+ index_dims = inputs.shape[1:-1]
+ indices = np.prod(index_dims)
+
+ # Flatten input features to a 1D index dimension if necessary.
+ if len(inputs.shape) > 3 and network_input_is_1d:
+ inputs = torch.reshape(inputs, [batch_size, indices, -1])
+
+ # Construct the position encoding.
+ if self.position_encoding_type == "trainable":
+ pos_enc = self.position_embeddings(batch_size, interpolate_pos_encoding, input_size)
+ elif self.position_encoding_type == "fourier":
+ pos_enc = self.position_embeddings(index_dims, batch_size, device=inputs.device, dtype=inputs.dtype)
+
+ # Optionally project them to a target dimension.
+ pos_enc = self.positions_projection(pos_enc)
+
+ if not network_input_is_1d:
+ # Reshape pos to match the input feature shape
+ # if the network takes non-1D inputs
+ sh = inputs.shape
+ pos_enc = torch.reshape(pos_enc, list(sh)[:-1] + [-1])
+ if self.concat_or_add_pos == "concat":
+ inputs_with_pos = torch.cat([inputs, pos_enc], dim=-1)
+ elif self.concat_or_add_pos == "add":
+ inputs_with_pos = inputs + pos_enc
+ return inputs_with_pos, inputs
+
+ def forward(
+ self,
+ inputs: torch.Tensor,
+ pos: torch.Tensor | None = None,
+ network_input_is_1d: bool = True,
+ interpolate_pos_encoding: bool = False,
+ ):
+ if self.prep_type == "conv":
+ # Convnet image featurization.
+ # Downsamples spatially by a factor of 4
+ inputs = self.convnet(inputs)
+
+ elif self.prep_type == "conv1x1":
+ # map inputs to self.out_channels
+ inputs = self.convnet_1x1(inputs)
+
+ elif self.prep_type == "pixels":
+ # if requested, downsamples in the crudest way
+ if inputs.ndim == 4:
+ inputs = inputs[:: self.spatial_downsample, :: self.spatial_downsample]
+ elif inputs.ndim == 5:
+ inputs = inputs[
+ :, :: self.temporal_downsample, :, :: self.spatial_downsample, :: self.spatial_downsample
+ ]
+ else:
+ raise ValueError("Unsupported data format for pixels.")
+
+ elif self.prep_type == "patches":
+ # Space2depth featurization.
+ # Video: B x T x C x H x W
+ inputs = space_to_depth(
+ inputs, temporal_block_size=self.temporal_downsample, spatial_block_size=self.spatial_downsample
+ )
+
+ if inputs.ndim == 5 and inputs.shape[1] == 1:
+ # for flow
+ inputs = inputs.squeeze(dim=1)
+
+ # Optionally apply conv layer.
+ inputs = self.conv_after_patches(inputs)
+
+ if self.prep_type != "patches":
+ # move channels to last dimension, as the _build_network_inputs method below expects this
+ if inputs.ndim == 4:
+ inputs = inputs.permute(0, 2, 3, 1)
+ elif inputs.ndim == 5:
+ inputs = inputs.permute(0, 1, 3, 4, 2)
+ else:
+ raise ValueError("Unsupported data format for conv1x1.")
+
+ inputs, inputs_without_pos = self._build_network_inputs(inputs, network_input_is_1d, interpolate_pos_encoding)
+ modality_sizes = None # Size for each modality, only needed for multimodal
+
+ return inputs, modality_sizes, inputs_without_pos
+
+
+class PerceiverOneHotPreprocessor(AbstractPreprocessor):
+ """
+ One-hot preprocessor for Perceiver Encoder. Can be used to add a dummy index dimension to the input.
+
+ Args:
+ config ([`PerceiverConfig`]):
+ Model configuration.
+ """
+
+ def __init__(self, config: PerceiverConfig) -> None:
+ super().__init__()
+ self.config: PerceiverConfig = config
+
+ @property
+ def num_channels(self) -> int:
+ return self.config.num_labels
+
+ def forward(self, inputs: torch.Tensor, pos: torch.Tensor | None = None, network_input_is_1d: bool = True):
+ # Add a dummy index dimension.
+ inputs = inputs[:, None, :]
+
+ # No position encodings, so the 1st (input) and 3rd (inputs_without_pos)
+ # outputs are identical.
+ return inputs, None, inputs
+
+
+class PerceiverAudioPreprocessor(AbstractPreprocessor):
+ """
+ Audio preprocessing for Perceiver Encoder.
+
+ Args:
+ config ([*PerceiverConfig*]):
+ Model configuration.
+ prep_type (`str`, *optional*, defaults to `"patches"`):
+ Preprocessor type to use. Only "patches" is supported.
+ samples_per_patch (`int`, *optional*, defaults to 96):
+ Number of samples per patch.
+ position_encoding_type (`str`, *optional*, defaults to `"fourier"`):
+ Type of position encoding to use. Can be "trainable" or "fourier".
+ concat_or_add_pos (`str`, *optional*, defaults to `"concat"`):
+ How to concatenate the position encoding to the input. Can be "concat" or "add".
+ out_channels (`int`, *optional*, defaults to 64):
+ Number of channels in the output.
+ project_pos_dim (`int`, *optional*, defaults to -1):
+ Dimension of the position encoding to project to. If -1, no projection is applied.
+ **position_encoding_kwargs (`Dict`, *optional*):
+ Keyword arguments for the position encoding.
+ """
+
+ def __init__(
+ self,
+ config,
+ prep_type: str = "patches",
+ samples_per_patch: int = 96,
+ position_encoding_type: str = "fourier",
+ concat_or_add_pos: str = "concat",
+ out_channels=64,
+ project_pos_dim=-1,
+ **position_encoding_kwargs,
+ ):
+ super().__init__()
+ self.config = config
+
+ if prep_type != "patches":
+ raise ValueError(f"Prep_type {prep_type} is invalid, can only be 'patches'.")
+
+ if concat_or_add_pos not in ["concat", "add"]:
+ raise ValueError(f"Concat_or_pos {concat_or_add_pos} is invalid, can only be 'concat' or 'add'.")
+
+ self.samples_per_patch = samples_per_patch
+ self.position_encoding_type = position_encoding_type
+ self.concat_or_add_pos = concat_or_add_pos
+ self.project_pos_dim = project_pos_dim
+
+ # Position embeddings
+ self.position_embeddings, self.positions_projection = build_position_encoding(
+ position_encoding_type=position_encoding_type,
+ out_channels=out_channels,
+ project_pos_dim=project_pos_dim,
+ **position_encoding_kwargs,
+ )
+
+ @property
+ def num_channels(self) -> int:
+ # position embedding
+ if self.project_pos_dim > 0:
+ pos_dim = self.project_pos_dim
+ else:
+ pos_dim = self.position_embeddings.output_size()
+ if self.concat_or_add_pos == "add":
+ return pos_dim
+ return self.samples_per_patch + pos_dim
+
+ def _build_network_inputs(self, inputs):
+ """Construct the final input, including position encoding."""
+ batch_size = inputs.shape[0]
+ index_dims = inputs.shape[1:-1]
+
+ # Construct the position encoding.
+ if self.position_encoding_type == "trainable":
+ pos_enc = self.position_embeddings(batch_size)
+ elif self.position_encoding_type == "fourier":
+ pos_enc = self.position_embeddings(index_dims, batch_size, device=inputs.device, dtype=inputs.dtype)
+
+ # Optionally project them to a target dimension.
+ pos_enc = self.positions_projection(pos_enc)
+
+ if self.concat_or_add_pos == "concat":
+ inputs_with_pos = torch.cat([inputs, pos_enc], dim=-1)
+ elif self.concat_or_add_pos == "add":
+ inputs_with_pos = inputs + pos_enc
+
+ return inputs_with_pos, inputs
+
+ def forward(
+ self,
+ inputs: torch.Tensor,
+ pos: torch.Tensor | None = None,
+ network_input_is_1d: bool = True,
+ interpolate_pos_encoding: bool = False,
+ ):
+ inputs = torch.reshape(inputs, [inputs.shape[0], -1, self.samples_per_patch])
+
+ inputs, inputs_without_pos = self._build_network_inputs(inputs)
+ modality_sizes = None # Size for each modality, only needed for multimodal
+
+ return inputs, modality_sizes, inputs_without_pos
+
+
+class PerceiverMultimodalPreprocessor(AbstractPreprocessor):
+ """
+ Multimodal preprocessing for Perceiver Encoder.
+
+ Inputs for each modality are preprocessed, then padded with trainable position embeddings to have the same number
+ of channels.
+
+ Args:
+ modalities (`Mapping[str, PreprocessorType]`):
+ Dict mapping modality name to preprocessor.
+ mask_probs (`dict[str, float]`):
+ Dict mapping modality name to masking probability of that modality.
+ min_padding_size (`int`, *optional*, defaults to 2):
+ The minimum padding size for all modalities. The final output will have num_channels equal to the maximum
+ channels across all modalities plus min_padding_size.
+ """
+
+ def __init__(
+ self,
+ modalities: Mapping[str, PreprocessorType],
+ mask_probs: Mapping[str, float] | None = None,
+ min_padding_size: int = 2,
+ ):
+ super().__init__()
+ self.modalities = nn.ModuleDict(modalities)
+ self.min_padding_size = min_padding_size
+ self.mask_probs = mask_probs if mask_probs is not None else {}
+ self.padding = nn.ParameterDict(
+ {
+ modality: nn.Parameter(torch.randn(1, self.num_channels - preprocessor.num_channels))
+ for modality, preprocessor in modalities.items()
+ }
+ )
+ self.mask = nn.ParameterDict(
+ {modality: nn.Parameter(torch.randn(1, self.num_channels)) for modality, _ in self.mask_probs.items()}
+ )
+
+ @property
+ def num_channels(self) -> int:
+ max_channel_size = max(processor.num_channels for _, processor in self.modalities.items())
+ common_channel_size = max_channel_size + self.min_padding_size
+ return common_channel_size
+
+ def forward(
+ self,
+ inputs: Mapping[str, torch.Tensor],
+ pos: torch.Tensor | None = None,
+ network_input_is_1d: bool = True,
+ interpolate_pos_encoding: bool = False,
+ ) -> PreprocessorOutputType:
+ padded = {}
+ modality_sizes = {}
+ inputs_without_pos = {}
+ for modality, preprocessor in self.modalities.items():
+ # preprocess each modality using the respective preprocessor.
+ output, _, inputs_without_pos[modality] = preprocessor(
+ inputs[modality], pos=pos, network_input_is_1d=network_input_is_1d
+ )
+
+ # pad to the same common_channel_size.
+ batch_size, num_samples, num_channels = output.shape
+ pos_enc = self.padding[modality].expand(batch_size, -1, -1)
+
+ padding = torch.broadcast_to(
+ pos_enc,
+ [batch_size, num_samples, self.num_channels - num_channels],
+ )
+ output_padded = torch.cat([output, padding], dim=2)
+
+ # mask if required
+ if modality in self.mask_probs:
+ mask_token = self.mask[modality].expand(batch_size, -1, -1)
+ mask_prob = self.mask_probs[modality]
+ mask = torch.bernoulli(torch.full([batch_size, num_samples], mask_prob))
+ mask = torch.unsqueeze(mask, dim=2).to(mask_token.device)
+ output_padded = (1 - mask) * output_padded + mask * mask_token
+
+ padded[modality] = output_padded
+ modality_sizes[modality] = output_padded.shape[1]
+
+ # Apply a predictable ordering to the modalities
+ padded_ls = [padded[k] for k in sorted(padded.keys())]
+
+ # Finally, concatenate along the time dimension
+ final_inputs = torch.cat(padded_ls, dim=1)
+
+ return final_inputs, modality_sizes, inputs_without_pos
+
+
+__all__ = [
+ "PerceiverForImageClassificationConvProcessing",
+ "PerceiverForImageClassificationFourier",
+ "PerceiverForImageClassificationLearned",
+ "PerceiverForMaskedLM",
+ "PerceiverForMultimodalAutoencoding",
+ "PerceiverForOpticalFlow",
+ "PerceiverForSequenceClassification",
+ "PerceiverLayer",
+ "PerceiverModel",
+ "PerceiverPreTrainedModel",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/tokenization_perceiver.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/tokenization_perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..6abaf31aed11e601ef288a1d12d2bbedab36abb1
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/perceiver/tokenization_perceiver.py
@@ -0,0 +1,197 @@
+# 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.
+"""Tokenization class for Perceiver."""
+
+from ...tokenization_python import AddedToken, PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class PerceiverTokenizer(PreTrainedTokenizer):
+ """
+ Construct a Perceiver tokenizer. The Perceiver 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:
+ pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ bos_token (`str`, *optional*, defaults to `"[BOS]"`):
+ The BOS token (reserved in the vocab, but not actually used).
+ eos_token (`str`, *optional*, defaults to `"[EOS]"`):
+ The end of sequence token (reserved in the vocab, but not actually used).
+
+
+
+ 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`.
+
+
+
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The MASK token, useful for masked language modeling.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The CLS token (reserved in the vocab, but not actually used).
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from two sequences.
+
+ """
+
+ model_input_names = ["input_ids", "attention_mask"]
+
+ def __init__(
+ self,
+ pad_token="[PAD]",
+ bos_token="[BOS]",
+ eos_token="[EOS]",
+ mask_token="[MASK]",
+ cls_token="[CLS]",
+ sep_token="[SEP]",
+ model_max_length=2048,
+ **kwargs,
+ ) -> None:
+ pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
+ bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
+ eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
+ mask_token = AddedToken(mask_token, lstrip=False, rstrip=False) if isinstance(mask_token, str) else mask_token
+ cls_token = AddedToken(cls_token, lstrip=False, rstrip=False) if isinstance(cls_token, str) else cls_token
+ sep_token = AddedToken(sep_token, lstrip=False, rstrip=False) if isinstance(sep_token, str) else sep_token
+
+ self._utf_vocab_size = 2**8 # utf is 8 bits
+
+ # Since these tokens are not part of the vocabulary, we manually add them
+ self._added_tokens_decoder: dict[str, int] = {
+ 0: pad_token,
+ 1: bos_token,
+ 2: eos_token,
+ 3: mask_token,
+ 4: cls_token,
+ 5: sep_token,
+ }
+ self._num_special_tokens = len(self._added_tokens_decoder)
+ super().__init__(
+ pad_token=pad_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ mask_token=mask_token,
+ cls_token=cls_token,
+ sep_token=sep_token,
+ model_max_length=model_max_length,
+ **kwargs,
+ )
+
+ def get_vocab(self) -> dict[str, int]:
+ vocab = {}
+ for i in range(self._utf_vocab_size):
+ token = chr(i)
+ vocab[token] = i + self._num_special_tokens
+ vocab.update(self.added_tokens_encoder)
+ return vocab
+
+ @property
+ def vocab_size(self):
+ return self._utf_vocab_size
+
+ 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 [1] + [0] * len(token_ids_0) + [1]
+ return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+
+ 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. A sequence has the
+ following format:
+
+ - single sequence: `[CLS] X [SEP]`
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+ Args:
+ token_ids_0 (`list[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`list[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `list[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+ else:
+ return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + token_ids_1 + [self.sep_token_id]
+
+ 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 = self.unk_token_id
+ else:
+ token_id = ord(token) + self._num_special_tokens
+ 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._num_special_tokens)
+ return token
+
+ # TODO @ArthurZ refactor this as well....
+ 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_encoder:
+ tok_string = str(token).encode("utf-8")
+ else:
+ tok_string = bytes([ord(token)])
+ bstring += tok_string
+ string = bstring.decode("utf-8", errors="replace")
+ return string
+
+ # PerceiverTokenizer has no vocab file
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ return ()
+
+
+__all__ = ["PerceiverTokenizer"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf351e817fdfa519d1fdc84cb0084214fb37bd4f
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/__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_swin import *
+ from .modeling_swin import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fc542f03a54c5d09c4c791d87b2cc109975c601
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py
@@ -0,0 +1,90 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Swin Transformer model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="microsoft/swin-tiny-patch4-window7-224")
+@strict
+class SwinConfig(BackboneConfigMixin, PreTrainedConfig):
+ r"""
+ depths (`list(int)`, *optional*, defaults to `[2, 2, 6, 2]`):
+ Depth of each layer in the Transformer encoder.
+ num_heads (`list(int)`, *optional*, defaults to `[3, 6, 12, 24]`):
+ Number of attention heads in each layer of the Transformer encoder.
+ window_size (`int`, *optional*, defaults to 7):
+ Size of windows.
+ encoder_stride (`int`, *optional*, defaults to 32):
+ Factor to increase the spatial resolution by in the decoder head for masked image modeling.
+
+ Example:
+
+ ```python
+ >>> from transformers import SwinConfig, SwinModel
+
+ >>> # Initializing a Swin microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> configuration = SwinConfig()
+
+ >>> # Initializing a model (with random weights) from the microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> model = SwinModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "swin"
+
+ attribute_map = {
+ "num_attention_heads": "num_heads",
+ "num_hidden_layers": "num_layers",
+ }
+
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 4
+ num_channels: int = 3
+ embed_dim: int = 96
+ depths: list[int] | tuple[int, ...] = (2, 2, 6, 2)
+ num_heads: list[int] | tuple[int, ...] = (3, 6, 12, 24)
+ window_size: int = 7
+ mlp_ratio: float | int = 4.0
+ qkv_bias: bool = True
+ hidden_dropout_prob: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ drop_path_rate: float | int = 0.1
+ hidden_act: str = "gelu"
+ use_absolute_embeddings: bool = False
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-5
+ encoder_stride: int = 32
+ _out_features: list[str] | None = None
+ _out_indices: list[int] | None = None
+
+ def __post_init__(self, **kwargs):
+ self.num_layers = len(self.depths)
+ # we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
+ # this indicates the channel dimension after the last stage of the model
+ self.hidden_size = int(self.embed_dim * 2 ** (len(self.depths) - 1))
+ self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+ self.set_output_features_output_indices(
+ out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+ )
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["SwinConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..931fdaff839886d54aee5ea30c7f6258c035d7f7
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py
@@ -0,0 +1,1163 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/swin/modular_swin.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_swin.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# 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.
+
+import collections.abc
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BackboneOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_int
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_swin import SwinConfig
+
+
+class SwinDropPath(nn.Module):
+ """Stochastic depth (DropPath) per sample, for residual blocks.
+
+ Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
+ `_.
+ """
+
+ def __init__(self, drop_prob: float = 0.0) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ if self.drop_prob == 0.0 or not self.training:
+ return hidden_states
+ keep_prob = 1 - self.drop_prob
+ shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
+ random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
+ random_tensor = torch.floor(random_tensor + keep_prob)
+ return hidden_states.div(keep_prob) * random_tensor
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin encoder's outputs, with potential hidden states and attentions.
+ """
+)
+@dataclass
+class SwinEncoderOutput(ModelOutput):
+ r"""
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin model's outputs that also contains a pooling of the last hidden states.
+ """
+)
+@dataclass
+class SwinModelOutput(ModelOutput):
+ r"""
+ pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
+ Average pooling of the last layer hidden-state.
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ pooler_output: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin masked image model outputs.
+ """
+)
+@dataclass
+class SwinMaskedImageModelingOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
+ Masked image modeling (MLM) loss.
+ reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Reconstructed pixel values.
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ loss: torch.FloatTensor | None = None
+ reconstruction: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin outputs for image classification.
+ """
+)
+@dataclass
+class SwinImageClassifierOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Classification (or regression if config.num_labels==1) loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+ Classification (or regression if config.num_labels==1) scores (before SoftMax).
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+class SwinEmbeddings(nn.Module):
+ """
+ Construct the patch and position embeddings. Optionally, also the mask token.
+ """
+
+ def __init__(self, config, use_mask_token=False):
+ super().__init__()
+
+ self.patch_embeddings = SwinPatchEmbeddings(config)
+ num_patches = self.patch_embeddings.num_patches
+ self.patch_grid = self.patch_embeddings.grid_size
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None
+
+ self.position_embeddings = (
+ nn.Parameter(torch.zeros(1, num_patches, config.embed_dim)) if config.use_absolute_embeddings else None
+ )
+
+ self.norm = nn.LayerNorm(config.embed_dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ Interpolate pre-trained position encodings to support higher-resolution images at inference.
+ Unlike ViT, Swin has no CLS token, so position embeddings cover patch positions only.
+ """
+ num_patches = embeddings.shape[1]
+ num_positions = self.position_embeddings.shape[1]
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ 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 = self.position_embeddings.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,
+ )
+
+ return patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> tuple[torch.Tensor]:
+ _, num_channels, height, width = pixel_values.shape
+ embeddings, output_dimensions = self.patch_embeddings(pixel_values)
+ embeddings = self.norm(embeddings)
+ batch_size, seq_len, _ = embeddings.size()
+
+ if bool_masked_pos is not None:
+ mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
+ # replace the masked visual tokens by mask_tokens
+ mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
+ embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
+
+ if self.position_embeddings is not None:
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embeddings
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings, output_dimensions
+
+
+class SwinPatchEmbeddings(nn.Module):
+ """
+ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+ `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+ Transformer.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ image_size, patch_size = config.image_size, config.patch_size
+ num_channels, hidden_size = config.num_channels, config.embed_dim
+ image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+ self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
+
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+ def maybe_pad(self, pixel_values, height, width):
+ """Pad pixel_values to be divisible by patch_size if needed."""
+ if width % self.patch_size[1] != 0:
+ pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
+ pixel_values = nn.functional.pad(pixel_values, pad_values)
+ if height % self.patch_size[0] != 0:
+ pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
+ pixel_values = nn.functional.pad(pixel_values, pad_values)
+ return pixel_values
+
+ def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]:
+ _, num_channels, height, width = pixel_values.shape
+ # pad the input to be divisible by self.patch_size, if needed
+ pixel_values = self.maybe_pad(pixel_values, height, width)
+ embeddings = self.projection(pixel_values)
+ _, _, height, width = embeddings.shape
+ output_dimensions = (height, width)
+ embeddings = embeddings.flatten(2).transpose(1, 2)
+
+ return embeddings, output_dimensions
+
+
+class SwinPatchMerging(nn.Module):
+ """
+ Patch Merging Layer.
+
+ Args:
+ dim (`int`):
+ Number of input channels.
+ """
+
+ def __init__(self, dim: int) -> None:
+ super().__init__()
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+ self.norm = nn.LayerNorm(4 * dim)
+
+ def maybe_pad(self, input_feature: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """Pad input feature map to be divisible by 2 in both spatial dimensions if needed."""
+ if (height % 2 == 1) or (width % 2 == 1):
+ input_feature = nn.functional.pad(input_feature, (0, 0, 0, width % 2, 0, height % 2))
+ return input_feature
+
+ def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor:
+ height, width = input_dimensions
+ # `dim` is height * width
+ batch_size, dim, num_channels = input_feature.shape
+
+ input_feature = input_feature.view(batch_size, height, width, num_channels)
+ # pad input to be divisible by width and height, if needed
+ input_feature = self.maybe_pad(input_feature, height, width)
+ # Interleave rows and columns to produce [batch_size, height/2*width/2, 4*num_channels]
+ input_feature = torch.cat(
+ [input_feature[:, row::2, col::2, :] for col in range(2) for row in range(2)], dim=-1
+ )
+ input_feature = input_feature.view(batch_size, -1, 4 * num_channels)
+
+ input_feature = self.norm(input_feature)
+ input_feature = self.reduction(input_feature)
+
+ return input_feature
+
+
+class SwinRelativePositionBias(nn.Module):
+ """
+ Relative position bias for Swin's window-based attention, following the style of BeitRelativePositionBias.
+
+ Unlike BeiT, Swin has no CLS token, so the table covers exactly (2*ws_h-1)*(2*ws_w-1) unique
+ relative positions. The lookup index is purely determined by window_size (static), so it is stored
+ as a non-persistent buffer (recomputed from config on load, never serialised). The table values
+ are learned parameters and must be re-read on every forward call.
+ """
+
+ def __init__(self, num_heads: int, window_size: tuple[int, int]):
+ super().__init__()
+ self.window_size = window_size
+ self.window_area = window_size[0] * window_size[1]
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+ )
+ # Non-persistent: fully determined by window_size, no need to serialise.
+ # Stored flat so forward avoids an extra .view() call.
+ self.register_buffer(
+ "relative_position_index",
+ self._create_relative_position_index().view(-1),
+ persistent=False,
+ )
+
+ def _create_relative_position_index(self) -> torch.Tensor:
+ coords_h = torch.arange(self.window_size[0])
+ coords_w = torch.arange(self.window_size[1])
+
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) # 2, Wh, Ww
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
+
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+
+ # shift to start from 0 and compute a unique flat index for each (dh, dw) pair
+ relative_coords[:, :, 0] += self.window_size[0] - 1
+ relative_coords[:, :, 1] += self.window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+
+ return relative_coords.sum(-1) # Wh*Ww, Wh*Ww
+
+ def forward(self) -> torch.Tensor:
+ relative_position_bias = self.relative_position_bias_table[self.relative_position_index]
+ relative_position_bias = relative_position_bias.view(self.window_area, self.window_area, -1)
+ return relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0) # 1, num_heads, Wh*Ww, Wh*Ww
+
+
+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, 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 SwinAttention(nn.Module):
+ def __init__(self, config: SwinConfig, hidden_size: int, num_attention_heads: int, window_size: int):
+ super().__init__()
+ self.config = config
+ self.num_attention_heads = num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.attention_dropout = config.attention_probs_dropout_prob
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.k_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.v_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.o_proj = nn.Linear(hidden_size, hidden_size)
+
+ self.relative_position_bias = SwinRelativePositionBias(num_attention_heads, (window_size, window_size))
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # hidden_states: (batch_size * num_windows, window_size * window_size, channels)
+ 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)
+
+ # Combine relative position bias with the cyclic-shift attention mask for SW-MSA
+ relative_position_bias = self.relative_position_bias() # 1, num_heads, ws*ws, ws*ws
+ if attention_mask is not None:
+ # attention_mask: (num_windows, ws*ws, ws*ws)
+ num_windows = attention_mask.shape[0]
+ batch_size = input_shape[0] // num_windows
+ seq_len = input_shape[1]
+ # Expand to (batch * num_windows, 1, ws*ws, ws*ws) for broadcasting
+ attention_mask = (
+ attention_mask.unsqueeze(1) # (num_windows, 1, ws*ws, ws*ws)
+ .unsqueeze(0) # (1, num_windows, 1, ws*ws, ws*ws)
+ .expand(batch_size, -1, -1, -1, -1) # (batch, num_windows, 1, ws*ws, ws*ws)
+ .reshape(-1, 1, seq_len, seq_len) # (batch * num_windows, 1, ws*ws, ws*ws)
+ )
+ combined_mask = relative_position_bias + attention_mask
+ else:
+ combined_mask = relative_position_bias
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ combined_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class SwinMLP(nn.Module):
+ def __init__(self, config: SwinConfig, dim: int):
+ super().__init__()
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(dim, int(config.mlp_ratio * dim))
+ self.fc2 = nn.Linear(int(config.mlp_ratio * dim), dim)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+
+ return hidden_states
+
+
+def window_partition(input_feature, window_size):
+ """
+ Partitions the given input into windows.
+ """
+ batch_size, height, width, num_channels = input_feature.shape
+ input_feature = input_feature.view(
+ batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
+ )
+ windows = input_feature.transpose(2, 3).contiguous().view(-1, window_size, window_size, num_channels)
+ return windows
+
+
+def window_reverse(windows, window_size, height, width):
+ """
+ Merges windows to produce higher resolution features.
+ """
+ num_channels = windows.shape[-1]
+ windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
+ windows = windows.transpose(2, 3).contiguous().view(-1, height, width, num_channels)
+ return windows
+
+
+class SwinLayer(GradientCheckpointingLayer):
+ def __init__(
+ self,
+ config: SwinConfig,
+ dim: int,
+ input_resolution: tuple[int, int],
+ num_heads: int,
+ drop_path_rate: float = 0.0,
+ shift_size: int = 0,
+ ):
+ super().__init__()
+ self.attention = SwinAttention(config, dim, num_heads, window_size=config.window_size)
+ self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.mlp = SwinMLP(config, dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.window_size = config.window_size
+ self.shift_size = shift_size
+ self.input_resolution = input_resolution
+ self.drop_path = SwinDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ if not always_partition:
+ self.set_shift_and_window_size(input_dimensions)
+ height, width = input_dimensions
+ batch_size, _, channels = hidden_states.size()
+ shortcut = hidden_states
+
+ hidden_states = self.layernorm_before(hidden_states)
+ hidden_states = hidden_states.view(batch_size, height, width, channels)
+
+ # pad hidden_states to multiples of window size
+ hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
+ _, height_pad, width_pad, _ = hidden_states.shape
+
+ hidden_states_windows = window_partition(self.cyclic_shift(hidden_states), self.window_size)
+ hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
+ attn_mask = self.get_attn_mask(
+ height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device
+ )
+
+ attention_output, attn_weights = self.attention(hidden_states_windows, attn_mask, **kwargs)
+ attention_output = self.dropout(attention_output)
+
+ attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
+ attention_windows = self.cyclic_shift(
+ window_reverse(attention_windows, self.window_size, height_pad, width_pad), reverse=True
+ )
+
+ if pad_values[3] > 0 or pad_values[5] > 0:
+ attention_windows = attention_windows[:, :height, :width, :].contiguous()
+
+ attention_windows = attention_windows.view(batch_size, height * width, channels)
+ hidden_states = shortcut + self.drop_path(attention_windows)
+
+ residual = hidden_states
+ hidden_states = self.layernorm_after(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.dropout(hidden_states) + residual
+
+ return hidden_states, attn_weights
+
+ def set_shift_and_window_size(self, input_resolution: tuple[int, int]) -> None:
+ """Clamp window and shift sizes when the window is larger than the input resolution."""
+ if min(input_resolution) <= self.window_size:
+ self.shift_size = torch_int(0)
+ self.window_size = (
+ torch.min(torch.tensor(input_resolution)) if torch.jit.is_tracing() else min(input_resolution)
+ )
+
+ def get_attn_mask(self, height: int, width: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor | None:
+ """Build the cyclic-shift attention mask for shifted-window MSA; returns None when shift_size is 0."""
+ if self.shift_size > 0:
+ img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)
+ height_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ width_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ count = 0
+ for height_slice in height_slices:
+ for width_slice in width_slices:
+ img_mask[:, height_slice, width_slice, :] = count
+ count += 1
+
+ mask_windows = window_partition(img_mask, self.window_size)
+ mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
+ else:
+ attn_mask = None
+ return attn_mask
+
+ def maybe_pad(self, hidden_states: torch.Tensor, height: int, width: int) -> tuple[torch.Tensor, tuple[int, ...]]:
+ """Pad feature map so both spatial dimensions are divisible by window_size."""
+ pad_right = (self.window_size - width % self.window_size) % self.window_size
+ pad_bottom = (self.window_size - height % self.window_size) % self.window_size
+ pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
+ hidden_states = nn.functional.pad(hidden_states, pad_values)
+ return hidden_states, pad_values
+
+ def cyclic_shift(self, hidden_states: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+ """Apply a cyclic shift along the spatial dimensions for shifted-window attention."""
+ if self.shift_size > 0:
+ direction = 1 if reverse else -1
+ hidden_states = torch.roll(
+ hidden_states,
+ shifts=(direction * self.shift_size, direction * self.shift_size),
+ dims=(1, 2),
+ )
+ return hidden_states
+
+
+class SwinStage(GradientCheckpointingLayer):
+ def __init__(
+ self,
+ config: SwinConfig,
+ dim: int,
+ input_resolution: tuple[int, int],
+ depth: int,
+ num_heads: int,
+ drop_path: list[float],
+ downsample,
+ ):
+ super().__init__()
+ self.config = config
+ self.blocks = nn.ModuleList(
+ [
+ SwinLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+ )
+ for i in range(depth)
+ ]
+ )
+
+ self.downsample = downsample(dim=dim) if downsample is not None else None
+
+ def get_reshaped_hidden_states(
+ self,
+ hidden_states: torch.Tensor,
+ hidden_states_before_downsampling: torch.Tensor,
+ height: int,
+ width: int,
+ output_hidden_states_before_downsampling: bool,
+ ) -> torch.Tensor:
+ """
+ Select the spatial hidden states for this stage and reshape from (B, L, C) to (B, C, H, W).
+
+ The chosen state and its resolution depend on output_hidden_states_before_downsampling:
+ - True → pre-downsampling states at (height, width) — used by the backbone.
+ - False → post-downsampling states at half the resolution (if a downsampler exists).
+ """
+ if output_hidden_states_before_downsampling:
+ spatial_state, h, w = hidden_states_before_downsampling, height, width
+ elif self.downsample is not None:
+ spatial_state, h, w = hidden_states, (height + 1) // 2, (width + 1) // 2
+ else:
+ spatial_state, h, w = hidden_states, height, width
+
+ batch_size, _, hidden_size = spatial_state.shape
+ return spatial_state.view(batch_size, h, w, hidden_size).permute(0, 3, 1, 2).contiguous()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ output_hidden_states_before_downsampling: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+ height, width = input_dimensions
+ last_attn_weights = None
+ for layer_module in self.blocks:
+ hidden_states, last_attn_weights = layer_module(
+ hidden_states, input_dimensions, always_partition=always_partition, **kwargs
+ )
+
+ hidden_states_before_downsampling = hidden_states
+ if self.downsample is not None:
+ hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)
+
+ reshaped_hidden_states = self.get_reshaped_hidden_states(
+ hidden_states, hidden_states_before_downsampling, height, width, output_hidden_states_before_downsampling
+ )
+
+ return hidden_states, reshaped_hidden_states, last_attn_weights
+
+
+@auto_docstring
+class SwinPreTrainedModel(PreTrainedModel):
+ config: SwinConfig
+ base_model_prefix = "swin"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["SwinStage"]
+ _supports_sdpa = True
+ _supports_flash_attn = False
+ _supports_flex_attn = False
+ _supports_attention_backend = True
+ _can_compile_fullgraph = True
+ _can_record_outputs = {
+ # capture_initial_hidden_state=True: prepend the embedding input (args[0] of SwinStage 0) so that
+ # hidden_states[0] has the same shape as the patch embeddings (num_patches, embed_dim).
+ "hidden_states": OutputRecorder(SwinStage, index=0, capture_initial_hidden_state=True),
+ # reshaped_hidden_states are collected explicitly by SwinEncoder (per stage) and the stem
+ # is prepended in SwinModel.forward, so they are NOT captured via hooks here.
+ # index=2: SwinStage returns (hidden_states, reshaped_hidden_states, last_attn_weights);
+ # capture the last block's attention weights at index 2, giving one entry per stage.
+ "attentions": OutputRecorder(SwinStage, index=2, capture_initial_hidden_state=False),
+ }
+ _input_embed_layer = "patch_embeddings"
+ # relative_position_index is now a non-persistent buffer (recomputed from window_size in __init__).
+ _keys_to_ignore_on_load_unexpected = [
+ r"attention\.self\.relative_position_index",
+ r"attention\.relative_position_bias\.relative_position_index",
+ ]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, SwinEmbeddings):
+ if module.mask_token is not None:
+ init.zeros_(module.mask_token)
+ if module.position_embeddings is not None:
+ init.zeros_(module.position_embeddings)
+ elif isinstance(module, SwinRelativePositionBias):
+ init.zeros_(module.relative_position_bias_table)
+ init.copy_(module.relative_position_index, module._create_relative_position_index().view(-1))
+
+
+class SwinEncoder(SwinPreTrainedModel):
+ def __init__(self, config: SwinConfig, grid_size: tuple[int, int]):
+ super().__init__(config)
+ self.num_layers = len(config.depths)
+ self.config = config
+ dpr = [config.drop_path_rate * i / max(sum(config.depths) - 1, 1) for i in range(sum(config.depths))]
+ self.layers = nn.ModuleList(
+ [
+ SwinStage(
+ config=config,
+ dim=int(config.embed_dim * 2**layer_idx),
+ input_resolution=(grid_size[0] // (2**layer_idx), grid_size[1] // (2**layer_idx)),
+ depth=config.depths[layer_idx],
+ num_heads=config.num_heads[layer_idx],
+ drop_path=dpr[sum(config.depths[:layer_idx]) : sum(config.depths[: layer_idx + 1])],
+ downsample=SwinPatchMerging if (layer_idx < self.num_layers - 1) else None,
+ )
+ for layer_idx in range(self.num_layers)
+ ]
+ )
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ output_hidden_states: bool = False,
+ output_hidden_states_before_downsampling: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinEncoderOutput:
+ r"""
+ input_dimensions (`tuple[int, int]`):
+ Spatial `(height, width)` of the patch grid entering the encoder.
+ always_partition (`bool`, *optional*, defaults to `False`):
+ If `True`, always apply window partitioning regardless of input resolution.
+ output_hidden_states_before_downsampling (`bool`, *optional*, defaults to `False`):
+ If `True`, `reshaped_hidden_states` contains pre-downsampling feature maps.
+ """
+ all_reshaped_hidden_states = None
+ if output_hidden_states:
+ # Prepend the stem: hidden_states is the patch embedding output (B, N, C),
+ # reshape it to spatial (B, C, H, W) as the first reshaped hidden state.
+ batch_size, _, hidden_size = hidden_states.shape
+ stem_spatial = (
+ hidden_states.view(batch_size, *input_dimensions, hidden_size).permute(0, 3, 1, 2).contiguous()
+ )
+ all_reshaped_hidden_states = (stem_spatial,)
+
+ for layer_module in self.layers:
+ hidden_states, reshaped_hidden_state, _ = layer_module(
+ hidden_states,
+ input_dimensions,
+ always_partition=always_partition,
+ output_hidden_states_before_downsampling=output_hidden_states_before_downsampling,
+ **kwargs,
+ )
+ if output_hidden_states:
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+ if layer_module.downsample is not None:
+ input_dimensions = ((input_dimensions[0] + 1) // 2, (input_dimensions[1] + 1) // 2)
+
+ return SwinEncoderOutput(
+ last_hidden_state=hidden_states,
+ reshaped_hidden_states=all_reshaped_hidden_states,
+ )
+
+
+@auto_docstring
+class SwinModel(SwinPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
+ r"""
+ add_pooling_layer (`bool`, *optional*, defaults to `True`):
+ Whether or not to apply pooling layer.
+ use_mask_token (`bool`, *optional*, defaults to `False`):
+ Whether or not to create and apply mask tokens in the embedding layer.
+ """
+ super().__init__(config)
+ self.config = config
+ self.num_layers = len(config.depths)
+ self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))
+
+ self.embeddings = SwinEmbeddings(config, use_mask_token=use_mask_token)
+ self.encoder = SwinEncoder(config, self.embeddings.patch_grid)
+
+ self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
+ self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinModelOutput:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+ """
+ # FIXME: output_hidden_states must be popped manually here because SwinEncoder takes it as an
+ # explicit argument (not via **kwargs), so it is not captured by the @capture_outputs decorator.
+ output_hidden_states = kwargs.pop("output_hidden_states", self.config.output_hidden_states)
+
+ embedding_output, input_dimensions = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ input_dimensions,
+ output_hidden_states=output_hidden_states,
+ **kwargs,
+ )
+
+ sequence_output = encoder_outputs.last_hidden_state
+ sequence_output = self.layernorm(sequence_output)
+
+ pooled_output = None
+ if self.pooler is not None:
+ pooled_output = self.pooler(sequence_output.transpose(1, 2))
+ pooled_output = torch.flatten(pooled_output, 1)
+
+ return SwinModelOutput(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin Model with a decoder on top for masked image modeling, as proposed in [SimMIM](https://huggingface.co/papers/2111.09886).
+
+
+
+ Note that we provide a script to pre-train this model on custom data in our [examples
+ directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining).
+
+
+ """
+)
+class SwinForMaskedImageModeling(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.swin = SwinModel(config, add_pooling_layer=False, use_mask_token=True)
+
+ num_features = int(config.embed_dim * 2 ** (config.num_layers - 1))
+ self.decoder = nn.Sequential(
+ nn.Conv2d(
+ in_channels=num_features, out_channels=config.encoder_stride**2 * config.num_channels, kernel_size=1
+ ),
+ nn.PixelShuffle(config.encoder_stride),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinMaskedImageModelingOutput:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+
+ Examples:
+ ```python
+ >>> from transformers import AutoImageProcessor, SwinForMaskedImageModeling
+ >>> 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()))
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/swin-base-simmim-window6-192")
+ >>> model = SwinForMaskedImageModeling.from_pretrained("microsoft/swin-base-simmim-window6-192")
+
+ >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
+ >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
+ >>> # create random boolean mask of shape (batch_size, num_patches)
+ >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
+
+ >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
+ >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction
+ >>> list(reconstructed_pixel_values.shape)
+ [1, 3, 192, 192]
+ ```"""
+ outputs = self.swin(
+ pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ sequence_output = outputs.last_hidden_state
+ # Reshape to (batch_size, num_channels, height, width)
+ sequence_output = sequence_output.transpose(1, 2)
+ batch_size, num_channels, sequence_length = sequence_output.shape
+ height = width = math.floor(sequence_length**0.5)
+ sequence_output = sequence_output.reshape(batch_size, num_channels, height, width)
+
+ # Reconstruct pixel values
+ reconstructed_pixel_values = self.decoder(sequence_output)
+
+ masked_im_loss = None
+ if bool_masked_pos is not None:
+ size = self.config.image_size // self.config.patch_size
+ bool_masked_pos = bool_masked_pos.reshape(-1, size, size)
+ mask = (
+ bool_masked_pos.repeat_interleave(self.config.patch_size, 1)
+ .repeat_interleave(self.config.patch_size, 2)
+ .unsqueeze(1)
+ .contiguous()
+ )
+ reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none")
+ masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels
+
+ return SwinMaskedImageModelingOutput(
+ loss=masked_im_loss,
+ reconstruction=reconstructed_pixel_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
+ the [CLS] token) e.g. for ImageNet.
+
+
+
+ Note that it's possible to fine-tune Swin on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """
+)
+class SwinForImageClassification(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.swin = SwinModel(config)
+
+ # Classifier head
+ self.classifier = (
+ nn.Linear(self.swin.num_features, 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.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinImageClassifierOutput:
+ 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 = self.swin(
+ pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ pooled_output = outputs.pooler_output
+
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config, **kwargs)
+
+ return SwinImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+class SwinBackbone(BackboneMixin, SwinPreTrainedModel):
+ _keys_to_ignore_on_load_missing = [r"swin.layernorm.*"]
+
+ def __init__(self, config: SwinConfig):
+ super().__init__(config)
+
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+ self.swin = SwinModel(config, add_pooling_layer=False)
+
+ # Add layer norms to hidden states of out_features
+ hidden_states_norms = {}
+ for stage, num_channels in zip(self.out_features, self.channels):
+ hidden_states_norms[stage] = nn.LayerNorm(num_channels)
+ self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224")
+ >>> model = AutoBackbone.from_pretrained(
+ ... "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"]
+ ... )
+
+ >>> inputs = processor(image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> feature_maps = outputs.feature_maps
+ >>> list(feature_maps[-1].shape)
+ [1, 768, 7, 7]
+ ```
+ """
+ kwargs["output_hidden_states"] = True # required to extract layers for the stages
+ # always_partition=True preserves shifted-window attention at all resolutions.
+ # output_hidden_states_before_downsampling=True captures pre-downsampling feature maps per stage.
+ outputs = self.swin(
+ pixel_values,
+ always_partition=True,
+ output_hidden_states_before_downsampling=True,
+ **kwargs,
+ )
+
+ feature_maps = ()
+ for stage, hidden_state in zip(self.stage_names, outputs.reshaped_hidden_states):
+ if stage in self.out_features:
+ batch_size, num_channels, height, width = hidden_state.shape
+ hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
+ hidden_state = hidden_state.view(batch_size, height * width, num_channels)
+ hidden_state = self.hidden_states_norms[stage](hidden_state)
+ hidden_state = hidden_state.view(batch_size, height, width, num_channels)
+ hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
+ feature_maps += (hidden_state,)
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.reshaped_hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "SwinForImageClassification",
+ "SwinForMaskedImageModeling",
+ "SwinModel",
+ "SwinPreTrainedModel",
+ "SwinBackbone",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modular_swin.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modular_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..94325cd8ae06b15e250c1f9361e8c97efc250dd5
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin/modular_swin.py
@@ -0,0 +1,1122 @@
+# 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.
+"""PyTorch Swin Transformer model."""
+
+import collections.abc
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BackboneOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+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 OutputRecorder, capture_outputs
+from ..vit.modeling_vit import (
+ PreTrainedModel,
+ ViTAttention,
+ ViTLayer,
+ ViTMLP,
+ ViTPreTrainedModel,
+ eager_attention_forward,
+)
+from .configuration_swin import SwinConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class SwinDropPath(nn.Module):
+ """Stochastic depth (DropPath) per sample, for residual blocks.
+
+ Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
+ `_.
+ """
+
+ def __init__(self, drop_prob: float = 0.0) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ if self.drop_prob == 0.0 or not self.training:
+ return hidden_states
+ keep_prob = 1 - self.drop_prob
+ shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
+ random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
+ random_tensor = torch.floor(random_tensor + keep_prob)
+ return hidden_states.div(keep_prob) * random_tensor
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin encoder's outputs, with potential hidden states and attentions.
+ """
+)
+@dataclass
+class SwinEncoderOutput(ModelOutput):
+ r"""
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin model's outputs that also contains a pooling of the last hidden states.
+ """
+)
+@dataclass
+class SwinModelOutput(ModelOutput):
+ r"""
+ pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
+ Average pooling of the last layer hidden-state.
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ pooler_output: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin masked image model outputs.
+ """
+)
+@dataclass
+class SwinMaskedImageModelingOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
+ Masked image modeling (MLM) loss.
+ reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Reconstructed pixel values.
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ loss: torch.FloatTensor | None = None
+ reconstruction: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin outputs for image classification.
+ """
+)
+@dataclass
+class SwinImageClassifierOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Classification (or regression if config.num_labels==1) loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+ Classification (or regression if config.num_labels==1) scores (before SoftMax).
+ reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+def window_partition(input_feature, window_size):
+ """
+ Partitions the given input into windows.
+ """
+ batch_size, height, width, num_channels = input_feature.shape
+ input_feature = input_feature.view(
+ batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
+ )
+ windows = input_feature.transpose(2, 3).contiguous().view(-1, window_size, window_size, num_channels)
+ return windows
+
+
+def window_reverse(windows, window_size, height, width):
+ """
+ Merges windows to produce higher resolution features.
+ """
+ num_channels = windows.shape[-1]
+ windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
+ windows = windows.transpose(2, 3).contiguous().view(-1, height, width, num_channels)
+ return windows
+
+
+class SwinEmbeddings(nn.Module):
+ """
+ Construct the patch and position embeddings. Optionally, also the mask token.
+ """
+
+ def __init__(self, config, use_mask_token=False):
+ super().__init__()
+
+ self.patch_embeddings = SwinPatchEmbeddings(config)
+ num_patches = self.patch_embeddings.num_patches
+ self.patch_grid = self.patch_embeddings.grid_size
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None
+
+ self.position_embeddings = (
+ nn.Parameter(torch.zeros(1, num_patches, config.embed_dim)) if config.use_absolute_embeddings else None
+ )
+
+ self.norm = nn.LayerNorm(config.embed_dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ Interpolate pre-trained position encodings to support higher-resolution images at inference.
+ Unlike ViT, Swin has no CLS token, so position embeddings cover patch positions only.
+ """
+ num_patches = embeddings.shape[1]
+ num_positions = self.position_embeddings.shape[1]
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ 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 = self.position_embeddings.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,
+ )
+
+ return patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> tuple[torch.Tensor]:
+ _, num_channels, height, width = pixel_values.shape
+ embeddings, output_dimensions = self.patch_embeddings(pixel_values)
+ embeddings = self.norm(embeddings)
+ batch_size, seq_len, _ = embeddings.size()
+
+ if bool_masked_pos is not None:
+ mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
+ # replace the masked visual tokens by mask_tokens
+ mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
+ embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
+
+ if self.position_embeddings is not None:
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embeddings
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings, output_dimensions
+
+
+class SwinPatchEmbeddings(nn.Module):
+ """
+ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+ `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+ Transformer.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ image_size, patch_size = config.image_size, config.patch_size
+ num_channels, hidden_size = config.num_channels, config.embed_dim
+ image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+ self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
+
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+ def maybe_pad(self, pixel_values, height, width):
+ """Pad pixel_values to be divisible by patch_size if needed."""
+ if width % self.patch_size[1] != 0:
+ pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
+ pixel_values = nn.functional.pad(pixel_values, pad_values)
+ if height % self.patch_size[0] != 0:
+ pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
+ pixel_values = nn.functional.pad(pixel_values, pad_values)
+ return pixel_values
+
+ def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]:
+ _, num_channels, height, width = pixel_values.shape
+ # pad the input to be divisible by self.patch_size, if needed
+ pixel_values = self.maybe_pad(pixel_values, height, width)
+ embeddings = self.projection(pixel_values)
+ _, _, height, width = embeddings.shape
+ output_dimensions = (height, width)
+ embeddings = embeddings.flatten(2).transpose(1, 2)
+
+ return embeddings, output_dimensions
+
+
+class SwinPatchMerging(nn.Module):
+ """
+ Patch Merging Layer.
+
+ Args:
+ dim (`int`):
+ Number of input channels.
+ """
+
+ def __init__(self, dim: int) -> None:
+ super().__init__()
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+ self.norm = nn.LayerNorm(4 * dim)
+
+ def maybe_pad(self, input_feature: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """Pad input feature map to be divisible by 2 in both spatial dimensions if needed."""
+ if (height % 2 == 1) or (width % 2 == 1):
+ input_feature = nn.functional.pad(input_feature, (0, 0, 0, width % 2, 0, height % 2))
+ return input_feature
+
+ def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor:
+ height, width = input_dimensions
+ # `dim` is height * width
+ batch_size, dim, num_channels = input_feature.shape
+
+ input_feature = input_feature.view(batch_size, height, width, num_channels)
+ # pad input to be divisible by width and height, if needed
+ input_feature = self.maybe_pad(input_feature, height, width)
+ # Interleave rows and columns to produce [batch_size, height/2*width/2, 4*num_channels]
+ input_feature = torch.cat(
+ [input_feature[:, row::2, col::2, :] for col in range(2) for row in range(2)], dim=-1
+ )
+ input_feature = input_feature.view(batch_size, -1, 4 * num_channels)
+
+ input_feature = self.norm(input_feature)
+ input_feature = self.reduction(input_feature)
+
+ return input_feature
+
+
+class SwinRelativePositionBias(nn.Module):
+ """
+ Relative position bias for Swin's window-based attention, following the style of BeitRelativePositionBias.
+
+ Unlike BeiT, Swin has no CLS token, so the table covers exactly (2*ws_h-1)*(2*ws_w-1) unique
+ relative positions. The lookup index is purely determined by window_size (static), so it is stored
+ as a non-persistent buffer (recomputed from config on load, never serialised). The table values
+ are learned parameters and must be re-read on every forward call.
+ """
+
+ def __init__(self, num_heads: int, window_size: tuple[int, int]):
+ super().__init__()
+ self.window_size = window_size
+ self.window_area = window_size[0] * window_size[1]
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+ )
+ # Non-persistent: fully determined by window_size, no need to serialise.
+ # Stored flat so forward avoids an extra .view() call.
+ self.register_buffer(
+ "relative_position_index",
+ self._create_relative_position_index().view(-1),
+ persistent=False,
+ )
+
+ def _create_relative_position_index(self) -> torch.Tensor:
+ coords_h = torch.arange(self.window_size[0])
+ coords_w = torch.arange(self.window_size[1])
+
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) # 2, Wh, Ww
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
+
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+
+ # shift to start from 0 and compute a unique flat index for each (dh, dw) pair
+ relative_coords[:, :, 0] += self.window_size[0] - 1
+ relative_coords[:, :, 1] += self.window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+
+ return relative_coords.sum(-1) # Wh*Ww, Wh*Ww
+
+ def forward(self) -> torch.Tensor:
+ relative_position_bias = self.relative_position_bias_table[self.relative_position_index]
+ relative_position_bias = relative_position_bias.view(self.window_area, self.window_area, -1)
+ return relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0) # 1, num_heads, Wh*Ww, Wh*Ww
+
+
+class SwinAttention(ViTAttention):
+ def __init__(self, config: SwinConfig, hidden_size: int, num_attention_heads: int, window_size: int):
+ super().__init__(config)
+ self.num_attention_heads = num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.scaling = self.head_dim**-0.5
+
+ self.q_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.k_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.v_proj = nn.Linear(hidden_size, hidden_size, bias=config.qkv_bias)
+ self.o_proj = nn.Linear(hidden_size, hidden_size)
+
+ self.relative_position_bias = SwinRelativePositionBias(num_attention_heads, (window_size, window_size))
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # hidden_states: (batch_size * num_windows, window_size * window_size, channels)
+ 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)
+
+ # Combine relative position bias with the cyclic-shift attention mask for SW-MSA
+ relative_position_bias = self.relative_position_bias() # 1, num_heads, ws*ws, ws*ws
+ if attention_mask is not None:
+ # attention_mask: (num_windows, ws*ws, ws*ws)
+ num_windows = attention_mask.shape[0]
+ batch_size = input_shape[0] // num_windows
+ seq_len = input_shape[1]
+ # Expand to (batch * num_windows, 1, ws*ws, ws*ws) for broadcasting
+ attention_mask = (
+ attention_mask.unsqueeze(1) # (num_windows, 1, ws*ws, ws*ws)
+ .unsqueeze(0) # (1, num_windows, 1, ws*ws, ws*ws)
+ .expand(batch_size, -1, -1, -1, -1) # (batch, num_windows, 1, ws*ws, ws*ws)
+ .reshape(-1, 1, seq_len, seq_len) # (batch * num_windows, 1, ws*ws, ws*ws)
+ )
+ combined_mask = relative_position_bias + attention_mask
+ else:
+ combined_mask = relative_position_bias
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ combined_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class SwinMLP(ViTMLP):
+ def __init__(self, config: SwinConfig, dim: int):
+ nn.Module.__init__(self)
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(dim, int(config.mlp_ratio * dim))
+ self.fc2 = nn.Linear(int(config.mlp_ratio * dim), dim)
+
+
+class SwinLayer(ViTLayer):
+ def __init__(
+ self,
+ config: SwinConfig,
+ dim: int,
+ input_resolution: tuple[int, int],
+ num_heads: int,
+ drop_path_rate: float = 0.0,
+ shift_size: int = 0,
+ ):
+ super().__init__()
+ self.window_size = config.window_size
+ self.attention = SwinAttention(config, dim, num_heads, window_size=config.window_size)
+ self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.mlp = SwinMLP(config, dim)
+ self.shift_size = shift_size
+ self.input_resolution = input_resolution
+ self.drop_path = SwinDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+
+ def set_shift_and_window_size(self, input_resolution: tuple[int, int]) -> None:
+ """Clamp window and shift sizes when the window is larger than the input resolution."""
+ if min(input_resolution) <= self.window_size:
+ self.shift_size = torch_int(0)
+ self.window_size = (
+ torch.min(torch.tensor(input_resolution)) if torch.jit.is_tracing() else min(input_resolution)
+ )
+
+ def get_attn_mask(self, height: int, width: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor | None:
+ """Build the cyclic-shift attention mask for shifted-window MSA; returns None when shift_size is 0."""
+ if self.shift_size > 0:
+ img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)
+ height_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ width_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ count = 0
+ for height_slice in height_slices:
+ for width_slice in width_slices:
+ img_mask[:, height_slice, width_slice, :] = count
+ count += 1
+
+ mask_windows = window_partition(img_mask, self.window_size)
+ mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
+ else:
+ attn_mask = None
+ return attn_mask
+
+ def maybe_pad(self, hidden_states: torch.Tensor, height: int, width: int) -> tuple[torch.Tensor, tuple[int, ...]]:
+ """Pad feature map so both spatial dimensions are divisible by window_size."""
+ pad_right = (self.window_size - width % self.window_size) % self.window_size
+ pad_bottom = (self.window_size - height % self.window_size) % self.window_size
+ pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
+ hidden_states = nn.functional.pad(hidden_states, pad_values)
+ return hidden_states, pad_values
+
+ def cyclic_shift(self, hidden_states: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+ """Apply a cyclic shift along the spatial dimensions for shifted-window attention."""
+ if self.shift_size > 0:
+ direction = 1 if reverse else -1
+ hidden_states = torch.roll(
+ hidden_states,
+ shifts=(direction * self.shift_size, direction * self.shift_size),
+ dims=(1, 2),
+ )
+ return hidden_states
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ if not always_partition:
+ self.set_shift_and_window_size(input_dimensions)
+ height, width = input_dimensions
+ batch_size, _, channels = hidden_states.size()
+ shortcut = hidden_states
+
+ hidden_states = self.layernorm_before(hidden_states)
+ hidden_states = hidden_states.view(batch_size, height, width, channels)
+
+ # pad hidden_states to multiples of window size
+ hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
+ _, height_pad, width_pad, _ = hidden_states.shape
+
+ hidden_states_windows = window_partition(self.cyclic_shift(hidden_states), self.window_size)
+ hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
+ attn_mask = self.get_attn_mask(
+ height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device
+ )
+
+ attention_output, attn_weights = self.attention(hidden_states_windows, attn_mask, **kwargs)
+ attention_output = self.dropout(attention_output)
+
+ attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
+ attention_windows = self.cyclic_shift(
+ window_reverse(attention_windows, self.window_size, height_pad, width_pad), reverse=True
+ )
+
+ if pad_values[3] > 0 or pad_values[5] > 0:
+ attention_windows = attention_windows[:, :height, :width, :].contiguous()
+
+ attention_windows = attention_windows.view(batch_size, height * width, channels)
+ hidden_states = shortcut + self.drop_path(attention_windows)
+
+ residual = hidden_states
+ hidden_states = self.layernorm_after(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.dropout(hidden_states) + residual
+
+ return hidden_states, attn_weights
+
+
+class SwinStage(GradientCheckpointingLayer):
+ def __init__(
+ self,
+ config: SwinConfig,
+ dim: int,
+ input_resolution: tuple[int, int],
+ depth: int,
+ num_heads: int,
+ drop_path: list[float],
+ downsample,
+ ):
+ super().__init__()
+ self.config = config
+ self.blocks = nn.ModuleList(
+ [
+ SwinLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+ )
+ for i in range(depth)
+ ]
+ )
+
+ self.downsample = downsample(dim=dim) if downsample is not None else None
+
+ def get_reshaped_hidden_states(
+ self,
+ hidden_states: torch.Tensor,
+ hidden_states_before_downsampling: torch.Tensor,
+ height: int,
+ width: int,
+ output_hidden_states_before_downsampling: bool,
+ ) -> torch.Tensor:
+ """
+ Select the spatial hidden states for this stage and reshape from (B, L, C) to (B, C, H, W).
+
+ The chosen state and its resolution depend on output_hidden_states_before_downsampling:
+ - True → pre-downsampling states at (height, width) — used by the backbone.
+ - False → post-downsampling states at half the resolution (if a downsampler exists).
+ """
+ if output_hidden_states_before_downsampling:
+ spatial_state, h, w = hidden_states_before_downsampling, height, width
+ elif self.downsample is not None:
+ spatial_state, h, w = hidden_states, (height + 1) // 2, (width + 1) // 2
+ else:
+ spatial_state, h, w = hidden_states, height, width
+
+ batch_size, _, hidden_size = spatial_state.shape
+ return spatial_state.view(batch_size, h, w, hidden_size).permute(0, 3, 1, 2).contiguous()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ output_hidden_states_before_downsampling: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+ height, width = input_dimensions
+ last_attn_weights = None
+ for layer_module in self.blocks:
+ hidden_states, last_attn_weights = layer_module(
+ hidden_states, input_dimensions, always_partition=always_partition, **kwargs
+ )
+
+ hidden_states_before_downsampling = hidden_states
+ if self.downsample is not None:
+ hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)
+
+ reshaped_hidden_states = self.get_reshaped_hidden_states(
+ hidden_states, hidden_states_before_downsampling, height, width, output_hidden_states_before_downsampling
+ )
+
+ return hidden_states, reshaped_hidden_states, last_attn_weights
+
+
+@auto_docstring
+class SwinPreTrainedModel(ViTPreTrainedModel):
+ config: SwinConfig
+ _no_split_modules = ["SwinStage"]
+ _supports_flash_attn = False
+ _supports_flex_attn = False
+ # relative_position_index is now a non-persistent buffer (recomputed from window_size in __init__).
+ _keys_to_ignore_on_load_unexpected = [
+ r"attention\.self\.relative_position_index",
+ r"attention\.relative_position_bias\.relative_position_index",
+ ]
+ _can_record_outputs = {
+ # capture_initial_hidden_state=True: prepend the embedding input (args[0] of SwinStage 0) so that
+ # hidden_states[0] has the same shape as the patch embeddings (num_patches, embed_dim).
+ "hidden_states": OutputRecorder(SwinStage, index=0, capture_initial_hidden_state=True),
+ # reshaped_hidden_states are collected explicitly by SwinEncoder (per stage) and the stem
+ # is prepended in SwinModel.forward, so they are NOT captured via hooks here.
+ # index=2: SwinStage returns (hidden_states, reshaped_hidden_states, last_attn_weights);
+ # capture the last block's attention weights at index 2, giving one entry per stage.
+ "attentions": OutputRecorder(SwinStage, index=2, capture_initial_hidden_state=False),
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ PreTrainedModel._init_weights(self, module)
+ if isinstance(module, SwinEmbeddings):
+ if module.mask_token is not None:
+ init.zeros_(module.mask_token)
+ if module.position_embeddings is not None:
+ init.zeros_(module.position_embeddings)
+ elif isinstance(module, SwinRelativePositionBias):
+ init.zeros_(module.relative_position_bias_table)
+ init.copy_(module.relative_position_index, module._create_relative_position_index().view(-1))
+
+
+class SwinEncoder(SwinPreTrainedModel):
+ def __init__(self, config: SwinConfig, grid_size: tuple[int, int]):
+ super().__init__(config)
+ self.num_layers = len(config.depths)
+ self.config = config
+ dpr = [config.drop_path_rate * i / max(sum(config.depths) - 1, 1) for i in range(sum(config.depths))]
+ self.layers = nn.ModuleList(
+ [
+ SwinStage(
+ config=config,
+ dim=int(config.embed_dim * 2**layer_idx),
+ input_resolution=(grid_size[0] // (2**layer_idx), grid_size[1] // (2**layer_idx)),
+ depth=config.depths[layer_idx],
+ num_heads=config.num_heads[layer_idx],
+ drop_path=dpr[sum(config.depths[:layer_idx]) : sum(config.depths[: layer_idx + 1])],
+ downsample=SwinPatchMerging if (layer_idx < self.num_layers - 1) else None,
+ )
+ for layer_idx in range(self.num_layers)
+ ]
+ )
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ always_partition: bool = False,
+ output_hidden_states: bool = False,
+ output_hidden_states_before_downsampling: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinEncoderOutput:
+ r"""
+ input_dimensions (`tuple[int, int]`):
+ Spatial `(height, width)` of the patch grid entering the encoder.
+ always_partition (`bool`, *optional*, defaults to `False`):
+ If `True`, always apply window partitioning regardless of input resolution.
+ output_hidden_states_before_downsampling (`bool`, *optional*, defaults to `False`):
+ If `True`, `reshaped_hidden_states` contains pre-downsampling feature maps.
+ """
+ all_reshaped_hidden_states = None
+ if output_hidden_states:
+ # Prepend the stem: hidden_states is the patch embedding output (B, N, C),
+ # reshape it to spatial (B, C, H, W) as the first reshaped hidden state.
+ batch_size, _, hidden_size = hidden_states.shape
+ stem_spatial = (
+ hidden_states.view(batch_size, *input_dimensions, hidden_size).permute(0, 3, 1, 2).contiguous()
+ )
+ all_reshaped_hidden_states = (stem_spatial,)
+
+ for layer_module in self.layers:
+ hidden_states, reshaped_hidden_state, _ = layer_module(
+ hidden_states,
+ input_dimensions,
+ always_partition=always_partition,
+ output_hidden_states_before_downsampling=output_hidden_states_before_downsampling,
+ **kwargs,
+ )
+ if output_hidden_states:
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+ if layer_module.downsample is not None:
+ input_dimensions = ((input_dimensions[0] + 1) // 2, (input_dimensions[1] + 1) // 2)
+
+ return SwinEncoderOutput(
+ last_hidden_state=hidden_states,
+ reshaped_hidden_states=all_reshaped_hidden_states,
+ )
+
+
+@auto_docstring
+class SwinModel(SwinPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
+ r"""
+ add_pooling_layer (`bool`, *optional*, defaults to `True`):
+ Whether or not to apply pooling layer.
+ use_mask_token (`bool`, *optional*, defaults to `False`):
+ Whether or not to create and apply mask tokens in the embedding layer.
+ """
+ super().__init__(config)
+ self.config = config
+ self.num_layers = len(config.depths)
+ self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))
+
+ self.embeddings = SwinEmbeddings(config, use_mask_token=use_mask_token)
+ self.encoder = SwinEncoder(config, self.embeddings.patch_grid)
+
+ self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
+ self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinModelOutput:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+ """
+ # FIXME: output_hidden_states must be popped manually here because SwinEncoder takes it as an
+ # explicit argument (not via **kwargs), so it is not captured by the @capture_outputs decorator.
+ output_hidden_states = kwargs.pop("output_hidden_states", self.config.output_hidden_states)
+
+ embedding_output, input_dimensions = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ input_dimensions,
+ output_hidden_states=output_hidden_states,
+ **kwargs,
+ )
+
+ sequence_output = encoder_outputs.last_hidden_state
+ sequence_output = self.layernorm(sequence_output)
+
+ pooled_output = None
+ if self.pooler is not None:
+ pooled_output = self.pooler(sequence_output.transpose(1, 2))
+ pooled_output = torch.flatten(pooled_output, 1)
+
+ return SwinModelOutput(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin Model with a decoder on top for masked image modeling, as proposed in [SimMIM](https://huggingface.co/papers/2111.09886).
+
+
+
+ Note that we provide a script to pre-train this model on custom data in our [examples
+ directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining).
+
+
+ """
+)
+class SwinForMaskedImageModeling(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.swin = SwinModel(config, add_pooling_layer=False, use_mask_token=True)
+
+ num_features = int(config.embed_dim * 2 ** (config.num_layers - 1))
+ self.decoder = nn.Sequential(
+ nn.Conv2d(
+ in_channels=num_features, out_channels=config.encoder_stride**2 * config.num_channels, kernel_size=1
+ ),
+ nn.PixelShuffle(config.encoder_stride),
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinMaskedImageModelingOutput:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+
+ Examples:
+ ```python
+ >>> from transformers import AutoImageProcessor, SwinForMaskedImageModeling
+ >>> 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()))
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/swin-base-simmim-window6-192")
+ >>> model = SwinForMaskedImageModeling.from_pretrained("microsoft/swin-base-simmim-window6-192")
+
+ >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
+ >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
+ >>> # create random boolean mask of shape (batch_size, num_patches)
+ >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
+
+ >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
+ >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction
+ >>> list(reconstructed_pixel_values.shape)
+ [1, 3, 192, 192]
+ ```"""
+ outputs = self.swin(
+ pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ sequence_output = outputs.last_hidden_state
+ # Reshape to (batch_size, num_channels, height, width)
+ sequence_output = sequence_output.transpose(1, 2)
+ batch_size, num_channels, sequence_length = sequence_output.shape
+ height = width = math.floor(sequence_length**0.5)
+ sequence_output = sequence_output.reshape(batch_size, num_channels, height, width)
+
+ # Reconstruct pixel values
+ reconstructed_pixel_values = self.decoder(sequence_output)
+
+ masked_im_loss = None
+ if bool_masked_pos is not None:
+ size = self.config.image_size // self.config.patch_size
+ bool_masked_pos = bool_masked_pos.reshape(-1, size, size)
+ mask = (
+ bool_masked_pos.repeat_interleave(self.config.patch_size, 1)
+ .repeat_interleave(self.config.patch_size, 2)
+ .unsqueeze(1)
+ .contiguous()
+ )
+ reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none")
+ masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels
+
+ return SwinMaskedImageModelingOutput(
+ loss=masked_im_loss,
+ reconstruction=reconstructed_pixel_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
+ the [CLS] token) e.g. for ImageNet.
+
+
+
+ Note that it's possible to fine-tune Swin on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """
+)
+class SwinForImageClassification(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.swin = SwinModel(config)
+
+ # Classifier head
+ self.classifier = (
+ nn.Linear(self.swin.num_features, 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.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SwinImageClassifierOutput:
+ 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 = self.swin(
+ pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ pooled_output = outputs.pooler_output
+
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config, **kwargs)
+
+ return SwinImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Swin backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+class SwinBackbone(BackboneMixin, SwinPreTrainedModel):
+ _keys_to_ignore_on_load_missing = [r"swin.layernorm.*"]
+
+ def __init__(self, config: SwinConfig):
+ super().__init__(config)
+
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+ self.swin = SwinModel(config, add_pooling_layer=False)
+
+ # Add layer norms to hidden states of out_features
+ hidden_states_norms = {}
+ for stage, num_channels in zip(self.out_features, self.channels):
+ hidden_states_norms[stage] = nn.LayerNorm(num_channels)
+ self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224")
+ >>> model = AutoBackbone.from_pretrained(
+ ... "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"]
+ ... )
+
+ >>> inputs = processor(image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> feature_maps = outputs.feature_maps
+ >>> list(feature_maps[-1].shape)
+ [1, 768, 7, 7]
+ ```
+ """
+ kwargs["output_hidden_states"] = True # required to extract layers for the stages
+ # always_partition=True preserves shifted-window attention at all resolutions.
+ # output_hidden_states_before_downsampling=True captures pre-downsampling feature maps per stage.
+ outputs = self.swin(
+ pixel_values,
+ always_partition=True,
+ output_hidden_states_before_downsampling=True,
+ **kwargs,
+ )
+
+ feature_maps = ()
+ for stage, hidden_state in zip(self.stage_names, outputs.reshaped_hidden_states):
+ if stage in self.out_features:
+ batch_size, num_channels, height, width = hidden_state.shape
+ hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
+ hidden_state = hidden_state.view(batch_size, height * width, num_channels)
+ hidden_state = self.hidden_states_norms[stage](hidden_state)
+ hidden_state = hidden_state.view(batch_size, height, width, num_channels)
+ hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
+ feature_maps += (hidden_state,)
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.reshaped_hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "SwinForImageClassification",
+ "SwinForMaskedImageModeling",
+ "SwinModel",
+ "SwinPreTrainedModel",
+ "SwinBackbone",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa3a5c4c82f8cce587c5de36f8b79cd55e8af3bb
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/__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_wav2vec2 import *
+ from .feature_extraction_wav2vec2 import *
+ from .modeling_wav2vec2 import *
+ from .processing_wav2vec2 import *
+ from .tokenization_wav2vec2 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/configuration_wav2vec2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/configuration_wav2vec2.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a0fd126d362a2562e5ece8a63568d9676e23e49
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/configuration_wav2vec2.py
@@ -0,0 +1,245 @@
+# 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.
+"""Wav2Vec2 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="facebook/wav2vec2-base-960h")
+@strict
+class Wav2Vec2Config(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 [`Wav2Vec2ForCTC`].
+ 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, 3, 3)`):
+ 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_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_loss_reduction (`str`, *optional*, defaults to `"sum"`):
+ Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an
+ instance of [`Wav2Vec2ForCTC`].
+ 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 [`Wav2Vec2ForCTC`].
+ 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.
+ 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 network should be stacked on top of the Wav2Vec2 Encoder. Can be very useful for
+ warm-starting Wav2Vec2 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 3):
+ Number of convolutional layers that should be used in 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`.
+ adapter_attn_dim (`int`, *optional*):
+ Dimension of the attention adapter weights to be used in each attention block. An example of a model using
+ attention adapters is [facebook/mms-1b-all](https://huggingface.co/facebook/mms-1b-all).
+
+ Example:
+
+ ```python
+ >>> from transformers import Wav2Vec2Config, Wav2Vec2Model
+
+ >>> # Initializing a Wav2Vec2 facebook/wav2vec2-base-960h style configuration
+ >>> configuration = Wav2Vec2Config()
+
+ >>> # Initializing a model (with random weights) from the facebook/wav2vec2-base-960h style configuration
+ >>> model = Wav2Vec2Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "wav2vec2"
+
+ vocab_size: int | None = 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 = "sum"
+ 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 | 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 = 3
+ output_hidden_size: int | None = None
+ adapter_attn_dim: int | None = None
+
+ def __post_init__(self, **kwargs):
+ self.num_feat_extract_layers = len(self.conv_dim)
+ 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 (
+ (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__ = ["Wav2Vec2Config"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/feature_extraction_wav2vec2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/feature_extraction_wav2vec2.py
new file mode 100644
index 0000000000000000000000000000000000000000..dea2f3af5b486e245c7e26669e177352079499bf
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/feature_extraction_wav2vec2.py
@@ -0,0 +1,239 @@
+# 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 Wav2Vec2
+"""
+
+import numpy as np
+
+from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
+from ...feature_extraction_utils import BatchFeature
+from ...utils import PaddingStrategy, TensorType, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Wav2Vec2FeatureExtractor(SequenceFeatureExtractor):
+ r"""
+ Constructs a Wav2Vec2 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.
+
+ Args:
+ feature_size (`int`, *optional*, defaults to 1):
+ 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).
+ padding_value (`float`, *optional*, defaults to 0.0):
+ The value that is used to fill the padding values.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly
+ improve the performance for some models, *e.g.*,
+ [wav2vec2-lv60](https://huggingface.co/models?search=lv60).
+ return_attention_mask (`bool`, *optional*, defaults to `False`):
+ Whether or not [`~Wav2Vec2FeatureExtractor.__call__`] should return `attention_mask`.
+
+
+
+ Wav2Vec2 models that have set `config.feat_extract_norm == "group"`, such as
+ [wav2vec2-base](https://huggingface.co/facebook/wav2vec2-base-960h), have **not** been trained using
+ `attention_mask`. For such models, `input_values` should simply be padded with 0 and no `attention_mask`
+ should be passed.
+
+ For Wav2Vec2 models that have set `config.feat_extract_norm == "layer"`, such as
+ [wav2vec2-lv60](https://huggingface.co/facebook/wav2vec2-large-960h-lv60-self), `attention_mask` should be
+ passed for batched inference.
+
+ """
+
+ model_input_names = ["input_values", "attention_mask"]
+
+ def __init__(
+ self,
+ feature_size=1,
+ sampling_rate=16000,
+ padding_value=0.0,
+ return_attention_mask=False,
+ do_normalize=True,
+ **kwargs,
+ ):
+ super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
+ self.return_attention_mask = return_attention_mask
+ self.do_normalize = do_normalize
+
+ @staticmethod
+ def zero_mean_unit_var_norm(
+ input_values: list[np.ndarray], attention_mask: list[np.ndarray], padding_value: float = 0.0
+ ) -> list[np.ndarray]:
+ """
+ Every array in the list is normalized to have zero mean and unit variance
+ """
+ if attention_mask is not None:
+ attention_mask = np.array(attention_mask, np.int32)
+ normed_input_values = []
+
+ for vector, length in zip(input_values, attention_mask.sum(-1)):
+ normed_slice = (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)
+ if length < normed_slice.shape[0]:
+ normed_slice[length:] = padding_value
+
+ normed_input_values.append(normed_slice)
+ else:
+ normed_input_values = [(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]
+
+ return normed_input_values
+
+ 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_attention_mask: bool | None = None,
+ return_tensors: str | TensorType | None = None,
+ sampling_rate: int | 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 `False`):
+ 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)
+
+
+
+ Wav2Vec2 models that have set `config.feat_extract_norm == "group"`, such as
+ [wav2vec2-base](https://huggingface.co/facebook/wav2vec2-base-960h), have **not** been trained using
+ `attention_mask`. For such models, `input_values` should simply be padded with 0 and no
+ `attention_mask` should be passed.
+
+ For Wav2Vec2 models that have set `config.feat_extract_norm == "layer"`, such as
+ [wav2vec2-lv60](https://huggingface.co/facebook/wav2vec2-large-960h-lv60-self), `attention_mask` should
+ be passed for batched inference.
+
+
+
+ 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):
+ """
+
+ 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)))
+ )
+
+ # always return batch
+ if not is_batched:
+ raw_speech = [raw_speech]
+
+ # convert into correct format for padding
+ encoded_inputs = BatchFeature({"input_values": raw_speech})
+
+ 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,
+ )
+
+ # convert input values to correct format
+ input_values = padded_inputs["input_values"]
+ if not isinstance(input_values[0], np.ndarray):
+ padded_inputs["input_values"] = [np.asarray(array, dtype=np.float32) for array in input_values]
+ elif (
+ not isinstance(input_values, np.ndarray)
+ and isinstance(input_values[0], np.ndarray)
+ and input_values[0].dtype is np.dtype(np.float64)
+ ):
+ padded_inputs["input_values"] = [array.astype(np.float32) for array in input_values]
+ elif isinstance(input_values, np.ndarray) and input_values.dtype is np.dtype(np.float64):
+ padded_inputs["input_values"] = input_values.astype(np.float32)
+
+ # convert attention_mask to correct format
+ 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]
+
+ # zero-mean and unit-variance normalization
+ if self.do_normalize:
+ attention_mask = (
+ attention_mask
+ if self._get_padding_strategies(padding, max_length=max_length) is not PaddingStrategy.DO_NOT_PAD
+ else None
+ )
+ padded_inputs["input_values"] = self.zero_mean_unit_var_norm(
+ padded_inputs["input_values"], attention_mask=attention_mask, padding_value=self.padding_value
+ )
+
+ if return_tensors is not None:
+ padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
+
+ return padded_inputs
+
+
+__all__ = ["Wav2Vec2FeatureExtractor"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py
new file mode 100644
index 0000000000000000000000000000000000000000..274a033657107e2a43cf36b12258d5ff501c913b
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py
@@ -0,0 +1,2153 @@
+# 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 Wav2Vec2 model."""
+
+import math
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from safetensors.torch import load_file as safe_load_file
+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_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ CausalLMOutput,
+ 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 (
+ ModelOutput,
+ TransformersKwargs,
+ auto_docstring,
+ cached_file,
+ check_torch_load_is_safe,
+ is_peft_available,
+ logging,
+)
+from .configuration_wav2vec2 import Wav2Vec2Config
+
+
+WAV2VEC2_ADAPTER_PT_FILE = "adapter.{}.bin"
+WAV2VEC2_ADAPTER_SAFE_FILE = "adapter.{}.safetensors"
+
+
+logger = logging.get_logger(__name__)
+
+
+_HIDDEN_STATES_START_POSITION = 2
+
+
+@auto_docstring(
+ custom_intro="""
+ Output type of [`Wav2Vec2ForPreTraining`], with potential hidden states and attentions.
+ """
+)
+@dataclass
+class Wav2Vec2ForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (*optional*, returned when `sample_negative_indices` are passed, `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).
+ 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.
+ contrastive_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):
+ The contrastive loss (L_m) as stated in the [official paper](https://huggingface.co/papers/2006.11477).
+ diversity_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):
+ The diversity loss (L_d) as stated in the [official paper](https://huggingface.co/papers/2006.11477).
+ """
+
+ loss: 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
+ contrastive_loss: torch.FloatTensor | None = None
+ diversity_loss: torch.FloatTensor | None = None
+
+
+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
+
+
+def _sample_negative_indices(features_shape: tuple, num_negatives: int, mask_time_indices: np.ndarray | None = None):
+ """
+ Sample `num_negatives` vectors from feature vectors.
+ """
+ batch_size, sequence_length = features_shape
+
+ # generate indices of the positive vectors themselves, repeat them `num_negatives` times
+ sequence_length_range = np.arange(sequence_length)
+
+ # get `num_negatives` random vector indices from the same utterance
+ sampled_negative_indices = np.zeros(shape=(batch_size, sequence_length, num_negatives), dtype=np.int32)
+
+ mask_time_indices = (
+ mask_time_indices.astype(bool) if mask_time_indices is not None else np.ones(features_shape, dtype=bool)
+ )
+
+ for batch_idx in range(batch_size):
+ high = mask_time_indices[batch_idx].sum() - 1
+ mapped_masked_indices = sequence_length_range[mask_time_indices[batch_idx]]
+
+ feature_indices = np.broadcast_to(np.arange(high + 1)[:, None], (high + 1, num_negatives))
+ sampled_indices = np.random.randint(0, high, size=(high + 1, num_negatives))
+ # avoid sampling the same positive vector, but keep the distribution uniform
+ sampled_indices[sampled_indices >= feature_indices] += 1
+
+ # remap to actual indices
+ sampled_negative_indices[batch_idx][mask_time_indices[batch_idx]] = mapped_masked_indices[sampled_indices]
+
+ # correct for batch size
+ sampled_negative_indices[batch_idx] += batch_idx * sequence_length
+
+ return sampled_negative_indices
+
+
+class Wav2Vec2NoLayerNormConvLayer(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 Wav2Vec2LayerNormConvLayer(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 Wav2Vec2GroupNormConvLayer(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 Wav2Vec2PositionalConvEmbedding(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 = Wav2Vec2SamePadLayer(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 Wav2Vec2SamePadLayer(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 Wav2Vec2FeatureEncoder(nn.Module):
+ """Construct the features from raw audio waveform"""
+
+ def __init__(self, config):
+ super().__init__()
+
+ if config.feat_extract_norm == "group":
+ conv_layers = [Wav2Vec2GroupNormConvLayer(config, layer_id=0)] + [
+ Wav2Vec2NoLayerNormConvLayer(config, layer_id=i + 1) for i in range(config.num_feat_extract_layers - 1)
+ ]
+ elif config.feat_extract_norm == "layer":
+ conv_layers = [
+ Wav2Vec2LayerNormConvLayer(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 Wav2Vec2FeatureProjection(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
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class Wav2Vec2Attention(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: Wav2Vec2Config | 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
+ input_shape = hidden_states.shape[:-1]
+
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ kv_shape = (*current_states.shape[:-1], -1, self.head_dim)
+ key_states = self.k_proj(current_states).view(kv_shape).transpose(1, 2)
+ value_states = self.v_proj(current_states).view(kv_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(*input_shape, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights, None
+
+
+class Wav2Vec2FeedForward(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 Wav2Vec2EncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config):
+ super().__init__()
+ self.attention = Wav2Vec2Attention(
+ 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 = Wav2Vec2FeedForward(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 Wav2Vec2EncoderLayerStableLayerNorm(GradientCheckpointingLayer):
+ def __init__(self, config):
+ super().__init__()
+ self.attention = Wav2Vec2Attention(
+ 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 = Wav2Vec2FeedForward(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 = Wav2Vec2AttnAdapterLayer(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 Wav2Vec2Encoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.pos_conv_embed = Wav2Vec2PositionalConvEmbedding(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([Wav2Vec2EncoderLayer(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 Wav2Vec2EncoderStableLayerNorm(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.pos_conv_embed = Wav2Vec2PositionalConvEmbedding(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(
+ [Wav2Vec2EncoderLayerStableLayerNorm(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 Wav2Vec2GumbelVectorQuantizer(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.conv_dim[-1], self.num_groups * self.num_vars)
+
+ # can be decayed for training
+ self.temperature = 2
+
+ @staticmethod
+ def _compute_perplexity(probs, mask=None):
+ if mask is not None:
+ mask_extended = mask.flatten()[:, None, None].expand(probs.shape)
+ probs = torch.where(mask_extended, probs, torch.zeros_like(probs))
+ marginal_probs = probs.sum(dim=0) / mask.sum()
+ else:
+ 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, mask_time_indices=None):
+ 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, mask_time_indices)
+ 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, mask_time_indices)
+
+ 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
+
+
+class Wav2Vec2Adapter(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ # feature dim might need to be down-projected
+ if config.output_hidden_size != config.hidden_size:
+ self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)
+ self.proj_layer_norm = nn.LayerNorm(config.output_hidden_size)
+ else:
+ self.proj = self.proj_layer_norm = None
+
+ self.layers = nn.ModuleList(Wav2Vec2AdapterLayer(config) for _ in range(config.num_adapter_layers))
+ self.layerdrop = config.layerdrop
+
+ def forward(self, hidden_states):
+ # down project hidden_states if necessary
+ if self.proj is not None and self.proj_layer_norm is not None:
+ hidden_states = self.proj(hidden_states)
+ hidden_states = self.proj_layer_norm(hidden_states)
+
+ hidden_states = hidden_states.transpose(1, 2)
+
+ for layer in self.layers:
+ layerdrop_prob = np.random.random()
+ if not self.training or (layerdrop_prob > self.layerdrop):
+ hidden_states = layer(hidden_states)
+
+ hidden_states = hidden_states.transpose(1, 2)
+ return hidden_states
+
+
+class Wav2Vec2AdapterLayer(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ config.output_hidden_size,
+ 2 * config.output_hidden_size,
+ config.adapter_kernel_size,
+ stride=config.adapter_stride,
+ padding=1,
+ )
+
+ def forward(self, hidden_states):
+ hidden_states = self.conv(hidden_states)
+ hidden_states = nn.functional.glu(hidden_states, dim=1)
+
+ return hidden_states
+
+
+class Wav2Vec2AttnAdapterLayer(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
+
+
+@auto_docstring
+class Wav2Vec2PreTrainedModel(PreTrainedModel):
+ config: Wav2Vec2Config
+ base_model_prefix = "wav2vec2"
+ 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"""
+ # Wav2Vec2ForPreTraining last 2 linear layers need standard Linear init.
+ if isinstance(module, Wav2Vec2ForPreTraining):
+ module.project_hid.reset_parameters()
+ module.project_q.reset_parameters()
+ # gumbel softmax requires special init
+ elif isinstance(module, Wav2Vec2GumbelVectorQuantizer):
+ init.normal_(module.weight_proj.weight, mean=0.0, std=1)
+ init.zeros_(module.weight_proj.bias)
+ init.uniform_(module.codevectors)
+ elif isinstance(module, Wav2Vec2PositionalConvEmbedding):
+ 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, Wav2Vec2FeatureProjection):
+ 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, add_adapter: bool | None = None):
+ """
+ Computes the output length of the convolutional layers
+ """
+
+ add_adapter = self.config.add_adapter if add_adapter is None else add_adapter
+
+ 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)
+
+ if add_adapter:
+ for _ in range(self.config.num_adapter_layers):
+ input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)
+
+ return input_lengths
+
+ def _get_feature_vector_attention_mask(
+ self, feature_vector_length: int, attention_mask: torch.LongTensor, add_adapter=None
+ ):
+ # 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, add_adapter=add_adapter)
+ output_lengths = output_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 _get_adapters(self):
+ if self.config.adapter_attn_dim is None:
+ raise ValueError(f"{self.__class__} has no adapter layers. Make sure to define `config.adapter_attn_dim`.")
+
+ adapter_weights = {}
+ for name, module in self.named_modules():
+ if isinstance(module, Wav2Vec2AttnAdapterLayer):
+ for param_name, param in module.named_parameters():
+ adapter_weights[".".join([name, param_name])] = param
+
+ if isinstance(self, Wav2Vec2ForCTC):
+ for name, param in self.lm_head.named_parameters():
+ adapter_weights[".".join(["lm_head", name])] = param
+
+ return adapter_weights
+
+ def init_adapter_layers(self):
+ """
+ (Re-)initialize attention adapter layers and lm head for adapter-only fine-tuning
+ """
+ # init attention adapters
+ for module in self.modules():
+ if isinstance(module, Wav2Vec2AttnAdapterLayer):
+ self._init_weights(module)
+
+ # init lm head
+ if isinstance(self, Wav2Vec2ForCTC):
+ self._init_weights(self.lm_head)
+
+ def load_adapter(self, target_lang: str, force_load=True, **kwargs):
+ r"""
+ Load a language adapter model from a pre-trained adapter model.
+
+ Parameters:
+ target_lang (`str`):
+ Has to be a language id of an existing adapter weight. Adapter weights are stored in the format
+ adapter..safetensors or adapter..bin
+ force_load (`bool`, defaults to `True`):
+ Whether the weights shall be loaded even if `target_lang` matches `self.target_lang`.
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+ cached versions if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ local_files_only(`bool`, *optional*, defaults to `False`):
+ Whether or not to only look at local files (i.e., do not try to download the model).
+ token (`str` or `bool`, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
+ the token generated when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+
+
+
+ To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`.
+
+
+
+ mirror (`str`, *optional*):
+ Mirror source to accelerate downloads in China. If you are from China and have an accessibility
+ problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
+ Please refer to the mirror site for more information.
+
+
+
+ Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to
+ use this method in a firewalled environment.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import Wav2Vec2ForCTC, AutoProcessor
+
+ >>> ckpt = "facebook/mms-1b-all"
+ >>> processor = AutoProcessor.from_pretrained(ckpt)
+ >>> model = Wav2Vec2ForCTC.from_pretrained(ckpt, target_lang="eng")
+ >>> # set specific language
+ >>> processor.tokenizer.set_target_lang("spa")
+ >>> model.load_adapter("spa")
+ ```
+ """
+ if self.config.adapter_attn_dim is None:
+ raise ValueError(f"Cannot load_adapter for {target_lang} if `config.adapter_attn_dim` is not defined.")
+
+ if target_lang == self.target_lang and not force_load:
+ logger.warning(f"Adapter weights are already set to {target_lang}.")
+ return
+
+ cache_dir = kwargs.pop("cache_dir", None)
+ force_download = kwargs.pop("force_download", False)
+ proxies = kwargs.pop("proxies", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ token = kwargs.pop("token", None)
+ revision = kwargs.pop("revision", None)
+ use_safetensors = kwargs.pop("use_safetensors", None)
+ model_path_or_id = self.config._name_or_path
+ state_dict = None
+
+ # 1. Let's first try loading a safetensors adapter weight
+ if use_safetensors is not False:
+ filepath = WAV2VEC2_ADAPTER_SAFE_FILE.format(target_lang)
+
+ try:
+ weight_path = cached_file(
+ model_path_or_id,
+ filename=filepath,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ cache_dir=cache_dir,
+ )
+
+ state_dict = safe_load_file(weight_path)
+
+ except OSError:
+ if use_safetensors:
+ # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
+ # to the original exception.
+ raise
+
+ except Exception:
+ # For any other exception, we throw a generic error.
+ if use_safetensors:
+ raise OSError(
+ f"Can't load the model for '{model_path_or_id}'. If you were trying to load it"
+ " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
+ f" same name. Otherwise, make sure '{model_path_or_id}' is the correct path to a"
+ f" directory containing a file named {filepath}."
+ )
+
+ # 2. If this didn't work let's try loading a PyTorch adapter weight
+ if state_dict is None:
+ filepath = WAV2VEC2_ADAPTER_PT_FILE.format(target_lang)
+
+ try:
+ weight_path = cached_file(
+ model_path_or_id,
+ filename=filepath,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ cache_dir=cache_dir,
+ )
+
+ check_torch_load_is_safe()
+ state_dict = torch.load(
+ weight_path,
+ map_location="cpu",
+ weights_only=True,
+ )
+
+ except OSError:
+ # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
+ # to the original exception.
+ raise
+
+ except ValueError:
+ raise
+
+ except Exception:
+ # For any other exception, we throw a generic error.
+ raise OSError(
+ f"Can't load the model for '{model_path_or_id}'. If you were trying to load it"
+ " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
+ f" same name. Otherwise, make sure '{model_path_or_id}' is the correct path to a"
+ f" directory containing a file named {filepath}."
+ )
+
+ adapter_weights = self._get_adapters()
+ unexpected_keys = set(state_dict.keys()) - set(adapter_weights.keys())
+ missing_keys = set(adapter_weights.keys()) - set(state_dict.keys())
+
+ if len(unexpected_keys) > 0:
+ raise ValueError(f"The adapter weights {weight_path} has unexpected keys: {', '.join(unexpected_keys)}.")
+ elif len(missing_keys) > 0:
+ raise ValueError(f"The adapter weights {weight_path} has missing keys: {', '.join(missing_keys)}.")
+
+ # make sure now vocab size is correct
+ target_vocab_size = state_dict["lm_head.weight"].shape[0]
+ if target_vocab_size != self.config.vocab_size:
+ self.lm_head = nn.Linear(
+ self.config.output_hidden_size, target_vocab_size, device=self.device, dtype=self.dtype
+ )
+ self.config.vocab_size = target_vocab_size
+
+ # make sure that adapter weights are put in exactly the same precision and device placement and overwritten adapter weights
+ state_dict = {k: v.to(adapter_weights[k]) for k, v in state_dict.items()}
+ self.load_state_dict(state_dict, strict=False)
+
+ # set target language correctly
+ self.target_lang = target_lang
+
+
+@auto_docstring
+class Wav2Vec2Model(Wav2Vec2PreTrainedModel):
+ def __init__(self, config: Wav2Vec2Config):
+ super().__init__(config)
+ self.config = config
+ self.feature_extractor = Wav2Vec2FeatureEncoder(config)
+ self.feature_projection = Wav2Vec2FeatureProjection(config)
+
+ # model only needs masking vector if mask prob is > 0.0
+ 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_())
+
+ if config.do_stable_layer_norm:
+ self.encoder = Wav2Vec2EncoderStableLayerNorm(config)
+ else:
+ self.encoder = Wav2Vec2Encoder(config)
+
+ self.adapter = Wav2Vec2Adapter(config) if config.add_adapter else None
+
+ # 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.feature_extractor._freeze_parameters()
+
+ 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 | Wav2Vec2BaseModelOutput:
+ 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, add_adapter=False
+ )
+
+ 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 self.adapter is not None:
+ hidden_states = self.adapter(hidden_states)
+
+ if not return_dict:
+ return (hidden_states, extract_features) + encoder_outputs[1:]
+
+ return Wav2Vec2BaseModelOutput(
+ last_hidden_state=hidden_states,
+ extract_features=extract_features,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Wav2Vec2 Model with a quantizer and `VQ` head on top.
+ """
+)
+class Wav2Vec2ForPreTraining(Wav2Vec2PreTrainedModel):
+ def __init__(self, config: Wav2Vec2Config):
+ super().__init__(config)
+ self.wav2vec2 = Wav2Vec2Model(config)
+ self.dropout_features = nn.Dropout(config.feat_quantizer_dropout)
+
+ self.quantizer = Wav2Vec2GumbelVectorQuantizer(config)
+
+ self.project_hid = nn.Linear(config.hidden_size, config.proj_codevector_dim)
+ self.project_q = nn.Linear(config.codevector_dim, config.proj_codevector_dim)
+
+ # 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.wav2vec2.feature_extractor._freeze_parameters()
+
+ @staticmethod
+ def compute_contrastive_logits(
+ target_features: torch.FloatTensor,
+ negative_features: torch.FloatTensor,
+ predicted_features: torch.FloatTensor,
+ temperature: float = 0.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).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,
+ mask_time_indices: torch.BoolTensor | None = None,
+ sampled_negative_indices: torch.BoolTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | Wav2Vec2ForPreTrainingOutput:
+ 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.
+ sampled_negative_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_negatives)`, *optional*):
+ Indices indicating which quantized target vectors are used as negative sampled vectors in contrastive loss.
+ Required input for pre-training.
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoFeatureExtractor, Wav2Vec2ForPreTraining
+ >>> from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices, _sample_negative_indices
+ >>> from datasets import load_dataset
+
+ >>> feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
+ >>> model = Wav2Vec2ForPreTraining.from_pretrained("facebook/wav2vec2-base")
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> input_values = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt").input_values # Batch size 1
+
+ >>> # compute masked indices
+ >>> batch_size, raw_sequence_length = input_values.shape
+ >>> sequence_length = model._get_feat_extract_output_lengths(raw_sequence_length).item()
+ >>> mask_time_indices = _compute_mask_indices(
+ ... shape=(batch_size, sequence_length), mask_prob=0.2, mask_length=2
+ ... )
+ >>> sampled_negative_indices = _sample_negative_indices(
+ ... features_shape=(batch_size, sequence_length),
+ ... num_negatives=model.config.num_negatives,
+ ... mask_time_indices=mask_time_indices,
+ ... )
+ >>> mask_time_indices = torch.tensor(data=mask_time_indices, device=input_values.device, dtype=torch.long)
+ >>> sampled_negative_indices = torch.tensor(
+ ... data=sampled_negative_indices, device=input_values.device, dtype=torch.long
+ ... )
+
+ >>> with torch.no_grad():
+ ... outputs = model(input_values, mask_time_indices=mask_time_indices)
+
+ >>> # compute cosine similarity between predicted (=projected_states) and target (=projected_quantized_states)
+ >>> cosine_sim = torch.cosine_similarity(outputs.projected_states, outputs.projected_quantized_states, dim=-1)
+
+ >>> # show that cosine similarity is much higher than random
+ >>> cosine_sim[mask_time_indices.to(torch.bool)].mean() > 0.5
+ tensor(True)
+
+ >>> # for contrastive loss training model should be put into train mode
+ >>> model = model.train()
+ >>> loss = model(
+ ... input_values, mask_time_indices=mask_time_indices, sampled_negative_indices=sampled_negative_indices
+ ... ).loss
+ ```"""
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ if mask_time_indices is not None:
+ mask_time_indices = mask_time_indices.to(torch.bool)
+
+ outputs = self.wav2vec2(
+ input_values,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ mask_time_indices=mask_time_indices,
+ return_dict=return_dict,
+ )
+
+ # 1. project all transformed features (including masked) to final vq dim
+ transformer_features = self.project_hid(outputs[0])
+
+ # 2. quantize all (unmasked) extracted features and project to final vq dim
+ extract_features = self.dropout_features(outputs[1])
+
+ 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, add_adapter=False
+ )
+
+ quantized_features, codevector_perplexity = self.quantizer(
+ extract_features, mask_time_indices=mask_time_indices
+ )
+
+ quantized_features = quantized_features.to(self.project_q.weight.dtype)
+ quantized_features = self.project_q(quantized_features)
+
+ loss = contrastive_loss = diversity_loss = None
+ if sampled_negative_indices is not None:
+ batch_size, sequence_length, hidden_size = quantized_features.shape
+
+ # for training, we sample negatives
+ # 3. sample K negatives (distractors) quantized states for contrastive loss
+ # if attention_mask is passed, make sure that padded feature vectors cannot be sampled
+ # sample negative quantized vectors BTC => (BxT)C
+ negative_quantized_features = quantized_features.view(-1, hidden_size)[
+ sampled_negative_indices.long().view(-1)
+ ]
+ negative_quantized_features = negative_quantized_features.view(
+ batch_size, sequence_length, -1, hidden_size
+ ).permute(2, 0, 1, 3)
+
+ # 4. compute logits, corresponding to `logs = sim(c_t, [q_t, \sim{q}_t]) / \kappa`
+ # of equation (3) in https://huggingface.co/papers/2006.11477
+ logits = self.compute_contrastive_logits(
+ quantized_features[None, :],
+ negative_quantized_features,
+ transformer_features,
+ self.config.contrastive_logits_temperature,
+ )
+
+ # 5. if a negative vector is identical to the positive (i.e. when codebook utilization is low),
+ # its cosine similarity will be masked
+ neg_is_pos = (quantized_features == negative_quantized_features).all(-1)
+
+ if neg_is_pos.any():
+ logits[1:][neg_is_pos] = float("-inf")
+
+ # 6. compute contrastive loss \mathbf{L}_m = cross_entropy(logs) =
+ # -log(exp(sim(c_t, q_t)/\kappa) / \sum_{\sim{q}} exp(sim(c_t, \sim{q})/\kappa))
+ logits = logits.transpose(0, 2).reshape(-1, logits.size(0))
+ target = ((1 - mask_time_indices.long()) * -100).transpose(0, 1).flatten()
+
+ contrastive_loss = nn.functional.cross_entropy(logits.float(), target, reduction="sum")
+ # 7. compute diversity loss: \mathbf{L}_d
+ num_codevectors = self.config.num_codevectors_per_group * self.config.num_codevector_groups
+ diversity_loss = ((num_codevectors - codevector_perplexity) / num_codevectors) * mask_time_indices.sum()
+
+ # 8. \mathbf{L} = \mathbf{L}_m + \alpha * \mathbf{L}_d
+ loss = contrastive_loss + self.config.diversity_loss_weight * diversity_loss
+
+ if not return_dict:
+ if loss is not None:
+ return (loss, transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
+ return (transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
+
+ return Wav2Vec2ForPreTrainingOutput(
+ loss=loss,
+ projected_states=transformer_features,
+ projected_quantized_states=quantized_features,
+ codevector_perplexity=codevector_perplexity,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ contrastive_loss=contrastive_loss,
+ diversity_loss=diversity_loss,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Wav2Vec2 Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).
+ """
+)
+class Wav2Vec2ForCTC(Wav2Vec2PreTrainedModel):
+ 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 [`Wav2Vec2ForCTC`] with adapters. Uses 'eng' by
+ default.
+ """
+ super().__init__(config)
+
+ self.wav2vec2 = Wav2Vec2Model(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: `Wav2Vec2ForCTC.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 Wav2Vec2 so that we do not have to introduce a new API to
+ # [`PreTrainedModel`]. While slightly hacky, Wav2Vec2 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.wav2vec2.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.wav2vec2.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.wav2vec2(
+ 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="""
+ Wav2Vec2 Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like
+ SUPERB Keyword Spotting.
+ """
+)
+class Wav2Vec2ForSequenceClassification(Wav2Vec2PreTrainedModel):
+ 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 Wav2Vec2 adapters (config.add_adapter=True)"
+ )
+ self.wav2vec2 = Wav2Vec2Model(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.wav2vec2.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.wav2vec2.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 [`Wav2Vec2Processor.__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.wav2vec2(
+ 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 Wav2Vec2ForAudioFrameClassification(Wav2Vec2PreTrainedModel):
+ 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 Wav2Vec2 adapters (config.add_adapter=True)"
+ )
+ self.wav2vec2 = Wav2Vec2Model(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.wav2vec2.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.wav2vec2.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 [`Wav2Vec2Processor.__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.wav2vec2(
+ 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="""
+ Wav2Vec2 Model with an XVector feature extraction head on top for tasks like Speaker Verification.
+ """
+)
+class Wav2Vec2ForXVector(Wav2Vec2PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.wav2vec2 = Wav2Vec2Model(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.wav2vec2.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.wav2vec2.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 [`Wav2Vec2Processor.__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.wav2vec2(
+ 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__ = [
+ "Wav2Vec2ForAudioFrameClassification",
+ "Wav2Vec2ForCTC",
+ "Wav2Vec2ForPreTraining",
+ "Wav2Vec2ForSequenceClassification",
+ "Wav2Vec2ForXVector",
+ "Wav2Vec2Model",
+ "Wav2Vec2PreTrainedModel",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py
new file mode 100644
index 0000000000000000000000000000000000000000..10f1959dbaca3b03d34e196741deeefeac364bd3
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py
@@ -0,0 +1,105 @@
+# 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 Wav2Vec2
+"""
+
+from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import AudioInput, PreTokenizedInput, TextInput
+from ...utils import auto_docstring
+
+
+class Wav2Vec2ProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {}
+
+
+@auto_docstring
+class Wav2Vec2Processor(ProcessorMixin):
+ def __init__(self, feature_extractor, tokenizer):
+ super().__init__(feature_extractor, tokenizer)
+
+ @auto_docstring
+ def __call__(
+ self,
+ audio: AudioInput | None = None,
+ text: str | list[str] | TextInput | PreTokenizedInput | None = None,
+ **kwargs: Unpack[Wav2Vec2ProcessorKwargs],
+ ):
+ r"""
+ Returns:
+ This method returns the results of each `call` method. If both are used, the output is a dictionary containing the results of both.
+ """
+ if audio is None and text is None:
+ raise ValueError("You need to specify either an `audio` or `text` input to process.")
+
+ output_kwargs = self._merge_kwargs(
+ Wav2Vec2ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ 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
+ elif audio is None:
+ return encodings
+ else:
+ inputs["labels"] = encodings["input_ids"]
+ return inputs
+
+ def pad(self, *args, **kwargs):
+ """
+ This method operates on batches of extracted features and/or tokenized text. It forwards all arguments to
+ [`Wav2Vec2FeatureExtractor.pad`] and/or [`PreTrainedTokenizer.pad`] depending on the input modality and returns their outputs. If both modalities are passed, [`Wav2Vec2FeatureExtractor.pad`] and [`PreTrainedTokenizer.pad`] are called.
+
+ Args:
+ input_features:
+ When the first argument is a dictionary containing a batch of tensors, or the `input_features` argument is present, it is passed to [`Wav2Vec2FeatureExtractor.pad`].
+ labels:
+ When the `label` argument is present, it is passed to [`PreTrainedTokenizer.pad`].
+
+ Returns:
+ This method returns the results of each `pad` method. If both are used, the output is a dictionary containing the results of both.
+ """
+ input_features = kwargs.pop("input_features", None)
+ labels = kwargs.pop("labels", None)
+ if len(args) > 0:
+ input_features = args[0]
+ args = args[1:]
+
+ if input_features is not None:
+ input_features = self.feature_extractor.pad(input_features, *args, **kwargs)
+ if labels is not None:
+ labels = self.tokenizer.pad(labels, **kwargs)
+
+ if labels is None:
+ return input_features
+ elif input_features is None:
+ return labels
+ else:
+ input_features["labels"] = labels["input_ids"]
+ return input_features
+
+ @property
+ def model_input_names(self):
+ # The processor doesn't return text ids and the model seems to not need them
+ feature_extractor_input_names = self.feature_extractor.model_input_names
+ return feature_extractor_input_names + ["labels"]
+
+
+__all__ = ["Wav2Vec2Processor"]