diff --git a/third_party/transformers/src/transformers/models/align/__init__.py b/third_party/transformers/src/transformers/models/align/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aaa64dfb6064b10e820fca01e9e632aa7ecd68c6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/align/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_align import * + from .modeling_align import * + from .processing_align import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/align/configuration_align.py b/third_party/transformers/src/transformers/models/align/configuration_align.py new file mode 100644 index 0000000000000000000000000000000000000000..91760eeae750e07dc166bbd125df2f8f70cdb42e --- /dev/null +++ b/third_party/transformers/src/transformers/models/align/configuration_align.py @@ -0,0 +1,211 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ALIGN model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="kakaobrain/align-base") +@strict +class AlignTextConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import AlignTextConfig, AlignTextModel + + >>> # Initializing a AlignTextConfig with kakaobrain/align-base style configuration + >>> configuration = AlignTextConfig() + + >>> # Initializing a AlignTextModel (with random weights) from the kakaobrain/align-base style configuration + >>> model = AlignTextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "align_text_model" + base_config_key = "text_config" + + vocab_size: int = 30522 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.1 + attention_probs_dropout_prob: float | int = 0.1 + max_position_embeddings: int = 512 + type_vocab_size: int = 2 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + pad_token_id: int | None = 0 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + + +@auto_docstring(checkpoint="kakaobrain/align-base") +@strict +class AlignVisionConfig(PreTrainedConfig): + r""" + width_coefficient (`float`, *optional*, defaults to 2.0): + Scaling coefficient for network width at each stage. + depth_coefficient (`float`, *optional*, defaults to 3.1): + Scaling coefficient for network depth at each stage. + depth_divisor (`int`, *optional*, defaults to 8): + A unit of network width. + kernel_sizes (`list[int]`, *optional*, defaults to `[3, 3, 5, 3, 5, 5, 3]`): + List of kernel sizes to be used in each block. + in_channels (`list[int]`, *optional*, defaults to `[32, 16, 24, 40, 80, 112, 192]`): + List of input channel sizes to be used in each block for convolutional layers. + out_channels (`list[int]`, *optional*, defaults to `[16, 24, 40, 80, 112, 192, 320]`): + List of output channel sizes to be used in each block for convolutional layers. + depthwise_padding (`list[int]`, *optional*, defaults to `[]`): + List of block indices with square padding. + strides (`list[int]`, *optional*, defaults to `[1, 2, 2, 2, 1, 2, 1]`): + List of stride sizes to be used in each block for convolutional layers. + num_block_repeats (`list[int]`, *optional*, defaults to `[1, 2, 2, 3, 3, 4, 1]`): + List of the number of times each block is to repeated. + expand_ratios (`list[int]`, *optional*, defaults to `[1, 6, 6, 6, 6, 6, 6]`): + List of scaling coefficient of each block. + squeeze_expansion_ratio (`float`, *optional*, defaults to 0.25): + Squeeze expansion ratio. + hidden_dim (`int`, *optional*, defaults to 1280): + The hidden dimension of the layer before the classification head. + pooling_type (`str` or `function`, *optional*, defaults to `"mean"`): + Type of final pooling to be applied before the dense classification head. Available options are [`"mean"`, + `"max"`] + batch_norm_momentum (`float`, *optional*, defaults to 0.99): + The momentum used by the batch normalization layers. + drop_connect_rate (`float`, *optional*, defaults to 0.2): + The drop rate for skip connections. + + Example: + + ```python + >>> from transformers import AlignVisionConfig, AlignVisionModel + + >>> # Initializing a AlignVisionConfig with kakaobrain/align-base style configuration + >>> configuration = AlignVisionConfig() + + >>> # Initializing a AlignVisionModel (with random weights) from the kakaobrain/align-base style configuration + >>> model = AlignVisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "align_vision_model" + base_config_key = "vision_config" + + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 600 + width_coefficient: float = 2.0 + depth_coefficient: float = 3.1 + depth_divisor: int = 8 + kernel_sizes: list[int] | tuple[int, ...] = (3, 3, 5, 3, 5, 5, 3) + in_channels: list[int] | tuple[int, ...] = (32, 16, 24, 40, 80, 112, 192) + out_channels: list[int] | tuple[int, ...] = (16, 24, 40, 80, 112, 192, 320) + depthwise_padding: list | tuple[int, ...] = () + strides: list[int] | tuple[int, ...] = (1, 2, 2, 2, 1, 2, 1) + num_block_repeats: list[int] | tuple[int, ...] = (1, 2, 2, 3, 3, 4, 1) + expand_ratios: list[int] | tuple[int, ...] = (1, 6, 6, 6, 6, 6, 6) + squeeze_expansion_ratio: float = 0.25 + hidden_act: str = "swish" + hidden_dim: int = 2560 + pooling_type: str = "mean" + initializer_range: float = 0.02 + batch_norm_eps: float = 0.001 + batch_norm_momentum: float = 0.99 + drop_connect_rate: float | int = 0.2 + + def __post_init__(self, **kwargs): + self.num_hidden_layers = sum(self.num_block_repeats) * 4 + for attr in [ + "kernel_sizes", + "in_channels", + "out_channels", + "depthwise_padding", + "strides", + "num_block_repeats", + "expand_ratios", + ]: + # cast tuple so it can be JSON-ized when saving + setattr(self, attr, list(getattr(self, attr))) + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="kakaobrain/align-base") +@strict +class AlignConfig(PreTrainedConfig): + r""" + temperature_init_value (`float`, *optional*, defaults to 1.0): + The initial value of the *temperature* parameter. Default is used as per the original ALIGN implementation. + + Example: + + ```python + >>> from transformers import AlignConfig, AlignModel + + >>> # Initializing a AlignConfig with kakaobrain/align-base style configuration + >>> configuration = AlignConfig() + + >>> # Initializing a AlignModel (with random weights) from the kakaobrain/align-base style configuration + >>> model = AlignModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a AlignConfig from a AlignTextConfig and a AlignVisionConfig + >>> from transformers import AlignTextConfig, AlignVisionConfig + + >>> # Initializing ALIGN Text and Vision configurations + >>> config_text = AlignTextConfig() + >>> config_vision = AlignVisionConfig() + + >>> config = AlignConfig(text_config=config_text, vision_config=config_vision) + ```""" + + model_type = "align" + sub_configs = {"text_config": AlignTextConfig, "vision_config": AlignVisionConfig} + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + projection_dim: int = 640 + temperature_init_value: float = 1.0 + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + if self.text_config is None: + self.text_config = AlignTextConfig() + logger.info("`text_config` is `None`. Initializing the `AlignTextConfig` with default values.") + elif isinstance(self.text_config, dict): + self.text_config = AlignTextConfig(**self.text_config) + + if self.vision_config is None: + self.vision_config = AlignVisionConfig() + logger.info("`vision_config` is `None`. initializing the `AlignVisionConfig` with default values.") + elif isinstance(self.vision_config, dict): + self.vision_config = AlignVisionConfig(**self.vision_config) + + super().__post_init__(**kwargs) + + +__all__ = ["AlignTextConfig", "AlignVisionConfig", "AlignConfig"] diff --git a/third_party/transformers/src/transformers/models/align/convert_align_tf_to_hf.py b/third_party/transformers/src/transformers/models/align/convert_align_tf_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..9ba9352f4240748f4e72d71773e7a85fc221c37a --- /dev/null +++ b/third_party/transformers/src/transformers/models/align/convert_align_tf_to_hf.py @@ -0,0 +1,388 @@ +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert ALIGN checkpoints from the original repository.""" + +import argparse +import os +from io import BytesIO + +import align +import httpx +import numpy as np +import tensorflow as tf +import torch +from PIL import Image +from tokenizer import Tokenizer + +from transformers import ( + AlignConfig, + AlignModel, + AlignProcessor, + BertConfig, + BertTokenizer, + EfficientNetConfig, + EfficientNetImageProcessor, +) +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def preprocess(image): + image = tf.image.resize(image, (346, 346)) + image = tf.image.crop_to_bounding_box(image, (346 - 289) // 2, (346 - 289) // 2, 289, 289) + return image + + +def get_align_config(): + vision_config = EfficientNetConfig.from_pretrained("google/efficientnet-b7") + vision_config.image_size = 289 + vision_config.hidden_dim = 640 + vision_config.id2label = {"0": "LABEL_0", "1": "LABEL_1"} + vision_config.label2id = {"LABEL_0": 0, "LABEL_1": 1} + vision_config.depthwise_padding = [] + + text_config = BertConfig() + config = AlignConfig(text_config=text_config, vision_config=vision_config, projection_dim=640) + return config + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +def get_processor(): + image_processor = EfficientNetImageProcessor( + do_center_crop=True, + rescale_factor=1 / 127.5, + rescale_offset=True, + do_normalize=False, + include_top=False, + resample=Image.BILINEAR, + ) + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + tokenizer.model_max_length = 64 + processor = AlignProcessor(image_processor=image_processor, tokenizer=tokenizer) + return processor + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def rename_keys(original_param_names): + # EfficientNet image encoder + block_names = [v.split("_")[0].split("block")[1] for v in original_param_names if v.startswith("block")] + block_names = list(set(block_names)) + block_names = sorted(block_names) + num_blocks = len(block_names) + block_name_mapping = {b: str(i) for b, i in zip(block_names, range(num_blocks))} + + rename_keys = [] + rename_keys.append(("stem_conv/kernel:0", "embeddings.convolution.weight")) + rename_keys.append(("stem_bn/gamma:0", "embeddings.batchnorm.weight")) + rename_keys.append(("stem_bn/beta:0", "embeddings.batchnorm.bias")) + rename_keys.append(("stem_bn/moving_mean:0", "embeddings.batchnorm.running_mean")) + rename_keys.append(("stem_bn/moving_variance:0", "embeddings.batchnorm.running_var")) + + for b in block_names: + hf_b = block_name_mapping[b] + rename_keys.append((f"block{b}_expand_conv/kernel:0", f"encoder.blocks.{hf_b}.expansion.expand_conv.weight")) + rename_keys.append((f"block{b}_expand_bn/gamma:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.weight")) + rename_keys.append((f"block{b}_expand_bn/beta:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.bias")) + rename_keys.append( + (f"block{b}_expand_bn/moving_mean:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_mean") + ) + rename_keys.append( + (f"block{b}_expand_bn/moving_variance:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_var") + ) + rename_keys.append( + (f"block{b}_dwconv/depthwise_kernel:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight") + ) + rename_keys.append((f"block{b}_bn/gamma:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight")) + rename_keys.append((f"block{b}_bn/beta:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias")) + rename_keys.append( + (f"block{b}_bn/moving_mean:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean") + ) + rename_keys.append( + (f"block{b}_bn/moving_variance:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var") + ) + + rename_keys.append((f"block{b}_se_reduce/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.weight")) + rename_keys.append((f"block{b}_se_reduce/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.bias")) + rename_keys.append((f"block{b}_se_expand/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.weight")) + rename_keys.append((f"block{b}_se_expand/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.bias")) + rename_keys.append( + (f"block{b}_project_conv/kernel:0", f"encoder.blocks.{hf_b}.projection.project_conv.weight") + ) + rename_keys.append((f"block{b}_project_bn/gamma:0", f"encoder.blocks.{hf_b}.projection.project_bn.weight")) + rename_keys.append((f"block{b}_project_bn/beta:0", f"encoder.blocks.{hf_b}.projection.project_bn.bias")) + rename_keys.append( + (f"block{b}_project_bn/moving_mean:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_mean") + ) + rename_keys.append( + (f"block{b}_project_bn/moving_variance:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_var") + ) + + key_mapping = {} + for item in rename_keys: + if item[0] in original_param_names: + key_mapping[item[0]] = "vision_model." + item[1] + + # BERT text encoder + rename_keys = [] + old = "tf_bert_model/bert" + new = "text_model" + for i in range(12): + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/query/kernel:0", + f"{new}.encoder.layer.{i}.attention.self.query.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/query/bias:0", + f"{new}.encoder.layer.{i}.attention.self.query.bias", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/key/kernel:0", + f"{new}.encoder.layer.{i}.attention.self.key.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/key/bias:0", + f"{new}.encoder.layer.{i}.attention.self.key.bias", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/value/kernel:0", + f"{new}.encoder.layer.{i}.attention.self.value.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/self/value/bias:0", + f"{new}.encoder.layer.{i}.attention.self.value.bias", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/output/dense/kernel:0", + f"{new}.encoder.layer.{i}.attention.output.dense.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/output/dense/bias:0", + f"{new}.encoder.layer.{i}.attention.output.dense.bias", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0", + f"{new}.encoder.layer.{i}.attention.output.LayerNorm.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0", + f"{new}.encoder.layer.{i}.attention.output.LayerNorm.bias", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/intermediate/dense/kernel:0", + f"{new}.encoder.layer.{i}.intermediate.dense.weight", + ) + ) + rename_keys.append( + ( + f"{old}/encoder/layer_._{i}/intermediate/dense/bias:0", + f"{new}.encoder.layer.{i}.intermediate.dense.bias", + ) + ) + rename_keys.append( + (f"{old}/encoder/layer_._{i}/output/dense/kernel:0", f"{new}.encoder.layer.{i}.output.dense.weight") + ) + rename_keys.append( + (f"{old}/encoder/layer_._{i}/output/dense/bias:0", f"{new}.encoder.layer.{i}.output.dense.bias") + ) + rename_keys.append( + (f"{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0", f"{new}.encoder.layer.{i}.output.LayerNorm.weight") + ) + rename_keys.append( + (f"{old}/encoder/layer_._{i}/output/LayerNorm/beta:0", f"{new}.encoder.layer.{i}.output.LayerNorm.bias") + ) + + rename_keys.append((f"{old}/embeddings/word_embeddings/weight:0", f"{new}.embeddings.word_embeddings.weight")) + rename_keys.append( + (f"{old}/embeddings/position_embeddings/embeddings:0", f"{new}.embeddings.position_embeddings.weight") + ) + rename_keys.append( + (f"{old}/embeddings/token_type_embeddings/embeddings:0", f"{new}.embeddings.token_type_embeddings.weight") + ) + rename_keys.append((f"{old}/embeddings/LayerNorm/gamma:0", f"{new}.embeddings.LayerNorm.weight")) + rename_keys.append((f"{old}/embeddings/LayerNorm/beta:0", f"{new}.embeddings.LayerNorm.bias")) + + rename_keys.append((f"{old}/pooler/dense/kernel:0", f"{new}.pooler.dense.weight")) + rename_keys.append((f"{old}/pooler/dense/bias:0", f"{new}.pooler.dense.bias")) + rename_keys.append(("dense/kernel:0", "text_projection.weight")) + rename_keys.append(("dense/bias:0", "text_projection.bias")) + rename_keys.append(("dense/bias:0", "text_projection.bias")) + rename_keys.append(("temperature:0", "temperature")) + + for item in rename_keys: + if item[0] in original_param_names: + key_mapping[item[0]] = item[1] + return key_mapping + + +def replace_params(hf_params, tf_params, key_mapping): + list(hf_params.keys()) + + for key, value in tf_params.items(): + if key not in key_mapping: + continue + + hf_key = key_mapping[key] + if "_conv" in key and "kernel" in key: + new_hf_value = torch.from_numpy(value).permute(3, 2, 0, 1) + elif "embeddings" in key: + new_hf_value = torch.from_numpy(value) + elif "depthwise_kernel" in key: + new_hf_value = torch.from_numpy(value).permute(2, 3, 0, 1) + elif "kernel" in key: + new_hf_value = torch.from_numpy(np.transpose(value)) + elif "temperature" in key: + new_hf_value = value + elif "bn/gamma" in key or "bn/beta" in key: + new_hf_value = torch.from_numpy(np.transpose(value)).squeeze() + else: + new_hf_value = torch.from_numpy(value) + + # Replace HF parameters with original TF model parameters + hf_params[hf_key].copy_(new_hf_value) + + +@torch.no_grad() +def convert_align_checkpoint(checkpoint_path, pytorch_dump_folder_path, save_model, push_to_hub): + """ + Copy/paste/tweak model's weights to our ALIGN structure. + """ + # Load original model + seq_length = 64 + tok = Tokenizer(seq_length) + original_model = align.Align("efficientnet-b7", "bert-base", 640, seq_length, tok.get_vocab_size()) + original_model.compile() + original_model.load_weights(checkpoint_path) + + tf_params = original_model.trainable_variables + tf_non_train_params = original_model.non_trainable_variables + tf_params = {param.name: param.numpy() for param in tf_params} + for param in tf_non_train_params: + tf_params[param.name] = param.numpy() + tf_param_names = list(tf_params.keys()) + + # Load HuggingFace model + config = get_align_config() + hf_model = AlignModel(config).eval() + hf_params = hf_model.state_dict() + + # Create src-to-dst parameter name mapping dictionary + print("Converting parameters...") + key_mapping = rename_keys(tf_param_names) + replace_params(hf_params, tf_params, key_mapping) + + # Initialize processor + processor = get_processor() + inputs = processor( + images=prepare_img(), text="A picture of a cat", padding="max_length", max_length=64, return_tensors="pt" + ) + + # HF model inference + hf_model.eval() + with torch.no_grad(): + outputs = hf_model(**inputs) + + hf_image_features = outputs.image_embeds.detach().numpy() + hf_text_features = outputs.text_embeds.detach().numpy() + + # Original model inference + original_model.trainable = False + tf_image_processor = EfficientNetImageProcessor( + do_center_crop=True, + do_rescale=False, + do_normalize=False, + include_top=False, + resample=Image.BILINEAR, + ) + image = tf_image_processor(images=prepare_img(), return_tensors="tf", data_format="channels_last")["pixel_values"] + text = tok(tf.constant(["A picture of a cat"])) + + image_features = original_model.image_encoder(image, training=False) + text_features = original_model.text_encoder(text, training=False) + + image_features = tf.nn.l2_normalize(image_features, axis=-1) + text_features = tf.nn.l2_normalize(text_features, axis=-1) + + # Check whether original and HF model outputs match -> np.allclose + if not np.allclose(image_features, hf_image_features, atol=1e-3): + raise ValueError("The predicted image features are not the same.") + if not np.allclose(text_features, hf_text_features, atol=1e-3): + raise ValueError("The predicted text features are not the same.") + print("Model outputs match!") + + if save_model: + # Create folder to save model + if not os.path.isdir(pytorch_dump_folder_path): + os.mkdir(pytorch_dump_folder_path) + # Save converted model and image processor + hf_model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + # Push model and image processor to hub + print("Pushing converted ALIGN to the hub...") + processor.push_to_hub("align-base") + hf_model.push_to_hub("align-base") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--checkpoint_path", + default="./weights/model-weights", + type=str, + help="Path to the pretrained TF ALIGN checkpoint.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default="hf_model", + type=str, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument("--save_model", action="store_true", help="Save model to local") + parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub") + + args = parser.parse_args() + convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/align/modeling_align.py b/third_party/transformers/src/transformers/models/align/modeling_align.py new file mode 100644 index 0000000000000000000000000000000000000000..62abb9438ea9b6f6ab6bd8f15177513635b47591 --- /dev/null +++ b/third_party/transformers/src/transformers/models/align/modeling_align.py @@ -0,0 +1,1180 @@ +# Copyright 2023 The Google Research Team Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch ALIGN model.""" + +import math +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithNoAttention, + BaseModelOutputWithPooling, + BaseModelOutputWithPoolingAndNoAttention, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_align import AlignConfig, AlignTextConfig, AlignVisionConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + """ +) +class AlignVisionModelOutput(ModelOutput): + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + """ + + image_embeds: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for text model's outputs that also contains a pooling of the last hidden states. + """ +) +class AlignTextModelOutput(ModelOutput): + r""" + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + """ + + text_embeds: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring +class AlignOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`AlignTextModel`]. + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The output of [`AlignVisionModel`]. + text_model_output (`BaseModelOutputWithPooling`): + The output of the [`AlignTextModel`]. + vision_model_output (`BaseModelOutputWithPoolingAndNoAttention`): + The output of the [`AlignVisionModel`]. + """ + + loss: torch.FloatTensor | None = None + logits_per_image: torch.FloatTensor | None = None + logits_per_text: torch.FloatTensor | None = None + text_embeds: torch.FloatTensor | None = None + image_embeds: torch.FloatTensor | None = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPoolingAndNoAttention = None + + def to_tuple(self) -> tuple[Any]: + return tuple( + self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple() + for k in self.keys() + ) + + +# contrastive loss function, adapted from +# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html +def contrastive_loss(logits: torch.Tensor) -> torch.Tensor: + return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device), label_smoothing=0.1) + + +def align_loss(similarity: torch.Tensor) -> torch.Tensor: + caption_loss = contrastive_loss(similarity) + image_loss = contrastive_loss(similarity.t()) + return (caption_loss + image_loss) / 2.0 + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.round_filters with EfficientNet->AlignVision +def round_filters(config: AlignVisionConfig, num_channels: int): + r""" + Round number of filters based on depth multiplier. + """ + divisor = config.depth_divisor + num_channels *= config.width_coefficient + new_dim = max(divisor, int(num_channels + divisor / 2) // divisor * divisor) + + # Make sure that round down does not go down by more than 10%. + if new_dim < 0.9 * num_channels: + new_dim += divisor + + return int(new_dim) + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.correct_pad +def correct_pad(kernel_size: int | tuple, adjust: bool = True): + r""" + Utility function to get the tuple padding value for the depthwise convolution. + + Args: + kernel_size (`int` or `tuple`): + Kernel size of the convolution layers. + adjust (`bool`, *optional*, defaults to `True`): + Adjusts padding value to apply to right and bottom sides of the input. + """ + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + + correct = (kernel_size[0] // 2, kernel_size[1] // 2) + if adjust: + return (correct[1] - 1, correct[1], correct[0] - 1, correct[0]) + else: + return (correct[1], correct[1], correct[0], correct[0]) + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.EfficientNetEmbeddings with EfficientNet->AlignVision +class AlignVisionEmbeddings(nn.Module): + r""" + A module that corresponds to the stem module of the original work. + """ + + def __init__(self, config: AlignVisionConfig): + super().__init__() + + self.out_dim = round_filters(config, 32) + self.padding = nn.ZeroPad2d(padding=(0, 1, 0, 1)) + self.convolution = nn.Conv2d( + config.num_channels, self.out_dim, kernel_size=3, stride=2, padding="valid", bias=False + ) + self.batchnorm = nn.BatchNorm2d(self.out_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum) + self.activation = ACT2FN[config.hidden_act] + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + features = self.padding(pixel_values) + features = self.convolution(features) + features = self.batchnorm(features) + features = self.activation(features) + + return features + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.EfficientNetDepthwiseConv2d with EfficientNet->AlignVision +class AlignVisionDepthwiseConv2d(nn.Conv2d): + def __init__( + self, + in_channels, + depth_multiplier=1, + kernel_size=3, + stride=1, + padding=0, + dilation=1, + bias=True, + padding_mode="zeros", + ): + out_channels = in_channels * depth_multiplier + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=in_channels, + bias=bias, + padding_mode=padding_mode, + ) + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.EfficientNetExpansionLayer with EfficientNet->AlignVision +class AlignVisionExpansionLayer(nn.Module): + r""" + This corresponds to the expansion phase of each block in the original implementation. + """ + + def __init__(self, config: AlignVisionConfig, in_dim: int, out_dim: int, stride: int): + super().__init__() + self.expand_conv = nn.Conv2d( + in_channels=in_dim, + out_channels=out_dim, + kernel_size=1, + padding="same", + bias=False, + ) + self.expand_bn = nn.BatchNorm2d(num_features=out_dim, eps=config.batch_norm_eps) + self.expand_act = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor: + # Expand phase + hidden_states = self.expand_conv(hidden_states) + hidden_states = self.expand_bn(hidden_states) + hidden_states = self.expand_act(hidden_states) + + return hidden_states + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.EfficientNetDepthwiseLayer with EfficientNet->AlignVision +class AlignVisionDepthwiseLayer(nn.Module): + r""" + This corresponds to the depthwise convolution phase of each block in the original implementation. + """ + + def __init__( + self, + config: AlignVisionConfig, + in_dim: int, + stride: int, + kernel_size: int, + adjust_padding: bool, + ): + super().__init__() + self.stride = stride + conv_pad = "valid" if self.stride == 2 else "same" + padding = correct_pad(kernel_size, adjust=adjust_padding) + + self.depthwise_conv_pad = nn.ZeroPad2d(padding=padding) + self.depthwise_conv = AlignVisionDepthwiseConv2d( + in_dim, kernel_size=kernel_size, stride=stride, padding=conv_pad, bias=False + ) + self.depthwise_norm = nn.BatchNorm2d( + num_features=in_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum + ) + self.depthwise_act = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor: + # Depthwise convolution + if self.stride == 2: + hidden_states = self.depthwise_conv_pad(hidden_states) + + hidden_states = self.depthwise_conv(hidden_states) + hidden_states = self.depthwise_norm(hidden_states) + hidden_states = self.depthwise_act(hidden_states) + + return hidden_states + + +# Copied from transformers.models.efficientnet.modeling_efficientnet.EfficientNetSqueezeExciteLayer with EfficientNet->AlignVision +class AlignVisionSqueezeExciteLayer(nn.Module): + r""" + This corresponds to the Squeeze and Excitement phase of each block in the original implementation. + """ + + def __init__(self, config: AlignVisionConfig, in_dim: int, expand_dim: int, expand: bool = False): + super().__init__() + self.dim = expand_dim if expand else in_dim + self.dim_se = max(1, int(in_dim * config.squeeze_expansion_ratio)) + + self.squeeze = nn.AdaptiveAvgPool2d(output_size=1) + self.reduce = nn.Conv2d( + in_channels=self.dim, + out_channels=self.dim_se, + kernel_size=1, + padding="same", + ) + self.expand = nn.Conv2d( + in_channels=self.dim_se, + out_channels=self.dim, + kernel_size=1, + padding="same", + ) + self.act_reduce = ACT2FN[config.hidden_act] + self.act_expand = nn.Sigmoid() + + def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor: + inputs = hidden_states + hidden_states = self.squeeze(hidden_states) + hidden_states = self.reduce(hidden_states) + hidden_states = self.act_reduce(hidden_states) + + hidden_states = self.expand(hidden_states) + hidden_states = self.act_expand(hidden_states) + hidden_states = torch.mul(inputs, hidden_states) + + return hidden_states + + +class AlignVisionFinalBlockLayer(nn.Module): + r""" + This corresponds to the final phase of each block in the original implementation. + """ + + def __init__( + self, config: AlignVisionConfig, in_dim: int, out_dim: int, stride: int, drop_rate: float, id_skip: bool + ): + super().__init__() + self.apply_dropout = stride == 1 and not id_skip + self.project_conv = nn.Conv2d( + in_channels=in_dim, + out_channels=out_dim, + kernel_size=1, + padding="same", + bias=False, + ) + self.project_bn = nn.BatchNorm2d( + num_features=out_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum + ) + self.dropout = nn.Dropout(p=drop_rate) + + def forward(self, embeddings: torch.FloatTensor, hidden_states: torch.FloatTensor) -> torch.Tensor: + hidden_states = self.project_conv(hidden_states) + hidden_states = self.project_bn(hidden_states) + + if self.apply_dropout: + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + embeddings + + return hidden_states + + +class AlignVisionBlock(nn.Module): + r""" + This corresponds to the block module of original the EfficientNet vision encoder implementation. + + Args: + config ([`AlignVisionConfig`]): + Model configuration class. + in_dim (`int`): + Number of input channels. + out_dim (`int`): + Number of output channels. + stride (`int`): + Stride size to be used in convolution layers. + expand_ratio (`int`): + Expand ratio to set the output dimensions for the expansion and squeeze-excite layers. + kernel_size (`int`): + Kernel size for the depthwise convolution layer. + drop_rate (`float`): + Dropout rate to be used in the final phase of each block. + id_skip (`bool`): + Whether to apply dropout and sum the final hidden states with the input embeddings during the final phase + of each block. Set to `True` for the first block of each stage. + adjust_padding (`bool`): + Whether to apply padding to only right and bottom side of the input kernel before the depthwise convolution + operation, set to `True` for inputs with odd input sizes. + """ + + def __init__( + self, + config: AlignVisionConfig, + in_dim: int, + out_dim: int, + stride: int, + expand_ratio: int, + kernel_size: int, + drop_rate: float, + id_skip: bool, + adjust_padding: bool, + ): + super().__init__() + self.expand_ratio = expand_ratio + self.expand = self.expand_ratio != 1 + expand_in_dim = in_dim * expand_ratio + + if self.expand: + self.expansion = AlignVisionExpansionLayer( + config=config, in_dim=in_dim, out_dim=expand_in_dim, stride=stride + ) + + self.depthwise_conv = AlignVisionDepthwiseLayer( + config=config, + in_dim=expand_in_dim if self.expand else in_dim, + stride=stride, + kernel_size=kernel_size, + adjust_padding=adjust_padding, + ) + self.squeeze_excite = AlignVisionSqueezeExciteLayer( + config=config, in_dim=in_dim, expand_dim=expand_in_dim, expand=self.expand + ) + self.projection = AlignVisionFinalBlockLayer( + config=config, + in_dim=expand_in_dim if self.expand else in_dim, + out_dim=out_dim, + stride=stride, + drop_rate=drop_rate, + id_skip=id_skip, + ) + + def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor: + embeddings = hidden_states + # Expansion and depthwise convolution phase + if self.expand_ratio != 1: + hidden_states = self.expansion(hidden_states) + hidden_states = self.depthwise_conv(hidden_states) + + # Squeeze and excite phase + hidden_states = self.squeeze_excite(hidden_states) + hidden_states = self.projection(embeddings, hidden_states) + return hidden_states + + +class AlignVisionEncoder(nn.Module): + r""" + Forward propagates the embeddings through each vision encoder (EfficientNet) block. + + Args: + config ([`AlignVisionConfig`]): + Model configuration class. + """ + + def __init__(self, config: AlignVisionConfig): + super().__init__() + self.depth_coefficient = config.depth_coefficient + + def round_repeats(repeats): + # Round number of block repeats based on depth multiplier. + return int(math.ceil(self.depth_coefficient * repeats)) + + num_base_blocks = len(config.in_channels) + num_blocks = sum(round_repeats(n) for n in config.num_block_repeats) + + curr_block_num = 0 + blocks = [] + for i in range(num_base_blocks): + in_dim = round_filters(config, config.in_channels[i]) + out_dim = round_filters(config, config.out_channels[i]) + stride = config.strides[i] + kernel_size = config.kernel_sizes[i] + expand_ratio = config.expand_ratios[i] + + for j in range(round_repeats(config.num_block_repeats[i])): + id_skip = j == 0 + stride = 1 if j > 0 else stride + in_dim = out_dim if j > 0 else in_dim + adjust_padding = curr_block_num not in config.depthwise_padding + drop_rate = config.drop_connect_rate * curr_block_num / num_blocks + + block = AlignVisionBlock( + config=config, + in_dim=in_dim, + out_dim=out_dim, + stride=stride, + kernel_size=kernel_size, + expand_ratio=expand_ratio, + drop_rate=drop_rate, + id_skip=id_skip, + adjust_padding=adjust_padding, + ) + blocks.append(block) + curr_block_num += 1 + + self.blocks = nn.ModuleList(blocks) + + def forward( + self, + hidden_states: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithNoAttention: + for block in self.blocks: + hidden_states = block(hidden_states) + + return BaseModelOutputWithNoAttention( + last_hidden_state=hidden_states, + ) + + +class AlignTextEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + self.register_buffer( + "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs + # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves + # issue #5664 + if token_type_ids is None: + if hasattr(self, "token_type_ids"): + buffered_token_type_ids = self.token_type_ids[:, :seq_length] + buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length) + token_type_ids = buffered_token_type_ids_expanded + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = inputs_embeds + token_type_embeddings + + position_embeddings = self.position_embeddings(position_ids) + embeddings += position_embeddings + + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs, +): + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights + + +class AlignTextSelfAttention(nn.Module): + def __init__(self, config): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + + self.config = config + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + self.attention_dropout = config.attention_probs_dropout_prob + self.scaling = self.attention_head_size**-0.5 + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.attention_head_size) + + query_states = self.query(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.key(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.value(hidden_states).view(hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return attn_output, attn_weights + + +# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->AlignText +class AlignTextSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class AlignTextAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = AlignTextSelfAttention(config) + self.output = AlignTextSelfOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states, _ = self.self( + hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = self.output(hidden_states, residual) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->AlignText +class AlignTextIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->AlignText +class AlignTextOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class AlignTextLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = AlignTextAttention(config) + self.intermediate = AlignTextIntermediate(config) + self.output = AlignTextOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + hidden_states = self.attention( + hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + + hidden_states = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, hidden_states + ) + + return hidden_states + + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class AlignTextEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([AlignTextLayer(config) for i in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + for layer_module in self.layer: + hidden_states = layer_module( + hidden_states, + attention_mask, + **kwargs, + ) + + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert -> AlignText +class AlignTextPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +@auto_docstring +class AlignPreTrainedModel(PreTrainedModel): + config: AlignConfig + base_model_prefix = "align" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + """Initialize the weights""" + std = self.config.initializer_range + if isinstance(module, (nn.Linear, nn.Conv2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, AlignModel): + init.xavier_uniform_(module.text_projection.weight) + init.zeros_(module.text_projection.bias) + init.constant_(module.temperature, self.config.temperature_init_value) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + if isinstance(module, (nn.LayerNorm, nn.BatchNorm2d)): + init.zeros_(module.bias) + init.ones_(module.weight) + if getattr(module, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + elif isinstance(module, AlignTextEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + init.zeros_(module.token_type_ids) + + +@auto_docstring( + custom_intro=""" + The text model from ALIGN without any head or projection on top. + """ +) +class AlignTextModel(AlignPreTrainedModel): + config: AlignTextConfig + input_modalities = ("text",) + _no_split_modules = ["AlignTextEmbeddings"] + _can_record_outputs = { + "hidden_states": AlignTextLayer, + "attentions": AlignTextSelfAttention, + } + + def __init__(self, config: AlignTextConfig, add_pooling_layer: bool = True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + self.embeddings = AlignTextEmbeddings(config) + self.encoder = AlignTextEncoder(config) + + self.pooler = AlignTextPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from transformers import AutoTokenizer, AlignTextModel + + >>> model = AlignTextModel.from_pretrained("kakaobrain/align-base") + >>> tokenizer = AutoTokenizer.from_pretrained("kakaobrain/align-base") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = torch.ones(((batch_size, seq_length)), device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + **kwargs, + ) + sequence_output = encoder_outputs[0] + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + ) + + +@auto_docstring( + custom_intro=""" + The vision model from ALIGN without any head or projection on top. + """ +) +class AlignVisionModel(AlignPreTrainedModel): + config: AlignVisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = False + _input_embed_layer = "convolution" + _no_split_modules = ["AlignVisionBlock"] + _can_record_outputs = { + "hidden_states": AlignVisionBlock, + } + + def __init__(self, config: AlignVisionConfig): + super().__init__(config) + self.config = config + self.embeddings = AlignVisionEmbeddings(config) + self.encoder = AlignVisionEncoder(config) + + # Final pooling layer + if config.pooling_type == "mean": + self.pooler = nn.AvgPool2d(config.hidden_dim, ceil_mode=True) + elif config.pooling_type == "max": + self.pooler = nn.MaxPool2d(config.hidden_dim, ceil_mode=True) + else: + raise ValueError(f"config.pooling must be one of ['mean', 'max'] got {config.pooling}") + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPoolingAndNoAttention: + r""" + Examples: + + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, AlignVisionModel + + >>> model = AlignVisionModel.from_pretrained("kakaobrain/align-base") + >>> processor = AutoProcessor.from_pretrained("kakaobrain/align-base") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled CLS states + ```""" + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + embedding_output = self.embeddings(pixel_values) + encoder_outputs = self.encoder( + embedding_output, + **kwargs, + ) + last_hidden_state = encoder_outputs[0] + pooled_output = self.pooler(last_hidden_state) + pooled_output = pooled_output.reshape(pooled_output.shape[:2]) + + return BaseModelOutputWithPoolingAndNoAttention( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + ) + + +@auto_docstring +class AlignModel(AlignPreTrainedModel): + config: AlignConfig + + def __init__(self, config: AlignConfig): + super().__init__(config) + + if not isinstance(config.text_config, AlignTextConfig): + raise TypeError( + "config.text_config is expected to be of type AlignTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, AlignVisionConfig): + raise TypeError( + "config.vision_config is expected to be of type AlignVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + + self.text_model = AlignTextModel(text_config) + self.vision_model = AlignVisionModel(vision_config) + + self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim) + self.temperature = nn.Parameter(torch.tensor(self.config.temperature_init_value)) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, AlignModel + + >>> model = AlignModel.from_pretrained("kakaobrain/align-base") + >>> tokenizer = AutoTokenizer.from_pretrained("kakaobrain/align-base") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + >>> with torch.inference_mode(): + ... text_features = model.get_text_features(**inputs) + ```""" + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + last_hidden_state = text_outputs[0][:, 0, :] + text_outputs.pooler_output = self.text_projection(last_hidden_state) + + return text_outputs + + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, AlignModel + >>> from transformers.image_utils import load_image + + >>> model = AlignModel.from_pretrained("kakaobrain/align-base") + >>> processor = AutoProcessor.from_pretrained("kakaobrain/align-base") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = load_image(url) + + >>> inputs = processor(images=image, return_tensors="pt") + >>> with torch.inference_mode(): + ... image_features = model.get_image_features(**inputs) + ```""" + return self.vision_model(pixel_values=pixel_values, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + return_loss: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | AlignOutput: + r""" + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + + Examples: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, AlignModel + >>> from transformers.image_utils import load_image + + >>> model = AlignModel.from_pretrained("kakaobrain/align-base") + >>> processor = AutoProcessor.from_pretrained("kakaobrain/align-base") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = load_image(url) + + >>> inputs = processor( + ... images=image, text=["a photo of a cat", "a photo of a dog"], return_tensors="pt", padding=True + ... ) + + >>> with torch.inference_mode(): + ... outputs = model(**inputs) + >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score + >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities + ```""" + vision_outputs = self.vision_model( + pixel_values=pixel_values, + **kwargs, + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + image_embeds = vision_outputs[1] + text_embeds = text_outputs[0][:, 0, :] + text_embeds = self.text_projection(text_embeds) + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logits_per_text = torch.matmul(text_embeds, image_embeds.t()) / self.temperature + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = align_loss(logits_per_text) + + return AlignOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +__all__ = ["AlignPreTrainedModel", "AlignTextModel", "AlignVisionModel", "AlignModel"] diff --git a/third_party/transformers/src/transformers/models/align/processing_align.py b/third_party/transformers/src/transformers/models/align/processing_align.py new file mode 100644 index 0000000000000000000000000000000000000000..fa15fcce3de6cc52da5f37f541ffa0fa6cbcd651 --- /dev/null +++ b/third_party/transformers/src/transformers/models/align/processing_align.py @@ -0,0 +1,40 @@ +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Image/Text processor class for ALIGN +""" + +from ...processing_utils import ProcessingKwargs, ProcessorMixin +from ...utils import auto_docstring + + +class AlignProcessorKwargs(ProcessingKwargs, total=False): + # see processing_utils.ProcessingKwargs documentation for usage. + _defaults = { + "text_kwargs": { + "padding": "max_length", + "max_length": 64, + }, + } + + +@auto_docstring +class AlignProcessor(ProcessorMixin): + valid_processor_kwargs = AlignProcessorKwargs + + def __init__(self, image_processor, tokenizer): + super().__init__(image_processor, tokenizer) + + +__all__ = ["AlignProcessor"] diff --git a/third_party/transformers/src/transformers/models/barthez/__init__.py b/third_party/transformers/src/transformers/models/barthez/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9e11571fc6d49eb75f8a1779d522093af269a00 --- /dev/null +++ b/third_party/transformers/src/transformers/models/barthez/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .tokenization_barthez import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/barthez/tokenization_barthez.py b/third_party/transformers/src/transformers/models/barthez/tokenization_barthez.py new file mode 100644 index 0000000000000000000000000000000000000000..c356f25cf3106dcc8e4c5002e8fe595996e50073 --- /dev/null +++ b/third_party/transformers/src/transformers/models/barthez/tokenization_barthez.py @@ -0,0 +1,144 @@ +# Copyright 2020 Ecole Polytechnique and the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License +"""Tokenization classes for the BARThez model.""" + +from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers +from tokenizers.models import Unigram + +from ...tokenization_python import AddedToken +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"} + + +SPIECE_UNDERLINE = "▁" + + +class BarthezTokenizer(TokenizersBackend): + """ + Adapted from [`CamembertTokenizer`] and [`BartTokenizer`]. Construct a "fast" BARThez tokenizer. Based on + [SentencePiece](https://github.com/google/sentencepiece). + + This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + + + + When building a sequence using special tokens, this is not the token that is used for the beginning of + sequence. The token used is the `cls_token`. + + + + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + + + + When building a sequence using special tokens, this is not the token that is used for the end of sequence. + The token used is the `sep_token`. + + + + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + vocab_file (`str`, *optional*): + [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that + contains the vocabulary necessary to instantiate a tokenizer. + vocab (`str`, `dict` or `list`, *optional*): + Custom vocabulary dictionary. If not provided, vocabulary is loaded from vocab_file. + add_prefix_space (`bool`, *optional*, defaults to `True`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + slow_tokenizer_class = None + + def __init__( + self, + vocab: str | dict | list | None = None, + bos_token="", + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + add_prefix_space=True, + **kwargs, + ): + # Mask token behave like a normal word, i.e. include the space before it + mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token + self.add_prefix_space = add_prefix_space + + if vocab is not None: + self._vocab = vocab + else: + self._vocab = [ + (str(pad_token), 0.0), + (str(unk_token), 0.0), + (str(cls_token), 0.0), + (str(sep_token), 0.0), + (str(mask_token), 0.0), + ] + + self._tokenizer = Tokenizer(Unigram(self._vocab, unk_id=3, byte_fallback=False)) + + self._tokenizer.normalizer = normalizers.Sequence( + [ + normalizers.Replace(Regex(r"\s{2,}|[\n\r\t]"), " "), + normalizers.NFC(), + normalizers.Strip(left=False, right=True), + ] + ) + prepend_scheme = "always" if add_prefix_space else "never" + self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme=prepend_scheme) + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme=prepend_scheme) + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + sep_token=sep_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + add_prefix_space=add_prefix_space, + **kwargs, + ) + + +__all__ = ["BarthezTokenizer"] diff --git a/third_party/transformers/src/transformers/models/colpali/__init__.py b/third_party/transformers/src/transformers/models/colpali/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fa1b63fd009803e27cc25ebe591ab4fa2cd32cb9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/__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_colpali import * + from .modeling_colpali import * + from .processing_colpali import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/colpali/configuration_colpali.py b/third_party/transformers/src/transformers/models/colpali/configuration_colpali.py new file mode 100644 index 0000000000000000000000000000000000000000..5cff87c9559906d11e93b691390a2faa8232b88b --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/configuration_colpali.py @@ -0,0 +1,64 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ColPali model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging +from ..auto import CONFIG_MAPPING, AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="vidore/colpali-v1.2") +@strict +class ColPaliConfig(PreTrainedConfig): + r""" + Example: + + ```python + from transformers.models.colpali import ColPaliConfig, ColPaliForRetrieval + + config = ColPaliConfig() + model = ColPaliForRetrieval(config) + ``` + """ + + model_type = "colpali" + sub_configs = {"vlm_config": PreTrainedConfig, "text_config": AutoConfig} + + vlm_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + embedding_dim: int = 128 + + def __post_init__(self, **kwargs): + if self.vlm_config is None: + self.vlm_config = CONFIG_MAPPING["paligemma"]() + logger.info( + "`vlm_config` is `None`. Initializing `vlm_config` with the `PaliGemmaConfig` with default values." + ) + elif isinstance(self.vlm_config, dict): + self.vlm_config = CONFIG_MAPPING[self.vlm_config["model_type"]](**self.vlm_config) + + self.text_config = self.text_config if self.text_config is not None else self.vlm_config.text_config + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "gemma") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + + super().__post_init__(**kwargs) + + +__all__ = ["ColPaliConfig"] diff --git a/third_party/transformers/src/transformers/models/colpali/convert_colpali_weights_to_hf.py b/third_party/transformers/src/transformers/models/colpali/convert_colpali_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..48664ab143be404cd2822d99b5db443572f63345 --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/convert_colpali_weights_to_hf.py @@ -0,0 +1,221 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Convert ColPali weights from the original repository to the HF model format. + +Original repository: https://github.com/illuin-tech/colpali. + +NOTE: This script was originally run using `torch==2.5.1` and with: + +```bash +python src/transformers/models/colpali/convert_colpali_weights_to_hf.py \ + --model_id vidore/colpali-v1.2-merged \ + --revision 89fd9736194236a1ecb7a9ec9b04f537f6f896af \ + --original_vlm_name_or_path google/paligemma-3b-mix-448 \ + --output_dir vidore/colpali-v1.2-hf-internal \ + --push_to_hub + +python src/transformers/models/colpali/convert_colpali_weights_to_hf.py \ + --model_id vidore/colpali-v1.3-merged \ + --revision 5b955e3415a7c5468ab33119d98d6d45c3a5b2c3 \ + --original_vlm_name_or_path google/paligemma-3b-mix-448 \ + --output_dir vidore/colpali-v1.3-hf \ + --push_to_hub +``` +""" + +import argparse +import glob +from pathlib import Path +from typing import Any + +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open + +from transformers import AutoConfig +from transformers.models.colpali import ColPaliForRetrieval +from transformers.models.colpali.configuration_colpali import ColPaliConfig +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +ORIGINAL_DTYPE = torch.bfloat16 + + +def rename_state_dict_keys(state_dict: dict[str, Any]) -> dict[str, Any]: + new_state_dict = {} + for key, value in state_dict.items(): + new_key = key + if key.startswith("custom_text_proj"): + new_key = key.replace("custom_text_proj", "embedding_proj_layer") + if key.startswith("model."): + new_key = key.replace("model.", "vlm.", 1) + new_state_dict[new_key] = value + return new_state_dict + + +def load_original_state_dict(model_id: str, revision: str | None = None) -> dict[str, torch.Tensor]: + directory_path = snapshot_download( + repo_id=model_id, + revision=revision, + allow_patterns=["*.safetensors"], + ) + + original_state_dict = {} + for path in glob.glob(f"{directory_path}/*"): + if path.endswith(".safetensors"): + with safe_open(path, framework="pt", device="cpu") as f: + for key in f.keys(): + original_state_dict[key] = f.get_tensor(key) + + # Some weights are tied, so `lm.head`` is not saved. Let's clone to load state dict. + if "lm_head.weight" not in original_state_dict: + original_state_dict["vlm.language_model.lm_head.weight"] = original_state_dict[ + "model.language_model.model.embed_tokens.weight" + ].clone() + + return original_state_dict + + +@torch.no_grad() +def convert_colpali_weights_to_hf( + model_id: str, + output_dir: str, + push_to_hub: bool, + revision: str | None = None, + original_vlm_name_or_path: str | None = None, +): + # Load the original model data + original_config = AutoConfig.from_pretrained( + model_id, + revision=revision, + ) + if original_vlm_name_or_path is not None: + original_config._name_or_path = original_vlm_name_or_path + if hasattr(original_config, "architectures"): + delattr(original_config, "architectures") + + original_state_dict = load_original_state_dict(model_id, revision=revision) + + # Format the state_dict keys + original_state_dict = rename_state_dict_keys(original_state_dict) + + # Create the new config + config = ColPaliConfig( + vlm_config=original_config, + embedding_dim=128, # hardcoded in the original model + ) + config.model_type = "colpali" + config.is_composition = False + + # Load the untrained model + model = ColPaliForRetrieval(config=config).to("cpu").eval() + print("Created model with new config and randomly initialized weights") + + # NOTE: The model was initialized with float32 weights. We need to convert it to the desired precision. + # There are two ways to set the model's dtype: + # - Using `model.from_pretrained(..., dtype=dtype_precision)` doesn't convert the hyperparameters to the desired precision. + # - Using `model.to(dtype_precision)` converts all values - including the hyperparameters - to the desired precision. + # The following snippet allows a fine-grained control over the model's dtype, making sure that all + # the new weights' dtypes match the original model. + for param in model.parameters(): + param.data = param.data.to(ORIGINAL_DTYPE) + print(f"Converted the new model weights to `{ORIGINAL_DTYPE}`") + + # Load the original weights + model.load_state_dict(original_state_dict) + print("Loaded original model weights") + + # Tie the weights (following ColPali's `__init__`` step) + if model.vlm.language_model._tied_weights_keys is not None: + prefix = "vlm.language_model." + prefixed_mapping = { + f"{prefix}{target}": f"{prefix}{source}" + for target, source in model.vlm.language_model._tied_weights_keys.items() + } + if isinstance(model._tied_weights_keys, dict): + model._tied_weights_keys.update(prefixed_mapping) + else: + model._tied_weights_keys = prefixed_mapping + + # Sanity check: ensure all keys are the same + state_dict_keys_old = set(original_state_dict.keys()) + state_dict_keys_new = set(model.state_dict().keys()) + disjoint_keys = state_dict_keys_old.symmetric_difference(state_dict_keys_new) + if disjoint_keys: + raise ValueError(f"Incompatible keys: {disjoint_keys}") + + # Save the model + if push_to_hub: + model.push_to_hub(output_dir, private=True) + print(f"Model pushed to the hub at `{output_dir}`") + else: + Path(output_dir).mkdir(exist_ok=True, parents=True) + model.save_pretrained(output_dir) + print(f"Model saved to `{output_dir}`") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description=""" + This script converts the original ColPali model to the HF model format. + + Example usage: + ```bash + python src/transformers/models/colpali/convert_colpali_weights_to_hf.py \ + --model_id vidore/colpali-v1.2-merged \ + --revision 89fd9736194236a1ecb7a9ec9b04f537f6f896af \ + --original_vlm_name_or_path google/paligemma-3b-mix-448 \ + --output_dir vidore/colpali-v1.2-hf \ + --push_to_hub + ``` + """ + ) + parser.add_argument( + "--model_id", + help="Model ID of the original model to convert", + ) + parser.add_argument( + "--output_dir", + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--push_to_hub", + help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally", + action="store_true", + default=False, + ) + parser.add_argument( + "--revision", + help="Revision of the model to download", + default=None, + ) + parser.add_argument( + "--original_vlm_name_or_path", + help="Name or path of the original VLM backbone model", + default=None, + ) + args = parser.parse_args() + + convert_colpali_weights_to_hf( + model_id=args.model_id, + output_dir=args.output_dir, + push_to_hub=args.push_to_hub, + revision=args.revision, + original_vlm_name_or_path=args.original_vlm_name_or_path, + ) diff --git a/third_party/transformers/src/transformers/models/colpali/modeling_colpali.py b/third_party/transformers/src/transformers/models/colpali/modeling_colpali.py new file mode 100644 index 0000000000000000000000000000000000000000..dc7dcc306d335fbcf4c8771eb46eaddbeb284d6b --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/modeling_colpali.py @@ -0,0 +1,169 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch ColPali model""" + +from dataclasses import dataclass + +import torch +from torch import nn + +from transformers import AutoModel + +from ... import initialization as init +from ...cache_utils import Cache +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple +from .configuration_colpali import ColPaliConfig + + +@auto_docstring +class ColPaliPreTrainedModel(PreTrainedModel): + config: ColPaliConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + _no_split_modules = [] + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + + @torch.no_grad() + def _init_weights(self, module): + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.vlm_config.text_config.initializer_range + ) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for ColPali embeddings output. + """ +) +class ColPaliForRetrievalOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + The embeddings of the model. + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder after projecting last hidden state. + """ + + loss: torch.FloatTensor | None = None + embeddings: torch.Tensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + The ColPali architecture leverages VLMs to construct efficient multi-vector embeddings directly + from document images (“screenshots”) for document retrieval. The model is trained to maximize the similarity + between these document embeddings and the corresponding query embeddings, using the late interaction method + introduced in ColBERT. + + Using ColPali removes the need for potentially complex and brittle layout recognition and OCR pipelines with a + single model that can take into account both the textual and visual content (layout, charts, etc.) of a document. + + ColPali is part of the ColVision model family, which was first introduced in the following paper: + [*ColPali: Efficient Document Retrieval with Vision Language Models*](https://huggingface.co/papers/2407.01449). + """ +) +class ColPaliForRetrieval(ColPaliPreTrainedModel): + base_model_prefix = "vlm" + + def __init__(self, config: ColPaliConfig): + super().__init__(config) + self.config = config + self.vocab_size = config.vlm_config.text_config.vocab_size + + self.vlm = AutoModel.from_config(config.vlm_config) + + self.embedding_dim = self.config.embedding_dim + self.embedding_proj_layer = nn.Linear( + self.config.vlm_config.text_config.hidden_size, + self.embedding_dim, + ) + + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ColPaliForRetrievalOutput: + if pixel_values is not None: + pixel_values = pixel_values.to(dtype=self.dtype) + output_hidden_states = kwargs.pop("output_hidden_states", None) + if output_hidden_states is None: + output_hidden_states = self.config.output_hidden_states + + vlm_output = self.vlm( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + output_hidden_states=True, + **kwargs, + ) + vlm_hidden_states = vlm_output.hidden_states if output_hidden_states else None + vlm_image_hidden_states = vlm_output.image_hidden_states if pixel_values is not None else None + + last_hidden_states = vlm_output[0] # (batch_size, sequence_length, hidden_size) + proj_dtype = self.embedding_proj_layer.weight.dtype + embeddings = self.embedding_proj_layer(last_hidden_states.to(proj_dtype)) # (batch_size, sequence_length, dim) + + # L2 normalization + embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True) # (batch_size, sequence_length, dim) + + if attention_mask is not None: + embeddings = embeddings * attention_mask.unsqueeze(-1) # (batch_size, sequence_length, dim) + + return ColPaliForRetrievalOutput( + embeddings=embeddings, + past_key_values=vlm_output.past_key_values, + hidden_states=vlm_hidden_states, + attentions=vlm_output.attentions, + image_hidden_states=vlm_image_hidden_states, + ) + + +__all__ = [ + "ColPaliForRetrieval", + "ColPaliPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/colpali/modular_colpali.py b/third_party/transformers/src/transformers/models/colpali/modular_colpali.py new file mode 100644 index 0000000000000000000000000000000000000000..dec2fee64dea644741529271e62cf3444f9f7d22 --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/modular_colpali.py @@ -0,0 +1,296 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import Optional, Union + +from transformers.models.paligemma.processing_paligemma import IMAGE_TOKEN, PaliGemmaProcessor, build_string_from_input + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, make_flat_list_of_images +from ...processing_utils import ProcessingKwargs, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import is_torch_available, logging + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class ColPaliProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": "longest", + }, + "images_kwargs": { + "data_format": "channels_first", + "do_convert_rgb": True, + }, + "common_kwargs": {"return_tensors": "pt"}, + } + + +class ColPaliProcessor(PaliGemmaProcessor): + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + visual_prompt_prefix: str = "Describe the image.", + query_prefix: str = "Question: ", + ): + r""" + visual_prompt_prefix (`str`, *optional*, defaults to `"Describe the image."`): + A string that gets tokenized and prepended to the image tokens. + query_prefix (`str`, *optional*, defaults to `"Question: "`): + A prefix to be used for the query. + """ + self.visual_prompt_prefix = visual_prompt_prefix + self.query_prefix = query_prefix + super().__init__(image_processor=image_processor, tokenizer=tokenizer, chat_template=chat_template) + + @property + def query_augmentation_token(self) -> str: + """ + Return the query augmentation token. + + Query augmentation buffers are used as reasoning buffers during inference. + """ + return self.tokenizer.pad_token + + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + output_kwargs = self._merge_kwargs( + ColPaliProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + suffix = output_kwargs["text_kwargs"].pop("suffix", None) + + return_token_type_ids = True + + if text is None and images is None: + raise ValueError("Either text or images must be provided") + if text is not None and images is not None: + raise ValueError("Only one of text or images can be processed at a time") + + if images is not None: + images = self.image_processor.fetch_images(images) + images = make_flat_list_of_images(images) + texts_doc = [self.visual_prompt_prefix] * len(images) + images = [image.convert("RGB") for image in images] + + input_strings = [ + build_string_from_input( + prompt=prompt, + bos_token=self.tokenizer.bos_token, + image_seq_len=self.image_seq_length, + image_token=IMAGE_TOKEN, + num_images=len(image_list) if isinstance(image_list, list) else 1, + ) + for prompt, image_list in zip(texts_doc, images) + ] + pixel_values = self.image_processor(images, **output_kwargs["images_kwargs"])["pixel_values"] + + # max_length has to account for the image tokens + if output_kwargs["text_kwargs"].get("max_length", None) is not None: + output_kwargs["text_kwargs"]["max_length"] += self.image_seq_length + + inputs = self.tokenizer( + input_strings, + return_token_type_ids=return_token_type_ids, + **output_kwargs["text_kwargs"], + ) + + return_data = {**inputs, "pixel_values": pixel_values} + + if return_token_type_ids: + labels = inputs["input_ids"].masked_fill(inputs["token_type_ids"] == 0, -100) + return_data.update({"labels": labels}) + + return BatchFeature(data=return_data) + + elif text is not None: + if isinstance(text, str): + text = [text] + elif not (isinstance(text, list) and isinstance(text[0], str)): + raise ValueError("Text must be a string or a list of strings") + + if suffix is None: + suffix = self.query_augmentation_token * 10 + + texts_query: list[str] = [] + for query in text: + query = self.tokenizer.bos_token + self.query_prefix + query + suffix + "\n" + texts_query.append(query) + + output_kwargs["text_kwargs"]["max_length"] = output_kwargs["text_kwargs"].get("max_length", 50) + + batch_query = self.tokenizer( + texts_query, + return_token_type_ids=return_token_type_ids, + **output_kwargs["text_kwargs"], + ) + + return batch_query + + def process_images( + self, + images: ImageInput | None = None, + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + """ + Prepare for the model one or several image(s). This method is a wrapper around the `__call__` method of the ColPaliProcessor's + [`ColPaliProcessor.__call__`]. + + This method forwards the `images` and `kwargs` arguments to the image processor. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a + number of channels, H and W are image height and width. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + return self.__call__(images=images, **kwargs) + + def process_queries( + self, + text: TextInput | list[TextInput], + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + """ + Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColPaliProcessor's + [`ColPaliProcessor.__call__`]. + + This method forwards the `text` and `kwargs` arguments to the tokenizer. + + Args: + text (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + """ + return self.__call__(text=text, **kwargs) + + def score_retrieval( + self, + query_embeddings: Union["torch.Tensor", list["torch.Tensor"]], + passage_embeddings: Union["torch.Tensor", list["torch.Tensor"]], + batch_size: int = 128, + output_dtype: Optional["torch.dtype"] = None, + output_device: Union["torch.device", str] = "cpu", + ) -> "torch.Tensor": + """ + Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector + query embeddings (`qs`) and passage embeddings (`ps`). For ColPali, a passage is the + image of a document page. + + Because the embedding tensors are multi-vector and can thus have different shapes, they + should be fed as: + (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim) + (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually + obtained by padding the list of tensors. + + Args: + query_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Query embeddings. + passage_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Passage embeddings. + batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores. + output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor. + If `None`, the dtype of the input embeddings is used. + output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor. + + Returns: + `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score + tensor is saved on the "cpu" device. + """ + + if len(query_embeddings) == 0: + raise ValueError("No queries provided") + if len(passage_embeddings) == 0: + raise ValueError("No passages provided") + + if query_embeddings[0].device != passage_embeddings[0].device: + raise ValueError("Queries and passages must be on the same device") + + if query_embeddings[0].dtype != passage_embeddings[0].dtype: + raise ValueError("Queries and passages must have the same dtype") + + if output_dtype is None: + output_dtype = query_embeddings[0].dtype + + scores: list[torch.Tensor] = [] + + for i in range(0, len(query_embeddings), batch_size): + batch_scores: list[torch.Tensor] = [] + batch_queries = torch.nn.utils.rnn.pad_sequence( + query_embeddings[i : i + batch_size], batch_first=True, padding_value=0 + ) + for j in range(0, len(passage_embeddings), batch_size): + batch_passages = torch.nn.utils.rnn.pad_sequence( + passage_embeddings[j : j + batch_size], batch_first=True, padding_value=0 + ) + batch_scores.append( + torch.einsum("bnd,csd->bcns", batch_queries, batch_passages).max(dim=3)[0].sum(dim=2) + ) + scores.append(torch.cat(batch_scores, dim=1).to(output_dtype).to(output_device)) + + return torch.cat(scores, dim=0) + + +__all__ = [ + "ColPaliProcessor", +] diff --git a/third_party/transformers/src/transformers/models/colpali/processing_colpali.py b/third_party/transformers/src/transformers/models/colpali/processing_colpali.py new file mode 100644 index 0000000000000000000000000000000000000000..efdccaee4c1d80adf486557dd72420edaa0655e4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/colpali/processing_colpali.py @@ -0,0 +1,367 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/colpali/modular_colpali.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_colpali.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import Optional, Union + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, make_flat_list_of_images +from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import AddedToken, PreTokenizedInput, TextInput +from ...utils import auto_docstring, is_torch_available + + +if is_torch_available(): + import torch + + +class ColPaliProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": "longest", + }, + "images_kwargs": { + "data_format": "channels_first", + "do_convert_rgb": True, + }, + "common_kwargs": {"return_tensors": "pt"}, + } + + +IMAGE_TOKEN = "" +EXTRA_TOKENS = [f"4}>" for i in range(1024)] + [f"3}>" for i in range(128)] + + +def build_string_from_input(prompt, bos_token, image_seq_len, image_token, num_images): + """ + Builds a string from the input prompt and image tokens. + For example, for the call: + build_string_from_input( + prompt="Prefix str" + bos_token="", + image_seq_len=3, + image_token="", + ) + The output will be: + "Initial str" + Args: + prompt (`list[Union[str, ImageInput]]`): The input prompt. + bos_token (`str`): The beginning of sentence token. + image_seq_len (`int`): The length of the image sequence. + image_token (`str`): The image token. + num_images (`int`): Number of images in the prompt. + """ + return f"{image_token * image_seq_len * num_images}{bos_token}{prompt}\n" + + +@auto_docstring +class ColPaliProcessor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + visual_prompt_prefix: str = "Describe the image.", + query_prefix: str = "Question: ", + ): + r""" + visual_prompt_prefix (`str`, *optional*, defaults to `"Describe the image."`): + A string that gets tokenized and prepended to the image tokens. + query_prefix (`str`, *optional*, defaults to `"Question: "`): + A prefix to be used for the query. + """ + self.visual_prompt_prefix = visual_prompt_prefix + self.query_prefix = query_prefix + if not hasattr(image_processor, "image_seq_length"): + raise ValueError("Image processor is missing an `image_seq_length` attribute.") + + self.image_seq_length = image_processor.image_seq_length + + if not hasattr(tokenizer, "image_token"): + image_token = AddedToken(IMAGE_TOKEN, normalized=False, special=True) + tokens_to_add = {"additional_special_tokens": [image_token]} + tokenizer.add_special_tokens(tokens_to_add) + self.image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN) + self.image_token = IMAGE_TOKEN + else: + self.image_token_id = tokenizer.image_token_id + self.image_token = tokenizer.image_token + + tokenizer.add_tokens(EXTRA_TOKENS) + tokenizer.add_bos_token = False + tokenizer.add_eos_token = False + + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + output_kwargs = self._merge_kwargs( + ColPaliProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + suffix = output_kwargs["text_kwargs"].pop("suffix", None) + + return_token_type_ids = True + + if text is None and images is None: + raise ValueError("Either text or images must be provided") + if text is not None and images is not None: + raise ValueError("Only one of text or images can be processed at a time") + + if images is not None: + images = self.image_processor.fetch_images(images) + images = make_flat_list_of_images(images) + texts_doc = [self.visual_prompt_prefix] * len(images) + images = [image.convert("RGB") for image in images] + + input_strings = [ + build_string_from_input( + prompt=prompt, + bos_token=self.tokenizer.bos_token, + image_seq_len=self.image_seq_length, + image_token=IMAGE_TOKEN, + num_images=len(image_list) if isinstance(image_list, list) else 1, + ) + for prompt, image_list in zip(texts_doc, images) + ] + pixel_values = self.image_processor(images, **output_kwargs["images_kwargs"])["pixel_values"] + + # max_length has to account for the image tokens + if output_kwargs["text_kwargs"].get("max_length", None) is not None: + output_kwargs["text_kwargs"]["max_length"] += self.image_seq_length + + inputs = self.tokenizer( + input_strings, + return_token_type_ids=return_token_type_ids, + **output_kwargs["text_kwargs"], + ) + + return_data = {**inputs, "pixel_values": pixel_values} + + if return_token_type_ids: + labels = inputs["input_ids"].masked_fill(inputs["token_type_ids"] == 0, -100) + return_data.update({"labels": labels}) + + return BatchFeature(data=return_data) + + elif text is not None: + if isinstance(text, str): + text = [text] + elif not (isinstance(text, list) and isinstance(text[0], str)): + raise ValueError("Text must be a string or a list of strings") + + if suffix is None: + suffix = self.query_augmentation_token * 10 + + texts_query: list[str] = [] + for query in text: + query = self.tokenizer.bos_token + self.query_prefix + query + suffix + "\n" + texts_query.append(query) + + output_kwargs["text_kwargs"]["max_length"] = output_kwargs["text_kwargs"].get("max_length", 50) + + batch_query = self.tokenizer( + texts_query, + return_token_type_ids=return_token_type_ids, + **output_kwargs["text_kwargs"], + ) + + return batch_query + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (list[list[str]], *optional*): + The input sizes formatted as (height, width) per each image. + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + vision_data = {} + if image_sizes is not None: + num_image_tokens = [self.image_seq_length] * len(image_sizes) + num_image_patches = [1] * len(image_sizes) + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + return MultiModalData(**vision_data) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + ["token_type_ids", "labels"] + image_processor_input_names = self.image_processor.model_input_names + return list(tokenizer_input_names + image_processor_input_names) + + @property + def query_augmentation_token(self) -> str: + """ + Return the query augmentation token. + + Query augmentation buffers are used as reasoning buffers during inference. + """ + return self.tokenizer.pad_token + + def process_images( + self, + images: ImageInput | None = None, + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + """ + Prepare for the model one or several image(s). This method is a wrapper around the `__call__` method of the ColPaliProcessor's + [`ColPaliProcessor.__call__`]. + + This method forwards the `images` and `kwargs` arguments to the image processor. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a + number of channels, H and W are image height and width. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + return self.__call__(images=images, **kwargs) + + def process_queries( + self, + text: TextInput | list[TextInput], + **kwargs: Unpack[ColPaliProcessorKwargs], + ) -> BatchFeature: + """ + Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColPaliProcessor's + [`ColPaliProcessor.__call__`]. + + This method forwards the `text` and `kwargs` arguments to the tokenizer. + + Args: + text (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + """ + return self.__call__(text=text, **kwargs) + + def score_retrieval( + self, + query_embeddings: Union["torch.Tensor", list["torch.Tensor"]], + passage_embeddings: Union["torch.Tensor", list["torch.Tensor"]], + batch_size: int = 128, + output_dtype: Optional["torch.dtype"] = None, + output_device: Union["torch.device", str] = "cpu", + ) -> "torch.Tensor": + """ + Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector + query embeddings (`qs`) and passage embeddings (`ps`). For ColPali, a passage is the + image of a document page. + + Because the embedding tensors are multi-vector and can thus have different shapes, they + should be fed as: + (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim) + (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually + obtained by padding the list of tensors. + + Args: + query_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Query embeddings. + passage_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Passage embeddings. + batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores. + output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor. + If `None`, the dtype of the input embeddings is used. + output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor. + + Returns: + `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score + tensor is saved on the "cpu" device. + """ + + if len(query_embeddings) == 0: + raise ValueError("No queries provided") + if len(passage_embeddings) == 0: + raise ValueError("No passages provided") + + if query_embeddings[0].device != passage_embeddings[0].device: + raise ValueError("Queries and passages must be on the same device") + + if query_embeddings[0].dtype != passage_embeddings[0].dtype: + raise ValueError("Queries and passages must have the same dtype") + + if output_dtype is None: + output_dtype = query_embeddings[0].dtype + + scores: list[torch.Tensor] = [] + + for i in range(0, len(query_embeddings), batch_size): + batch_scores: list[torch.Tensor] = [] + batch_queries = torch.nn.utils.rnn.pad_sequence( + query_embeddings[i : i + batch_size], batch_first=True, padding_value=0 + ) + for j in range(0, len(passage_embeddings), batch_size): + batch_passages = torch.nn.utils.rnn.pad_sequence( + passage_embeddings[j : j + batch_size], batch_first=True, padding_value=0 + ) + batch_scores.append( + torch.einsum("bnd,csd->bcns", batch_queries, batch_passages).max(dim=3)[0].sum(dim=2) + ) + scores.append(torch.cat(batch_scores, dim=1).to(output_dtype).to(output_device)) + + return torch.cat(scores, dim=0) + + +__all__ = ["ColPaliProcessor"] diff --git a/third_party/transformers/src/transformers/models/eomt_dinov3/__init__.py b/third_party/transformers/src/transformers/models/eomt_dinov3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..10e4ae987f7dd7642a190d594be313a98f955713 --- /dev/null +++ b/third_party/transformers/src/transformers/models/eomt_dinov3/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_eomt_dinov3 import * + from .modeling_eomt_dinov3 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/eomt_dinov3/configuration_eomt_dinov3.py b/third_party/transformers/src/transformers/models/eomt_dinov3/configuration_eomt_dinov3.py new file mode 100644 index 0000000000000000000000000000000000000000..93f6350f26d2e08c5faf36ae5079b451bc3becc7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/eomt_dinov3/configuration_eomt_dinov3.py @@ -0,0 +1,107 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/eomt_dinov3/modular_eomt_dinov3.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_eomt_dinov3.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="tue-mps/coco_panoptic_eomt_large_640_dinov3") +@strict +class EomtDinov3Config(PreTrainedConfig): + r""" + layerscale_value (`float`, *optional*, defaults to 1.0): + Initial value for the LayerScale parameter. + num_upscale_blocks (`int`, *optional*, defaults to 2): + Number of upsampling blocks used in the decoder or segmentation head. + num_blocks (`int`, *optional*, defaults to 4): + Number of feature blocks or stages in the architecture. + no_object_weight (`float`, *optional*, defaults to 0.1): + Loss weight for the "no object" class in panoptic/instance segmentation. + train_num_points (`int`, *optional*, defaults to 12544): + Number of points to sample for mask loss computation during training. + oversample_ratio (`float`, *optional*, defaults to 3.0): + Oversampling ratio used in point sampling for mask training. + importance_sample_ratio (`float`, *optional*, defaults to 0.75): + Ratio of points to sample based on importance during training. + num_queries (`int`, *optional*, defaults to 200): + Number of object queries in the Transformer. + num_register_tokens (`int`, *optional*, defaults to 4): + Number of learnable register tokens added to the transformer input. + query_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in query projection. + key_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in key projection. + value_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in value projection. + proj_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in output projection. + use_gated_mlp (`bool`, *optional*, defaults to `False`): + Whether to use gated MLP layers. + pos_embed_shift (`float`, *optional*): + Shift value for position embeddings. + pos_embed_jitter (`float`, *optional*): + Jitter value for position embeddings. + pos_embed_rescale (`float`, *optional*, defaults to 2.0): + Rescale value for position embeddings. + """ + + model_type = "eomt_dinov3" + + hidden_size: int = 1024 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-6 + image_size: int | list[int] | tuple[int, int] = 640 + patch_size: int | list[int] | tuple[int, int] = 16 + num_channels: int = 3 + layerscale_value: float = 1.0 + drop_path_rate: float | int = 0.0 + num_upscale_blocks: int = 2 + attention_dropout: float | int = 0.0 + num_blocks: int = 4 + no_object_weight: float = 0.1 + class_weight: float = 2.0 + mask_weight: float = 5.0 + dice_weight: float = 5.0 + train_num_points: int = 12544 + oversample_ratio: float = 3.0 + importance_sample_ratio: float = 0.75 + num_queries: int = 200 + num_register_tokens: int = 4 + default_theta = 100.0 + intermediate_size: int = 4096 + rope_parameters: RopeParameters | dict | None = None + query_bias: bool = True + key_bias: bool = False + value_bias: bool = True + proj_bias: bool = True + mlp_bias: bool = True + use_gated_mlp: bool = False + pos_embed_shift: float | None = None + pos_embed_jitter: float | None = None + pos_embed_rescale: float | None = 2.0 + + +__all__ = ["EomtDinov3Config"] diff --git a/third_party/transformers/src/transformers/models/eomt_dinov3/convert_eomt_dinov3_to_hf.py b/third_party/transformers/src/transformers/models/eomt_dinov3/convert_eomt_dinov3_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..b4df176cadcbefccadac0df0043b5943b7caf622 --- /dev/null +++ b/third_party/transformers/src/transformers/models/eomt_dinov3/convert_eomt_dinov3_to_hf.py @@ -0,0 +1,670 @@ +# Copyright 2026 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for EoMT-DINOv3 checkpoints. + +To convert one of the official checkpoints directly from the Hugging Face Hub you can run: + +```bash +HF_TOKEN=your_token_here \ +python -m transformers.models.eomt_dinov3.convert_eomt_dinov3_to_hf \ + --model-id tue-mps/coco_panoptic_eomt_large_640_dinov3 \ + --output-dir /tmp/eomt_converted \ + --verify \ + --original-repo-path /tmp/eomt +``` + +Make sure the token used above has been granted access to the gated DINOv3 weights. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from collections.abc import Iterable +from pathlib import Path +from typing import NamedTuple + +import requests +import torch +from accelerate import init_empty_weights +from huggingface_hub import hf_hub_download +from PIL import Image +from safetensors.torch import load_file + +from transformers import EomtDinov3Config, EomtDinov3ForUniversalSegmentation, EomtImageProcessorFast + + +CAT_URL = "http://images.cocodataset.org/val2017/000000039769.jpg" + + +DEFAULT_BACKBONE_REPO_ID = "facebook/dinov3-vitl16-pretrain-lvd1689m" +DEFAULT_IMAGE_SIZE = 640 + + +class CheckpointSpec(NamedTuple): + """Metadata describing how to convert an official EoMT-DINOv3 checkpoint.""" + + model_id: str + backbone_repo_id: str + image_size: int + + +CHECKPOINT_CATALOG: tuple[CheckpointSpec, ...] = ( + CheckpointSpec( + model_id="tue-mps/coco_panoptic_eomt_small_640_dinov3", + backbone_repo_id="facebook/dinov3-vits16-pretrain-lvd1689m", + image_size=640, + ), + CheckpointSpec( + model_id="tue-mps/coco_panoptic_eomt_base_640_dinov3", + backbone_repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m", + image_size=640, + ), + CheckpointSpec( + model_id="tue-mps/coco_panoptic_eomt_large_640_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=640, + ), + CheckpointSpec( + model_id="tue-mps/coco_panoptic_eomt_large_1280_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=1280, + ), + CheckpointSpec( + model_id="tue-mps/ade20k_semantic_eomt_large_512_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=512, + ), + CheckpointSpec( + model_id="tue-mps/coco_instance_eomt_large_640_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=640, + ), + CheckpointSpec( + model_id="tue-mps/coco_instance_eomt_large_1280_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=1280, + ), + CheckpointSpec( + model_id="tue-mps/ade_semantic_eomt_large_512_dinov3", + backbone_repo_id="facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size=512, + ), +) + + +def _build_checkpoint_index() -> dict[str, CheckpointSpec]: + index: dict[str, CheckpointSpec] = {} + + for spec in CHECKPOINT_CATALOG: + keys = {spec.model_id.lower(), spec.model_id.split("/", maxsplit=1)[-1].lower()} + for key in keys: + index[key] = spec + + return index + + +CHECKPOINT_SPECS = _build_checkpoint_index() + + +def resolve_checkpoint_spec(model_id: str) -> CheckpointSpec: + key = model_id.lower() + if key not in CHECKPOINT_SPECS: + available = ", ".join(sorted(spec.model_id for spec in CHECKPOINT_CATALOG)) + raise ValueError(f"Unknown checkpoint '{model_id}'. Available options: {available}.") + return CHECKPOINT_SPECS[key] + + +def print_checkpoint_catalog() -> None: + print("Supported checkpoints:") + for spec in CHECKPOINT_CATALOG: + print(f"- {spec.model_id} (image_size={spec.image_size}, backbone={spec.backbone_repo_id})") + + +DELTA_KEY_REPLACEMENTS: tuple[tuple[str, str], ...] = ( + (r"^network\.encoder\.backbone\.patch_embed\.cls_token$", "embeddings.cls_token"), + (r"^network\.encoder\.backbone\.patch_embed\.register_tokens$", "embeddings.register_tokens"), + ( + r"^network\.encoder\.backbone\.patch_embed\.patch_embeddings\.", + "embeddings.patch_embeddings.", + ), + (r"^network\.encoder\.backbone\.blocks\.(\d+)\.", r"layers.\1."), + (r"^network\.encoder\.backbone\.norm\.", "layernorm."), + (r"^network\.q\.", "query."), + (r"^network\.class_head\.", "class_predictor."), + (r"^network\.mask_head\.0\.", "mask_head.fc1."), + (r"^network\.mask_head\.2\.", "mask_head.fc2."), + (r"^network\.mask_head\.4\.", "mask_head.fc3."), + (r"^network\.upscale\.(\d+)\.conv1\.", r"upscale_block.block.\1.conv1."), + (r"^network\.upscale\.(\d+)\.conv2\.", r"upscale_block.block.\1.conv2."), + (r"^network\.upscale\.(\d+)\.norm\.", r"upscale_block.block.\1.layernorm2d."), + (r"^network\.attn_mask_probs$", "attn_mask_probs"), + (r"^criterion\.", "criterion."), +) + +SKIP_KEYS = { + "network.encoder.pixel_mean", + "network.encoder.pixel_std", +} + + +def _rename_delta_key(key: str) -> tuple[str | None, bool]: + if key in SKIP_KEYS: + return None, False + + for pattern, replacement in DELTA_KEY_REPLACEMENTS: + if re.match(pattern, key): + new_key = re.sub(pattern, replacement, key) + return new_key, key.startswith("network.encoder.backbone") + + if key.startswith("network.encoder.backbone"): + raise KeyError(f"Unhandled backbone key: {key}") + + return None, False + + +def convert_delta_state_dict(state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], set[str]]: + converted: dict[str, torch.Tensor] = {} + backbone_keys: set[str] = set() + + for key, value in state_dict.items(): + new_key, is_backbone = _rename_delta_key(key) + if new_key is None: + continue + + converted[new_key] = value + if is_backbone: + backbone_keys.add(new_key) + + return converted, backbone_keys + + +def map_dinov3_state_to_eomt(base_state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + mapped: dict[str, torch.Tensor] = {} + + for key, tensor in base_state_dict.items(): + if key.startswith("layer."): + new_key = key.replace("layer.", "layers.", 1) + elif key == "norm.weight": + new_key = "layernorm.weight" + elif key == "norm.bias": + new_key = "layernorm.bias" + else: + new_key = key + + if new_key == "embeddings.mask_token": + continue + + mapped[new_key] = tensor + + return mapped + + +def merge_backbone_weights( + base_backbone: dict[str, torch.Tensor], + delta_backbone: dict[str, torch.Tensor], + backbone_delta_keys: Iterable[str], +) -> dict[str, torch.Tensor]: + merged = dict(base_backbone) + + for key, value in delta_backbone.items(): + if key in backbone_delta_keys: + merged[key] = merged[key] + value + else: + merged[key] = value + + return merged + + +def build_eomt_config( + *, + base_config: dict[str, object], + delta_state_dict: dict[str, torch.Tensor], + image_size: int, +) -> EomtDinov3Config: + num_queries = delta_state_dict["network.q.weight"].shape[0] + num_blocks = delta_state_dict["network.attn_mask_probs"].numel() + num_upscale_blocks = len({int(key.split(".")[2]) for key in delta_state_dict if key.startswith("network.upscale")}) + num_register_tokens = delta_state_dict["network.encoder.backbone.patch_embed.register_tokens"].shape[1] + num_labels = delta_state_dict["network.class_head.weight"].shape[0] - 1 + + config = EomtDinov3Config( + hidden_size=base_config["hidden_size"], + num_hidden_layers=base_config["num_hidden_layers"], + num_attention_heads=base_config["num_attention_heads"], + intermediate_size=base_config["intermediate_size"], + hidden_act=base_config["hidden_act"], + hidden_dropout_prob=0.0, + initializer_range=base_config["initializer_range"], + layer_norm_eps=base_config["layer_norm_eps"], + image_size=image_size, + patch_size=base_config["patch_size"], + num_channels=base_config.get("num_channels", 3), + layerscale_value=base_config.get("layerscale_value", 1.0), + drop_path_rate=base_config.get("drop_path_rate", 0.0), + attention_dropout=base_config.get("attention_dropout", 0.0), + num_upscale_blocks=num_upscale_blocks, + num_blocks=num_blocks, + num_queries=num_queries, + num_register_tokens=num_register_tokens, + rope_parameters={"rope_theta": base_config.get("rope_theta", 100.0)}, + query_bias=base_config.get("query_bias", True), + key_bias=base_config.get("key_bias", False), + value_bias=base_config.get("value_bias", True), + proj_bias=base_config.get("proj_bias", True), + mlp_bias=base_config.get("mlp_bias", True), + use_gated_mlp=base_config.get("use_gated_mlp", False), + pos_embed_shift=base_config.get("pos_embed_shift"), + pos_embed_jitter=base_config.get("pos_embed_jitter"), + pos_embed_rescale=base_config.get("pos_embed_rescale"), + num_labels=num_labels, + ) + + return config + + +def convert_checkpoint( + *, + delta_state_dict: dict[str, torch.Tensor], + backbone_repo_id: str, + image_size: int, +) -> tuple[EomtDinov3Config, dict[str, torch.Tensor]]: + # load model.safetensors + filepath = hf_hub_download(backbone_repo_id, filename="model.safetensors") + base_state_dict = load_file(filepath) + + # load config.json + filepath = hf_hub_download(backbone_repo_id, filename="config.json") + with open(filepath, "r") as f: + base_config = json.load(f) + + mapped_base = map_dinov3_state_to_eomt(base_state_dict) + converted_delta, backbone_delta_keys = convert_delta_state_dict(delta_state_dict) + merged_state_dict = merge_backbone_weights(mapped_base, converted_delta, backbone_delta_keys) + + config = build_eomt_config( + base_config=base_config, + delta_state_dict=delta_state_dict, + image_size=image_size, + ) + + return config, merged_state_dict + + +def ensure_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + if "state_dict" in state_dict: + return state_dict["state_dict"] + return state_dict + + +def convert_model( + *, + model_id: str, + output_dir: Path, + verify: bool, + original_repo_path: Path | None, + push_to_hub: bool = False, +) -> None: + # resolve checkpoint spec + spec = resolve_checkpoint_spec(model_id) + backbone_repo_id = spec.backbone_repo_id + image_size = spec.image_size + + # load delta state + delta_path = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin") + raw_delta_state = torch.load(delta_path, map_location="cpu") + delta_state_dict = ensure_state_dict(raw_delta_state) + + # convert checkpoint + config, merged_state_dict = convert_checkpoint( + delta_state_dict=delta_state_dict, + backbone_repo_id=backbone_repo_id, + image_size=image_size, + ) + + with init_empty_weights(): + model = EomtDinov3ForUniversalSegmentation(config) + + model.load_state_dict(merged_state_dict, strict=True, assign=True) + + processor = EomtImageProcessorFast( + size={"shortest_edge": image_size, "longest_edge": image_size}, + do_split_image=False, + do_pad=True, + ) + + if verify: + verify_conversion( + hf_model=model, + processor=processor, + delta_state_dict=delta_state_dict, + backbone_repo_id=backbone_repo_id, + image_size=image_size, + original_repo_path=original_repo_path, + ) + + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + model.save_pretrained(output_dir) + processor.save_pretrained(output_dir) + + if push_to_hub: + # Extract model name from model_id (e.g. "tue-mps/coco_panoptic_eomt_large_640_dinov3" -> "eomt-dinov3-coco-panoptic-large-640") + base_name = model_id.split("/")[-1] # e.g. "coco_panoptic_eomt_large_640_dinov3" + parts = base_name.replace("_dinov3", "").split("_") # ["coco", "panoptic", "eomt", "large", "640"] + # Reorder to: eomt-dinov3-{task}-{dataset}-{size}-{resolution} + if "eomt" in parts: + parts.remove("eomt") + repo_name = "eomt-dinov3-" + "-".join(parts) + model.push_to_hub(repo_id=f"nielsr/{repo_name}") + processor.push_to_hub(repo_id=f"nielsr/{repo_name}") + + +def _prepare_image(processor: EomtImageProcessorFast) -> torch.Tensor: + image = Image.open(requests.get(CAT_URL, stream=True).raw).convert("RGB") + inputs = processor(images=image, do_normalize=False, return_tensors="pt") + return inputs.pixel_values + + +def _load_original_model( + *, + original_repo_path: Path, + backbone_repo_id: str, + image_size: int, + num_labels: int, + num_queries: int, + num_blocks: int, + delta_state_dict: dict[str, torch.Tensor], +) -> torch.nn.Module: + sys.path.insert(0, str(original_repo_path)) + + from models.eomt import EoMT + from models.vit import ViT + + encoder = ViT((image_size, image_size), backbone_name=backbone_repo_id, ckpt_path=None) + model = EoMT(encoder=encoder, num_classes=num_labels, num_q=num_queries, num_blocks=num_blocks) + + state_dict = model.state_dict() + for key, value in delta_state_dict.items(): + if key in SKIP_KEYS or key.startswith("criterion"): + continue + + target_key = key + if key.startswith("network."): + target_key = key.replace("network.", "", 1) + + if target_key.startswith("encoder.backbone"): + state_dict[target_key] = state_dict[target_key] + value + elif target_key in state_dict: + state_dict[target_key] = value + + model.load_state_dict(state_dict) + model.eval() + return model + + +class BackboneVerificationOutputs(NamedTuple): + patch_embeddings: torch.Tensor + rope_embeddings: tuple[torch.Tensor, torch.Tensor] + hidden_states: list[torch.Tensor] + mask_logits: list[torch.Tensor] + class_logits: list[torch.Tensor] + sequence_output: torch.Tensor + + +def _collect_original_backbone_states(model, pixel_values: torch.Tensor) -> BackboneVerificationOutputs: + backbone = model.encoder.backbone + hidden_states = (pixel_values - model.encoder.pixel_mean) / model.encoder.pixel_std + + rope = None + if hasattr(backbone, "rope_embeddings"): + rope = backbone.rope_embeddings(hidden_states) + + hidden_states = backbone.patch_embed(hidden_states) + patch_embeddings = hidden_states.detach().clone() + if rope is None: + raise ValueError("Original model is missing rope embeddings") + outputs = [] + mask_logits_list = [] + class_logits_list = [] + attn_mask = None + + for idx, block in enumerate(backbone.blocks): + if idx == len(backbone.blocks) - model.num_blocks: + query = model.q.weight[None, :, :].expand(hidden_states.shape[0], -1, -1) + hidden_states = torch.cat((query, hidden_states), dim=1) + + if idx >= len(backbone.blocks) - model.num_blocks: + norm_hidden_states = backbone.norm(hidden_states) + mask_logits, class_logits = model._predict(norm_hidden_states) + mask_logits_list.append(mask_logits) + class_logits_list.append(class_logits) + attn_mask = model._attn_mask(hidden_states, mask_logits, idx) + + attn_module = block.attention if hasattr(block, "attention") else block.attn + attn_output = model._attn(attn_module, block.norm1(hidden_states), attn_mask, rope=rope) + if hasattr(block, "layer_scale1"): + hidden_states = hidden_states + block.layer_scale1(attn_output) + else: + hidden_states = hidden_states + block.ls1(attn_output) + + mlp_output = block.mlp(block.norm2(hidden_states)) + if hasattr(block, "layer_scale2"): + hidden_states = hidden_states + block.layer_scale2(mlp_output) + else: + hidden_states = hidden_states + block.ls2(mlp_output) + + outputs.append(hidden_states) + + sequence_output = backbone.norm(hidden_states) + mask_logits, class_logits = model._predict(sequence_output) + mask_logits_list.append(mask_logits) + class_logits_list.append(class_logits) + + return BackboneVerificationOutputs( + patch_embeddings=patch_embeddings, + rope_embeddings=rope, + hidden_states=outputs, + mask_logits=mask_logits_list, + class_logits=class_logits_list, + sequence_output=sequence_output, + ) + + +def _collect_hf_backbone_states( + model: EomtDinov3ForUniversalSegmentation, pixel_values: torch.Tensor +) -> BackboneVerificationOutputs: + position_embeddings = model.rope_embeddings(pixel_values) + hidden_states = model.embeddings(pixel_values) + patch_embeddings = hidden_states.detach().clone() + + outputs = [] + mask_logits_list = [] + class_logits_list = [] + attention_mask = None + + for idx, layer_module in enumerate(model.layers): + if idx == model.num_hidden_layers - model.config.num_blocks: + query = model.query.weight[None, :, :].expand(hidden_states.shape[0], -1, -1).to(hidden_states.device) + hidden_states = torch.cat((query, hidden_states), dim=1) + + if idx >= model.num_hidden_layers - model.config.num_blocks: + norm_hidden_states = model.layernorm(hidden_states) + mask_logits, class_logits = model.predict(norm_hidden_states) + mask_logits_list.append(mask_logits) + class_logits_list.append(class_logits) + + probs_index = idx - model.num_hidden_layers + model.config.num_blocks + if model.training or model.attn_mask_probs[probs_index] > 0: + attention_mask = torch.ones( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.shape[1], + device=hidden_states.device, + dtype=torch.bool, + ) + + interpolated_logits = torch.nn.functional.interpolate( + mask_logits, + size=model.grid_size, + mode="bilinear", + ).view(mask_logits.size(0), mask_logits.size(1), -1) + + num_query_tokens = model.config.num_queries + encoder_start_tokens = num_query_tokens + model.embeddings.num_prefix_tokens + attention_mask[:, :num_query_tokens, encoder_start_tokens:] = interpolated_logits > 0 + + attention_mask = model._disable_attention_mask( + attention_mask, + prob=model.attn_mask_probs[probs_index], + num_query_tokens=num_query_tokens, + encoder_start_tokens=encoder_start_tokens, + device=hidden_states.device, + ) + + attention_mask = attention_mask[:, None, ...].expand(-1, model.config.num_attention_heads, -1, -1) + + bool_attention_mask = attention_mask + attention_mask = attention_mask.float().masked_fill(~bool_attention_mask, -1e9) + if attention_mask.dtype != hidden_states.dtype: + attention_mask = attention_mask.to(dtype=hidden_states.dtype) + + hidden_states = layer_module( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + outputs.append(hidden_states) + + sequence_output = model.layernorm(hidden_states) + mask_logits, class_logits = model.predict(sequence_output) + mask_logits_list.append(mask_logits) + class_logits_list.append(class_logits) + + return BackboneVerificationOutputs( + patch_embeddings=patch_embeddings, + rope_embeddings=position_embeddings, + hidden_states=outputs, + mask_logits=mask_logits_list, + class_logits=class_logits_list, + sequence_output=sequence_output, + ) + + +def _assert_allclose(reference: Iterable[torch.Tensor], actual: Iterable[torch.Tensor], message: str) -> None: + for idx, (ref_tensor, act_tensor) in enumerate(zip(reference, actual)): + if not torch.allclose(ref_tensor, act_tensor, atol=1e-4, rtol=1e-4): + raise ValueError(f"Mismatch in {message} at index {idx}") + + +def verify_conversion( + *, + hf_model: EomtDinov3ForUniversalSegmentation, + processor: EomtImageProcessorFast, + delta_state_dict: dict[str, torch.Tensor], + backbone_repo_id: str, + image_size: int, + original_repo_path: Path | None, +) -> None: + if original_repo_path is None: + raise ValueError("Original repository path is required for verification") + + torch.manual_seed(0) + + pixel_values = _prepare_image(processor) + + image_mean = torch.tensor(processor.image_mean, dtype=pixel_values.dtype, device=pixel_values.device)[ + None, :, None, None + ] + image_std = torch.tensor(processor.image_std, dtype=pixel_values.dtype, device=pixel_values.device)[ + None, :, None, None + ] + normalized_pixel_values = (pixel_values - image_mean) / image_std + + original_model = _load_original_model( + original_repo_path=original_repo_path, + backbone_repo_id=backbone_repo_id, + image_size=image_size, + num_labels=hf_model.config.num_labels, + num_queries=hf_model.config.num_queries, + num_blocks=hf_model.config.num_blocks, + delta_state_dict=delta_state_dict, + ) + + hf_model.eval() + + with torch.no_grad(): + orig_outputs = _collect_original_backbone_states(original_model, pixel_values) + hf_outputs = _collect_hf_backbone_states(hf_model, normalized_pixel_values) + + patch_abs_diff = (orig_outputs.patch_embeddings - hf_outputs.patch_embeddings).abs() + print(f"Patch embedding max abs diff: {patch_abs_diff.max().item():.6e}") + + rope_abs_diffs = [(orig - hf).abs() for orig, hf in zip(orig_outputs.rope_embeddings, hf_outputs.rope_embeddings)] + rope_max_diffs = [diff.max().item() for diff in rope_abs_diffs] + print( + "RoPE embedding max abs diff: " + + ", ".join(f"component_{idx}={value:.6e}" for idx, value in enumerate(rope_max_diffs)) + ) + + if not torch.allclose(orig_outputs.patch_embeddings, hf_outputs.patch_embeddings, atol=1e-4, rtol=1e-4): + raise ValueError("Mismatch in patch embeddings") + + _assert_allclose(orig_outputs.rope_embeddings, hf_outputs.rope_embeddings, "rope embeddings") + _assert_allclose(orig_outputs.hidden_states, hf_outputs.hidden_states, "backbone hidden states") + _assert_allclose(orig_outputs.mask_logits, hf_outputs.mask_logits, "mask logits") + _assert_allclose(orig_outputs.class_logits, hf_outputs.class_logits, "class logits") + print("Looks good!") + + if not torch.allclose(orig_outputs.sequence_output, hf_outputs.sequence_output, atol=1e-4, rtol=1e-4): + raise ValueError("Mismatch in final sequence output") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Convert EoMT-DINOv3 checkpoints to 🤗 Transformers format") + parser.add_argument("--list-models", action="store_true", help="List supported checkpoint names and exit") + parser.add_argument( + "--model-id", + help="Name of an official EoMT-DINOv3 checkpoint to download and convert", + required=True, + ) + parser.add_argument("--output-dir", type=Path, help="Directory to save the converted model") + parser.add_argument("--verify", action="store_true") + parser.add_argument("--original-repo-path", type=Path, default=None, help="Path to the original EoMT repository") + parser.add_argument( + "--push-to-hub", action="store_true", help="Whether to push the converted model to the Hugging Face Hub." + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if args.list_models: + print_checkpoint_catalog() + return + + convert_model( + model_id=args.model_id, + output_dir=args.output_dir, + verify=args.verify, + original_repo_path=args.original_repo_path, + push_to_hub=args.push_to_hub, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/eomt_dinov3/modeling_eomt_dinov3.py b/third_party/transformers/src/transformers/models/eomt_dinov3/modeling_eomt_dinov3.py new file mode 100644 index 0000000000000000000000000000000000000000..fe7b05dc69c45396293ef11defc1a75be4bd5518 --- /dev/null +++ b/third_party/transformers/src/transformers/models/eomt_dinov3/modeling_eomt_dinov3.py @@ -0,0 +1,1379 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/eomt_dinov3/modular_eomt_dinov3.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_eomt_dinov3.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...file_utils import ModelOutput, is_scipy_available, requires_backends +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import TransformersKwargs, auto_docstring, is_accelerate_available +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_eomt_dinov3 import EomtDinov3Config + + +if is_scipy_available(): + from scipy.optimize import linear_sum_assignment + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import reduce + + +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 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 + + +def apply_rotary_pos_emb( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs +) -> tuple[torch.Tensor, torch.Tensor]: + """Applies Rotary Position Embedding to the query and key tensors, but only to the patch tokens, + ignoring the prefix tokens (cls token and register tokens). + + 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. + + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + + num_tokens = q.shape[-2] + num_patches = sin.shape[-2] + num_prefix_tokens = num_tokens - num_patches # cls token + register tokens + + q_prefix_tokens, q_patches = q.split((num_prefix_tokens, num_patches), dim=-2) + k_prefix_tokens, k_patches = k.split((num_prefix_tokens, num_patches), dim=-2) + + # apply rope only to patch tokens + q_patches = (q_patches * cos) + (rotate_half(q_patches) * sin) + k_patches = (k_patches * cos) + (rotate_half(k_patches) * sin) + + q = torch.cat((q_prefix_tokens, q_patches), dim=-2) + k = torch.cat((k_prefix_tokens, k_patches), dim=-2) + + return q, k + + +class EomtDinov3Attention(nn.Module): + """ + Multi-headed attention compatible with ALL_ATTENTION_FUNCTIONS. + """ + + def __init__(self, config: EomtDinov3Config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.is_causal = False + + self.scaling = self.head_dim**-0.5 + self.is_causal = False + + self.dropout = config.attention_dropout + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.key_bias) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.value_bias) + + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.query_bias) + self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.proj_bias) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + batch_size, patches, _ = 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(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(batch_size, patches, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class EomtDinov3Embeddings(nn.Module): + """ + Construct the CLS token, mask token, position and patch embeddings. + """ + + def __init__(self, config: EomtDinov3Config): + super().__init__() + self.config = config + self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.register_tokens = nn.Parameter(torch.empty(1, config.num_register_tokens, config.hidden_size)) + self.patch_embeddings = nn.Conv2d( + config.num_channels, config.hidden_size, kernel_size=config.patch_size, stride=config.patch_size + ) + self.num_prefix_tokens = 1 + config.num_register_tokens + + def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None) -> torch.Tensor: + batch_size = pixel_values.shape[0] + target_dtype = self.patch_embeddings.weight.dtype + + # (batch_size, num_channels, height, width) -> (batch_size, num_patches, hidden_size) + patch_embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype)) + patch_embeddings = patch_embeddings.flatten(2).transpose(1, 2) + + if bool_masked_pos is not None: + mask_token = self.mask_token.to(patch_embeddings.dtype) + patch_embeddings = torch.where(bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings) + + # Add CLS and register tokens + cls_token = self.cls_token.expand(batch_size, -1, -1) + register_tokens = self.register_tokens.expand(batch_size, -1, -1) + embeddings = torch.cat([cls_token, register_tokens, patch_embeddings], dim=1) + + return embeddings + + +def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +class EomtDinov3DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float | None = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return f"p={self.drop_prob}" + + +class EomtDinov3MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +class EomtDinov3GatedMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class EomtDinov3Layer(GradientCheckpointingLayer): + """This corresponds to the Block class in the original implementation.""" + + def __init__(self, config: EomtDinov3Config): + super().__init__() + + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attention = EomtDinov3Attention(config) + self.layer_scale1 = EomtDinov3LayerScale(config) + self.drop_path = EomtDinov3DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity() + + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + if config.use_gated_mlp: + self.mlp = EomtDinov3GatedMLP(config) + else: + self.mlp = EomtDinov3MLP(config) + self.layer_scale2 = EomtDinov3LayerScale(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + # Attention with residual connection + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states, _ = self.attention( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.layer_scale1(hidden_states) + hidden_states = self.drop_path(hidden_states) + residual + + # MLP with residual connection + residual = hidden_states + hidden_states = self.norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.layer_scale2(hidden_states) + hidden_states = self.drop_path(hidden_states) + residual + + return hidden_states + + +class EomtDinov3LayerScale(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size)) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + return hidden_state * self.lambda1 + + +@compile_compatible_method_lru_cache(maxsize=32) +def get_patches_center_coordinates( + num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """ + Computes the 2D coordinates of the centers of image patches, normalized to the range [-1, +1]. + The center of each patch is exactly halfway between its top-left and bottom-right corners. + + Args: + num_patches_h (int): Number of patches along the vertical (height) axis. + num_patches_w (int): Number of patches along the horizontal (width) axis. + dtype (torch.dtype): The desired data type of the returned tensor. + + Returns: + torch.Tensor: A tensor of shape (height * width, 2), where each row contains the (y, x) + coordinates of a patch center, normalized to [-1, +1]. + """ + coords_h = torch.arange(0.5, num_patches_h, dtype=dtype, device=device) + coords_w = torch.arange(0.5, num_patches_w, dtype=dtype, device=device) + coords_h = coords_h / num_patches_h + coords_w = coords_w / num_patches_w + # (height, width, 2) -> (height * width, 2) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) + coords = coords.flatten(0, 1) + # Shift range [0, 1] to [-1, +1] + coords = 2.0 * coords - 1.0 + return coords + + +def augment_patches_center_coordinates( + coords: torch.Tensor, + shift: float | None = None, + jitter: float | None = None, + rescale: float | None = None, +) -> torch.Tensor: + # Shift coords by adding a uniform value in [-shift, shift] + if shift is not None: + shift_hw = torch.empty((1, 2), device=coords.device, dtype=coords.dtype) + shift_hw = shift_hw.uniform_(-shift, shift) + coords = coords + shift_hw + + # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter] + if jitter is not None: + jitter_range = np.log(jitter) + jitter_hw = torch.empty((1, 2), device=coords.device, dtype=coords.dtype) + jitter_hw = jitter_hw.uniform_(-jitter_range, jitter_range).exp() + coords = coords * jitter_hw + + # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale] + if rescale is not None: + rescale_range = np.log(rescale) + rescale_hw = torch.empty(1, device=coords.device, dtype=coords.dtype) + rescale_hw = rescale_hw.uniform_(-rescale_range, rescale_range).exp() + coords = coords * rescale_hw + + return coords + + +class EomtDinov3RotaryEmbedding(nn.Module): + inv_freq: Tensor + + def __init__(self, config: EomtDinov3Config, device=None): + super().__init__() + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + raise ValueError("`EomtDinov3` only supports `default` RoPE! Please check your `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) + + def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + _, _, height, width = pixel_values.shape + num_patches_h = height // self.config.patch_size + num_patches_w = width // self.config.patch_size + + device = pixel_values.device + device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu" + + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + # Although we could precompute static patch_coords from image_size and patch_size in the config, + # the model was trained with random_scale, so it can process images of varying sizes. + # Therefore, it's better to compute patch_coords dynamically (with lru_cache). + patch_coords = get_patches_center_coordinates( + num_patches_h, num_patches_w, dtype=torch.float32, device=device + ) + if self.training: + patch_coords = augment_patches_center_coordinates( + patch_coords, + shift=self.config.pos_embed_shift, + jitter=self.config.pos_embed_jitter, + rescale=self.config.pos_embed_rescale, + ) + + # (height * width, 2, head_dim / 4) -> (height * width, head_dim / 2) -> (height * width, head_dim) + angles = 2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :] + angles = angles.flatten(1, 2) + angles = angles.tile(2) + + cos = torch.cos(angles) + sin = torch.sin(angles) + + dtype = pixel_values.dtype + return cos.to(dtype=dtype), sin.to(dtype=dtype) + + @staticmethod + def compute_default_rope_parameters( + config: EomtDinov3Config | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> torch.Tensor: + """ + 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"] + head_dim = config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1 / base ** torch.arange(0, 1, 4 / head_dim, dtype=torch.float32, device=device) + return inv_freq, attention_factor + + +# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py +def sample_point( + input_features: torch.Tensor, point_coordinates: torch.Tensor, add_dim=False, **kwargs +) -> torch.Tensor: + """ + A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors. + + Args: + input_features (`torch.Tensor` of shape (batch_size, channels, height, width)): + A tensor that contains features map on a height * width grid + point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,: + 2)): + A tensor that contains [0, 1] * [0, 1] normalized point coordinates + add_dim (`bool`): + boolean value to keep track of added dimension + + Returns: + point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels, + height_grid, width_grid): + A tensor that contains features for points in `point_coordinates`. + """ + if point_coordinates.dim() == 3: + add_dim = True + point_coordinates = point_coordinates.unsqueeze(2) + + # use nn.function.grid_sample to get features for points in `point_coordinates` via bilinear interpolation + point_features = torch.nn.functional.grid_sample(input_features, 2.0 * point_coordinates - 1.0, **kwargs) + if add_dim: + point_features = point_features.squeeze(3) + + return point_features + + +def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor: + """ + A pair wise version of the dice loss, see `dice_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + `torch.Tensor`: The computed loss between each pairs. + """ + inputs = inputs.sigmoid().flatten(1) + numerator = 2 * torch.matmul(inputs, labels.T) + # using broadcasting to get a [num_queries, NUM_CLASSES] matrix + denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :] + loss = 1 - (numerator + 1) / (denominator + 1) + return loss + + +def pair_wise_sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + r""" + A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss between each pairs. + """ + + height_and_width = inputs.shape[1] + + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) + cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) + + loss_pos = torch.matmul(cross_entropy_loss_pos / height_and_width, labels.T) + loss_neg = torch.matmul(cross_entropy_loss_neg / height_and_width, (1 - labels).T) + loss = loss_pos + loss_neg + return loss + + +# Adapted from https://github.com/facebookresearch/EomtDinov3/blob/main/eomt_dinov3/modeling/matcher.py +class EomtDinov3HungarianMatcher(nn.Module): + """This class computes an assignment between the labels and the predictions of the network. + + For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more + predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are + un-matched (and thus treated as non-objects). + """ + + def __init__( + self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0, num_points: int = 12544 + ): + """Creates the matcher + + Params: + cost_class (`float`, *optional*, defaults to 1.0): + Relative weight of the classification error in the matching cost. + cost_mask (`float`, *optional*, defaults to 1.0): + This is the relative weight of the focal loss of the binary mask in the matching cost. + cost_dice (`float`, *optional*, defaults to 1.0): + This is the relative weight of the dice loss of the binary mask in the matching cost. + num_points (`int`, *optional*, defaults to 12544): + No. of points to sample on which the mask loss will be calculated. The same set of K points are + uniformly sampled for all prediction and ground truth masks to construct the cost matrix for bipartite + matching. + """ + super().__init__() + if cost_class == 0 and cost_mask == 0 and cost_dice == 0: + raise ValueError("All costs can't be 0") + + self.num_points = num_points + self.cost_class = cost_class + self.cost_mask = cost_mask + self.cost_dice = cost_dice + + @torch.no_grad() + def forward( + self, + masks_queries_logits: torch.Tensor, + class_queries_logits: torch.Tensor, + mask_labels: torch.Tensor, + class_labels: torch.Tensor, + ) -> list[tuple[Tensor]]: + """ + Params: + masks_queries_logits (`torch.Tensor`): + A tensor of dim `batch_size, num_queries, num_labels` with the classification logits. + class_queries_logits (`torch.Tensor`): + A tensor of dim `batch_size, num_queries, height, width` with the predicted masks. + class_labels (`torch.Tensor`): + A tensor of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the + target) containing the class labels. + mask_labels (`torch.Tensor`): + A tensor of dim `num_target_boxes, height, width` containing the target masks. + + Returns: + matched_indices (`list[tuple[Tensor]]`): A list of size batch_size, containing tuples of (index_i, index_j) + where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected labels (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes). + """ + indices: list[tuple[np.array]] = [] + + # iterate through batch size + batch_size = masks_queries_logits.shape[0] + for i in range(batch_size): + pred_probs = class_queries_logits[i].softmax(-1) + pred_mask = masks_queries_logits[i] + + # Compute the classification cost. Contrary to the loss, we don't use the NLL, but approximate it in 1 - proba[target class]. The 1 is a constant that doesn't change the matching, it can be omitted. + cost_class = -pred_probs[:, class_labels[i]] + target_mask = mask_labels[i].to(pred_mask) + target_mask = target_mask[:, None] + pred_mask = pred_mask[:, None] + + # Sample ground truth and predicted masks + point_coordinates = torch.rand(1, self.num_points, 2, device=pred_mask.device) + + target_coordinates = point_coordinates.repeat(target_mask.shape[0], 1, 1) + target_mask = sample_point(target_mask, target_coordinates, align_corners=False).squeeze(1) + + pred_coordinates = point_coordinates.repeat(pred_mask.shape[0], 1, 1) + pred_mask = sample_point(pred_mask, pred_coordinates, align_corners=False).squeeze(1) + + # compute the cross entropy loss between each mask pairs -> shape (num_queries, num_labels) + cost_mask = pair_wise_sigmoid_cross_entropy_loss(pred_mask, target_mask) + # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels) + cost_dice = pair_wise_dice_loss(pred_mask, target_mask) + # final cost matrix + cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice + # eliminate infinite values in cost_matrix to avoid the error ``ValueError: cost matrix is infeasible`` + cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10)) + cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10)) + cost_matrix = torch.nan_to_num(cost_matrix, 0) + # do the assignment using the hungarian algorithm in scipy + assigned_indices: tuple[np.array] = linear_sum_assignment(cost_matrix.cpu()) + indices.append(assigned_indices) + + # It could be stacked in one tensor + matched_indices = [ + (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices + ] + return matched_indices + + +def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor: + r""" + Compute the DICE loss, similar to generalized IOU for masks as follows: + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$ + + In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$ + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + num_masks (`int`): + The number of masks present in the current batch, used for normalization. + + Returns: + `torch.Tensor`: The computed loss. + """ + probs = inputs.sigmoid().flatten(1) + numerator = 2 * (probs * labels).sum(-1) + denominator = probs.sum(-1) + labels.sum(-1) + loss = 1 - (numerator + 1) / (denominator + 1) + loss = loss.sum() / num_masks + return loss + + +def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor: + r""" + Args: + inputs (`torch.Tensor`): + A float tensor of arbitrary shape. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss. + """ + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss = criterion(inputs, labels) + + loss = cross_entropy_loss.mean(1).sum() / num_masks + return loss + + +# Adapted from https://github.com/facebookresearch/EomtDinov3/blob/main/eomt_dinov3/modeling/criterion.py +class EomtDinov3Loss(nn.Module): + def __init__(self, config: EomtDinov3Config, weight_dict: dict[str, float]): + """ + The EomtDinov3 Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we + compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair + of matched ground-truth / prediction (supervise class and mask) + + Args: + config (`EomtDinov3Config`): + The configuration for EomtDinov3 model also containing loss calculation specific parameters. + weight_dict (`dict[str, float]`): + A dictionary of weights to be applied to the different losses. + """ + super().__init__() + requires_backends(self, ["scipy"]) + self.num_labels = config.num_labels + self.weight_dict = weight_dict + + # Weight to apply to the null class + self.eos_coef = config.no_object_weight + empty_weight = torch.ones(self.num_labels + 1) + empty_weight[-1] = self.eos_coef + self.register_buffer("empty_weight", empty_weight) + + # pointwise mask loss parameters + self.num_points = config.train_num_points + self.oversample_ratio = config.oversample_ratio + self.importance_sample_ratio = config.importance_sample_ratio + + self.matcher = EomtDinov3HungarianMatcher( + cost_class=config.class_weight, + cost_dice=config.dice_weight, + cost_mask=config.mask_weight, + num_points=self.num_points, + ) + + def _max_by_axis(self, sizes: list[list[int]]) -> list[int]: + maxes = sizes[0] + for sublist in sizes[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + # Adapted from nested_tensor_from_tensor_list() in original implementation + def _pad_images_to_max_in_batch(self, tensors: list[Tensor]) -> tuple[Tensor, Tensor]: + # get the maximum size in the batch + max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors]) + # compute final size + batch_shape = [len(tensors)] + max_size + batch_size, _, height, width = batch_shape + dtype = tensors[0].dtype + device = tensors[0].device + padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device) + padding_masks = torch.ones((batch_size, height, width), dtype=torch.bool, device=device) + # pad the tensors to the size of the biggest one + for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks): + padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor) + padding_mask[: tensor.shape[1], : tensor.shape[2]] = False + + return padded_tensors, padding_masks + + def loss_labels( + self, class_queries_logits: Tensor, class_labels: list[Tensor], indices: tuple[np.array] + ) -> dict[str, Tensor]: + """Compute the losses related to the labels using cross entropy. + + Args: + class_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, num_labels` + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + """ + pred_logits = class_queries_logits + batch_size, num_queries, _ = pred_logits.shape + criterion = nn.CrossEntropyLoss(weight=self.empty_weight) + idx = self._get_predictions_permutation_indices(indices) # shape of (batch_size, num_queries) + target_classes_o = torch.cat( + [target[j] for target, (_, j) in zip(class_labels, indices)] + ) # shape of (batch_size, num_queries) + target_classes = torch.full( + (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device + ) + target_classes[idx] = target_classes_o + # Permute target_classes (batch_size, num_queries, num_labels) -> (batch_size, num_labels, num_queries) + pred_logits_transposed = pred_logits.transpose(1, 2) + loss_ce = criterion(pred_logits_transposed, target_classes) + losses = {"loss_cross_entropy": loss_ce} + return losses + + def loss_masks( + self, + masks_queries_logits: torch.Tensor, + mask_labels: list[torch.Tensor], + indices: tuple[np.array], + num_masks: int, + ) -> dict[str, torch.Tensor]: + """Compute the losses related to the masks using sigmoid_cross_entropy_loss and dice loss. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, height, width)`. + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + num_masks (`int)`: + The number of masks, used for normalization. + + Returns: + losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing two keys: + - **loss_mask** -- The loss computed using sigmoid cross entropy loss on the predicted and ground truth. + masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth, + masks. + """ + src_idx = self._get_predictions_permutation_indices(indices) + tgt_idx = self._get_targets_permutation_indices(indices) + # shape (batch_size * num_queries, height, width) + pred_masks = masks_queries_logits[src_idx] + # shape (batch_size, num_queries, height, width) + # pad all and stack the targets to the num_labels dimension + target_masks, _ = self._pad_images_to_max_in_batch(mask_labels) + target_masks = target_masks[tgt_idx] + + # No need to upsample predictions as we are using normalized coordinates + pred_masks = pred_masks[:, None] + target_masks = target_masks[:, None] + + # Sample point coordinates + with torch.no_grad(): + point_coordinates = self.sample_points_using_uncertainty( + pred_masks, + lambda logits: self.calculate_uncertainty(logits), + self.num_points, + self.oversample_ratio, + self.importance_sample_ratio, + ) + + point_labels = sample_point(target_masks, point_coordinates, align_corners=False).squeeze(1) + + point_logits = sample_point(pred_masks, point_coordinates, align_corners=False).squeeze(1) + + losses = { + "loss_mask": sigmoid_cross_entropy_loss(point_logits, point_labels, num_masks), + "loss_dice": dice_loss(point_logits, point_labels, num_masks), + } + + del pred_masks + del target_masks + return losses + + def _get_predictions_permutation_indices(self, indices): + # Permute predictions following indices + batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) + predictions_indices = torch.cat([src for (src, _) in indices]) + return batch_indices, predictions_indices + + def _get_targets_permutation_indices(self, indices): + # Permute labels following indices + batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) + target_indices = torch.cat([tgt for (_, tgt) in indices]) + return batch_indices, target_indices + + def calculate_uncertainty(self, logits: torch.Tensor) -> torch.Tensor: + """ + In EomtDinov3 paper, uncertainty is estimated as L1 distance between 0.0 and the logit prediction in 'logits' + for the foreground class in `classes`. + + Args: + logits (`torch.Tensor`): + A tensor of shape (R, 1, ...) for class-specific or class-agnostic, where R is the total number of predicted masks in all images and C is: + the number of foreground classes. The values are logits. + + Returns: + scores (`torch.Tensor`): A tensor of shape (R, 1, ...) that contains uncertainty scores with the most + uncertain locations having the highest uncertainty score. + """ + uncertainty_scores = -(torch.abs(logits)) + return uncertainty_scores + + def sample_points_using_uncertainty( + self, + logits: torch.Tensor, + uncertainty_function, + num_points: int, + oversample_ratio: int, + importance_sample_ratio: float, + ) -> torch.Tensor: + """ + This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The + uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit + prediction as input. + + Args: + logits (`float`): + Logit predictions for P points. + uncertainty_function: + A function that takes logit predictions for P points and returns their uncertainties. + num_points (`int`): + The number of points P to sample. + oversample_ratio (`int`): + Oversampling parameter. + importance_sample_ratio (`float`): + Ratio of points that are sampled via importance sampling. + + Returns: + point_coordinates (`torch.Tensor`): + Coordinates for P sampled points. + """ + + num_boxes = logits.shape[0] + num_points_sampled = int(num_points * oversample_ratio) + + # Get random point coordinates + point_coordinates = torch.rand(num_boxes, num_points_sampled, 2, device=logits.device) + # Get sampled prediction value for the point coordinates + point_logits = sample_point(logits, point_coordinates, align_corners=False) + # Calculate the uncertainties based on the sampled prediction values of the points + point_uncertainties = uncertainty_function(point_logits) + + num_uncertain_points = int(importance_sample_ratio * num_points) + num_random_points = num_points - num_uncertain_points + + idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] + shift = num_points_sampled * torch.arange(num_boxes, dtype=torch.long, device=logits.device) + idx += shift[:, None] + point_coordinates = point_coordinates.view(-1, 2)[idx.view(-1), :].view(num_boxes, num_uncertain_points, 2) + + if num_random_points > 0: + point_coordinates = torch.cat( + [point_coordinates, torch.rand(num_boxes, num_random_points, 2, device=logits.device)], + dim=1, + ) + return point_coordinates + + def forward( + self, + masks_queries_logits: torch.Tensor, + class_queries_logits: torch.Tensor, + mask_labels: list[torch.Tensor], + class_labels: list[torch.Tensor], + auxiliary_predictions: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: + """ + This performs the loss computation. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, height, width)`. + class_queries_logits (`torch.Tensor`): + A tensor of shape `(batch_size, num_queries, num_labels)`. + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*): + if `use_auxiliary_loss` was set to `true` in [`EomtDinov3Config`], then it contains the logits from + the inner layers of the EomtDinov3MaskedAttentionDecoder. + + Returns: + losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing three keys: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + - **loss_mask** -- The loss computed using sigmoid cross_entropy loss on the predicted and ground truth + masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth + masks. + if `use_auxiliary_loss` was set to `true` in [`EomtDinov3Config`], the dictionary contains additional + losses for each auxiliary predictions. + """ + + # retrieve the matching between the outputs of the last layer and the labels + indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + # compute the average number of target masks for normalization purposes + num_masks = self.get_num_masks(class_labels, device=class_labels[0].device) + # get all the losses + losses: dict[str, Tensor] = { + **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), + **self.loss_labels(class_queries_logits, class_labels, indices), + } + # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if auxiliary_predictions is not None: + for idx, aux_outputs in enumerate(auxiliary_predictions): + masks_queries_logits = aux_outputs["masks_queries_logits"] + class_queries_logits = aux_outputs["class_queries_logits"] + loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} + losses.update(loss_dict) + + return losses + + def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Computes the average number of target masks across the batch, for normalization purposes. + """ + num_masks = sum(len(classes) for classes in class_labels) + num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device) + world_size = 1 + if is_accelerate_available(): + if PartialState._shared_state != {}: + num_masks = reduce(num_masks) + world_size = PartialState().num_processes + + num_masks = torch.clamp(num_masks / world_size, min=1) + return num_masks + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`EomtDinov3ForUniversalSegmentationOutput`]. + + This output can be directly passed to [`~EomtDinov3ImageProcessor.post_process_semantic_segmentation`] or + [`~EomtDinov3ImageProcessor.post_process_instance_segmentation`] or + [`~EomtDinov3ImageProcessor.post_process_panoptic_segmentation`] to compute final segmentation maps. Please, see + [`~EomtDinov3ImageProcessor] for details regarding usage. + """ +) +class EomtDinov3ForUniversalSegmentationOutput(ModelOutput): + r""" + loss (`torch.Tensor`, *optional*): + The computed loss, returned when labels are present. + class_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each + query. Note the `+ 1` is needed because we incorporate the null class. + masks_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each + query. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last layer. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states all layers of the model. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Self and Cross Attentions weights from transformer decoder. + patch_offsets (`list[torch.Tensor]`, *optional*): + list of tuples indicating the image index and start and end positions of patches for semantic segmentation. + """ + + loss: torch.FloatTensor | None = None + class_queries_logits: torch.FloatTensor | None = None + masks_queries_logits: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + patch_offsets: list[torch.Tensor] | None = None + + +@auto_docstring +class EomtDinov3PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: EomtDinov3Config + base_model_prefix = "eomt_dinov3" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = False + _no_split_modules = ["EomtDinov3Layer"] + _supports_sdpa = True + _can_record_outputs = { + "hidden_states": EomtDinov3Layer, + "attentions": EomtDinov3Attention, + } + config_class = EomtDinov3Config + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + super()._init_weights(module) + std = self.config.initializer_range + if isinstance(module, EomtDinov3LayerScale): + if hasattr(module, "lambda1"): + init.constant_(module.lambda1, self.config.layerscale_value) + elif isinstance(module, EomtDinov3Embeddings): + init.trunc_normal_(module.cls_token, mean=0.0, std=std) + init.zeros_(module.register_tokens) + elif isinstance(module, EomtDinov3Loss): + empty_weight = torch.ones(module.num_labels + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + elif isinstance(module, EomtDinov3ForUniversalSegmentation): + init.ones_(module.attn_mask_probs) + + +class EomtDinov3LayerNorm2d(nn.LayerNorm): + def __init__(self, num_channels, eps=1e-6, affine=True): + super().__init__(num_channels, eps=eps, elementwise_affine=affine) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = hidden_state.permute(0, 2, 3, 1) + hidden_state = F.layer_norm(hidden_state, self.normalized_shape, self.weight, self.bias, self.eps) + hidden_state = hidden_state.permute(0, 3, 1, 2) + return hidden_state + + +class EomtDinov3ScaleLayer(nn.Module): + def __init__(self, config: EomtDinov3Config): + super().__init__() + hidden_size = config.hidden_size + self.conv1 = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2) + self.activation = ACT2FN[config.hidden_act] + self.conv2 = nn.Conv2d( + hidden_size, + hidden_size, + kernel_size=3, + padding=1, + groups=hidden_size, + bias=False, + ) + + self.layernorm2d = EomtDinov3LayerNorm2d(hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv1(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.layernorm2d(hidden_states) + return hidden_states + + +class EomtDinov3ScaleBlock(nn.Module): + def __init__(self, config: EomtDinov3Config): + super().__init__() + self.num_blocks = config.num_upscale_blocks + self.block = nn.ModuleList([EomtDinov3ScaleLayer(config) for _ in range(self.num_blocks)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for block in self.block: + hidden_states = block(hidden_states) + return hidden_states + + +class EomtDinov3MaskHead(nn.Module): + def __init__(self, config: EomtDinov3Config): + super().__init__() + + hidden_size = config.hidden_size + self.fc1 = nn.Linear(hidden_size, hidden_size) + self.fc2 = nn.Linear(hidden_size, hidden_size) + self.fc3 = nn.Linear(hidden_size, hidden_size) + self.activation = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.activation(self.fc1(hidden_states)) + hidden_states = self.activation(self.fc2(hidden_states)) + hidden_states = self.fc3(hidden_states) + return hidden_states + + +@auto_docstring( + custom_intro=""" + The EoMT-DINOv3 model with head on top for instance/semantic/panoptic segmentation. + """, +) +class EomtDinov3ForUniversalSegmentation(EomtDinov3PreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: EomtDinov3Config): + super().__init__(config) + self.config = config + self.num_hidden_layers = config.num_hidden_layers + self.embeddings = EomtDinov3Embeddings(config) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + self.query = nn.Embedding(config.num_queries, config.hidden_size) + self.layers = nn.ModuleList([EomtDinov3Layer(config) for _ in range(config.num_hidden_layers)]) + + self.upscale_block = EomtDinov3ScaleBlock(config) + self.mask_head = EomtDinov3MaskHead(config) + + self.class_predictor = nn.Linear(config.hidden_size, config.num_labels + 1) + + self.grid_size = (config.image_size // config.patch_size, config.image_size // config.patch_size) + self.weight_dict: dict[str, float] = { + "loss_cross_entropy": config.class_weight, + "loss_mask": config.mask_weight, + "loss_dice": config.dice_weight, + } + + self.criterion = EomtDinov3Loss(config=config, weight_dict=self.weight_dict) + + self.register_buffer("attn_mask_probs", torch.ones(config.num_blocks)) + + self.num_prefix_tokens = 1 + config.num_register_tokens + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.embeddings.register_parameter("mask_token", None) + + self.rope_embeddings = EomtDinov3RotaryEmbedding(config) + + self.post_init() + + def get_loss_dict( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: Tensor, + class_labels: Tensor, + auxiliary_predictions: dict[str, Tensor], + ) -> dict[str, Tensor]: + loss_dict: dict[str, Tensor] = self.criterion( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + auxiliary_predictions=auxiliary_predictions, + ) + + # weight each loss by `self.weight_dict[]` including auxiliary losses + for key, weight in self.weight_dict.items(): + for loss_key, loss in loss_dict.items(): + if key in loss_key: + loss *= weight + + return loss_dict + + def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: + return sum(loss_dict.values()) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: Tensor, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + patch_offsets: list[Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> EomtDinov3ForUniversalSegmentationOutput: + r""" + mask_labels (`list[torch.Tensor]`, *optional*): + list of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + patch_offsets (`list[torch.Tensor]`, *optional*): + list of tuples indicating the image index and start and end positions of patches for semantic segmentation. + """ + masks_queries_logits_per_layer, class_queries_logits_per_layer = (), () + + hidden_states = self.dropout(self.embeddings(pixel_values)) + position_embeddings = self.rope_embeddings(pixel_values.to(hidden_states.dtype)) + attention_mask = None + + for idx, layer_module in enumerate(self.layers): + if idx == self.num_hidden_layers - self.config.num_blocks: + query = self.query.weight[None, :, :].expand(hidden_states.shape[0], -1, -1).to(hidden_states.device) + hidden_states = torch.cat((query, hidden_states), dim=1) + + if idx >= self.num_hidden_layers - self.config.num_blocks and ( + self.training or self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks] > 0 + ): + norm_hidden_states = self.layernorm(hidden_states) + masks_queries_logits, class_queries_logits = self.predict(norm_hidden_states) + + masks_queries_logits_per_layer += (masks_queries_logits,) + class_queries_logits_per_layer += (class_queries_logits,) + + attention_mask = torch.ones( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.shape[1], + device=hidden_states.device, + dtype=torch.bool, + ) + + interpolated_logits = F.interpolate(masks_queries_logits, size=self.grid_size, mode="bilinear") + interpolated_logits = interpolated_logits.view( + interpolated_logits.size(0), interpolated_logits.size(1), -1 + ) + + num_query_tokens = self.config.num_queries + encoder_start_tokens = num_query_tokens + self.num_prefix_tokens + + # Set attention mask for queries to focus on encoder tokens based on interpolated logits + attention_mask[:, :num_query_tokens, encoder_start_tokens:] = interpolated_logits > 0 + + # Disable attention mask for random query tokens. + attention_mask = self._disable_attention_mask( + attention_mask, + prob=self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks], + num_query_tokens=num_query_tokens, + encoder_start_tokens=encoder_start_tokens, + device=attention_mask.device, + ) + + # Expand attention mask to 4d mask. + attention_mask = attention_mask[:, None, ...].expand(-1, self.config.num_attention_heads, -1, -1) + dtype_min = torch.finfo(hidden_states.dtype).min + attention_mask = attention_mask.to(hidden_states.dtype).masked_fill(~attention_mask, dtype_min) + + hidden_states = layer_module( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + + sequence_output = self.layernorm(hidden_states) + + masks_queries_logits, class_queries_logits = self.predict(sequence_output) + masks_queries_logits_per_layer += (masks_queries_logits,) + class_queries_logits_per_layer += (class_queries_logits,) + + loss = None + if mask_labels is not None and class_labels is not None: + loss = 0.0 + for masks_queries_logits, class_queries_logits in zip( + masks_queries_logits_per_layer, class_queries_logits_per_layer + ): + loss_dict = self.get_loss_dict( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + auxiliary_predictions=None, + ) + loss += self.get_loss(loss_dict) + + return EomtDinov3ForUniversalSegmentationOutput( + loss=loss, + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + last_hidden_state=sequence_output, + patch_offsets=patch_offsets, + ) + + def get_input_embeddings(self): + return self.embeddings.patch_embeddings + + def predict(self, logits: torch.Tensor): + query_tokens = logits[:, : self.config.num_queries, :] + class_logits = self.class_predictor(query_tokens) + + prefix_tokens = logits[:, self.config.num_queries + self.embeddings.num_prefix_tokens :, :] + prefix_tokens = prefix_tokens.transpose(1, 2) + + prefix_tokens = prefix_tokens.reshape(prefix_tokens.shape[0], -1, *self.grid_size) + + query_tokens = self.mask_head(query_tokens) + prefix_tokens = self.upscale_block(prefix_tokens) + + mask_logits = torch.einsum("bqc, bchw -> bqhw", query_tokens, prefix_tokens) + + return mask_logits, class_logits + + @staticmethod + def _disable_attention_mask(attn_mask, prob, num_query_tokens, encoder_start_tokens, device): + if prob < 1: + # Generate random queries to disable based on the probs + random_queries = torch.rand(attn_mask.shape[0], num_query_tokens, device=device) > prob + + # Disable attention to the query tokens, considering the prefix tokens + attn_mask[:, :num_query_tokens, encoder_start_tokens:][random_queries] = 1 + + return attn_mask + + +__all__ = ["EomtDinov3PreTrainedModel", "EomtDinov3ForUniversalSegmentation"] diff --git a/third_party/transformers/src/transformers/models/eomt_dinov3/modular_eomt_dinov3.py b/third_party/transformers/src/transformers/models/eomt_dinov3/modular_eomt_dinov3.py new file mode 100644 index 0000000000000000000000000000000000000000..2b6920a3b75838da9c9abc247d5f38782a365ed4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/eomt_dinov3/modular_eomt_dinov3.py @@ -0,0 +1,364 @@ +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch EoMT model backed by DINOv3.""" + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.functional as F +from huggingface_hub.dataclasses import strict +from torch import Tensor, nn + +from ... import initialization as init +from ...modeling_rope_utils import RopeParameters +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, +) +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..dinov3_vit.modeling_dinov3_vit import ( + DINOv3ViTAttention, + DINOv3ViTEmbeddings, + DINOv3ViTLayer, + DINOv3ViTLayerScale, + DINOv3ViTRopePositionEmbedding, +) +from ..eomt.configuration_eomt import EomtConfig +from ..eomt.modeling_eomt import ( + EomtForUniversalSegmentation, + EomtForUniversalSegmentationOutput, + EomtLoss, + EomtPreTrainedModel, +) + + +@auto_docstring(checkpoint="tue-mps/coco_panoptic_eomt_large_640_dinov3") +@strict +class EomtDinov3Config(EomtConfig): + r""" + layerscale_value (`float`, *optional*, defaults to 1.0): + Initial value for the LayerScale parameter. + num_upscale_blocks (`int`, *optional*, defaults to 2): + Number of upsampling blocks used in the decoder or segmentation head. + num_blocks (`int`, *optional*, defaults to 4): + Number of feature blocks or stages in the architecture. + no_object_weight (`float`, *optional*, defaults to 0.1): + Loss weight for the "no object" class in panoptic/instance segmentation. + train_num_points (`int`, *optional*, defaults to 12544): + Number of points to sample for mask loss computation during training. + oversample_ratio (`float`, *optional*, defaults to 3.0): + Oversampling ratio used in point sampling for mask training. + importance_sample_ratio (`float`, *optional*, defaults to 0.75): + Ratio of points to sample based on importance during training. + num_queries (`int`, *optional*, defaults to 200): + Number of object queries in the Transformer. + num_register_tokens (`int`, *optional*, defaults to 4): + Number of learnable register tokens added to the transformer input. + query_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in query projection. + key_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in key projection. + value_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in value projection. + proj_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in output projection. + use_gated_mlp (`bool`, *optional*, defaults to `False`): + Whether to use gated MLP layers. + pos_embed_shift (`float`, *optional*): + Shift value for position embeddings. + pos_embed_jitter (`float`, *optional*): + Jitter value for position embeddings. + pos_embed_rescale (`float`, *optional*, defaults to 2.0): + Rescale value for position embeddings. + """ + + model_type = "eomt_dinov3" + default_theta = 100.0 + + hidden_size: int = 1024 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + intermediate_size: int = 4096 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-6 + image_size: int | list[int] | tuple[int, int] = 640 + patch_size: int | list[int] | tuple[int, int] = 16 + num_channels: int = 3 + layerscale_value: float = 1.0 + drop_path_rate: float | int = 0.0 + num_upscale_blocks: int = 2 + attention_dropout: float | int = 0.0 + num_blocks: int = 4 + no_object_weight: float = 0.1 + class_weight: float = 2.0 + mask_weight: float = 5.0 + dice_weight: float = 5.0 + train_num_points: int = 12544 + oversample_ratio: float = 3.0 + importance_sample_ratio: float = 0.75 + num_queries: int = 200 + num_register_tokens: int = 4 + rope_parameters: RopeParameters | dict | None = None + query_bias: bool = True + key_bias: bool = False + value_bias: bool = True + proj_bias: bool = True + mlp_bias: bool = True + use_gated_mlp: bool = False + pos_embed_shift: float | None = None + pos_embed_jitter: float | None = None + pos_embed_rescale: float | None = 2.0 + + mlp_ratio = AttributeError() + use_swiglu_ffn = AttributeError() + + +class EomtDinov3Attention(DINOv3ViTAttention): + pass + + +class EomtDinov3Embeddings(DINOv3ViTEmbeddings): + def __init__(self, config: EomtDinov3Config): + super().__init__(config) + self.num_prefix_tokens = 1 + config.num_register_tokens + + +class EomtDinov3Layer(DINOv3ViTLayer): + pass + + +class EomtDinov3LayerScale(DINOv3ViTLayerScale): + pass + + +class EomtDinov3RotaryEmbedding(DINOv3ViTRopePositionEmbedding): + inv_freq: Tensor + + def __init__(self, config: EomtDinov3Config, device=None): + nn.Module.__init__(self) + 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": + raise ValueError("`EomtDinov3` only supports `default` RoPE! Please check your `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: EomtDinov3Config | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> torch.Tensor: + """ + 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"] + head_dim = config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1 / base ** torch.arange(0, 1, 4 / head_dim, dtype=torch.float32, device=device) + return inv_freq, attention_factor + + +class EomtDinov3Loss(EomtLoss): + pass + + +class EomtDinov3ForUniversalSegmentationOutput(EomtForUniversalSegmentationOutput): + pass + + +class EomtDinov3PreTrainedModel(EomtPreTrainedModel): + config_class = EomtDinov3Config + base_model_prefix = "eomt_dinov3" + _no_split_modules = ["EomtDinov3Layer"] + _can_record_outputs = { + "hidden_states": EomtDinov3Layer, + "attentions": EomtDinov3Attention, + } + + def _init_weights(self, module: nn.Module) -> None: + PreTrainedModel._init_weights(module) + std = self.config.initializer_range + if isinstance(module, EomtDinov3LayerScale): + if hasattr(module, "lambda1"): + init.constant_(module.lambda1, self.config.layerscale_value) + elif isinstance(module, EomtDinov3Embeddings): + init.trunc_normal_(module.cls_token, mean=0.0, std=std) + init.zeros_(module.register_tokens) + elif isinstance(module, EomtDinov3Loss): + empty_weight = torch.ones(module.num_labels + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + elif isinstance(module, EomtDinov3ForUniversalSegmentation): + init.ones_(module.attn_mask_probs) + + +@auto_docstring( + custom_intro=""" + The EoMT-DINOv3 model with head on top for instance/semantic/panoptic segmentation. + """, +) +class EomtDinov3ForUniversalSegmentation(EomtDinov3PreTrainedModel, EomtForUniversalSegmentation): + def __init__(self, config: EomtDinov3Config): + super().__init__(config) + + self.num_prefix_tokens = 1 + config.num_register_tokens + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.embeddings = EomtDinov3Embeddings(config) + self.embeddings.register_parameter("mask_token", None) + + self.rope_embeddings = EomtDinov3RotaryEmbedding(config) + self.layers = nn.ModuleList([EomtDinov3Layer(config) for _ in range(config.num_hidden_layers)]) + + self.post_init() + + # We redefine forward here because EoMT-DINOv3 uses DINOv3 backbone components (RoPE embeddings, layers) + # which require different integration than the base EoMT model that uses a separate encoder. + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + pixel_values: Tensor, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + patch_offsets: list[Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> EomtDinov3ForUniversalSegmentationOutput: + r""" + mask_labels (`list[torch.Tensor]`, *optional*): + list of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + patch_offsets (`list[torch.Tensor]`, *optional*): + list of tuples indicating the image index and start and end positions of patches for semantic segmentation. + """ + masks_queries_logits_per_layer, class_queries_logits_per_layer = (), () + + hidden_states = self.dropout(self.embeddings(pixel_values)) + position_embeddings = self.rope_embeddings(pixel_values.to(hidden_states.dtype)) + attention_mask = None + + for idx, layer_module in enumerate(self.layers): + if idx == self.num_hidden_layers - self.config.num_blocks: + query = self.query.weight[None, :, :].expand(hidden_states.shape[0], -1, -1).to(hidden_states.device) + hidden_states = torch.cat((query, hidden_states), dim=1) + + if idx >= self.num_hidden_layers - self.config.num_blocks and ( + self.training or self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks] > 0 + ): + norm_hidden_states = self.layernorm(hidden_states) + masks_queries_logits, class_queries_logits = self.predict(norm_hidden_states) + + masks_queries_logits_per_layer += (masks_queries_logits,) + class_queries_logits_per_layer += (class_queries_logits,) + + attention_mask = torch.ones( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.shape[1], + device=hidden_states.device, + dtype=torch.bool, + ) + + interpolated_logits = F.interpolate(masks_queries_logits, size=self.grid_size, mode="bilinear") + interpolated_logits = interpolated_logits.view( + interpolated_logits.size(0), interpolated_logits.size(1), -1 + ) + + num_query_tokens = self.config.num_queries + encoder_start_tokens = num_query_tokens + self.num_prefix_tokens + + # Set attention mask for queries to focus on encoder tokens based on interpolated logits + attention_mask[:, :num_query_tokens, encoder_start_tokens:] = interpolated_logits > 0 + + # Disable attention mask for random query tokens. + attention_mask = self._disable_attention_mask( + attention_mask, + prob=self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks], + num_query_tokens=num_query_tokens, + encoder_start_tokens=encoder_start_tokens, + device=attention_mask.device, + ) + + # Expand attention mask to 4d mask. + attention_mask = attention_mask[:, None, ...].expand(-1, self.config.num_attention_heads, -1, -1) + dtype_min = torch.finfo(hidden_states.dtype).min + attention_mask = attention_mask.to(hidden_states.dtype).masked_fill(~attention_mask, dtype_min) + + hidden_states = layer_module( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + + sequence_output = self.layernorm(hidden_states) + + masks_queries_logits, class_queries_logits = self.predict(sequence_output) + masks_queries_logits_per_layer += (masks_queries_logits,) + class_queries_logits_per_layer += (class_queries_logits,) + + loss = None + if mask_labels is not None and class_labels is not None: + loss = 0.0 + for masks_queries_logits, class_queries_logits in zip( + masks_queries_logits_per_layer, class_queries_logits_per_layer + ): + loss_dict = self.get_loss_dict( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + auxiliary_predictions=None, + ) + loss += self.get_loss(loss_dict) + + return EomtDinov3ForUniversalSegmentationOutput( + loss=loss, + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + last_hidden_state=sequence_output, + patch_offsets=patch_offsets, + ) + + +__all__ = [ + "EomtDinov3Config", + "EomtDinov3PreTrainedModel", + "EomtDinov3ForUniversalSegmentation", +] diff --git a/third_party/transformers/src/transformers/models/flaubert/__init__.py b/third_party/transformers/src/transformers/models/flaubert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e418a0f743810ffb168a48b9febff056b9a028d2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/flaubert/__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_flaubert import * + from .modeling_flaubert import * + from .tokenization_flaubert import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/flaubert/configuration_flaubert.py b/third_party/transformers/src/transformers/models/flaubert/configuration_flaubert.py new file mode 100644 index 0000000000000000000000000000000000000000..c205f3e4094b5c3cae62256b2ddd7bc170e31dde --- /dev/null +++ b/third_party/transformers/src/transformers/models/flaubert/configuration_flaubert.py @@ -0,0 +1,143 @@ +# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Flaubert configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="flaubert/flaubert_base_uncased") +@strict +class FlaubertConfig(PreTrainedConfig): + r""" + pre_norm (`bool`, *optional*, defaults to `False`): + Whether to apply the layer normalization before or after the feed forward layer following the attention in + each layer (Vaswani et al., Tensor2Tensor for Neural Machine Translation. 2018) + emb_dim (`int`, *optional*, defaults to 2048): + The dimensionality of embedding layer. + gelu_activation (`bool`, *optional*, defaults to True): + Whether to use GeLU activation function. + sinusoidal_embeddings (`bool`, *optional*, defaults to `False`): + Whether or not to use sinusoidal positional embeddings instead of absolute positional embeddings. + causal (`bool`, *optional*, defaults to `False`): + Whether or not the model should behave in a causal manner. Causal models use a triangular attention mask in + order to only attend to the left-side context instead if a bidirectional context. + asm (`bool`, *optional*, defaults to `False`): + Whether or not to use an adaptive log softmax projection layer instead of a linear layer for the prediction + layer. + n_langs (`int`, *optional*, defaults to 1): + The number of languages the model handles. Set to 1 for monolingual models. + use_lang_emb (`bool`, *optional*, defaults to `True`) + Whether to use language embeddings. Some models use additional language embeddings, see [the multilingual + models page](http://huggingface.co/transformers/multilingual.html#xlm-language-embeddings) for information + on how to use them. + embed_init_std (`float`, *optional*, defaults to 2048^-0.5): + The standard deviation of the truncated_normal_initializer for initializing the embedding matrices. + embed_init_std (`float`, *optional*, defaults to `2048**-0.5`): + Initializer std for embedding layers. + bos_index (`int`, *optional*, defaults to 0): + The index of the beginning of sentence token in the vocabulary. + eos_index (`int`, *optional*, defaults to 1): + The index of the end of sentence token in the vocabulary. + pad_index (`int`, *optional*, defaults to 2): + The index of the padding token in the vocabulary. + unk_index (`int`, *optional*, defaults to 3): + The index of the unknown token in the vocabulary. + mask_index (`int`, *optional*, defaults to 5): + The index of the masking token in the vocabulary. + is_encoder (`bool`, *optional*, defaults to True): + Whether the model is used as an encoder. + summary_type (`string`, *optional*, defaults to "first"): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Has to be one of the following options: + - `"last"`: Take the last token hidden state (like XLNet). + - `"first"`: Take the first token hidden state (like BERT). + - `"mean"`: Take the mean of all tokens hidden states. + - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2). + - `"attn"`: Not implemented now, use multi-head attention. + summary_use_proj (`bool`, *optional*, defaults to `True`): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Whether or not to add a projection after the vector extraction. + summary_activation (`str`, *optional*): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation. + summary_proj_to_labels (`bool`, *optional*, defaults to `True`): + Used in the sequence classification and multiple choice models. + Whether the projection outputs should have `config.num_labels` or `config.hidden_size` classes. + summary_first_dropout (`float`, *optional*, defaults to 0.1): + Used in the sequence classification and multiple choice models. + The dropout ratio to be used after the projection and activation. + start_n_top (`int`, *optional*, defaults to 5): + Used in the SQuAD evaluation script. + end_n_top (`int`, *optional*, defaults to 5): + Used in the SQuAD evaluation script. + mask_token_id (`int`, *optional*, defaults to 0): + Model agnostic parameter to identify masked tokens when generating text in an MLM context. + lang_id (`int`, *optional*, defaults to 1): + The ID of the language used by the model. This parameter is used when generating text in a given language. + """ + + model_type = "flaubert" + attribute_map = { + "hidden_size": "emb_dim", + "num_attention_heads": "n_heads", + "num_hidden_layers": "n_layers", + "n_words": "vocab_size", # For backward compatibility + "bos_index": "bos_token_id", + "eos_index": "eos_token_id", + "pad_index": "pad_token_id", + } + + pre_norm: bool = False + layerdrop: float | int = 0.0 + vocab_size: int = 30145 + emb_dim: int = 2048 + n_layers: int = 12 + n_heads: int = 16 + dropout: float | int = 0.1 + attention_dropout: float | int = 0.1 + gelu_activation: bool = True + sinusoidal_embeddings: bool = False + causal: bool = False + asm: bool = False + n_langs: int = 1 + use_lang_emb: bool = True + max_position_embeddings: int = 512 + embed_init_std: float = 2048**-0.5 + layer_norm_eps: float = 1e-12 + init_std: float = 0.02 + bos_index: int = 0 + eos_index: int = 1 + pad_index: int = 2 + unk_index: int = 3 + mask_index: int = 5 + is_encoder: bool = True + summary_type: str = "first" + summary_use_proj: bool = True + summary_activation: str | None = None + summary_proj_to_labels: bool = True + summary_first_dropout: float | int = 0.1 + start_n_top: int = 5 + end_n_top: int = 5 + mask_token_id: int = 0 + lang_id: int = 0 + pad_token_id: int | None = 2 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + tie_word_embeddings: bool = True + + +__all__ = ["FlaubertConfig"] diff --git a/third_party/transformers/src/transformers/models/flaubert/modeling_flaubert.py b/third_party/transformers/src/transformers/models/flaubert/modeling_flaubert.py new file mode 100644 index 0000000000000000000000000000000000000000..0a05370cda7aa3d0c14fab0e1994eb42330d5a5f --- /dev/null +++ b/third_party/transformers/src/transformers/models/flaubert/modeling_flaubert.py @@ -0,0 +1,1651 @@ +# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Flaubert model, based on XLM.""" + +import math +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import gelu, get_activation +from ...cache_utils import DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...modeling_outputs import ( + BaseModelOutput, + MaskedLMOutput, + MultipleChoiceModelOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_flaubert import FlaubertConfig + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.xlm.modeling_xlm.create_sinusoidal_embeddings +def create_sinusoidal_embeddings(n_pos, dim, out): + position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]) + out.requires_grad = False + out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2])) + out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2])) + out.detach_() + return out + + +# Copied from transformers.models.xlm.modeling_xlm.get_masks +def get_masks(slen, lengths, causal, padding_mask=None): + """ + Generate hidden states mask, and optionally an attention mask. + """ + alen = torch.arange(slen, dtype=torch.long, device=lengths.device) + if padding_mask is not None: + mask = padding_mask + else: + assert lengths.max().item() <= slen + mask = alen < lengths[:, None] + + # attention mask is the same as mask, or triangular inferior attention (causal) + bs = lengths.size(0) + if causal: + attn_mask = alen[None, None, :].repeat(bs, slen, 1) <= alen[None, :, None] + else: + attn_mask = mask + + # sanity check + assert mask.size() == (bs, slen) + assert causal is False or attn_mask.size() == (bs, slen, slen) + + return mask, attn_mask + + +# Copied from transformers.models.xlm.modeling_xlm.MultiHeadAttention +class MultiHeadAttention(nn.Module): + def __init__(self, n_heads, dim, config, layer_idx: int = 0): + super().__init__() + self.layer_id = layer_idx + self.dim = dim + self.n_heads = n_heads + self.head_dim = dim // n_heads + self.dropout = config.attention_dropout + assert self.dim % self.n_heads == 0 + + self.q_lin = nn.Linear(dim, dim) + self.k_lin = nn.Linear(dim, dim) + self.v_lin = nn.Linear(dim, dim) + self.out_lin = nn.Linear(dim, dim) + + def forward( + self, + input, + mask, + kv=None, + cache=None, + output_attentions=False, + **kwargs, + ): + """ + Self-attention (if kv is None) or attention over source sentence (provided by kv). + """ + # Input is (bs, qlen, dim) + # Mask is (bs, klen) (non-causal) or (bs, klen, klen) + bs, qlen, dim = input.size() + is_cross_attention = kv is not None + mask_reshape = (bs, 1, qlen, -1) if mask.dim() == 3 else (bs, 1, 1, -1) + + q = self.q_lin(input).view(bs, -1, self.n_heads, self.head_dim).transpose(1, 2) + if cache is not None: + if isinstance(cache, EncoderDecoderCache): + is_updated = cache.is_updated.get(self.layer_id) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = cache.cross_attention_cache + else: + curr_past_key_values = cache.self_attention_cache + else: + curr_past_key_values = cache + + current_states = kv if is_cross_attention else input + if is_cross_attention and cache is not None and is_updated: + # reuse k,v, cross_attentions + k = curr_past_key_values.key_cache[self.layer_id] + v = curr_past_key_values.value_cache[self.layer_id] + else: + k = self.k_lin(current_states) + v = self.v_lin(current_states) + k = k.view(bs, -1, self.n_heads, self.head_dim).transpose(1, 2) + v = v.view(bs, -1, self.n_heads, self.head_dim).transpose(1, 2) + + if cache is not None: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + k, v = curr_past_key_values.update(k, v, self.layer_id) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention: + cache.is_updated[self.layer_id] = True + + q = q / math.sqrt(self.head_dim) # (bs, n_heads, qlen, head_dim) + scores = torch.matmul(q, k.transpose(2, 3)) # (bs, n_heads, qlen, klen) + mask = (mask == 0).view(mask_reshape).expand_as(scores) # (bs, n_heads, qlen, klen) + scores.masked_fill_(mask, torch.finfo(scores.dtype).min) # (bs, n_heads, qlen, klen) + + weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) # (bs, n_heads, qlen, klen) + weights = nn.functional.dropout(weights, p=self.dropout, training=self.training) # (bs, n_heads, qlen, klen) + + context = torch.matmul(weights, v) # (bs, n_heads, qlen, head_dim) + context = context.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.head_dim) + + outputs = (self.out_lin(context),) + if output_attentions: + outputs = outputs + (weights,) + return outputs + + +# Copied from transformers.models.xlm.modeling_xlm.TransformerFFN +class TransformerFFN(nn.Module): + def __init__(self, in_dim, dim_hidden, out_dim, config): + super().__init__() + self.dropout = config.dropout + self.lin1 = nn.Linear(in_dim, dim_hidden) + self.lin2 = nn.Linear(dim_hidden, out_dim) + self.act = gelu if config.gelu_activation else nn.functional.relu + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + + def forward(self, input): + return apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input) + + def ff_chunk(self, input): + x = self.lin1(input) + x = self.act(x) + x = self.lin2(x) + x = nn.functional.dropout(x, p=self.dropout, training=self.training) + return x + + +@auto_docstring( + custom_intro=""" + The bare Flaubert Model transformer outputting raw hidden-states without any specific head on top. + """ +) +# Copied from transformers.models.xlm.modeling_xlm.XLMPredLayer with XLM->Flaubert +class FlaubertPredLayer(nn.Module): + """ + Prediction layer (cross_entropy or adaptive_softmax). + """ + + def __init__(self, config): + super().__init__() + self.asm = config.asm + self.n_words = config.n_words + self.pad_index = config.pad_index + dim = config.emb_dim + + if config.asm is False: + self.proj = nn.Linear(dim, config.n_words, bias=True) + else: + self.proj = nn.AdaptiveLogSoftmaxWithLoss( + in_features=dim, + n_classes=config.n_words, + cutoffs=config.asm_cutoffs, + div_value=config.asm_div_value, + head_bias=True, # default is False + ) + + def forward(self, x, y=None): + """Compute the loss, and optionally the scores.""" + outputs = () + if self.asm is False: + scores = self.proj(x) + outputs = (scores,) + outputs + if y is not None: + loss = nn.functional.cross_entropy(scores.view(-1, self.n_words), y.view(-1), reduction="mean") + outputs = (loss,) + outputs + else: + scores = self.proj.log_prob(x) + outputs = (scores,) + outputs + if y is not None: + _, loss = self.proj(x, y) + outputs = (loss,) + outputs + + return outputs + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of question answering models using a [`~modeling_utils.FlaubertSQuADHead`]. + """ +) +# Copied from transformers.models.xlm.modeling_xlm.XLMSquadHeadOutput with XLM->Flaubert +class FlaubertSquadHeadOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided): + Classification loss as the sum of start token, end token (and is_impossible if provided) classification + losses. + start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top config.start_n_top start token possibilities (beam-search). + start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top config.start_n_top start token possibilities (beam-search). + end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities + (beam-search). + end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search). + cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the `is_impossible` label of the answers. + """ + + loss: torch.FloatTensor | None = None + start_top_log_probs: torch.FloatTensor | None = None + start_top_index: torch.LongTensor | None = None + end_top_log_probs: torch.FloatTensor | None = None + end_top_index: torch.LongTensor | None = None + cls_logits: torch.FloatTensor | None = None + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerStartLogits with XLM->Flaubert +class FlaubertPoolerStartLogits(nn.Module): + """ + Compute SQuAD start logits from sequence hidden states. + + Args: + config ([`FlaubertConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: FlaubertConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, 1) + + def forward(self, hidden_states: torch.FloatTensor, p_mask: torch.FloatTensor | None = None) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + Returns: + `torch.FloatTensor`: The start logits for SQuAD. + """ + x = self.dense(hidden_states).squeeze(-1) + + if p_mask is not None: + if p_mask.dtype == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerEndLogits with XLM->Flaubert +class FlaubertPoolerEndLogits(nn.Module): + """ + Compute SQuAD end logits from sequence hidden states. + + Args: + config ([`FlaubertConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config: FlaubertConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dense_1 = nn.Linear(config.hidden_size, 1) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + p_mask: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The end logits for SQuAD. + """ + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + slen, hsz = hidden_states.shape[-2:] + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions) # shape (bsz, 1, hsz) + start_states = start_states.expand(-1, slen, -1) # shape (bsz, slen, hsz) + + x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1)) + x = self.activation(x) + x = self.LayerNorm(x) + x = self.dense_1(x).squeeze(-1) + + if p_mask is not None: + if p_mask.dtype == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerAnswerClass with XLM->Flaubert +class FlaubertPoolerAnswerClass(nn.Module): + """ + Compute SQuAD 2.0 answer class from classification and start tokens hidden states. + + Args: + config ([`FlaubertConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: FlaubertConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + cls_index: torch.LongTensor | None = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The SQuAD 2.0 answer class. + """ + # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample. + hsz = hidden_states.shape[-1] + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions).squeeze(-2) # shape (bsz, hsz) + + if cls_index is not None: + cls_index = cls_index[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, hsz) + else: + cls_token_state = hidden_states[:, -1, :] # shape (bsz, hsz) + + x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1)) + x = self.activation(x) + x = self.dense_1(x).squeeze(-1) + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMSQuADHead with XLM->Flaubert +class FlaubertSQuADHead(nn.Module): + r""" + A SQuAD head inspired by XLNet. + + Args: + config ([`FlaubertConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config: FlaubertConfig): + super().__init__() + self.start_n_top = config.start_n_top + self.end_n_top = config.end_n_top + + self.start_logits = FlaubertPoolerStartLogits(config) + self.end_logits = FlaubertPoolerEndLogits(config) + self.answer_class = FlaubertPoolerAnswerClass(config) + + @auto_docstring + def forward( + self, + hidden_states: torch.FloatTensor, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + cls_index: torch.LongTensor | None = None, + is_impossible: torch.LongTensor | None = None, + p_mask: torch.FloatTensor | None = None, + return_dict: bool = False, + ) -> FlaubertSquadHeadOutput | tuple[torch.FloatTensor]: + r""" + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + Final hidden states of the model on the sequence tokens. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the first token for the labeled span. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the last token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Whether the question has a possible answer in the paragraph or not. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + """ + start_logits = self.start_logits(hidden_states, p_mask=p_mask) + + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, let's remove the dimension added by batch splitting + for x in (start_positions, end_positions, cls_index, is_impossible): + if x is not None and x.dim() > 1: + x.squeeze_(-1) + + # during training, compute the end logits based on the ground truth of the start position + end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask) + + loss_fct = CrossEntropyLoss() + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if cls_index is not None and is_impossible is not None: + # Predict answerability from the representation of CLS and START + cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index) + loss_fct_cls = nn.BCEWithLogitsLoss() + cls_loss = loss_fct_cls(cls_logits, is_impossible) + + # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss + total_loss += cls_loss * 0.5 + + return FlaubertSquadHeadOutput(loss=total_loss) if return_dict else (total_loss,) + + else: + # during inference, compute the end logits based on beam search + bsz, slen, hsz = hidden_states.size() + start_log_probs = nn.functional.softmax(start_logits, dim=-1) # shape (bsz, slen) + + start_top_log_probs, start_top_index = torch.topk( + start_log_probs, self.start_n_top, dim=-1 + ) # shape (bsz, start_n_top) + start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz) + start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz) + start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz) + + hidden_states_expanded = hidden_states.unsqueeze(2).expand_as( + start_states + ) # shape (bsz, slen, start_n_top, hsz) + p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None + end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask) + end_log_probs = nn.functional.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top) + + end_top_log_probs, end_top_index = torch.topk( + end_log_probs, self.end_n_top, dim=1 + ) # shape (bsz, end_n_top, start_n_top) + end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top) + end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top) + + start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs) + cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index) + + if not return_dict: + return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + else: + return FlaubertSquadHeadOutput( + start_top_log_probs=start_top_log_probs, + start_top_index=start_top_index, + end_top_log_probs=end_top_log_probs, + end_top_index=end_top_index, + cls_logits=cls_logits, + ) + + +# Copied from transformers.models.xlm.modeling_xlm.XLMSequenceSummary with XLM->Flaubert +class FlaubertSequenceSummary(nn.Module): + r""" + Compute a single vector summary of a sequence hidden states. + + Args: + config ([`FlaubertConfig`]): + The config used by the model. Relevant arguments in the config class of the model are (refer to the actual + config class of your model for the default values it uses): + + - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are: + + - `"last"` -- Take the last token hidden state (like XLNet) + - `"first"` -- Take the first token hidden state (like Bert) + - `"mean"` -- Take the mean of all tokens hidden states + - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2) + - `"attn"` -- Not implemented now, use multi-head attention + + - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction. + - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes + (otherwise to `config.hidden_size`). + - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output, + another string or `None` will add no activation. + - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation. + - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation. + """ + + def __init__(self, config: FlaubertConfig): + super().__init__() + + self.summary_type = getattr(config, "summary_type", "last") + if self.summary_type == "attn": + # We should use a standard multi-head attention module with absolute positional embedding for that. + # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276 + # We can probably just use the multi-head attention module of PyTorch >=1.1.0 + raise NotImplementedError + + self.summary = nn.Identity() + if hasattr(config, "summary_use_proj") and config.summary_use_proj: + if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0: + num_classes = config.num_labels + else: + num_classes = config.hidden_size + self.summary = nn.Linear(config.hidden_size, num_classes) + + activation_string = getattr(config, "summary_activation", None) + self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity() + + self.first_dropout = nn.Identity() + if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0: + self.first_dropout = nn.Dropout(config.summary_first_dropout) + + self.last_dropout = nn.Identity() + if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0: + self.last_dropout = nn.Dropout(config.summary_last_dropout) + + def forward( + self, hidden_states: torch.FloatTensor, cls_index: torch.LongTensor | None = None + ) -> torch.FloatTensor: + """ + Compute a single vector summary of a sequence hidden states. + + Args: + hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`): + The hidden states of the last layer. + cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*): + Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token. + + Returns: + `torch.FloatTensor`: The summary of the sequence hidden states. + """ + if self.summary_type == "last": + output = hidden_states[:, -1] + elif self.summary_type == "first": + output = hidden_states[:, 0] + elif self.summary_type == "mean": + output = hidden_states.mean(dim=1) + elif self.summary_type == "cls_index": + if cls_index is None: + cls_index = torch.full_like( + hidden_states[..., :1, :], + hidden_states.shape[-2] - 1, + dtype=torch.long, + ) + else: + cls_index = cls_index.unsqueeze(-1).unsqueeze(-1) + cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),)) + # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states + output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size) + elif self.summary_type == "attn": + raise NotImplementedError + + output = self.first_dropout(output) + output = self.summary(output) + output = self.activation(output) + output = self.last_dropout(output) + + return output + + +@auto_docstring +# Copied from transformers.models.xlm.modeling_xlm.XLMPreTrainedModel with XLM->Flaubert +class FlaubertPreTrainedModel(PreTrainedModel): + config: FlaubertConfig + base_model_prefix = "transformer" + + @property + def dummy_inputs(self): + inputs_list = torch.tensor([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]) + attns_list = torch.tensor([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]]) + if self.config.use_lang_emb and self.config.n_langs > 1: + langs_list = torch.tensor([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]]) + else: + langs_list = None + return {"input_ids": inputs_list, "attention_mask": attns_list, "langs": langs_list} + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, nn.Embedding): + if self.config is not None and self.config.embed_init_std is not None: + init.normal_(module.weight, mean=0, std=self.config.embed_init_std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + if isinstance(module, nn.Linear): + if self.config is not None and self.config.init_std is not None: + init.normal_(module.weight, mean=0, std=self.config.init_std) + if module.bias is not None: + init.constant_(module.bias, 0.0) + if isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + if isinstance(module, FlaubertModel): + if self.config.sinusoidal_embeddings: + init.copy_( + module.position_embeddings.weight, + create_sinusoidal_embeddings( + self.config.max_position_embeddings, + self.config.emb_dim, + out=torch.empty_like(module.position_embeddings.weight), + ), + ) + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@auto_docstring +class FlaubertModel(FlaubertPreTrainedModel): + def __init__(self, config): # , dico, is_encoder, with_output): + super().__init__(config) + + # encoder / decoder, output layer + self.is_encoder = config.is_encoder + self.is_decoder = not config.is_encoder + if self.is_decoder: + raise NotImplementedError("Currently Flaubert can only be used as an encoder") + # self.with_output = with_output + self.causal = config.causal + + # dictionary / languages + self.n_langs = config.n_langs + self.use_lang_emb = config.use_lang_emb + self.n_words = config.n_words + self.eos_index = config.eos_index + self.pad_index = config.pad_index + # self.dico = dico + # self.id2lang = config.id2lang + # self.lang2id = config.lang2id + # assert len(self.dico) == self.n_words + # assert len(self.id2lang) == len(self.lang2id) == self.n_langs + + # model parameters + self.dim = config.emb_dim # 512 by default + self.hidden_dim = self.dim * 4 # 2048 by default + self.n_heads = config.n_heads # 8 by default + self.n_layers = config.n_layers + self.dropout = config.dropout + self.attention_dropout = config.attention_dropout + assert self.dim % self.n_heads == 0, "transformer dim must be a multiple of n_heads" + + # embeddings + self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.dim) + if config.n_langs > 1 and config.use_lang_emb: + self.lang_embeddings = nn.Embedding(self.n_langs, self.dim) + self.embeddings = nn.Embedding(self.n_words, self.dim, padding_idx=self.pad_index) + self.layer_norm_emb = nn.LayerNorm(self.dim, eps=config.layer_norm_eps) + + # transformer layers + self.attentions = nn.ModuleList() + self.layer_norm1 = nn.ModuleList() + self.ffns = nn.ModuleList() + self.layer_norm2 = nn.ModuleList() + # if self.is_decoder: + # self.layer_norm15 = nn.ModuleList() + # self.encoder_attn = nn.ModuleList() + + for i in range(self.n_layers): + self.attentions.append(MultiHeadAttention(self.n_heads, self.dim, config=config, layer_idx=i)) + self.layer_norm1.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps)) + # if self.is_decoder: + # self.layer_norm15.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps)) + # self.encoder_attn.append(MultiHeadAttention(self.n_heads, self.dim, dropout=self.attention_dropout)) + self.ffns.append(TransformerFFN(self.dim, self.hidden_dim, self.dim, config=config)) + self.layer_norm2.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps)) + + self.layerdrop = getattr(config, "layerdrop", 0.0) + self.pre_norm = getattr(config, "pre_norm", False) + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + # Initialize weights and apply final processing + self.post_init() + + # Copied from transformers.models.xlm.modeling_xlm.XLMModel.get_input_embeddings + def get_input_embeddings(self): + return self.embeddings + + # Copied from transformers.models.xlm.modeling_xlm.XLMModel.set_input_embeddings + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + lengths: torch.LongTensor | None = None, + cache: dict[str, torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use `attention_mask` for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`: + cache (`dict[str, torch.FloatTensor]`, *optional*): + Dictionary strings to `torch.FloatTensor` that contains precomputed hidden-states (key and values in the + attention blocks) as computed by the model (see `cache` output below). Can be used to speed up sequential + decoding. The dictionary object will be modified in-place during the forward pass to add newly computed + hidden-states. + """ + 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 + + # removed: src_enc=None, src_len=None + if input_ids is not None: + bs, slen = input_ids.size() + else: + bs, slen = inputs_embeds.size()[:-1] + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if cache is None: + cache = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + + if lengths is None: + if input_ids is not None: + lengths = (input_ids != self.pad_index).sum(dim=1).long() + else: + lengths = torch.full((bs,), slen, device=device, dtype=torch.long) + # mask = input_ids != self.pad_index + + # check inputs + assert lengths.size(0) == bs + assert lengths.max().item() <= slen + # input_ids = input_ids.transpose(0, 1) # batch size as dimension 0 + # assert (src_enc is None) == (src_len is None) + # if src_enc is not None: + # assert self.is_decoder + # assert src_enc.size(0) == bs + + # generate masks + mask, attn_mask = get_masks(slen, lengths, self.causal, padding_mask=attention_mask) + # if self.is_decoder and src_enc is not None: + # src_mask = torch.arange(src_len.max(), dtype=torch.long, device=lengths.device) < src_len[:, None] + + # Setting the position-ids to the registered buffer in constructor, it helps + # when tracing the model without passing position-ids, solves + # issues similar to issue #5664 + if position_ids is None: + if hasattr(self, "position_ids"): + position_ids = self.position_ids[:, :slen] + position_ids = position_ids.expand((bs, slen)) + else: + position_ids = torch.arange(slen, dtype=torch.long, device=device) + position_ids = position_ids.unsqueeze(0).expand((bs, slen)) + else: + assert position_ids.size() == (bs, slen) # (slen, bs) + # position_ids = position_ids.transpose(0, 1) + + # langs + if langs is not None: + assert langs.size() == (bs, slen) # (slen, bs) + # langs = langs.transpose(0, 1) + + # do not recompute cached elements + if cache is not None and input_ids is not None: + _slen = slen - cache.get_seq_length() + input_ids = input_ids[:, -_slen:] + position_ids = position_ids[:, -_slen:] + if langs is not None: + langs = langs[:, -_slen:] + mask = mask[:, -_slen:] + attn_mask = attn_mask[:, -_slen:] + + # embeddings + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + tensor = inputs_embeds + self.position_embeddings(position_ids).expand_as(inputs_embeds) + if langs is not None and self.use_lang_emb and self.config.n_langs > 1: + tensor = tensor + self.lang_embeddings(langs) + if token_type_ids is not None: + tensor = tensor + self.embeddings(token_type_ids) + tensor = self.layer_norm_emb(tensor) + tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training) + tensor *= mask.unsqueeze(-1).to(tensor.dtype) + + # transformer layers + hidden_states = () if output_hidden_states else None + attentions = () if output_attentions else None + for i in range(self.n_layers): + # LayerDrop + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + if output_hidden_states: + hidden_states = hidden_states + (tensor,) + + # self attention + if not self.pre_norm: + attn_outputs = self.attentions[i]( + tensor, + attn_mask, + cache=cache, + output_attentions=output_attentions, + ) + attn = attn_outputs[0] + if output_attentions: + attentions = attentions + (attn_outputs[1],) + attn = nn.functional.dropout(attn, p=self.dropout, training=self.training) + tensor = tensor + attn + tensor = self.layer_norm1[i](tensor) + else: + tensor_normalized = self.layer_norm1[i](tensor) + attn_outputs = self.attentions[i](tensor_normalized, attn_mask, cache=cache[i]) + attn = attn_outputs[0] + if output_attentions: + attentions = attentions + (attn_outputs[1],) + attn = nn.functional.dropout(attn, p=self.dropout, training=self.training) + tensor = tensor + attn + + # FFN + if not self.pre_norm: + tensor = tensor + self.ffns[i](tensor) + tensor = self.layer_norm2[i](tensor) + else: + tensor_normalized = self.layer_norm2[i](tensor) + tensor = tensor + self.ffns[i](tensor_normalized) + + tensor *= mask.unsqueeze(-1).to(tensor.dtype) + + # Add last hidden state + if output_hidden_states: + hidden_states = hidden_states + (tensor,) + + if not return_dict: + return tuple(v for v in [tensor, hidden_states, attentions] if v is not None) + + return BaseModelOutput(last_hidden_state=tensor, hidden_states=hidden_states, attentions=attentions) + + +@auto_docstring( + custom_intro=""" + The Flaubert Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """ +) +class FlaubertWithLMHeadModel(FlaubertPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"pred_layer.proj.weight": "transformer.embeddings.weight"} + + def __init__(self, config): + super().__init__(config) + self.transformer = FlaubertModel(config) + self.pred_layer = FlaubertPredLayer(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.pred_layer.proj + + def set_output_embeddings(self, new_embeddings): + self.pred_layer.proj = new_embeddings + + def prepare_inputs_for_generation(self, input_ids, **kwargs): + # Overwritten -- uses a language id + + mask_token_id = self.config.mask_token_id + lang_id = self.config.lang_id + + effective_batch_size = input_ids.shape[0] + mask_token = torch.full((effective_batch_size, 1), mask_token_id, dtype=torch.long, device=input_ids.device) + input_ids = torch.cat([input_ids, mask_token], dim=1) + if lang_id is not None: + langs = torch.full_like(input_ids, lang_id) + else: + langs = None + return {"input_ids": input_ids, "langs": langs} + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MaskedLMOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use `attention_mask` for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`: + cache (`dict[str, torch.FloatTensor]`, *optional*): + Dictionary strings to `torch.FloatTensor` that contains precomputed hidden-states (key and values in the + attention blocks) as computed by the model (see `cache` output below). Can be used to speed up sequential + decoding. The dictionary object will be modified in-place during the forward pass to add newly computed + hidden-states. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + output = transformer_outputs[0] + outputs = self.pred_layer(output, labels) # (loss, logits) or (logits,) depending on if labels are provided. + + if not return_dict: + return outputs + transformer_outputs[1:] + + return MaskedLMOutput( + loss=outputs[0] if labels is not None else None, + logits=outputs[0] if labels is None else outputs[1], + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + Flaubert Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) + e.g. for GLUE tasks. + """ +) +# Copied from transformers.models.xlm.modeling_xlm.XLMForSequenceClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert +class FlaubertForSequenceClassification(FlaubertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + + self.transformer = FlaubertModel(config) + self.sequence_summary = FlaubertSequenceSummary(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use *attention_mask* for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`. + cache (`dict[str, torch.FloatTensor]`, *optional*): + Instance of `EncoderDecoderCache` that contains precomputed KV states. Can be used to speed up sequential + decoding. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + output = transformer_outputs[0] + logits = self.sequence_summary(output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +# Copied from transformers.models.xlm.modeling_xlm.XLMForTokenClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert +class FlaubertForTokenClassification(FlaubertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = FlaubertModel(config) + self.dropout = nn.Dropout(config.dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | TokenClassifierOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use *attention_mask* for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`. + cache (`dict[str, torch.FloatTensor]`, *optional*): + Instance of `EncoderDecoderCache` that contains precomputed KV states. Can be used to speed up sequential + decoding. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + Flaubert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear + layers on top of the hidden-states output to compute `span start logits` and `span end logits`). + """ +) +# Copied from transformers.models.xlm.modeling_xlm.XLMForQuestionAnsweringSimple with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert +class FlaubertForQuestionAnsweringSimple(FlaubertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.transformer = FlaubertModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.Tensor | None = None, + end_positions: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | QuestionAnsweringModelOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use *attention_mask* for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`. + cache (`dict[str, torch.FloatTensor]`, *optional*): + Instance of `EncoderDecoderCache` that contains precomputed KV states. Can be used to speed up sequential + decoding. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = transformer_outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + transformer_outputs[1:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of question answering models using a `SquadHead`. + """ +) +# Copied from transformer.models.xlm.modeling_xlm.XLMForQuestionAnsweringOutput with XLM->Flaubert +class FlaubertForQuestionAnsweringOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided): + Classification loss as the sum of start token, end token (and is_impossible if provided) classification + losses. + start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top config.start_n_top start token possibilities (beam-search). + start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top config.start_n_top start token possibilities (beam-search). + end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities + (beam-search). + end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search). + cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the `is_impossible` label of the answers. + """ + + loss: torch.FloatTensor | None = None + start_top_log_probs: torch.FloatTensor | None = None + start_top_index: torch.LongTensor | None = None + end_top_log_probs: torch.FloatTensor | None = None + end_top_index: torch.LongTensor | None = None + cls_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@auto_docstring +# Copied from transformers.models.xlm.modeling_xlm.XLMForQuestionAnswering with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert +class FlaubertForQuestionAnswering(FlaubertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.transformer = FlaubertModel(config) + self.qa_outputs = FlaubertSQuADHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.Tensor | None = None, + end_positions: torch.Tensor | None = None, + is_impossible: torch.Tensor | None = None, + cls_index: torch.Tensor | None = None, + p_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | FlaubertForQuestionAnsweringOutput: + r""" + langs (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use *attention_mask* for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`. + cache (`dict[str, torch.FloatTensor]`, *optional*): + Instance of `EncoderDecoderCache` that contains precomputed KV states. Can be used to speed up sequential + decoding. + is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels whether a question has an answer or no answer (SQuAD 2.0) + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the classification token to use as input for computing plausibility of the + answer. + p_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Optional mask of tokens which can't be in answers (e.g. [CLS], [PAD], ...). 1.0 means token should be + masked. 0.0 mean token is not masked. + + Example: + + ```python + >>> from transformers import AutoTokenizer, FlaubertForQuestionAnswering + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-mlm-en-2048") + >>> model = FlaubertForQuestionAnswering.from_pretrained("FacebookAI/xlm-mlm-en-2048") + + >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze( + ... 0 + ... ) # Batch size 1 + >>> start_positions = torch.tensor([1]) + >>> end_positions = torch.tensor([3]) + + >>> outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions) + >>> loss = outputs.loss + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + output = transformer_outputs[0] + + outputs = self.qa_outputs( + output, + start_positions=start_positions, + end_positions=end_positions, + cls_index=cls_index, + is_impossible=is_impossible, + p_mask=p_mask, + return_dict=return_dict, + ) + + if not return_dict: + return outputs + transformer_outputs[1:] + + return FlaubertForQuestionAnsweringOutput( + loss=outputs.loss, + start_top_log_probs=outputs.start_top_log_probs, + start_top_index=outputs.start_top_index, + end_top_log_probs=outputs.end_top_log_probs, + end_top_index=outputs.end_top_index, + cls_logits=outputs.cls_logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +# Copied from transformers.models.xlm.modeling_xlm.XLMForMultipleChoice with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert +class FlaubertForMultipleChoice(FlaubertPreTrainedModel): + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + + self.transformer = FlaubertModel(config) + self.sequence_summary = FlaubertSequenceSummary(config) + self.logits_proj = nn.Linear(config.num_labels, 1) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + langs: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, + cache: dict[str, torch.Tensor] | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MultipleChoiceModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + langs (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are + languages ids which can be obtained from the language names by using two conversion mappings provided in + the configuration of the model (only provided for multilingual models). More precisely, the *language name + to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the + *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string). + + See usage examples detailed in the [multilingual documentation](../multilingual). + token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each sentence that can be used to avoid performing attention on padding token indices. You can + also use *attention_mask* for the same result (see above), kept here for compatibility. Indices selected in + `[0, ..., input_ids.size(-1)]`. + cache (`dict[str, torch.FloatTensor]`, *optional*): + Instance of `EncoderDecoderCache` that contains precomputed KV states. Can be used to speed up sequential + decoding. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None + langs = langs.view(-1, langs.size(-1)) if langs is not None else None + inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + if lengths is not None: + logger.warning( + "The `lengths` parameter cannot be used with the Flaubert multiple choice models. Please use the " + "attention mask instead." + ) + lengths = None + + transformer_outputs = self.transformer( + input_ids=input_ids, + attention_mask=attention_mask, + langs=langs, + token_type_ids=token_type_ids, + position_ids=position_ids, + lengths=lengths, + cache=cache, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + output = transformer_outputs[0] + logits = self.sequence_summary(output) + logits = self.logits_proj(logits) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + output = (reshaped_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return MultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +__all__ = [ + "FlaubertForMultipleChoice", + "FlaubertForQuestionAnswering", + "FlaubertForQuestionAnsweringSimple", + "FlaubertForSequenceClassification", + "FlaubertForTokenClassification", + "FlaubertModel", + "FlaubertWithLMHeadModel", + "FlaubertPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/flaubert/tokenization_flaubert.py b/third_party/transformers/src/transformers/models/flaubert/tokenization_flaubert.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbadd34db3eaa7b93b579df9a474fcaf528a06c --- /dev/null +++ b/third_party/transformers/src/transformers/models/flaubert/tokenization_flaubert.py @@ -0,0 +1,536 @@ +# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for Flaubert.""" + +import json +import os +import re +import unicodedata + +from ...tokenization_python import PreTrainedTokenizer +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", +} + + +def convert_to_unicode(text): + """ + Converts `text` to Unicode (if it's not already), assuming UTF-8 input. + """ + + def ensure_text(s, encoding="utf-8", errors="strict"): + if isinstance(s, bytes): + return s.decode(encoding, errors) + elif isinstance(s, str): + return s + else: + raise TypeError(f"not expecting type '{type(s)}'") + + return ensure_text(text, encoding="utf-8", errors="ignore") + + +# Copied from transformers.models.xlm.tokenization_xlm.get_pairs +def get_pairs(word): + """ + Return set of symbol pairs in a word. word is represented as tuple of symbols (symbols being variable-length + strings) + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +# Copied from transformers.models.xlm.tokenization_xlm.replace_unicode_punct +def replace_unicode_punct(text): + """ + Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/replace-unicode-punctuation.perl + """ + text = text.replace("，", ",") + text = re.sub(r"。\s*", ". ", text) + text = text.replace("、", ",") + text = text.replace("”", '"') + text = text.replace("“", '"') + text = text.replace("∶", ":") + text = text.replace("：", ":") + text = text.replace("？", "?") + text = text.replace("《", '"') + text = text.replace("》", '"') + text = text.replace("）", ")") + text = text.replace("！", "!") + text = text.replace("（", "(") + text = text.replace("；", ";") + text = text.replace("１", "1") + text = text.replace("」", '"') + text = text.replace("「", '"') + text = text.replace("０", "0") + text = text.replace("３", "3") + text = text.replace("２", "2") + text = text.replace("５", "5") + text = text.replace("６", "6") + text = text.replace("９", "9") + text = text.replace("７", "7") + text = text.replace("８", "8") + text = text.replace("４", "4") + text = re.sub(r"．\s*", ". ", text) + text = text.replace("～", "~") + text = text.replace("’", "'") + text = text.replace("…", "...") + text = text.replace("━", "-") + text = text.replace("〈", "<") + text = text.replace("〉", ">") + text = text.replace("【", "[") + text = text.replace("】", "]") + text = text.replace("％", "%") + return text + + +# Copied from transformers.models.xlm.tokenization_xlm.remove_non_printing_char +def remove_non_printing_char(text): + """ + Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/remove-non-printing-char.perl + """ + output = [] + for char in text: + cat = unicodedata.category(char) + if cat.startswith("C"): + continue + output.append(char) + return "".join(output) + + +class FlaubertTokenizer(PreTrainedTokenizer): + """ + Construct a Flaubert tokenizer. Based on Byte-Pair Encoding. The tokenization process is the following: + + - Moses preprocessing and tokenization. + - Normalizing all inputs text. + - The arguments `special_tokens` and the function `set_special_tokens`, can be used to add additional symbols (like + "__classify__") to a vocabulary. + - The argument `do_lowercase` controls lower casing (automatically set for pretrained vocabularies). + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + Vocabulary file. + merges_file (`str`): + Merges file. + do_lowercase (`bool`, *optional*, defaults to `False`): + Controls lower casing. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + + + + When building a sequence using special tokens, this is not the token that is used for the beginning of + sequence. The token used is the `cls_token`. + + + + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + additional_special_tokens (`List[str]`, *optional*, defaults to `['', '', '', '', '', '', '', '', '', '']`): + List of additional special tokens. + lang2id (`Dict[str, int]`, *optional*): + Dictionary mapping languages string identifiers to their IDs. + id2lang (`Dict[int, str]`, *optional*): + Dictionary mapping language IDs to their string identifiers. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + def __init__( + self, + vocab_file, + merges_file, + do_lowercase=False, + unk_token="", + bos_token="", + sep_token="", + pad_token="", + cls_token="", + mask_token="", + additional_special_tokens=[ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + ], + lang2id=None, + id2lang=None, + **kwargs, + ): + do_lowercase_and_remove_accent = kwargs.pop("do_lowercase_and_remove_accent", None) + if do_lowercase_and_remove_accent is not None: + logger.warning( + "`do_lowercase_and_remove_accent` is passed as a keyword argument, but this won't do anything." + " `FlaubertTokenizer` will always set it to `False`." + ) + # always `False` + self.do_lowercase_and_remove_accent = False + + self.do_lowercase = do_lowercase + + try: + import sacremoses + except ImportError: + raise ImportError( + "You need to install sacremoses to use FlaubertTokenizer. " + "See https://pypi.org/project/sacremoses/ for installation." + ) + + self.sm = sacremoses + + # cache of sm.MosesPunctNormalizer instance + self.cache_moses_punct_normalizer = {} + # cache of sm.MosesTokenizer instance + self.cache_moses_tokenizer = {} + self.lang_with_custom_tokenizer = {"zh", "th", "ja"} + self.lang2id = lang2id + self.id2lang = id2lang + if lang2id is not None and id2lang is not None: + assert len(lang2id) == len(id2lang) + + self.ja_word_tokenizer = None + self.zh_word_tokenizer = None + + with open(vocab_file, encoding="utf-8") as vocab_handle: + self.encoder = json.load(vocab_handle) + self.decoder = {v: k for k, v in self.encoder.items()} + with open(merges_file, encoding="utf-8") as merges_handle: + merges = merges_handle.read().split("\n")[:-1] + merges = [tuple(merge.split()[:2]) for merge in merges] + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {} + + super().__init__( + do_lowercase=do_lowercase, + unk_token=unk_token, + bos_token=bos_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + additional_special_tokens=additional_special_tokens, + lang2id=lang2id, + id2lang=id2lang, + **kwargs, + ) + + @property + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.do_lower_case + def do_lower_case(self): + return self.do_lowercase_and_remove_accent + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_punct_norm + def moses_punct_norm(self, text, lang): + if lang not in self.cache_moses_punct_normalizer: + punct_normalizer = self.sm.MosesPunctNormalizer(lang=lang) + self.cache_moses_punct_normalizer[lang] = punct_normalizer + else: + punct_normalizer = self.cache_moses_punct_normalizer[lang] + return punct_normalizer.normalize(text) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_tokenize + def moses_tokenize(self, text, lang): + if lang not in self.cache_moses_tokenizer: + moses_tokenizer = self.sm.MosesTokenizer(lang=lang) + self.cache_moses_tokenizer[lang] = moses_tokenizer + else: + moses_tokenizer = self.cache_moses_tokenizer[lang] + return moses_tokenizer.tokenize(text, return_str=False, escape=False) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_pipeline + def moses_pipeline(self, text, lang): + text = replace_unicode_punct(text) + text = self.moses_punct_norm(text, lang) + text = remove_non_printing_char(text) + return text + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.ja_tokenize + def ja_tokenize(self, text): + if self.ja_word_tokenizer is None: + try: + import Mykytea + + self.ja_word_tokenizer = Mykytea.Mykytea( + f"-model {os.path.expanduser('~')}/local/share/kytea/model.bin" + ) + except (AttributeError, ImportError): + logger.error( + "Make sure you install KyTea (https://github.com/neubig/kytea) and it's python wrapper" + " (https://github.com/chezou/Mykytea-python) with the following steps" + ) + logger.error("1. git clone git@github.com:neubig/kytea.git && cd kytea") + logger.error("2. autoreconf -i") + logger.error("3. ./configure --prefix=$HOME/local") + logger.error("4. make && make install") + logger.error("5. pip install kytea") + raise + return list(self.ja_word_tokenizer.getWS(text)) + + @property + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.vocab_size + def vocab_size(self): + return len(self.encoder) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.get_vocab + def get_vocab(self): + return dict(self.encoder, **self.added_tokens_encoder) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.bpe + def bpe(self, token): + word = tuple(token[:-1]) + (token[-1] + "",) + if token in self.cache: + return self.cache[token] + pairs = get_pairs(word) + + if not pairs: + return token + "" + + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + else: + new_word.extend(word[i:j]) + i = j + + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = " ".join(word) + if word == "\n ": + word = "\n" + self.cache[token] = word + return word + + def preprocess_text(self, text): + text = text.replace("``", '"').replace("''", '"') + text = convert_to_unicode(text) + text = unicodedata.normalize("NFC", text) + + if self.do_lowercase: + text = text.lower() + + return text + + def _tokenize(self, text, bypass_tokenizer=False): + """ + Tokenize a string given language code using Moses. + + Details of tokenization: + + - [sacremoses](https://github.com/alvations/sacremoses): port of Moses + - Install with `pip install sacremoses` + + Args: + - bypass_tokenizer: Allow users to preprocess and tokenize the sentences externally (default = False) + (bool). If True, we only apply BPE. + + Returns: + List of tokens. + """ + lang = "fr" + if lang and self.lang2id and lang not in self.lang2id: + logger.error( + "Supplied language code not found in lang2id mapping. Please check that your language is supported by" + " the loaded pretrained model." + ) + + if bypass_tokenizer: + text = text.split() + else: + text = self.preprocess_text(text) + text = self.moses_pipeline(text, lang=lang) + text = self.moses_tokenize(text, lang=lang) + + split_tokens = [] + for token in text: + if token: + split_tokens.extend(list(self.bpe(token).split(" "))) + + return split_tokens + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer._convert_token_to_id + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.encoder.get(token, self.encoder.get(self.unk_token)) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer._convert_id_to_token + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.decoder.get(index, self.unk_token) + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.convert_tokens_to_string + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + out_string = "".join(tokens).replace("", " ").strip() + return out_string + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.build_inputs_with_special_tokens + def build_inputs_with_special_tokens( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. An XLM sequence has the following format: + + - single sequence: ` X ` + - pair of sequences: ` A B ` + + Args: + token_ids_0 (`List[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + + """ + bos = [self.bos_token_id] + sep = [self.sep_token_id] + + if token_ids_1 is None: + return bos + token_ids_0 + sep + return bos + token_ids_0 + sep + token_ids_1 + sep + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.get_special_tokens_mask + def get_special_tokens_mask( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False + ) -> list[int]: + """ + Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer `prepare_for_model` method. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True + ) + + if token_ids_1 is not None: + return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1] + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.save_vocabulary + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + if not os.path.isdir(save_directory): + logger.error(f"Vocabulary path ({save_directory}) should be a directory") + return + vocab_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] + ) + merge_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"] + ) + + with open(vocab_file, "w", encoding="utf-8") as f: + f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n") + + index = 0 + with open(merge_file, "w", encoding="utf-8") as writer: + for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]): + if index != token_index: + logger.warning( + f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive." + " Please check that the tokenizer is not corrupted!" + ) + index = token_index + writer.write(" ".join(bpe_tokens) + "\n") + index += 1 + + return vocab_file, merge_file + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.__getstate__ + def __getstate__(self): + state = self.__dict__.copy() + state["sm"] = None + return state + + # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.__setstate__ + def __setstate__(self, d): + self.__dict__ = d + + try: + import sacremoses + except ImportError: + raise ImportError( + "You need to install sacremoses to use XLMTokenizer. " + "See https://pypi.org/project/sacremoses/ for installation." + ) + + self.sm = sacremoses + + +__all__ = ["FlaubertTokenizer"] diff --git a/third_party/transformers/src/transformers/models/gemma/__init__.py b/third_party/transformers/src/transformers/models/gemma/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..80c8d30760c45526b97bfffd557a6d39b8e3e826 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_gemma import * + from .modeling_gemma import * + from .tokenization_gemma import * + from .tokenization_gemma_fast import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/gemma/configuration_gemma.py b/third_party/transformers/src/transformers/models/gemma/configuration_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..0ef2e8b31b8d2d73783c7cc22fe9761f18eef36d --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/configuration_gemma.py @@ -0,0 +1,86 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.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_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/gemma-7b") +@strict +class GemmaConfig(PreTrainedConfig): + r""" + use_bidirectional_attention (`bool`, *optional*): + If True, the model will attend to all text tokens instead of using a causal mask. + + ```python + >>> from transformers import GemmaModel, GemmaConfig + >>> # Initializing a Gemma gemma-7b style configuration + >>> configuration = GemmaConfig() + >>> # Initializing a model from the gemma-7b style configuration + >>> model = GemmaModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gemma" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 256000 + hidden_size: int = 3072 + intermediate_size: int = 24576 + num_hidden_layers: int = 28 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 256 + hidden_act: str = "gelu_pytorch_tanh" + max_position_embeddings: int = 8192 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + bos_token_id: int | None = 2 + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + use_bidirectional_attention: bool | None = None + + +__all__ = ["GemmaConfig"] diff --git a/third_party/transformers/src/transformers/models/gemma/convert_gemma_weights_to_hf.py b/third_party/transformers/src/transformers/models/gemma/convert_gemma_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..ae06b9db2fdb2cc644d46b1f211320905d87939d --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/convert_gemma_weights_to_hf.py @@ -0,0 +1,194 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os + +import torch + +from transformers import GemmaConfig, GemmaForCausalLM, GemmaTokenizer +from transformers.tokenization_utils_sentencepiece import SentencePieceExtractor + + +""" +Sample usage: + +``` +python src/transformers/models/gemma/convert_gemma_weights_to_hf.py \ + --input_dir /path/to/downloaded/gemma/weights --model_size 7B --output_dir /output/path +``` + +Thereafter, models can be loaded via: + +```py +from transformers import GemmaForCausalLM, GemmaTokenizerFast + +model = GemmaForCausalLM.from_pretrained("/output/path") +tokenizer = GemmaTokenizerFast.from_pretrained("/output/path") +``` + +Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions +come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM). +""" + +gemma_2b_config = GemmaConfig( + num_hidden_layers=18, + num_attention_heads=8, + num_key_value_heads=1, + hidden_size=2048, + intermediate_size=16384, +) + +gemma_7b_config = GemmaConfig() + +CONFIG_MAPPING = {"2B": gemma_2b_config, "7B": gemma_7b_config} +LAYER_NAME_MAPPING = {"embedder.weight": "model.embed_tokens.weight"} + + +def write_model(save_path, input_base_path, config, push_to_hub=False, dtype=torch.float32): + num_attn_heads = config.num_attention_heads + hidden_size = config.hidden_size + num_kv_heads = config.num_key_value_heads + head_dim = config.head_dim + + print(f"Fetching all parameters from the checkpoint at '{input_base_path}'") + model_state_dict = torch.load(input_base_path, map_location="cpu", weights_only=True)["model_state_dict"] + model_state_dict.pop("freqs_cis") + + state_dict = {} + for k, v in model_state_dict.items(): + if "qkv_proj" in k: + if num_kv_heads == 1: + v = v.reshape(num_attn_heads + num_kv_heads * 2, head_dim, hidden_size) + q_proj = v[:num_attn_heads, ...] + k_proj = v[num_attn_heads : num_attn_heads + num_kv_heads, ...].repeat(num_kv_heads, 1, 1) + v_proj = v[-num_kv_heads:, ...].repeat(num_kv_heads, 1, 1) + + state_dict[k.replace("qkv_proj", "q_proj")] = q_proj.reshape( + num_attn_heads * head_dim, hidden_size + ).clone() + state_dict[k.replace("qkv_proj", "k_proj")] = k_proj.reshape( + num_kv_heads * head_dim, hidden_size + ).clone() + state_dict[k.replace("qkv_proj", "v_proj")] = v_proj[0].clone() + else: + q_proj, k_proj, v_proj = torch.split(v, v.shape[0] // 3, 0) + state_dict[k.replace("qkv_proj", "q_proj")] = q_proj.reshape( + num_attn_heads * head_dim, hidden_size + ).clone() + state_dict[k.replace("qkv_proj", "k_proj")] = k_proj.reshape( + num_kv_heads * head_dim, hidden_size + ).clone() + state_dict[k.replace("qkv_proj", "v_proj")] = v_proj.clone() + + elif k == "embedder.weight": + state_dict[LAYER_NAME_MAPPING[k]] = v + state_dict["lm_head.weight"] = v + else: + state_dict[k] = v + + torch.set_default_dtype(dtype) + + print("Loading the checkpoint in a Gemma model.") + with torch.device("meta"): + model = GemmaForCausalLM(config) + model.load_state_dict(state_dict, assign=True, strict=False) + + model.config.dtype = torch.float32 + del model.config._name_or_path + print("Saving in the Transformers format.") + + if push_to_hub: + print(f"pushing the model to {save_path}") + model.push_to_hub(save_path, private=True) + else: + model.save_pretrained(save_path) + + +def write_tokenizer(input_tokenizer_path, save_path, push_to_hub=False): + # Initialize the tokenizer based on the `spm` model + print(f"Saving a GemmaTokenizer to {save_path}.") + sentencepiece_extractor = SentencePieceExtractor(input_tokenizer_path) + vocab, _, merges = sentencepiece_extractor.extract() + tokenizer = GemmaTokenizer( + vocab=vocab, + merges=merges, + padding_side="left", + ) + if push_to_hub: + tokenizer.push_to_hub(save_path) + else: + tokenizer.save_pretrained(save_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_checkpoint", + help="Absolute path to the target Gemma weights.", + required=True, + ) + parser.add_argument( + "--tokenizer_checkpoint", + help="Location of Gemma tokenizer model", + ) + parser.add_argument( + "--model_size", + default="7B", + choices=["2B", "7B", "tokenizer_only"], + help="'f' models correspond to the finetuned versions, and are specific to the Gemma2 official release. For more details on Gemma2, check out the original repo: https://huggingface.co/google/gemma-7b", + ) + parser.add_argument( + "--output_dir", + default="google/gemma-7b", + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--convert_tokenizer", + help="Whether or not to convert the tokenizer as well.", + action="store_true", + default=False, + ) + parser.add_argument( + "--push_to_hub", + help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally.", + action="store_true", + default=False, + ) + parser.add_argument( + "--dtype", + default="float32", + help="Target dtype of the converted model", + ) + args = parser.parse_args() + + if args.convert_tokenizer: + if args.tokenizer_checkpoint is None: + raise ValueError("Path to the tokenizer is required when passing --convert_tokenizer") + + spm_path = os.path.join(args.tokenizer_checkpoint) + write_tokenizer(spm_path, args.output_dir, args.push_to_hub) + + config = CONFIG_MAPPING[args.model_size] + dtype = getattr(torch, args.dtype) + write_model( + config=config, + input_base_path=args.input_checkpoint, + save_path=args.output_dir, + push_to_hub=args.push_to_hub, + dtype=dtype, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/gemma/modeling_gemma.py b/third_party/transformers/src/transformers/models/gemma/modeling_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..c6c5a55b87900447b98456c3b17cb51cbc1620b7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/modeling_gemma.py @@ -0,0 +1,538 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.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_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import ( + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_gemma import GemmaConfig + + +class GemmaTextScaledWordEmbedding(nn.Embedding): + """ + This module overrides nn.Embeddings' forward by multiplying with embeddings scale. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float = 1.0): + super().__init__(num_embeddings, embedding_dim, padding_idx) + self.scalar_embed_scale = embed_scale + self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False) + + def forward(self, input_ids: torch.Tensor): + return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) + + +class GemmaRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +class GemmaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class GemmaRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: GemmaConfig, 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: GemmaConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernelized_func(apply_rotary_pos_emb) +class GemmaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = not getattr(config, "use_bidirectional_attention", False) + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class GemmaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx) + + self.mlp = GemmaMLP(config) + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaRMSNorm(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 GemmaPreTrainedModel(PreTrainedModel): + config: GemmaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GemmaDecoderLayer"] + _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": GemmaDecoderLayer, + "attentions": GemmaAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight) + if "RMSNorm" in module.__class__.__name__: + init.zeros_(module.weight) + elif isinstance(module, GemmaTextScaledWordEmbedding): + init.constant_(module.embed_scale, module.scalar_embed_scale) + + +@auto_docstring +class GemmaModel(GemmaPreTrainedModel): + def __init__(self, config: GemmaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + # Gemma3 downcasts the below to bfloat16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402 + self.embed_tokens = GemmaTextScaledWordEmbedding( + config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5 + ) + self.layers = nn.ModuleList( + [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = GemmaRotaryEmbedding(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 = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + # embed positions + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +@auto_docstring +class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = GemmaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, GemmaForCausalLM + + >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class GemmaForSequenceClassification(GenericForSequenceClassification, GemmaPreTrainedModel): + pass + + +class GemmaForTokenClassification(GenericForTokenClassification, GemmaPreTrainedModel): + pass + + +__all__ = [ + "GemmaModel", + "GemmaForCausalLM", + "GemmaForSequenceClassification", + "GemmaForTokenClassification", + "GemmaPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/gemma/modular_gemma.py b/third_party/transformers/src/transformers/models/gemma/modular_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..25f436473fbef06a4015ac29a0f0d7e95b59669d --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/modular_gemma.py @@ -0,0 +1,268 @@ +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from huggingface_hub.dataclasses import strict +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_rope_utils import RopeParameters +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaForTokenClassification, + LlamaMLP, + LlamaModel, + LlamaPreTrainedModel, + LlamaRotaryEmbedding, +) + + +VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"} + +SPIECE_UNDERLINE = "▁" + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="google/gemma-7b") +@strict +class GemmaConfig(PreTrainedConfig): + r""" + use_bidirectional_attention (`bool`, *optional*): + If True, the model will attend to all text tokens instead of using a causal mask. + + ```python + >>> from transformers import GemmaModel, GemmaConfig + >>> # Initializing a Gemma gemma-7b style configuration + >>> configuration = GemmaConfig() + >>> # Initializing a model from the gemma-7b style configuration + >>> model = GemmaModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gemma" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 256000 + hidden_size: int = 3072 + intermediate_size: int = 24576 + num_hidden_layers: int = 28 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 256 + hidden_act: str = "gelu_pytorch_tanh" + max_position_embeddings: int = 8192 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + bos_token_id: int | None = 2 + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + use_bidirectional_attention: bool | None = None + + +class GemmaTextScaledWordEmbedding(nn.Embedding): + """ + This module overrides nn.Embeddings' forward by multiplying with embeddings scale. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float = 1.0): + super().__init__(num_embeddings, embedding_dim, padding_idx) + self.scalar_embed_scale = embed_scale + self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False) + + def forward(self, input_ids: torch.Tensor): + return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) + + +class GemmaRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +class GemmaMLP(LlamaMLP): + def __init__(self, config): + super().__init__(config) + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + +class GemmaRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class GemmaAttention(LlamaAttention): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.is_causal = not getattr(config, "use_bidirectional_attention", False) + + +class GemmaPreTrainedModel(LlamaPreTrainedModel): + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight) + if "RMSNorm" in module.__class__.__name__: + init.zeros_(module.weight) + elif isinstance(module, GemmaTextScaledWordEmbedding): + init.constant_(module.embed_scale, module.scalar_embed_scale) + + +class GemmaModel(LlamaModel): + def __init__(self, config: GemmaConfig): + super().__init__(config) + # Gemma3 downcasts the below to bfloat16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402 + self.embed_tokens = GemmaTextScaledWordEmbedding( + config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5 + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + # embed positions + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +class GemmaForCausalLM(LlamaForCausalLM): + def forward(**super_kwargs): + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, GemmaForCausalLM + + >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + return super().forward(**super_kwargs) + + +class GemmaForSequenceClassification(LlamaForSequenceClassification): + pass + + +class GemmaForTokenClassification(LlamaForTokenClassification): + pass + + +__all__ = [ + "GemmaConfig", + "GemmaModel", + "GemmaForCausalLM", + "GemmaForSequenceClassification", + "GemmaForTokenClassification", + "GemmaPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/gemma/tokenization_gemma.py b/third_party/transformers/src/transformers/models/gemma/tokenization_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..85f0a74bb4c8db81defb4af2f53dc826773ff531 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gemma/tokenization_gemma.py @@ -0,0 +1,109 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers +from tokenizers.models import BPE + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) +VOCAB_FILES_NAMES = {"tokenizer_file": "tokenizer.json"} + + +class GemmaTokenizer(TokenizersBackend): + """ + Construct a fast Gemma tokenizer (backed by HuggingFace's tokenizers library). + + This tokenizer uses a BPE model with byte fallback, no prefix space, and a normalizer that replaces + spaces with "▁". + + Args: + tokenizer_file (`str`, optional): + A tokenizers JSON file containing the serialization of a tokenizer. + unk_token (`str`, optional, defaults to ""): + The unknown token. + bos_token (`str`, optional, defaults to ""): + The beginning of sequence token. + eos_token (`str`, optional, defaults to ""): + The end of sequence token. + pad_token (`str`, optional, defaults to ""): + The padding token. + mask_token (`str`, optional, defaults to ""): + The mask token. + add_bos_token (`bool`, optional, defaults to True): + Whether or not to add a `bos_token` at the start of sequences. + add_eos_token (`bool`, optional, defaults to False): + Whether or not to add an `eos_token` at the end of sequences. + vocab (`str` or `dict[str, int]`, optional): + Custom vocabulary dict. If not provided, a minimal vocabulary is created using the special tokens. + """ + + vocab_files_names = VOCAB_FILES_NAMES + padding_side = "left" + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + merges: str | list[str] | None = None, + unk_token: str = "", + bos_token: str = "", + eos_token: str = "", + pad_token: str = "", + mask_token: str = "", + **kwargs, + ): + if vocab is None: + vocab = { + str(pad_token): 0, + str(eos_token): 1, + str(bos_token): 2, + str(unk_token): 3, + str(mask_token): 4, + } + self._vocab = vocab + self._merges = merges or [] + + self._tokenizer = Tokenizer( + BPE( + vocab=self._vocab, + merges=self._merges, + fuse_unk=True, + unk_token=str(unk_token), + dropout=None, + byte_fallback=True, + ) + ) + self._tokenizer.pre_tokenizer = pre_tokenizers.Split( + pattern=" ", behavior="merged_with_previous", invert=False + ) + + self._tokenizer.decoder = decoders.Sequence( + [decoders.Replace("▁", " "), decoders.ByteFallback(), decoders.Fuse()] + ) + self._tokenizer.normalizer = normalizers.Replace(" ", "▁") + super().__init__( + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + mask_token=mask_token, + **kwargs, + ) + + +__all__ = ["GemmaTokenizer"] diff --git a/third_party/transformers/src/transformers/models/gpt_bigcode/__init__.py b/third_party/transformers/src/transformers/models/gpt_bigcode/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92e985d92734550a5b0635941294669386d35749 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_bigcode/__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_gpt_bigcode import * + from .modeling_gpt_bigcode import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py b/third_party/transformers/src/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py new file mode 100644 index 0000000000000000000000000000000000000000..8b8bafbce06763b107c26356cd539d74358047b7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py @@ -0,0 +1,87 @@ +# Copyright 2023 The BigCode team and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPTBigCode configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/gpt_bigcode") +@strict +class GPTBigCodeConfig(PreTrainedConfig): + r""" + scale_attn_weights (`bool`, *optional*, defaults to `True`): + Scale attention weights by dividing by sqrt(hidden_size).. + attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether to call the fused softmax in float32. + scale_attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether to scale the attention softmax in float32. + multi_query (`bool`, *optional*, defaults to `True`): + Whether to use multi-query attention in the decoder. Ignored when `new_decoder_architecture` is `True`. + + Example: + + ```python + >>> from transformers import GPTBigCodeConfig, GPTBigCodeModel + + >>> # Initializing a GPTBigCode configuration + >>> configuration = GPTBigCodeConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = GPTBigCodeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gpt_bigcode" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "n_embd", + "max_position_embeddings": "n_positions", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + vocab_size: int = 50257 + n_positions: int = 1024 + n_embd: int = 768 + n_layer: int = 12 + n_head: int = 12 + n_inner: int | None = None + activation_function: str = "gelu_pytorch_tanh" + resid_pdrop: float | int = 0.1 + embd_pdrop: float | int = 0.1 + attn_pdrop: float | int = 0.1 + layer_norm_epsilon: float = 1e-5 + initializer_range: float = 0.02 + scale_attn_weights: bool = True + use_cache: bool = True + bos_token_id: int | None = 50256 + eos_token_id: int | list[int] | None = 50256 + pad_token_id: int | None = None + attention_softmax_in_fp32: bool = True + scale_attention_softmax_in_fp32: bool = True + multi_query: bool = True + add_cross_attention: bool = False + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + self.num_key_value_heads = 1 if self.multi_query else self.n_head + super().__post_init__(**kwargs) + + +__all__ = ["GPTBigCodeConfig"] diff --git a/third_party/transformers/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/third_party/transformers/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py new file mode 100644 index 0000000000000000000000000000000000000000..bd3f2bbf5417ec8f1a1ab859449cbaa4fa3d00fc --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -0,0 +1,798 @@ +# Copyright 2023 The Bigcode team and HuggingFace Inc. team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch GPTBigCode model.""" + +import math +from collections.abc import Callable + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, +) +from ...utils.generic import is_flash_attention_requested, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_gpt_bigcode import GPTBigCodeConfig + + +logger = logging.get_logger(__name__) + + +# Fused kernels +# Use separate functions for each case because conditionals prevent kernel fusion. +# TODO: Could have better fused kernels depending on scaling and dropout. +# Is it doable without writing 32 functions? +@torch.jit.script +def upcast_masked_softmax( + x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor, scale: float, softmax_dtype: torch.dtype +): + input_dtype = x.dtype + x = x.to(softmax_dtype) * scale + x = torch.where(mask, x, mask_value) + x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype) + return x + + +@torch.jit.script +def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype): + input_dtype = x.dtype + x = x.to(softmax_dtype) * scale + x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype) + return x + + +@torch.jit.script +def masked_softmax(x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor): + x = torch.where(mask, x, mask_value) + x = torch.nn.functional.softmax(x, dim=-1) + return x + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class GPTBigCodeAttention(nn.Module): + def __init__(self, config, is_cross_attention=False, layer_idx=None): + super().__init__() + self.config = config + + self.mask_value = None + self.multi_query = config.multi_query + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.kv_heads = 1 if self.multi_query else self.num_heads + self.kv_dim = self.kv_heads * self.head_dim + self.num_key_value_groups = self.num_heads // self.kv_heads + self.split_size = self.embed_dim + self.is_causal = True + + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + + self.scale_attn_weights = config.scale_attn_weights + self.scaling = self.head_dim**-0.5 if config.scale_attn_weights else 1.0 + self.is_cross_attention = is_cross_attention + + self.layer_idx = layer_idx + self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 + self.scale_attention_softmax_in_fp32 = ( + config.scale_attention_softmax_in_fp32 and config.attention_softmax_in_fp32 + ) + self.attn_pdrop = config.attn_pdrop + + if self.is_cross_attention: + if self.multi_query: + raise NotImplementedError("Multi-Query Attention not supported for cross_attention") + + self.c_attn = nn.Linear(self.embed_dim, 2 * self.embed_dim) + self.q_attn = nn.Linear(self.embed_dim, self.embed_dim) + else: + self.c_attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.kv_dim) + + self.c_proj = nn.Linear(self.embed_dim, self.embed_dim) + + self.attn_dropout = config.attn_pdrop + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + def forward( + self, + hidden_states: torch.Tensor, + layer_past: Cache | None = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None] | tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor, ...]]: + input_shape = hidden_states.shape[:-1] + + if layer_past is not None: + if isinstance(layer_past, EncoderDecoderCache): + is_updated = layer_past.is_updated.get(self.layer_idx) + if self.is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = layer_past.cross_attention_cache + else: + curr_past_key_values = layer_past.self_attention_cache + else: + curr_past_key_values = layer_past + + if self.is_cross_attention: + if not hasattr(self, "q_attn") or not self.is_cross_attention: + raise ValueError( + "If class is used as cross attention, the weights `q_attn` have to be defined. " + "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`." + ) + if layer_past is not None and is_updated: + # reuse k,v, cross_attentions + key = curr_past_key_values.layers[self.layer_idx].keys + value = curr_past_key_values.layers[self.layer_idx].values + else: + query = self.q_attn(hidden_states).view(*input_shape, -1, self.head_dim).transpose(1, 2) + key, value = self.c_attn(encoder_hidden_states).split((self.head_dim, self.head_dim), dim=-1) + else: + if self.multi_query: + query, key, value = ( + self.c_attn(hidden_states).unsqueeze(1).split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3) + ) + query = query.view(*input_shape, -1, self.head_dim).transpose(1, 2) + else: + query, key, value = ( + self.c_attn(hidden_states) + .view(*hidden_states.shape[:2], self.num_heads, 3 * self.head_dim) + .transpose(1, 2) + .split(3 * [self.head_dim], dim=3) + ) + + if layer_past is not None: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + key, value = curr_past_key_values.update(key, value, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if self.is_cross_attention: + layer_past.is_updated[self.layer_idx] = True + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query, + key, + value, + attention_mask, + dropout=0.0 if not self.training else self.attn_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.c_proj(attn_output) + attn_output = self.resid_dropout(attn_output) + return attn_output, attn_weights + + +class GPTBigCodeMLP(nn.Module): + def __init__(self, intermediate_size, config): + super().__init__() + embed_dim = config.hidden_size + self.c_fc = nn.Linear(embed_dim, intermediate_size) + self.c_proj = nn.Linear(intermediate_size, embed_dim) + self.act = ACT2FN[config.activation_function] + self.dropout = nn.Dropout(config.resid_pdrop) + + # Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP.forward + def forward(self, hidden_states: tuple[torch.FloatTensor] | None) -> torch.FloatTensor: + hidden_states = self.c_fc(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.c_proj(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class GPTBigCodeBlock(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + hidden_size = config.hidden_size + self.inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size + + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + self.attn = GPTBigCodeAttention(config, layer_idx=layer_idx) + + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + if config.add_cross_attention: + if config.multi_query: + raise NotImplementedError("Cross-attention not implemented for MQA") + + self.crossattention = GPTBigCodeAttention(config, is_cross_attention=True, layer_idx=layer_idx) + + self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + self.mlp = GPTBigCodeMLP(self.inner_dim, config) + + def forward( + self, + hidden_states: tuple[torch.Tensor] | None, + encoder_hidden_states: torch.Tensor | None = None, + layer_past: Cache | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_output, _ = self.attn( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + hidden_states = attn_output + residual + + if encoder_hidden_states is not None: + # add one self-attention block for cross-attention + if not hasattr(self, "crossattention"): + raise ValueError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with " + "cross-attention layers by setting `config.add_cross_attention=True`" + ) + residual = hidden_states + hidden_states = self.ln_cross_attn(hidden_states) + attn_output, _ = self.crossattention( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + **kwargs, + ) + # residual connection + hidden_states = residual + attn_output + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states) + hidden_states = residual + feed_forward_hidden_states + return hidden_states + + +@auto_docstring +class GPTBigCodePreTrainedModel(PreTrainedModel): + config: GPTBigCodeConfig + base_model_prefix = "transformer" + supports_gradient_checkpointing = True + _no_split_modules = ["GPTBigCodeBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + _can_record_outputs = { + "hidden_states": GPTBigCodeBlock, + "attentions": OutputRecorder(GPTBigCodeAttention, index=1, layer_name="attn"), + "cross_attentions": OutputRecorder(GPTBigCodeAttention, index=1, layer_name="crossattention"), + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, (GPTBigCodeMLP, GPTBigCodeAttention)): + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + init.normal_( + module.c_proj.weight, mean=0.0, std=self.config.initializer_range / math.sqrt(2 * self.config.n_layer) + ) + elif isinstance(module, GPTBigCodeModel): + max_positions = module.config.max_position_embeddings + init.copy_(module.bias, torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool))) + + +@auto_docstring +class GPTBigCodeModel(GPTBigCodePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.multi_query = config.multi_query + self.embed_dim = config.hidden_size + + self.wte = nn.Embedding(config.vocab_size, self.embed_dim) + self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) + + self.drop = nn.Dropout(config.embd_pdrop) + self.h = nn.ModuleList([GPTBigCodeBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)]) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + max_positions = config.max_position_embeddings + self.register_buffer( + "bias", torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)), persistent=False + ) + + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.wte + + def set_input_embeddings(self, new_embeddings): + self.wte = new_embeddings + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: + r""" + input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + + if is_flash_attention_requested(self.config): + encoder_attention_mask = ( + encoder_attention_mask.bool() + if (encoder_attention_mask is not None and 0 in encoder_attention_mask) + else None + ) + else: + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if ( + self.config.add_cross_attention + and encoder_hidden_states is not None + and encoder_attention_mask is not None + ): + if encoder_attention_mask.dim() == 2: + encoder_attention_mask.unsqueeze(1) + assert encoder_attention_mask.dim() == 3 + encoder_attention_mask = encoder_attention_mask.bool().unsqueeze(2 if self.multi_query else 1) + else: + encoder_attention_mask = None + + position_embeds = self.wpe(position_ids) + hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device) + + if token_type_ids is not None: + token_type_ids = token_type_ids.view(-1, inputs_embeds.shape[1]) + token_type_embeds = self.wte(token_type_ids) + hidden_states = hidden_states + token_type_embeds + + hidden_states = self.drop(hidden_states) + output_shape = inputs_embeds.shape[:-1] + (hidden_states.size(-1),) + + for block in self.h: + hidden_states = block( + hidden_states, # as a positional argument for gradient checkpointing + encoder_hidden_states, # as a positional argument for gradient checkpointing + layer_past=past_key_values, # as keyword argument so it can be removed by GradientCheckpointingLayer + attention_mask=causal_mask, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.ln_f(hidden_states) + + hidden_states = hidden_states.view(output_shape) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring( + custom_intro=""" + The GPT_BIGCODE Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """ +) +class GPTBigCodeForCausalLM(GPTBigCodePreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} + + def __init__(self, config): + super().__init__(config) + self.transformer = GPTBigCodeModel(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithCrossAttentions: + r""" + input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + labels (`torch.Tensor` of shape `(batch_size, input_ids_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = transformer_outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + cross_attentions=transformer_outputs.cross_attentions, + ) + + +@auto_docstring( + custom_intro=""" + The GPTBigCode Model transformer with a sequence classification head on top (linear layer). + + [`GPTBigCodeForSequenceClassification`] uses the last token in order to do the classification, as other causal + models (e.g. GPT-1) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """ +) +class GPTBigCodeForSequenceClassification(GPTBigCodePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = GPTBigCodeModel(config) + self.score = nn.Linear(config.n_embd, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SequenceClassifierOutputWithPast: + r""" + input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + labels (`torch.Tensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size, sequence_length = input_ids.shape[:2] + else: + batch_size, sequence_length = inputs_embeds.shape[:2] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class GPTBigCodeForTokenClassification(GPTBigCodePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = GPTBigCodeModel(config) + if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None: + classifier_dropout = config.classifier_dropout + elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | TokenClassifierOutput: + r""" + input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = transformer_outputs[0] + hidden_states = self.dropout(hidden_states) + logits = self.classifier(hidden_states) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1).to(logits.device)) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +__all__ = [ + "GPTBigCodeForSequenceClassification", + "GPTBigCodeForTokenClassification", + "GPTBigCodeForCausalLM", + "GPTBigCodeModel", + "GPTBigCodePreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/gpt_neox_japanese/__init__.py b/third_party/transformers/src/transformers/models/gpt_neox_japanese/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94ba39d69ad638c706f6ac8491e2dea80e269929 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_neox_japanese/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_gpt_neox_japanese import * + from .modeling_gpt_neox_japanese import * + from .tokenization_gpt_neox_japanese import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py b/third_party/transformers/src/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py new file mode 100644 index 0000000000000000000000000000000000000000..d073d55923af02af5ef8f5ba8675e87f1cdfecb1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py @@ -0,0 +1,80 @@ +# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPTNeoX Japanese model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="abeja/gpt-neox-japanese-2.7b") +@strict +class GPTNeoXJapaneseConfig(PreTrainedConfig): + r""" + intermediate_multiple_size (`int`, *optional*, defaults to 4): + Dimension of the "intermediate" layer in the Transformer encoder is calculated by hidden_size * + intermediate_multiple_size. + + Example: + + ```python + >>> from transformers import GPTNeoXJapaneseConfig, GPTNeoXJapaneseModel + + >>> # Initializing a GPTNeoXJapanese gpt-neox-japanese-2.7b style configuration + >>> configuration = GPTNeoXJapaneseConfig() + + >>> # Initializing a model (with random weights) from the gpt-neox-japanese-2.7b style configuration + >>> model = GPTNeoXJapaneseModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gpt_neox_japanese" + + vocab_size: int = 32000 + hidden_size: int = 2560 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + intermediate_multiple_size: int = 4 + hidden_act: str = "gelu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + use_cache: bool = True + bos_token_id: int | None = 31996 + eos_token_id: int | list[int] | None = 31999 + rope_parameters: RopeParameters | dict | None = None + attention_dropout: float | int = 0.1 + hidden_dropout: float | int = 0.0 + is_decoder: bool = False + pad_token_id: int | None = None + tie_word_embeddings: bool = True + + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + # Model uses non-standard naming for rope params, overwrite! + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rotary_emb_base", self.default_theta)) + self.rope_parameters.setdefault("partial_rotary_factor", kwargs.pop("rotary_pct", 1.0)) + self.standardize_rope_params() + return kwargs + + +__all__ = ["GPTNeoXJapaneseConfig"] diff --git a/third_party/transformers/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/third_party/transformers/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py new file mode 100644 index 0000000000000000000000000000000000000000..e334ce023d67629d23cfebb88e949364a6d22bde --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -0,0 +1,608 @@ +# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch GPTNeoX model.""" + +import math +from collections.abc import Callable +from typing import Optional + +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from ...utils.generic import maybe_autocast +from .configuration_gpt_neox_japanese import GPTNeoXJapaneseConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring +class GPTNeoXJapanesePreTrainedModel(PreTrainedModel): + config: GPTNeoXJapaneseConfig + base_model_prefix = "gpt_neox_japanese" + _no_split_modules = ["GPTNeoXJapaneseLayer"] + _skip_keys_device_placement = "past_key_values" + _can_compile_fullgraph = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, GPTNeoXJapaneseAttention): + if module.dense_bias is not None: + init.zeros_(module.dense_bias) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->GPTNeoXJapanese +class GPTNeoXJapaneseRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: GPTNeoXJapaneseConfig, 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: GPTNeoXJapaneseConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class GPTNeoXJapaneseAttention(nn.Module): + def __init__(self, config, use_bias=False, layer_idx=None): + super().__init__() + self.num_attention_heads = config.num_attention_heads + self.hidden_size = config.hidden_size + self.head_size = self.hidden_size // self.num_attention_heads + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.layer_idx = layer_idx + partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) + self.rotary_ndims = int(self.head_size * partial_rotary_factor) + self.attention_dropout = nn.Dropout(config.attention_dropout) + self.norm_factor = math.sqrt(self.head_size) + + self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False) + self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + # Activate bias if the last layer + self.use_bias = use_bias + self.dense_bias = nn.Parameter(torch.zeros(config.hidden_size)) if use_bias else None + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: torch.FloatTensor, + position_ids: torch.LongTensor, + layer_past: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ): + # Compute QKV + # Attention heads [batch, seq_len, hidden_size] + # --> [batch, seq_len, (np * 3 * head_size)] + qkv = self.query_key_value(hidden_states) + + # [batch, seq_len, (num_heads * 3 * head_size)] + # --> [batch, seq_len, num_heads, 3 * head_size] + new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size) + qkv = qkv.view(*new_qkv_shape) + + # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size] + query = qkv[..., : self.head_size].permute(0, 2, 1, 3) + key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3) + value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3) + + # Compute rotary embeddings on rotary_ndims + query_rot = query[..., : self.rotary_ndims] + query_pass = query[..., self.rotary_ndims :] + key_rot = key[..., : self.rotary_ndims] + key_pass = key[..., self.rotary_ndims :] + + cos, sin = position_embeddings + query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin) + query = torch.cat((query, query_pass), dim=-1).contiguous() + key = torch.cat((key, key_pass), dim=-1).contiguous() + + # Cache QKV values + if layer_past is not None: + key, value = layer_past.update(key, value, self.layer_idx) + + # Compute attention + attn_output, attn_weights = self._attn(query, key, value, attention_mask) + + # Reshape outputs + attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size) + attn_output = self.dense(attn_output) + + return attn_output, attn_weights, self.dense_bias + + @classmethod + def _split_heads(cls, tensor, num_attention_heads, attn_head_size): + """ + Splits hidden dim into attn_head_size and num_attention_heads + """ + # tensor: [bs, seq_len, hidden_size] + new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size) + # -> [bs, seq_len, num_attention_heads, attn_head_size] + tensor = tensor.view(new_shape) + # -> [bs, num_attention_heads, seq_len, attn_head_size] + tensor = tensor.permute(0, 2, 1, 3) + return tensor + + @classmethod + def _merge_heads(cls, tensor, num_attention_heads, attn_head_size): + """ + Merges attn_head_size dim and num_attn_heads dim into hidden dim + """ + # tensor [bs, num_attention_heads, seq_len, attn_head_size] + tensor = tensor.permute(0, 2, 1, 3).contiguous() + # -> [bs, seq_len, num_attention_heads, attn_head_size] + tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size) + # -> [bs, seq_len, hidden_size] + return tensor + + def _attn(self, query, key, value, attention_mask=None): + # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size] + # compute causal mask from causal mask buffer + batch_size, num_attention_heads, query_length, attn_head_size = query.size() + key_length = key.size(-2) + + query = query.view(batch_size * num_attention_heads, query_length, attn_head_size) + key = key.view(batch_size * num_attention_heads, key_length, attn_head_size) + + # [batch_size * num_heads, q_length, kv_length] + attn_scores = torch.zeros( + batch_size * num_attention_heads, + query_length, + key_length, + dtype=query.dtype, + device=key.device, + ) + attention_scores = torch.baddbmm( + attn_scores, + query, + key.transpose(1, 2), + beta=1.0, + alpha=1.0 / self.norm_factor, + ) + + attention_scores = attention_scores.view(batch_size, num_attention_heads, query_length, -1) + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + + attn_weights = nn.functional.softmax(attention_scores, dim=-1) + attn_weights = self.attention_dropout(attn_weights) + attn_weights = attn_weights.to(value.dtype) + + attn_output = torch.matmul(attn_weights, value) + return attn_output, attn_weights + + +def bias_dropout_add(x: Tensor, bias: Tensor, residual: Tensor | None, prob: float, training: bool) -> Tensor: + """add bias to x, apply dropout and residual connection + + Args: + x (Tensor): main path of output + bias (Tensor): None or attn_bias of the last attention layer + residual (Optional[Tensor]): residual value + prob (float): dropout probability + training (bool): whether in training mode or not + + Returns: + Tensor: dropout(x + bias) + residual + """ + if bias is not None: + x = x + bias + out = torch.nn.functional.dropout(x, p=prob, training=training) + if residual is not None: + out = residual + out + return out + + +class GPTNeoXJapaneseMLP(nn.Module): + def __init__(self, config): + super().__init__() + intermediate_size = int(config.hidden_size * config.intermediate_multiple_size) + self.dense_h_to_4h = nn.Linear(config.hidden_size, intermediate_size, bias=False) + # Project back to h. + self.dense_4h_to_h = nn.Linear(intermediate_size, config.hidden_size, bias=False) + self.act = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + intermediate = self.dense_h_to_4h(hidden_states) + intermediate = self.act(intermediate) + output = self.dense_4h_to_h(intermediate) + return output + + +class GPTNeoXJapaneseLayer(nn.Module): + def __init__(self, config, layer_number): + super().__init__() + self.layer_number = layer_number + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + # activate bias only last layer + self.attention = GPTNeoXJapaneseAttention( + config=config, use_bias=layer_number == config.num_hidden_layers - 1, layer_idx=layer_number + ) + self.mlp = GPTNeoXJapaneseMLP(config) + self.hidden_dropout = config.hidden_dropout + + def forward( + self, + hidden_states: torch.FloatTensor | None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + layer_past: Cache | None = None, + output_attentions: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ): + residual = hidden_states + ln_out = self.input_layernorm(hidden_states) + attn_output, attn_weights, attn_bias = self.attention( + ln_out, + attention_mask=attention_mask, + layer_past=layer_past, + use_cache=use_cache, + output_attentions=output_attentions, + position_ids=position_ids, + position_embeddings=position_embeddings, + ) + + # attn_output = (atten_output + bias) + residual + attn_output = bias_dropout_add( + attn_output, + bias=attn_bias.expand_as(residual) if attn_bias is not None else attn_bias, + residual=residual, + prob=self.hidden_dropout, + training=self.training, + ) + mlp_output = self.mlp(self.post_attention_layernorm(attn_output)) + + # attn_output = (mlp_output + mlp_bias) + atten_output + attn_output = bias_dropout_add( + mlp_output, bias=None, residual=attn_output, prob=self.hidden_dropout, training=self.training + ) + + return attn_output, attn_weights + + +@auto_docstring +class GPTNeoXJapaneseModel(GPTNeoXJapanesePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.config = config + + self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [GPTNeoXJapaneseLayer(config=config, layer_number=i) for i in range(config.num_hidden_layers)] + ) + self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.rotary_emb = GPTNeoXJapaneseRotaryEmbedding(config=config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_in + + def set_input_embeddings(self, value): + self.embed_in = value + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, GPTNeoXJapaneseModel + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b") + >>> model = GPTNeoXJapaneseModel.from_pretrained("abeja/gpt-neox-japanese-2.7b") + + >>> inputs = tokenizer("日本語のGPT-neoxがHugging Faceで使えます😀", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_in(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + all_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + outputs = layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + layer_past=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + ) + hidden_states = outputs[0] + if output_attentions: + all_attentions = all_attentions + (outputs[1],) + + hidden_states = self.final_layer_norm(hidden_states) + # Add last hidden state + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v for v in [hidden_states, past_key_values, all_hidden_states, all_attentions] if v is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attentions, + ) + + +@auto_docstring( + custom_intro=""" + GPTNeoXJapanese Model with a `language modeling` head on top for Classifier Model fine-tuning. + """ +) +class GPTNeoXJapaneseForCausalLM(GPTNeoXJapanesePreTrainedModel, GenerationMixin): + _tied_weights_keys = {"embed_out.weight": "gpt_neox_japanese.embed_in.weight"} + + def __init__(self, config): + super().__init__(config) + self.config = config + + self.gpt_neox_japanese = GPTNeoXJapaneseModel(config) + self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.embed_out + + def set_output_embeddings(self, new_embeddings): + self.embed_out = new_embeddings + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseConfig + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b") + >>> config = GPTNeoXJapaneseConfig.from_pretrained("abeja/gpt-neox-japanese-2.7b") + >>> config.is_decoder = True + >>> model = GPTNeoXJapaneseForCausalLM.from_pretrained("abeja/gpt-neox-japanese-2.7b", config=config) + + >>> inputs = tokenizer("日本語のGPT-neoxがHugging Faceで使えます😀", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> prediction_logits = outputs.logits + ``` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.gpt_neox_japanese( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.embed_out(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "GPTNeoXJapaneseForCausalLM", + "GPTNeoXJapaneseLayer", + "GPTNeoXJapaneseModel", + "GPTNeoXJapanesePreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py b/third_party/transformers/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py new file mode 100644 index 0000000000000000000000000000000000000000..68f7ca8c4058eaa82c9dbf7e8b63217fcae5eedd --- /dev/null +++ b/third_party/transformers/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py @@ -0,0 +1,368 @@ +# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for GPTNeoXJapanese.""" + +import collections +import json +import os +import re +import sys + +import numpy as np + +from ...tokenization_python import PreTrainedTokenizer +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "emoji_file": "emoji.json"} + + +def load_vocab_and_emoji(vocab_file, emoji_file): + """Loads a vocabulary file and emoji file into a dictionary.""" + with open(emoji_file, "r", encoding="utf-8") as f: + emoji = json.loads(f.read()) + + vocab = collections.OrderedDict() + raw_vocab = collections.OrderedDict() + ids_to_tokens = collections.OrderedDict() + with open(vocab_file, "r", encoding="utf-8") as f: + token = f.readlines() + token = [[t.rstrip("\n")] if (t == "," or "," not in t) else t.rstrip("\n").split(",") for t in token] + for idx, b in enumerate(token): + ids_to_tokens[idx] = b + raw_vocab[",".join(b)] = idx + for wd in b: + vocab[wd] = idx + + return vocab, raw_vocab, ids_to_tokens, emoji + + +class GPTNeoXJapaneseTokenizer(PreTrainedTokenizer): + """ + This tokenizer inherits from [`PreTrainedTokenizer`] and is based on Japanese special Sub-Word-Encoding that is + used in this repository (https://github.com/tanreinama/Japanese-BPEEncoder_V2). Check the repository for details. + Japanese has a relatively large vocabulary and there is no separation between words. Furthermore, the language is a + combination of hiragana, katakana, and kanji, and variants such as "1" and "①" are often used. In order to cope + with these, this tokenizer has the following features + - Subword-by-subword segmentation, which is intermediate between byte strings and morphological analysis. + - BPEs are created for each Kanji, Hiragana, and Katakana character, and there are no BPEs that cross character + types, such as Kanji + Hiragana or Hiragana + Katakana. + - All-byte encoding that does not require . + - Independent of UTF codes such as 2-byte and 3-byte characters + - Conversion of heterographs to the same token_id + - Emoji and Emoticon are grouped into 12 types as special tags. + + Example: + + ```python + >>> from transformers import GPTNeoXJapaneseTokenizer + + >>> tokenizer = GPTNeoXJapaneseTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b") + >>> # You can confirm both 慶応 and 慶應 are encoded to 17749 + >>> tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"] + [30014, 26883, 26638, 27228, 25, 26650, 31732, 31679, 27809, 26638, 17749, 31592, 17749, 31593, 321, 1281] + + >>> # Both 慶応 and 慶應 are decoded to 慶応 + >>> tokenizer.decode(tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"]) + '吾輩は猫である🐯。実は慶応(慶応)大学出身' + ``` + + Args: + vocab_file (`str`): + File containing the vocabulary. + emoji_file (`str`): + File containing the emoji. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The token used for padding + bos_token (`str`, *optional*, defaults to `"<|startoftext|>"`): + The beginning of sequence token. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + do_clean_text (`bool`, *optional*, defaults to `False`): + Whether or not to clean text for URL, EMAIL, TEL, Japanese DATE and Japanese PRICE. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + emoji_file, + unk_token="<|endoftext|>", + pad_token="<|endoftext|>", + bos_token="<|startoftext|>", + eos_token="<|endoftext|>", + do_clean_text=False, + **kwargs, + ): + if not os.path.isfile(vocab_file): + raise ValueError( + f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained" + " model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + if not os.path.isfile(emoji_file): + raise ValueError( + f"Can't find a emoji file at path '{emoji_file}'. To load the emoji information from a Google" + " pretrained model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + self.do_clean_text = do_clean_text + self.vocab, self.raw_vocab, self.ids_to_tokens, self.emoji = load_vocab_and_emoji(vocab_file, emoji_file) + self.subword_tokenizer = SubWordJapaneseTokenizer( + vocab=self.vocab, ids_to_tokens=self.ids_to_tokens, emoji=self.emoji + ) + super().__init__( + unk_token=unk_token, + pad_token=pad_token, + bos_token=bos_token, + eos_token=eos_token, + do_clean_text=do_clean_text, + special_tokens_pattern="none", + **kwargs, + ) + + @property + def vocab_size(self): + # self.vocab contains support for character fluctuation unique to Japanese, and has a large number of vocab + return len(self.raw_vocab) + + def get_vocab(self): + return dict(self.raw_vocab, **self.added_tokens_encoder) + + def _tokenize(self, text): + return self.subword_tokenizer.tokenize(text, clean=self.do_clean_text) + + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.vocab.get(token, self.vocab.get(self.unk_token)) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.subword_tokenizer.convert_id_to_token(index) + + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + out_string = "".join(tokens).strip() + return out_string + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + index = 0 + if os.path.isdir(save_directory): + vocab_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] + ) + emoji_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["emoji_file"] + ) + else: + vocab_file = ( + (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["vocab_file"] + ) + emoji_file = ( + (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["emoji_file"] + ) + with open(vocab_file, "w", encoding="utf-8") as writer: + for token_index, token in self.ids_to_tokens.items(): + if index != token_index: + logger.warning( + f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive." + " Please check that the vocabulary is not corrupted!" + ) + index = token_index + writer.write(",".join(token) + "\n") + index += 1 + with open(emoji_file, "w", encoding="utf-8") as writer: + json.dump(self.emoji, writer) + return vocab_file, emoji_file + + +class SubWordJapaneseTokenizer: + """ + https://github.com/tanreinama/Japanese-BPEEncoder_V2 This tokenizer class is under MIT License according to the + original repository. + + MIT License + + Copyright (c) 2020 tanreinama + + Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated + documentation files (the "Software"), to deal in the Software without restriction, including without limitation the + rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to + permit persons to whom the Software is furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all copies or substantial portions of + the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO + THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + """ + + def __init__(self, vocab, ids_to_tokens, emoji): + self.vocab = vocab # same as swe + self.ids_to_tokens = ids_to_tokens # same as bpe + self.emoji = emoji + self.maxlen = np.max([len(w) for w in self.vocab]) + self.content_repatter1 = re.compile(r"(https?|ftp)(:\/\/[-_\.!~*\'()a-zA-Z0-9;\/?:\@&=\+$,%#]+)") + self.content_repatter2 = re.compile(r"[A-Za-z0-9\._+]*@[\-_0-9A-Za-z]+(\.[A-Za-z]+)*") + self.content_repatter3 = re.compile(r"[\(]{0,1}[0-9]{2,4}[\)\-\(]{0,1}[0-9]{2,4}[\)\-]{0,1}[0-9]{3,4}") + self.content_repatter4 = re.compile( + r"([12]\d{3}[/\-年])*(0?[1-9]|1[0-2])[/\-月]((0?[1-9]|[12][0-9]|3[01])日?)*(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*" + ) + self.content_repatter5 = re.compile( + r"(明治|大正|昭和|平成|令和|㍾|㍽|㍼|㍻|\u32ff)\d{1,2}年(0?[1-9]|1[0-2])月(0?[1-9]|[12][0-9]|3[01])日(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*" + ) + # The original version of this regex displays catastrophic backtracking behaviour. We avoid this using + # possessive quantifiers in Py >= 3.11. In versions below this, we avoid the vulnerability using a slightly + # different regex that should generally have the same behaviour in most non-pathological cases. + if sys.version_info >= (3, 11): + self.content_repatter6 = re.compile( + r"(?:\d,\d{3}|[\d億])*+" + r"(?:\d,\d{3}|[\d万])*+" + r"(?:\d,\d{3}|[\d千])*+" + r"(?:千円|万円|千万円|円|千ドル|万ドル|千万ドル|ドル|千ユーロ|万ユーロ|千万ユーロ|ユーロ)+" + r"(?:\(税込\)|\(税抜\)|\+tax)*" + ) + else: + self.content_repatter6 = re.compile( + r"(?:\d,\d{3}|[\d億万千])*" + r"(?:千円|万円|千万円|円|千ドル|万ドル|千万ドル|ドル|千ユーロ|万ユーロ|千万ユーロ|ユーロ)+" + r"(?:\(税込\)|\(税抜\)|\+tax)*" + ) + keisen = "─━│┃┄┅┆┇┈┉┊┋┌┍┎┏┐┑┒┓└┕┖┗┘┙┚┛├┝┞┟┠┡┢┣┤┥┦┧┨┩┪┫┬┭┮┯┰┱┲┳┴┵┶┷┸┹┺┻┼┽┾┿╀╁╂╃╄╅╆╇╈╉╊╋╌╍╎╏═║╒╓╔╕╖╗╘╙╚╛╜╝╞╟╠╡╢╣╤╥╦╧╨╩╪╫╬╭╮╯╰╱╲╳╴╵╶╷╸╹╺╻╼╽╾╿" + blocks = "▀▁▂▃▄▅▆▇█▉▊▋▌▍▎▏▐░▒▓▔▕▖▗▘▙▚▛▜▝▞▟" + self.content_trans1 = str.maketrans(dict.fromkeys(keisen + blocks, "")) + + def __len__(self): + return len(self.ids_to_tokens) + + def clean_text(self, content): + content = self.content_repatter1.sub("", content) + content = self.content_repatter2.sub("", content) + content = self.content_repatter3.sub("", content) + content = self.content_repatter4.sub("", content) + content = self.content_repatter5.sub("", content) + content = self.content_repatter6.sub("", content) + content = content.translate(self.content_trans1) + while "" in content: + content = content.replace("", "") + return content + + def tokenize(self, text, clean=False): + text = text.replace(" ", "") + text = text.replace("　", "") + text = text.replace("\r\n", "
") + text = text.replace("\n", "
") + text = text.replace("\r", "
") + text = text.replace("\t", "") + text = text.replace("—", "ー") + text = text.replace("−", "ー") + for k, v in self.emoji["emoji"].items(): + if k in text: + text = text.replace(k, v) + if clean: + text = self.clean_text(text) + + def check_simbol(x): + e = x.encode() + if len(x) == 1 and len(e) == 2: + c = (int(e[0]) << 8) + int(e[1]) + if ( + (c >= 0xC2A1 and c <= 0xC2BF) + or (c >= 0xC780 and c <= 0xC783) + or (c >= 0xCAB9 and c <= 0xCBBF) + or (c >= 0xCC80 and c <= 0xCDA2) + ): + return True + return False + + def checku2e(x): + e = x.encode() + if len(x) == 1 and len(e) == 3: + c = (int(e[0]) << 16) + (int(e[1]) << 8) + int(e[2]) + if c >= 0xE28080 and c <= 0xE2B07F: + return True + return False + + pos = 0 + result = [] + while pos < len(text): + end = min(len(text), pos + self.maxlen + 1) if text[pos] == "<" else pos + 3 + candidates = [] # (token_id, token, pos) + for e in range(end, pos, -1): + wd = text[pos:e] + if wd in self.vocab: + if wd[0] == "<" and len(wd) > 2: + candidates = [(self.vocab[wd], wd, e)] + break + else: + candidates.append((self.vocab[wd], wd, e)) + if len(candidates) > 0: + # the smallest token_id is adopted + _, wd, e = min(candidates, key=lambda x: x[0]) + result.append(wd) + pos = e + else: + end = pos + 1 + wd = text[pos:end] + if check_simbol(wd): + result.append("") + elif checku2e(wd): + result.append("") + else: + for i in wd.encode("utf-8"): + result.append("<|byte%d|>" % i) + pos = end + return result + + def convert_id_to_token(self, index, breakline="\n"): + words = [] + byte_tokens = [] + word = self.ids_to_tokens[index][0] + if word[:6] == "<|byte" and word[-2:] == "|>": + byte_tokens.append(int(word[6:-2])) + else: + if len(byte_tokens) > 0: + words.append(bytearray(byte_tokens).decode("utf-8", errors="replace")) + byte_tokens = [] + if word[:7] == "<|emoji" and word[-2:] == "|>": + words.append(self.emoji["emoji_inv"][word]) + elif word == "": + words.append(" ") + elif word == "
": + words.append(breakline) + elif word == "": + words.append("\t") + elif word == "": + words.append("▀") + elif word == "": + words.append("ǀ") + elif word == "": + words.append("‖") + else: + words.append(word) + if len(byte_tokens) > 0: + words.append(bytearray(byte_tokens).decode("utf-8", errors="replace")) + text = "".join(words) + return text + + +__all__ = ["GPTNeoXJapaneseTokenizer"] diff --git a/third_party/transformers/src/transformers/models/idefics2/__init__.py b/third_party/transformers/src/transformers/models/idefics2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..54c9d8bd56d804b1c22962b9cd8a0d2134bce9a8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/__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_idefics2 import * + from .image_processing_idefics2 import * + from .image_processing_pil_idefics2 import * + from .modeling_idefics2 import * + from .processing_idefics2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/idefics2/configuration_idefics2.py b/third_party/transformers/src/transformers/models/idefics2/configuration_idefics2.py new file mode 100644 index 0000000000000000000000000000000000000000..0b170712b1c96851fac0690f197a9badf3041f9c --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/configuration_idefics2.py @@ -0,0 +1,165 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Idefics2 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging +from ..auto import CONFIG_MAPPING, AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="HuggingFaceM4/idefics2-8b") +@strict +class Idefics2VisionConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer + >>> from transformers.models.idefics2.configuration_idefics2 import Idefics2VisionConfig + + >>> # Initializing a Idefics2VisionConfig with google/siglip-base-patch16-224 style configuration + >>> configuration = Idefics2VisionConfig() + + >>> # Initializing a Idefics2VisionTransformer (with random weights) from the google/siglip-base-patch16-224 style configuration + >>> model = Idefics2VisionTransformer(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "idefics2_vision" + base_config_key = "vision_config" + + hidden_size: int = 768 + intermediate_size: int = 3072 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 224 + patch_size: int | list[int] | tuple[int, int] = 32 + hidden_act: str = "gelu_pytorch_tanh" + layer_norm_eps: float = 1e-6 + attention_dropout: float | int = 0.0 + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="HuggingFaceM4/idefics2-8b") +@strict +class Idefics2PerceiverConfig(PreTrainedConfig): + r""" + resampler_n_latents (`int`, *optional*, defaults to 64): + Number of latent embeddings to resample ("compress") the input sequence to (usually < 128). + resampler_depth (`int`, *optional*, defaults to 3): + Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (<= 3). + resampler_n_heads (`int`, *optional*, defaults to 16): + Number of heads in each Transformer block (for multi-headed self-attention). + resampler_head_dim (`int`, *optional*, defaults to 96): + Dimensionality of each head projection in the Transformer block. + """ + + model_type = "idefics2_perceiver" + + hidden_act: str = "silu" + hidden_size: int = 4096 + rms_norm_eps: float = 1e-06 + resampler_n_latents: int = 64 + resampler_depth: int = 3 + resampler_n_heads: int = 16 + resampler_head_dim: int = 96 + num_key_value_heads: int = 4 + attention_dropout: float | int = 0.0 + initializer_range: float = 0.02 + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.num_key_value_heads > self.resampler_n_heads: + raise ValueError( + f"num_key_value_heads={self.num_key_value_heads} must be less than or equal to" + f" resampler_n_heads={self.resampler_n_heads}" + ) + + +@auto_docstring(checkpoint="HuggingFaceM4/idefics2-8b") +@strict +class Idefics2Config(PreTrainedConfig): + r""" + perceiver_config (`IdeficsPerceiverConfig` or `dict`, *optional*): + Custom perceiver config or dict + + Example: + ```python + >>> from transformers import Idefics2Model, Idefics2Config + >>> # Initializing configuration + >>> configuration = Idefics2Config() + >>> # Initializing a model from the configuration + >>> model = Idefics2Model(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "idefics2" + sub_configs = { + "text_config": AutoConfig, + "perceiver_config": Idefics2PerceiverConfig, + "vision_config": Idefics2VisionConfig, + } + + use_cache: bool = True + image_token_id: int = 32_001 + tie_word_embeddings: bool = False + vision_config: dict | PreTrainedConfig | None = None + perceiver_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.perceiver_config is None: + self.perceiver_config = Idefics2PerceiverConfig() + logger.info("perciver_config is None, using default perceiver config") + elif isinstance(self.perceiver_config, dict): + self.perceiver_config = Idefics2PerceiverConfig(**self.perceiver_config) + + if self.vision_config is None: + self.vision_config = Idefics2VisionConfig() + logger.info("vision_config is None, using default vision config") + elif isinstance(self.vision_config, dict): + self.vision_config = Idefics2VisionConfig(**self.vision_config) + + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "mistral") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + logger.info("text_config is None, using default text config") + self.text_config = CONFIG_MAPPING["mistral"]( + max_position_embeddings=4096 * 8, + rms_norm_eps=1e-5, + # None in the original configuration_mistral, we set it to the unk_token_id + pad_token_id=0, + ) + + if self.text_config.hidden_size != self.perceiver_config.hidden_size: + self.perceiver_config.hidden_size = self.text_config.hidden_size + self.perceiver_config.rms_norm_eps = self.text_config.rms_norm_eps + logger.warning_once( + "Perceiver config has a different `hidden_size` than text config, which means default values were used. " + "In your model's config on the hub, add `hidden_size` and `rms_norm_eps` keys under the `perceiver_config` dict. " + ) + + super().__post_init__(**kwargs) + + +__all__ = ["Idefics2Config"] diff --git a/third_party/transformers/src/transformers/models/idefics2/convert_idefics2_weights_to_hf.py b/third_party/transformers/src/transformers/models/idefics2/convert_idefics2_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..1aee7410a62c35e57f996411eb5a461945d6de84 --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/convert_idefics2_weights_to_hf.py @@ -0,0 +1,184 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import copy + +import torch + +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoTokenizer, + Idefics2Config, + Idefics2ForConditionalGeneration, + Idefics2ImageProcessor, + Idefics2Processor, + MistralConfig, +) + + +EPILOG_TXT = """Example: + python transformers/src/transformers/models/idefics2/convert_idefics2_weights_to_hf.py --original_model_id HuggingFaceM4/idefics2-8b --output_hub_path org/idefics2 +""" + + +KEYS_TO_MODIFY_MAPPING = { + "lm_head.weight": "lm_head.linear.weight", + "model.layers": "model.text_model.layers", + "model.norm": "model.text_model.norm", + "model.perceiver_resampler": "model.connector.perceiver_resampler", + "model.modality_projection": "model.connector.modality_projection", +} + + +WEIGHTS_TO_MERGE_MAPPING = ( + # (weights to merge in merging order), (new weight name) + ( + ("model.embed_tokens.weight", "model.embed_tokens.additional_embedding.weight"), + "model.text_model.embed_tokens.weight", + ), + (("lm_head.linear.weight", "additional_fc.weight"), "lm_head.weight"), +) + + +def convert_state_dict_to_hf(state_dict): + new_state_dict = {} + for key, value in state_dict.items(): + if key.endswith(".inv_freq"): + continue + for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items(): + if key_to_modify in key: + key = key.replace(key_to_modify, new_key) + + new_state_dict[key] = value + return new_state_dict + + +def merge_weights(state_dict): + new_state_dict = copy.deepcopy(state_dict) + + # Merge the weights + for weights_to_merge, new_weight_name in WEIGHTS_TO_MERGE_MAPPING: + for weight in weights_to_merge: + assert weight in state_dict, f"Weight {weight} is missing in the state dict" + if new_weight_name not in new_state_dict: + new_state_dict[new_weight_name] = [state_dict[weight]] + else: + new_state_dict[new_weight_name].append(state_dict[weight]) + new_state_dict[new_weight_name] = torch.cat(new_state_dict[new_weight_name], dim=0) + + # Remove the weights that were merged + for weights_to_merge, new_weight_name in WEIGHTS_TO_MERGE_MAPPING: + for weight in weights_to_merge: + if weight in new_state_dict and weight != new_weight_name: + new_state_dict.pop(weight) + + return new_state_dict + + +def get_config(checkpoint): + if checkpoint == "HuggingFaceM4/idefics2": + # We load the config then recreate to use the text_config + config = AutoConfig.from_pretrained(checkpoint) + text_config = MistralConfig( + vocab_size=config.vocab_size + config.additional_vocab_size, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + num_hidden_layers=config.num_hidden_layers, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + hidden_act=config.hidden_act, + max_position_embeddings=config.max_position_embeddings, + initializer_range=config.initializer_range, + rms_norm_eps=config.rms_norm_eps, + tie_word_embeddings=config.tie_word_embeddings, + rope_theta=config.rope_theta, + sliding_window=config.sliding_window, + attention_dropout=config.attention_dropout, + pad_token_id=config.pad_token_id, + bos_token_id=config.bos_token_id, + eos_token_id=config.eos_token_id, + ) + perceiver_config = config.perceiver_config.to_dict() + config = Idefics2Config( + text_config=text_config.to_dict(), + vision_config=config.vision_config, + perceiver_config=perceiver_config, + use_cache=config.use_cache, + image_token_id=config.image_token_id, + tie_word_embeddings=config.tie_word_embeddings, + ) + return config + + return AutoConfig.from_pretrained(checkpoint) + + +def convert_idefics2_hub_to_hf(original_model_id, output_hub_path, push_to_hub): + # The original model maps to AutoModelForCausalLM, converted we map to Idefics2ForConditionalGeneration + original_model = AutoModelForCausalLM.from_pretrained(original_model_id, trust_remote_code=True) + # The original model doesn't use the idefics2 processing objects + image_seq_len = original_model.config.perceiver_config.resampler_n_latents + image_processor = Idefics2ImageProcessor() + tokenizer = AutoTokenizer.from_pretrained(original_model_id) + processor = Idefics2Processor( + image_processor=image_processor, + tokenizer=tokenizer, + image_seq_len=image_seq_len, + ) + state_dict = original_model.state_dict() + state_dict = convert_state_dict_to_hf(state_dict) + + # Merge weights + state_dict = merge_weights(state_dict) + + config = get_config(original_model_id) + + with torch.device("meta"): + model = Idefics2ForConditionalGeneration(config) + + model.load_state_dict(state_dict, strict=True, assign=True) + + model.save_pretrained(output_hub_path) + processor.save_pretrained(output_hub_path) + + if push_to_hub: + model.push_to_hub(output_hub_path, private=True) + processor.push_to_hub(output_hub_path, private=True) + + +def main(): + parser = argparse.ArgumentParser( + epilog=EPILOG_TXT, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--original_model_id", + help="Hub location of the text model", + ) + parser.add_argument( + "--output_hub_path", + help="Location on the hub of the converted model", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="If set, the model will be pushed to the hub after conversion.", + ) + args = parser.parse_args() + convert_idefics2_hub_to_hf(args.original_model_id, args.output_hub_path, args.push_to_hub) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/idefics2/image_processing_idefics2.py b/third_party/transformers/src/transformers/models/idefics2/image_processing_idefics2.py new file mode 100644 index 0000000000000000000000000000000000000000..61f2037ce276d529d2107b4737db18788649da4c --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/image_processing_idefics2.py @@ -0,0 +1,289 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for Idefics2.""" + +import numpy as np +import torch + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_vision_available + + +if is_vision_available(): + from PIL import Image + +from torchvision.transforms.v2 import functional as tvF + + +def get_resize_output_image_size(image, size: SizeDict) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Images are always channels-first (CHW). + """ + height, width = image.shape[-2:] + + min_len = size.shortest_edge + max_len = size.longest_edge + aspect_ratio = width / height + + if width >= height and width > max_len: + width = max_len + height = int(width / aspect_ratio) + elif height > width and height > max_len: + height = max_len + width = int(height * aspect_ratio) + height = max(height, min_len) + width = max(width, min_len) + return height, width + + +def convert_to_rgb(image: ImageInput) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + """ + if not is_vision_available() or not isinstance(image, Image.Image): + return image + + if image.mode == "RGB": + return image + + image_rgba = image.convert("RGBA") + background = Image.new("RGBA", image_rgba.size, (255, 255, 255)) + alpha_composite = Image.alpha_composite(background, image_rgba) + alpha_composite = alpha_composite.convert("RGB") + return alpha_composite + + +class Idefics2ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_image_splitting (`bool`, *optional*, defaults to `self.do_image_splitting`): + Whether to split the image into a sequence 4 equal sub-images concatenated with the original image. + """ + + do_image_splitting: bool + + +def get_max_height_width(images_list: list[list["torch.Tensor|np.ndarray"]]) -> tuple[int, int]: + """ + Get the maximum height and width across all images in a batch. + """ + image_sizes = [] + for images in images_list: + for image in images: + image_sizes.append(image.shape[-2:]) + + max_height = max(size[0] for size in image_sizes) + max_width = max(size[1] for size in image_sizes) + return (max_height, max_width) + + +def make_pixel_mask(image: "torch.Tensor", output_size: tuple[int, int]) -> "torch.Tensor": + """ + Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + """ + input_height, input_width = image.shape[-2:] + mask = torch.zeros(output_size, dtype=torch.int64, device=image.device) + mask[:input_height, :input_width] = 1 + return mask + + +@auto_docstring +class Idefics2ImageProcessor(TorchvisionBackend): + valid_kwargs = Idefics2ImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + do_convert_rgb = True + do_image_splitting = False + default_to_square = False + size = {"shortest_edge": 378, "longest_edge": 980} + model_input_names = ["pixel_values", "pixel_attention_mask"] + + def __init__(self, **kwargs: Unpack[Idefics2ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[Idefics2ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def convert_to_rgb(self, image: ImageInput) -> ImageInput: + """Convert an image to RGB format.""" + return convert_to_rgb(image) + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + **kwargs, + ) -> "torch.Tensor": + """Resize using Idefics2 shortest_edge/longest_edge logic.""" + if size.shortest_edge and size.longest_edge: + new_size = get_resize_output_image_size(image, size) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError("Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys.") + + return super().resize(image, SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def _prepare_images_structure(self, images: ImageInput, expected_ndims: int = 3) -> ImageInput: + """Prepare a nested images structure for processing.""" + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def split_images(self, images: "torch.Tensor") -> list[list["torch.Tensor"]]: + """ + Split a batch of images into 4 equal sub-images, and concatenate that sequence with the original image. + """ + height, width = images.shape[-2:] + + mid_width = width // 2 + mid_height = height // 2 + + batch_split_images = [ + images[..., :mid_height, :mid_width], + images[..., :mid_height, mid_width:], + images[..., mid_height:, :mid_width], + images[..., mid_height:, mid_width:], + images, + ] + + batch_split_images = [[image[i] for image in batch_split_images] for i in range(len(batch_split_images[0]))] + return batch_split_images + + def pad( + self, image: "torch.Tensor", padded_size: tuple[int, int], fill: int = 0 + ) -> tuple["torch.Tensor", "torch.Tensor"]: + """ + Pad an image to the specified size and create the corresponding pixel mask. + """ + original_size = image.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + + if original_size != padded_size: + padding = (0, 0, padding_right, padding_bottom) + image = tvF.pad(image, padding, fill=fill, padding_mode="constant") + + pixel_mask = torch.zeros(padded_size, dtype=torch.int64, device=image.device) + pixel_mask[: original_size[0], : original_size[1]] = 1 + + return image, pixel_mask + + def _preprocess( + self, + images: list[list["torch.Tensor"]], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + do_image_splitting: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape( + images, disable_grouping=disable_grouping, is_nested=True + ) + split_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_image_splitting: + stacked_images = self.split_images(stacked_images) + split_images_grouped[shape] = stacked_images + split_images = reorder_images(split_images_grouped, grouped_images_index, is_nested=True) + if do_image_splitting: + for i, group_images in enumerate(split_images): + split_images[i] = [image for sublist in group_images for image in sublist] + + grouped_images, grouped_images_index = group_images_by_shape( + split_images, disable_grouping=disable_grouping, is_nested=True + ) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size, resample=resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index, is_nested=True) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping, is_nested=True + ) + 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, is_nested=True) + + if do_pad: + max_num_images = max(len(images_) for images_ in processed_images) + max_height, max_width = get_max_height_width(processed_images) + + processed_images_padded = torch.zeros( + len(processed_images), + max_num_images, + *(processed_images[0][0].shape[0], max_height, max_width), + device=processed_images[0][0].device, + ) + pixel_attention_masks = torch.zeros( + len(processed_images), + max_num_images, + *(max_height, max_width), + device=processed_images[0][0].device, + ) + for i, images in enumerate(processed_images): + for j, image in enumerate(images): + processed_images_padded[i, j], pixel_attention_masks[i, j] = self.pad( + image, (max_height, max_width) + ) + processed_images = processed_images_padded + if do_pad: + data = {"pixel_values": processed_images, "pixel_attention_mask": pixel_attention_masks} + elif return_tensors == "pt": + data = {"pixel_values": torch.stack([torch.stack(images) for images in processed_images])} + else: + data = {"pixel_values": processed_images} + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["Idefics2ImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/idefics2/image_processing_pil_idefics2.py b/third_party/transformers/src/transformers/models/idefics2/image_processing_pil_idefics2.py new file mode 100644 index 0000000000000000000000000000000000000000..641bdf0d2682b76a6ba262ab34db048c515f7d63 --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/image_processing_pil_idefics2.py @@ -0,0 +1,278 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PIL Image processor class for Idefics2.""" + +from typing import TYPE_CHECKING + +import numpy as np +from PIL import Image + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import PaddingMode, pad +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_vision_available + + +if TYPE_CHECKING: + import torch + + +def _make_pixel_mask(image: np.ndarray, output_size: tuple[int, int]) -> np.ndarray: + """Make pixel mask: 1=valid, 0=padding. Images are CHW.""" + h, w = image.shape[-2:] + mask = np.zeros(output_size, dtype=np.int64) + mask[:h, :w] = 1 + return mask + + +# Adapted from transformers.models.idefics2.image_processing_idefics2.Idefics2ImageProcessorKwargs +class Idefics2ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_image_splitting (`bool`, *optional*, defaults to `self.do_image_splitting`): + Whether to split the image into a sequence 4 equal sub-images concatenated with the original image. + """ + + do_image_splitting: bool + + +# Adapted from transformers.models.idefics2.image_processing_idefics2.convert_to_rgb +def convert_to_rgb(image: ImageInput) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + """ + if not is_vision_available() or not isinstance(image, Image.Image): + return image + + if image.mode == "RGB": + return image + + image_rgba = image.convert("RGBA") + background = Image.new("RGBA", image_rgba.size, (255, 255, 255)) + alpha_composite = Image.alpha_composite(background, image_rgba) + alpha_composite = alpha_composite.convert("RGB") + return alpha_composite + + +# Adapted from transformers.models.idefics2.image_processing_idefics2.get_max_height_width +def get_max_height_width(images_list: list[list["torch.Tensor|np.ndarray"]]) -> tuple[int, int]: + """ + Get the maximum height and width across all images in a batch. + """ + image_sizes = [] + for images in images_list: + for image in images: + image_sizes.append(image.shape[-2:]) + + max_height = max(size[0] for size in image_sizes) + max_width = max(size[1] for size in image_sizes) + return (max_height, max_width) + + +# Adapted from transformers.models.idefics2.image_processing_idefics2.get_resize_output_image_size +def get_resize_output_image_size(image, size: SizeDict) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Images are always channels-first (CHW). + """ + height, width = image.shape[-2:] + + min_len = size.shortest_edge + max_len = size.longest_edge + aspect_ratio = width / height + + if width >= height and width > max_len: + width = max_len + height = int(width / aspect_ratio) + elif height > width and height > max_len: + height = max_len + width = int(height * aspect_ratio) + height = max(height, min_len) + width = max(width, min_len) + return height, width + + +@auto_docstring +class Idefics2ImageProcessorPil(PilBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + do_convert_rgb = True + do_image_splitting = False + size = {"shortest_edge": 378, "longest_edge": 980} + default_to_square = False + model_input_names = ["pixel_values", "pixel_attention_mask"] + valid_kwargs = Idefics2ImageProcessorKwargs + + def __init__(self, **kwargs: Unpack[Idefics2ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[Idefics2ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def convert_to_rgb(self, image: ImageInput) -> ImageInput: + """Convert an image to RGB format.""" + return convert_to_rgb(image) + + def _prepare_images_structure(self, images: ImageInput, expected_ndims: int = 3) -> ImageInput: + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: PILImageResampling = PILImageResampling.BILINEAR, + **kwargs, + ) -> np.ndarray: + """Resize using Idefics2 shortest_edge/longest_edge logic. Images are always CHW.""" + if size.shortest_edge and size.longest_edge: + new_size = get_resize_output_image_size(image, size) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError("Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys.") + return super().resize(image, SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def split_images(self, image: np.ndarray) -> list[np.ndarray]: + """Split an image into 4 equal sub-images, and concatenate that sequence with the original image. Images are CHW.""" + height, width = image.shape[-2:] + mid_width = width // 2 + mid_height = height // 2 + + return [ + image[:, :mid_height, :mid_width], + image[:, :mid_height, mid_width:], + image[:, mid_height:, :mid_width], + image[:, mid_height:, mid_width:], + image, + ] + + def pad( + self, + image: np.ndarray, + padded_size: tuple[int, int], + fill: int = 0, + return_pixel_mask: bool = True, + ): + """Pad image to padded_size. Mirrors TorchvisionBackend.pad. Images are always CHW.""" + original_size = image.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + + pixel_mask = _make_pixel_mask(image, output_size=padded_size) if return_pixel_mask else None + + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + image = pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=fill, + data_format="channels_first", + input_data_format="channels_first", + ) + + return image, pixel_mask + + def _preprocess( + self, + images: list[list[np.ndarray]], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + do_image_splitting: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + """Process a batch of images. Mirrors TorchvisionBackend._preprocess with per-image loops.""" + if do_image_splitting: + new_images = [] + for batch_images in images: + new_batch = [] + for image in batch_images: + new_batch.extend(self.split_images(image)) + new_images.append(new_batch) + images = new_images + + if do_resize: + images = [ + [self.resize(image=img, size=size, resample=resample) for img in batch_images] + for batch_images in images + ] + + if do_rescale: + images = [[self.rescale(img, rescale_factor) for img in batch_images] for batch_images in images] + if do_normalize: + images = [[self.normalize(img, image_mean, image_std) for img in batch_images] for batch_images in images] + + if do_pad: + max_num_images = max(len(images_) for images_ in images) + max_height, max_width = get_max_height_width(images) + num_channels = images[0][0].shape[0] + + padded_images_list = [ + [np.zeros((num_channels, max_height, max_width), dtype=np.float32) for _ in range(max_num_images)] + for _ in range(len(images)) + ] + pixel_attention_masks = [ + [np.zeros((max_height, max_width), dtype=np.int64) for _ in range(max_num_images)] + for _ in range(len(images)) + ] + + for i, batch_images in enumerate(images): + for j, image in enumerate(batch_images): + padded_images_list[i][j], pixel_attention_masks[i][j] = self.pad(image, (max_height, max_width)) + images = padded_images_list + + if do_pad: + data = { + "pixel_values": np.array(images), + "pixel_attention_mask": np.array(pixel_attention_masks), + } + elif return_tensors == "pt": + data = {"pixel_values": np.asarray(images)} + else: + data = {"pixel_values": images} + + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["Idefics2ImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/idefics2/modeling_idefics2.py b/third_party/transformers/src/transformers/models/idefics2/modeling_idefics2.py new file mode 100644 index 0000000000000000000000000000000000000000..0fead94e2dfd355cee8d85d05f8b29dcc4bf6d01 --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/modeling_idefics2.py @@ -0,0 +1,1147 @@ +# Copyright 2024 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Idefics2 model.""" + +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 ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_bidirectional_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import AutoModel +from .configuration_idefics2 import Idefics2Config, Idefics2PerceiverConfig, Idefics2VisionConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Idefics2 model's outputs that may also contain a past key/values (to speed up sequential decoding). + """ +) +class Idefics2BaseModelOutputWithPast(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver + """ + + last_hidden_state: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Idefics2 causal language model (or autoregressive) outputs. + """ +) +# Copied from transformers.models.idefics.modeling_idefics.IdeficsCausalLMOutputWithPast with Idefics->Idefics2 +class Idefics2CausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + + image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: tuple[torch.FloatTensor] | None = None + + +class Idefics2VisionEmbeddings(nn.Module): + """ + This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings` to enable images of variable + resolution. + + The modifications are adapted from [Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution](https://huggingface.co/papers/2307.06304) + which allows treating images in their native aspect ratio and without the need to resize them to the same + fixed size. In particular, we start from the original pre-trained SigLIP model + (which uses images of fixed-size square images) and adapt it by training on images of variable resolutions. + """ + + def __init__(self, config: Idefics2VisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + ) + + self.num_patches_per_side = self.image_size // self.patch_size + self.num_patches = self.num_patches_per_side**2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + + def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor: + batch_size, _, max_im_h, max_im_w = pixel_values.shape + + patch_embeds = self.patch_embedding(pixel_values) + embeddings = patch_embeds.flatten(2).transpose(1, 2) + + max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size + boundaries = torch.arange( + 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side, device=pixel_values.device + ) + position_ids = torch.full( + size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0, device=pixel_values.device + ) + + nb_patches_h = patch_attention_mask[:, :, 0].sum(dim=1) # (batch_size,) + nb_patches_w = patch_attention_mask[:, 0, :].sum(dim=1) # (batch_size,) + + step_h = 1.0 / nb_patches_h # (batch_size,) + step_w = 1.0 / nb_patches_w # (batch_size,) + + max_patches_h = patch_attention_mask.size(1) + max_patches_w = patch_attention_mask.size(2) + h_indices = torch.arange(max_patches_h, device=position_ids.device, dtype=torch.float32) + w_indices = torch.arange(max_patches_w, device=position_ids.device, dtype=torch.float32) + + fractional_coords_h = h_indices[None, :] * step_h[:, None] + fractional_coords_w = w_indices[None, :] * step_w[:, None] + + fractional_coords_h = torch.clamp(fractional_coords_h, max=(1.0 - 1e-6)) + fractional_coords_w = torch.clamp(fractional_coords_w, max=(1.0 - 1e-6)) + + fractional_coords_h = fractional_coords_h.to(pixel_values.dtype) + fractional_coords_w = fractional_coords_w.to(pixel_values.dtype) + + bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) + bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) + + pos_ids = bucket_coords_h[:, :, None] * self.num_patches_per_side + bucket_coords_w[:, None, :] + pos_ids = pos_ids.reshape(batch_size, -1) + + position_ids[patch_attention_mask.view(batch_size, -1)] = pos_ids[patch_attention_mask.view(batch_size, -1)] + + embeddings = embeddings + self.position_embedding(position_ids) + return embeddings + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if hasattr(module, "num_key_value_groups"): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# Copied from transformers.models.siglip.modeling_siglip.SiglipAttention with Siglip->Idefics2Vision +class Idefics2VisionAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__ + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + # Ignore copy + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + is_causal=self.is_causal, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + ) + + attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +# Copied from transformers.models.siglip.modeling_siglip.SiglipMLP with Siglip->Idefics2Vision +class Idefics2VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class Idefics2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + output_size: int, + hidden_act: str, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, output_size, bias=False) + self.act_fn = ACT2FN[hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +# Copied from transformers.models.siglip.modeling_siglip.SiglipMultiheadAttentionPoolingHead with Siglip->Idefics2 +class Idefics2MultiheadAttentionPoolingHead(nn.Module): + """Multihead Attention Pooling.""" + + def __init__(self, config: Idefics2VisionConfig): + super().__init__() + + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + # Ignore copy + self.mlp = Idefics2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + output_size=config.hidden_size, + ) + + def forward(self, hidden_state): + batch_size = hidden_state.shape[0] + probe = self.probe.repeat(batch_size, 1, 1) + + hidden_state = self.attention(probe, hidden_state, hidden_state)[0] + + residual = hidden_state + hidden_state = self.layernorm(hidden_state) + hidden_state = residual + self.mlp(hidden_state) + + return hidden_state[:, 0] + + +class Idefics2EncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Idefics2VisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = Idefics2VisionAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = Idefics2VisionMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + @auto_docstring + # Copied from transformers.models.siglip.modeling_siglip.SiglipEncoderLayer.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +# Copied from transformers.models.siglip.modeling_siglip.SiglipEncoder with Siglip->Idefics2 +class Idefics2Encoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`Idefics2EncoderLayer`]. + + Args: + config: Idefics2Config + """ + + def __init__(self, config: Idefics2Config): + super().__init__() + self.config = config + self.layers = nn.ModuleList([Idefics2EncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + # Ignore copy + @auto_docstring + def forward( + self, + inputs_embeds, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + hidden_states = inputs_embeds + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_states, + attention_mask, + **kwargs, + ) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +@auto_docstring +class Idefics2PreTrainedModel(PreTrainedModel): + config: Idefics2Config + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["Idefics2VisionAttention", "Idefics2MLP", "Idefics2PerceiverLayer", "Idefics2DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _supports_attention_backend = True + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Idefics2MultiheadAttentionPoolingHead): + init.normal_(module.probe) + elif isinstance(module, Idefics2PerceiverResampler): + init.ones_(module.latents) + + +@auto_docstring( + custom_intro=""" + Idefics2 vision encoder model that returnss raw image embeddings. + """ +) +class Idefics2VisionTransformer(Idefics2PreTrainedModel): + config: Idefics2VisionConfig + input_modalities = ("image",) + _can_record_outputs = { + "hidden_states": Idefics2EncoderLayer, + "attentions": Idefics2VisionAttention, + } + + def __init__(self, config: Idefics2VisionConfig): + super().__init__(config) + embed_dim = config.hidden_size + + self.config = config + self.embeddings = Idefics2VisionEmbeddings(config) + self.encoder = Idefics2Encoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + pixel_values, + patch_attention_mask: torch.BoolTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + r""" + patch_attention_mask (`torch.BoolTensor` of shape `(batch_size, num_patches_height, num_patches_width)`, *optional*): + The attention mask for the patches. + """ + batch_size = pixel_values.size(0) + if patch_attention_mask is None: + patch_size = self.config.patch_size + patch_attention_mask = torch.ones( + ( + batch_size, + pixel_values.size(2) // patch_size, + pixel_values.size(3) // patch_size, + ) + ) + patch_attention_mask = patch_attention_mask.to(dtype=torch.bool, device=pixel_values.device) + + hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) + + patch_attention_mask = patch_attention_mask.view(batch_size, -1) + patch_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + attention_mask=patch_attention_mask, + **kwargs, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.post_layernorm(last_hidden_state) + + return BaseModelOutput(last_hidden_state=last_hidden_state) + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Idefics2 +class Idefics2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Idefics2RMSNorm 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 Idefics2PerceiverAttention(nn.Module): + def __init__(self, config, layer_idx: int | None = None) -> None: + """Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`""" + super().__init__() + self.config = config + self.layer_idx = None + self.hidden_size = config.hidden_size + self.num_heads = config.resampler_n_heads + self.head_dim = config.resampler_head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.attention_dropout = config.attention_dropout + self.scaling = self.head_dim**-0.5 + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.is_causal = False + + def forward( + self, + latents: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension! + + Args: + latents (`torch.Tensor`): Tensor of shape [bsz, n_latents, embed_dim] representing fixed length latents to compress to. + context (`torch.Tensor`): Tensor of shape [bsz, seq, embed_dim] representing long-form context to resample. + attention_mask (`torch.Tensor`, *optional*): Tensor of shape [bsz, 1, seq, n_latents] representing attention mask. + position_ids (`torch.LongTensor`, *optional*): Tensor of shape [bsz, seq] representing position indices of each input token. + past_key_values (`Cache`, *optional*): Tuple of tensors containing cached key and value states. + output_attentions (`bool`, *optional*, defaults to `False`): Whether to return attention weights. + use_cache (`bool`, *optional*, defaults to `False`): Whether to use past_key_values for caching. + """ + bsz, q_len, _ = latents.size() + kv_seq_len = q_len + context.size()[1] + + hidden_states = torch.concat([context, latents], dim=-2) + + queries = self.q_proj(latents) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + keys = keys.view(bsz, kv_seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + values = values.view(bsz, kv_seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + past_key_values = getattr(self, "past_key_values", past_key_values) + + if past_key_values is not None: + keys, values = past_key_values.update(keys, values, 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, + queries, + keys, + values, + attention_mask, + is_causal=self.is_causal, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class Idefics2PerceiverLayer(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.n_latents = config.resampler_n_latents + self.depth = config.resampler_depth + self.rms_norm_eps = config.rms_norm_eps + + self.input_latents_norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.input_context_norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.self_attn = Idefics2PerceiverAttention(config, layer_idx=layer_idx) + self.post_attention_layernorm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.mlp = Idefics2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.hidden_size * 4, + output_size=config.hidden_size, + hidden_act=config.hidden_act, + ) + + def forward( + self, + latents: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + """ + Args: + latents (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + context (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_values (`Cache`, *optional*): cached past key and value projection states + """ + residual = latents + + latents = self.input_latents_norm(latents) + context = self.input_context_norm(context) + + latents, _ = self.self_attn( + latents=latents, + context=context, + attention_mask=attention_mask, + **kwargs, + ) + latents = residual + latents + residual = latents + + latents = self.post_attention_layernorm(latents) + latents = self.mlp(latents) + latents = residual + latents + + return latents + + +@auto_docstring( + custom_intro=""" + Idefics2 perceiver resampler model that performs `depth` blocks of cross-attention with a fixed + """ +) +class Idefics2PerceiverResampler(Idefics2PreTrainedModel): + config: Idefics2PerceiverConfig + input_modalities = ("image",) + _supports_sdpa = True + _supports_flash_attn = True + _supports_flex_attn = True + + def __init__(self, config) -> None: + super().__init__(config) + self.hidden_size = config.hidden_size + self.hidden_act = config.hidden_act + self.n_latents = config.resampler_n_latents + self.depth = config.resampler_depth + self.rms_norm_eps = config.rms_norm_eps + + # Create Latents for Perceiver + self.latents = nn.Parameter(torch.ones(self.n_latents, self.hidden_size)) + + # Create Transformer Blocks + self.layers = nn.ModuleList([Idefics2PerceiverLayer(config, idx) for idx in range(self.depth)]) + self.norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + + self.post_init() + + @auto_docstring + def forward( + self, + context: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + r""" + context (`torch.FloatTensor` of shape `(batch, seq_len, embed_dim)`): + Input to the layer. + """ + # seq embed -> bsz seq embed + latents = self.latents.unsqueeze(0).expand((context.shape[0], *self.latents.size())) + + latent_attention_mask = torch.ones( + (attention_mask.size(0), latents.size(1)), dtype=attention_mask.dtype, device=attention_mask.device + ) + attention_mask = torch.cat([attention_mask, latent_attention_mask], dim=-1) + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=latents, + attention_mask=attention_mask, + ) + + compressed_context = latents + for perceiver_layer in self.layers: + compressed_context = perceiver_layer( + compressed_context, + context, + attention_mask=attention_mask, + position_ids=None, + **kwargs, + ) + + compressed_context = self.norm(compressed_context) + + return compressed_context + + +class Idefics2Connector(nn.Module): + def __init__(self, config): + super().__init__() + self.modality_projection = Idefics2MLP( + hidden_size=config.vision_config.hidden_size, + intermediate_size=config.text_config.intermediate_size, + output_size=config.text_config.hidden_size, + hidden_act=config.text_config.hidden_act, + ) + self.perceiver_resampler = Idefics2PerceiverResampler._from_config(config.perceiver_config) + + def forward(self, image_hidden_states, attention_mask): + image_hidden_states = self.modality_projection(image_hidden_states) + image_hidden_states = self.perceiver_resampler(context=image_hidden_states, attention_mask=attention_mask) + return image_hidden_states + + +@auto_docstring( + custom_intro=""" + Idefics2 model consisting of a SIGLIP vision encoder and Mistral language decoder + """ +) +class Idefics2Model(Idefics2PreTrainedModel): + def __init__(self, config: Idefics2Config): + super().__init__(config) + self.padding_idx = self.config.text_config.pad_token_id + self.vocab_size = self.config.text_config.vocab_size + + self.vision_model = Idefics2VisionTransformer._from_config(config.vision_config) + self.connector = Idefics2Connector(config) + self.text_model = AutoModel.from_config(config.text_config) + + self.image_seq_len = config.perceiver_config.resampler_n_latents + self.image_token_id = self.config.image_token_id + + self.post_init() + + def get_input_embeddings(self): + return self.text_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.text_model.set_input_embeddings(value) + + def inputs_merger( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.Tensor | None, + image_hidden_states: torch.Tensor | None, + ): + """ + This method aims at merging the token embeddings with the image hidden states into one single sequence of vectors that are fed to the transformer LM. + The merging happens as follows: + - The text token sequence is: `tok_1 tok_2 tok_3 ... tok_4`. + - We get the image hidden states for the image through the vision encoder (and potentially the perceiver), and that hidden state is then projected into the text embedding space. + We thus have a sequence of image hidden states of size (1, image_seq_len, hidden_dim), where 1 is for batch_size of 1 image and hidden_dim is the hidden_dim of the LM transformer. + - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM. + - To fit the format of that sequence, `input_ids`, `inputs_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + image_hidden_states = image_hidden_states.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_hidden_states) + return inputs_embeds + + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + batch_size, num_images, num_channels, height, width = pixel_values.shape + pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility + pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:]) + + # Remove padding images - padding images are full 0. + nb_values_per_image = pixel_values.shape[1:].numel() + real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image + pixel_values = pixel_values[real_images_inds].contiguous() + + # Handle the vision attention mask + if pixel_attention_mask is None: + pixel_attention_mask = torch.ones( + size=(pixel_values.size(0), pixel_values.size(2), pixel_values.size(3)), + dtype=torch.bool, + device=pixel_values.device, + ) + else: + # Remove padding images from the mask/pP p + pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:]) + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + + patch_size = self.config.vision_config.patch_size + patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size) + patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size) + patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) == patch_size * patch_size).bool() + # Get sequence from the vision encoder + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state + + # Modality projection & resampling + image_features = self.connector( + image_hidden_states, attention_mask=patch_attention_mask.view(pixel_values.size(0), -1) + ) + image_outputs.pooler_output = image_features.view(-1, image_features.shape[-1]) + + return image_outputs + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """ + ) + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | Idefics2BaseModelOutputWithPast: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection and perceiver resampling. + """ + + if self.training and self.text_model.gradient_checkpointing and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # retrieve input_ids and inputs_embeds + if input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.text_model.get_input_embeddings()(input_ids) + + # START VISUAL INPUTS INTEGRATION + if pixel_values is not None and image_hidden_states is not None: + raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") + elif pixel_values is not None: + image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask, **kwargs).pooler_output + elif image_hidden_states is not None: + image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device) + + if image_hidden_states is not None: + # When we generate, we don't want to replace the potential image_token_id that we generated by images + # that simply don't exist + inputs_embeds = self.inputs_merger( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + image_hidden_states=image_hidden_states, + ) + + kwargs["return_dict"] = True + outputs = self.text_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + return Idefics2BaseModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + The Idefics2 Model with a language modeling head. It is made up a SigLIP vision encoder, with a language modeling head on top. + """ +) +class Idefics2ForConditionalGeneration(Idefics2PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.text_model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = Idefics2Model(config) + self.image_token_id = self.config.image_token_id + + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.vocab_size = config.text_config.vocab_size + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.text_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.text_model.set_input_embeddings(value) + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + return self.model.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, **kwargs + ) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Idefics2CausalLMOutputWithPast: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection and perceiver resampling. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or `model.image_token_id` (where `model` is your instance of `Idefics2ForConditionalGeneration`). + Tokens with indices set to `model.image_token_id` are ignored (masked), the loss is only + computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> import torch + >>> from PIL import Image + >>> from io import BytesIO + + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + >>> from transformers.image_utils import load_image + + >>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible + >>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg") + >>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg") + >>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg") + + >>> processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b-base") + >>> model = AutoModelForImageTextToText.from_pretrained("HuggingFaceM4/idefics2-8b-base", device_map="auto") + + >>> BAD_WORDS_IDS = processor.tokenizer(["", ""], add_special_tokens=False).input_ids + >>> EOS_WORDS_IDS = [processor.tokenizer.eos_token_id] + + >>> # Create inputs + >>> prompts = [ + ... "In this image, we can see the city of New York, and more specifically the Statue of Liberty.In this image,", + ... "In which city is that bridge located?", + ... ] + >>> images = [[image1, image2], [image3]] + >>> inputs = processor(images=images, text=prompts, padding=True, return_tensors="pt").to("cuda") + + >>> # Generate + >>> generated_ids = model.generate(**inputs, bad_words_ids=BAD_WORDS_IDS, max_new_tokens=20) + >>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True) + + >>> print(generated_texts) + ['In this image, we can see the city of New York, and more specifically the Statue of Liberty. In this image, we can see the city of New York, and more specifically the Statue of Liberty.\n\n', 'In which city is that bridge located?\n\nThe bridge is located in the city of Pittsburgh, Pennsylvania.\n\n\nThe bridge is'] + ```""" + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs + ) + + return Idefics2CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + pixel_values=None, + pixel_attention_mask=None, + image_hidden_states=None, + logits_to_keep=None, + is_first_iteration=False, + use_cache=False, + **kwargs, + ): + # Overwritten -- there are mutually exclusive inputs (if the logic to make `image_hidden_states` take + # precedence is moved to the model, we can remove this fn) + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + logits_to_keep=logits_to_keep, + is_first_iteration=is_first_iteration, + use_cache=use_cache, + **kwargs, + ) + + if image_hidden_states is not None or (use_cache and not is_first_iteration): + model_inputs["pixel_values"] = None + model_inputs["pixel_attention_mask"] = None + + return model_inputs + + +__all__ = ["Idefics2ForConditionalGeneration", "Idefics2PreTrainedModel", "Idefics2Model"] diff --git a/third_party/transformers/src/transformers/models/idefics2/processing_idefics2.py b/third_party/transformers/src/transformers/models/idefics2/processing_idefics2.py new file mode 100644 index 0000000000000000000000000000000000000000..dd87290838ff4d847c31d23f2ac5fe9affe512af --- /dev/null +++ b/third_party/transformers/src/transformers/models/idefics2/processing_idefics2.py @@ -0,0 +1,190 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for IDEFICS2. +""" + +import re +from itertools import accumulate +from typing import TYPE_CHECKING, Union + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, is_valid_image, load_image +from ...processing_utils import ( + ProcessingKwargs, + ProcessorMixin, + Unpack, +) +from ...tokenization_utils_base import AddedToken, TextInput +from ...utils import auto_docstring, logging + + +if TYPE_CHECKING: + from ...tokenization_utils_base import PreTokenizedInput + + +logger = logging.get_logger(__name__) + + +def is_url(val) -> bool: + return isinstance(val, str) and val.startswith("http") + + +def is_image_or_image_url(elem): + return is_url(elem) or is_valid_image(elem) + + +class Idefics2ProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "add_special_tokens": True, + "padding": False, + "is_split_into_words": False, + }, + } + + +@auto_docstring +class Idefics2Processor(ProcessorMixin): + def __init__( + self, image_processor, tokenizer=None, image_seq_len: int = 64, chat_template: str | None = None, **kwargs + ): + r""" + image_seq_len (`int`, *optional*, defaults to 64): + The length of the image sequence i.e. the number of tokens per image in the input. + This parameter is used to build the string from the input prompt and image tokens and should match the + config.perceiver_config.resampler_n_latents value for the model used. + """ + if not hasattr(tokenizer, "image_token"): + self.fake_image_token = AddedToken("", normalized=False, special=True).content + self.image_token = AddedToken("", normalized=False, special=True).content + tokens_to_add = {"additional_special_tokens": [self.fake_image_token, self.image_token]} + tokenizer.add_special_tokens(tokens_to_add) + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + else: + self.fake_image_token = tokenizer.image_boundary_token + self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id + + self.end_of_utterance_token = AddedToken("", normalized=False, special=True) + tokenizer.add_special_tokens({"additional_special_tokens": [self.end_of_utterance_token]}) + self.image_seq_len = image_seq_len + + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + def _extract_images_from_prompts(self, prompts): + prompt_images = [] + for prompt in prompts: + images = [] + for elem in prompt: + if is_valid_image(elem): + images.append(elem) + elif is_url(elem): + images.append(load_image(elem)) + prompt_images.append(images) + return prompt_images + + @auto_docstring + def __call__( + self, + images: ImageInput | list[ImageInput] | list[list[ImageInput]] = None, + text: Union[TextInput, "PreTokenizedInput", list[TextInput], list["PreTokenizedInput"]] = None, + **kwargs: Unpack[Idefics2ProcessorKwargs], + ) -> BatchFeature: + if text is None and images is None: + raise ValueError("You must provide either `text` or `images`.") + + output_kwargs = self._merge_kwargs( + Idefics2ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + + n_images_in_text = [] + inputs = {} + + if text is not None: + if isinstance(text, str): + text = [text] + elif not isinstance(text, list) and not isinstance(text[0], str): + raise ValueError("Invalid input text. Please provide a string, or a list of strings") + + # Replace the image token with fake tokens around the expanded image token sequence of length `image_seq_len` + fake_image_token = self.fake_image_token + image_token = self.image_token + image_str = f"{fake_image_token}{image_token * self.image_seq_len}{fake_image_token}" + + if self.image_processor.do_image_splitting: + # A single image token is split into 4 patches + 1 original image + image_str = image_str * 5 + + prompt_strings = [] + closing_fake_pattern = re.compile(rf"{re.escape(fake_image_token)}(?=[^\s<])") + for sample in text: + n_images_in_text.append(sample.count(image_token)) + sample = sample.replace(image_token, image_str) + # Remove any double fake tokens if images are adjacent + sample = sample.replace(f"{fake_image_token}{fake_image_token}", f"{fake_image_token}") + # Ensure words attached directly after the closing fake token remain word-boundary aligned + sample = closing_fake_pattern.sub(f"{fake_image_token} ", sample) + prompt_strings.append(sample) + + text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) + inputs.update(text_inputs) + + if images is not None: + if is_image_or_image_url(images): + images = [[images]] + elif isinstance(images, (list, tuple)) and is_image_or_image_url(images[0]): + if text is not None: + if sum(n_images_in_text) != len(images): + raise ValueError( + f"The total number of {image_token} tokens in the prompts should be the same as the number of images passed." + f" Found {sum(n_images_in_text)} {image_token} tokens and {len(images)} images." + ) + # Reorganize the images to match the prompts + cumsum_images_in_text = [0] + list(accumulate(n_images_in_text)) + images = [ + images[cumsum_images_in_text[i] : cumsum_images_in_text[i + 1]] + for i in range(len(n_images_in_text)) + ] + else: + images = [images] + + elif ( + not isinstance(images, (list, tuple)) + and not isinstance(images[0], (list, tuple)) + and not is_image_or_image_url(images[0][0]) + ): + raise ValueError( + "Invalid input images. Please provide a single image or a list of images or a list of list of images." + ) + + n_images_in_images = [len(sample) for sample in images] + if text is not None and not n_images_in_images == n_images_in_text: + raise ValueError( + f"The number of images in the text {n_images_in_text} and images {n_images_in_images} should be the same." + ) + + # Load images if they are URLs + images = [[load_image(im) for im in sample] for sample in images] + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + inputs.update(image_inputs) + + return BatchFeature(inputs, tensor_type=return_tensors) + + +__all__ = ["Idefics2Processor"] diff --git a/third_party/transformers/src/transformers/models/jais2/__init__.py b/third_party/transformers/src/transformers/models/jais2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..137ee4e32e1425730fe8d6346776ebbf64719a0f --- /dev/null +++ b/third_party/transformers/src/transformers/models/jais2/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_jais2 import * + from .modeling_jais2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/jais2/configuration_jais2.py b/third_party/transformers/src/transformers/models/jais2/configuration_jais2.py new file mode 100644 index 0000000000000000000000000000000000000000..7139661d757573d73c67d00f7ad5647a6c514073 --- /dev/null +++ b/third_party/transformers/src/transformers/models/jais2/configuration_jais2.py @@ -0,0 +1,101 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/jais2/modular_jais2.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_jais2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring +from ...utils.type_validators import interval + + +@auto_docstring(checkpoint="inceptionai/Jais-2-8B-Chat") +@strict +class Jais2Config(PreTrainedConfig): + r""" + ```python + >>> from transformers import Jais2Model, Jais2Config + + >>> # Initializing a Jais2 jais2-7b style configuration + >>> configuration = Jais2Config() + + >>> # Initializing a model from the jais2-7b style configuration + >>> model = Jais2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "jais2" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 150272 + hidden_size: int = 3328 + intermediate_size: int = 26624 + num_hidden_layers: int = 32 + num_attention_heads: int = 26 + num_key_value_heads: int | None = None + hidden_act: str = "relu2" + max_position_embeddings: int = 8192 + initializer_range: float = interval(min=0.0, max=1.0)(default=0.02) + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 150024 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = True + attention_dropout: int | float | None = 0.0 + mlp_bias: bool = True + head_dim: int | None = None + layer_norm_eps: float = 1e-5 + + def __post_init__(self, **kwargs): + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +__all__ = ["Jais2Config"] diff --git a/third_party/transformers/src/transformers/models/jais2/modeling_jais2.py b/third_party/transformers/src/transformers/models/jais2/modeling_jais2.py new file mode 100644 index 0000000000000000000000000000000000000000..5e6a37c0172ddfd7b273cab52b69ca437b24c351 --- /dev/null +++ b/third_party/transformers/src/transformers/models/jais2/modeling_jais2.py @@ -0,0 +1,471 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/jais2/modular_jais2.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_jais2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn as nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_jais2 import Jais2Config + + +class Jais2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernelized_func(apply_rotary_pos_emb) +class Jais2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Jais2Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Jais2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Jais2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Jais2Attention(config=config, layer_idx=layer_idx) + + self.mlp = Jais2MLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + 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 Jais2PreTrainedModel(PreTrainedModel): + config: Jais2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Jais2DecoderLayer"] + _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": Jais2DecoderLayer, + "attentions": Jais2Attention, + } + + +class Jais2RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Jais2Config, 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: Jais2Config | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@auto_docstring +class Jais2Model(Jais2PreTrainedModel): + def __init__(self, config: Jais2Config): + 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( + [Jais2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.rotary_emb = Jais2RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = Jais2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, Jais2ForCausalLM + + >>> model = Jais2ForCausalLM.from_pretrained("inceptionai/Jais-2-8B-Chat") + >>> tokenizer = AutoTokenizer.from_pretrained("inceptionai/Jais-2-8B-Chat") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = ["Jais2Model", "Jais2ForCausalLM", "Jais2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/jais2/modular_jais2.py b/third_party/transformers/src/transformers/models/jais2/modular_jais2.py new file mode 100644 index 0000000000000000000000000000000000000000..3fbdf2c8cd4647dda6a5380277c56074cbb63f3a --- /dev/null +++ b/third_party/transformers/src/transformers/models/jais2/modular_jais2.py @@ -0,0 +1,107 @@ +# Copyright 2025 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch.nn as nn +from huggingface_hub.dataclasses import strict + +from ...utils import auto_docstring, can_return_tuple +from ..llama.configuration_llama import LlamaConfig +from ..llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaModel, + LlamaPreTrainedModel, +) +from ..nemotron.modeling_nemotron import NemotronMLP + + +@auto_docstring(checkpoint="inceptionai/Jais-2-8B-Chat") +@strict +class Jais2Config(LlamaConfig): + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + + vocab_size: int = 150272 + hidden_size: int = 3328 + intermediate_size: int = 26624 + num_attention_heads: int = 26 + hidden_act: str = "relu2" + max_position_embeddings: int = 8192 + layer_norm_eps: float = 1e-5 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 150024 + attention_bias: bool = True + mlp_bias: bool = True + rms_norm_eps = AttributeError() + pretraining_tp = AttributeError() + + +class Jais2MLP(NemotronMLP): + pass + + +class Jais2DecoderLayer(LlamaDecoderLayer): + def __init__(self, config: Jais2Config, layer_idx: int): + super().__init__(config, layer_idx) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + +class Jais2PreTrainedModel(LlamaPreTrainedModel): + pass + + +class Jais2Model(LlamaModel): + def __init__(self, config: Jais2Config): + super().__init__(config) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + +class Jais2ForCausalLM(LlamaForCausalLM): + @can_return_tuple + @auto_docstring + def forward(self, **super_kwargs): + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, Jais2ForCausalLM + + >>> model = Jais2ForCausalLM.from_pretrained("inceptionai/Jais-2-8B-Chat") + >>> tokenizer = AutoTokenizer.from_pretrained("inceptionai/Jais-2-8B-Chat") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + return super().forward(**super_kwargs) + + +__all__ = [ + "Jais2Config", + "Jais2Model", + "Jais2ForCausalLM", + "Jais2PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/__init__.py b/third_party/transformers/src/transformers/models/layoutlmv2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5e12b2f207e79e68e780f326f355a55554d665dc --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_layoutlmv2 import * + from .feature_extraction_layoutlmv2 import * + from .image_processing_layoutlmv2 import * + from .image_processing_pil_layoutlmv2 import * + from .modeling_layoutlmv2 import * + from .processing_layoutlmv2 import * + from .tokenization_layoutlmv2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/configuration_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/configuration_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..2fd4221d8eebff3530a502114abe01cc9e388eb8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/configuration_layoutlmv2.py @@ -0,0 +1,159 @@ +# Copyright 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. +"""LayoutLMv2 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, is_detectron2_available + + +# soft dependency +if is_detectron2_available(): + import detectron2 + + +@auto_docstring(checkpoint="microsoft/layoutlmv2-base-uncased") +@strict +class LayoutLMv2Config(PreTrainedConfig): + r""" + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum value that the 2D position embedding might ever be used with. Typically set this to something + large just in case (e.g., 1024). + max_rel_pos (`int`, *optional*, defaults to 128): + The maximum number of relative positions to be used in the self-attention mechanism. + rel_pos_bins (`int`, *optional*, defaults to 32): + The number of relative position bins to be used in the self-attention mechanism. + fast_qkv (`bool`, *optional*, defaults to `True`): + Whether or not to use a single matrix for the queries, keys, values in the self-attention layers. + max_rel_2d_pos (`int`, *optional*, defaults to 256): + The maximum number of relative 2D positions in the self-attention mechanism. + rel_2d_pos_bins (`int`, *optional*, defaults to 64): + The number of 2D relative position bins in the self-attention mechanism. + convert_sync_batchnorm (`bool`, *optional*, defaults to `True`): + Whether or not to convert batch normalization layers to synchronized batch normalization layers. + image_feature_pool_shape (`list[int]`, *optional*, defaults to `[7, 7, 256]`): + The shape of the average-pooled feature map. + coordinate_size (`int`, *optional*, defaults to 128): + Dimension of the coordinate embeddings. + shape_size (`int`, *optional*, defaults to 128): + Dimension of the width and height embeddings. + has_relative_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a relative attention bias in the self-attention mechanism. + has_spatial_attention_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use a spatial attention bias in the self-attention mechanism. + has_visual_segment_embedding (`bool`, *optional*, defaults to `False`): + Whether or not to add visual segment embeddings. + detectron2_config_args (`dict`, *optional*): + Dictionary containing the configuration arguments of the Detectron2 visual backbone. Refer to [this + file](https://github.com/microsoft/unilm/blob/master/layoutlmft/layoutlmft/models/layoutlmv2/detectron2_config.py) + for details regarding default values. + + Example: + + ```python + >>> from transformers import LayoutLMv2Config, LayoutLMv2Model + + >>> # Initializing a LayoutLMv2 microsoft/layoutlmv2-base-uncased style configuration + >>> configuration = LayoutLMv2Config() + + >>> # Initializing a model (with random weights) from the microsoft/layoutlmv2-base-uncased style configuration + >>> model = LayoutLMv2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "layoutlmv2" + + vocab_size: int = 30522 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.1 + attention_probs_dropout_prob: float | int = 0.1 + max_position_embeddings: int = 512 + type_vocab_size: int = 2 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + pad_token_id: int | None = 0 + max_2d_position_embeddings: int = 1024 + max_rel_pos: int = 128 + rel_pos_bins: int = 32 + fast_qkv: bool = True + max_rel_2d_pos: int = 256 + rel_2d_pos_bins: int = 64 + convert_sync_batchnorm: bool = True + image_feature_pool_shape: list[int] | tuple[int, ...] = (7, 7, 256) + coordinate_size: int = 128 + shape_size: int = 128 + has_relative_attention_bias: bool = True + has_spatial_attention_bias: bool = True + has_visual_segment_embedding: bool = False + detectron2_config_args: dict | None = None + + def __post_init__(self, **kwargs): + super().__post_init__(**kwargs) + self.detectron2_config_args = ( + self.detectron2_config_args + if self.detectron2_config_args is not None + else self.get_default_detectron2_config() + ) + + @classmethod + def get_default_detectron2_config(cls): + return { + "MODEL.MASK_ON": True, + "MODEL.PIXEL_STD": [57.375, 57.120, 58.395], + "MODEL.BACKBONE.NAME": "build_resnet_fpn_backbone", + "MODEL.FPN.IN_FEATURES": ["res2", "res3", "res4", "res5"], + "MODEL.ANCHOR_GENERATOR.SIZES": [[32], [64], [128], [256], [512]], + "MODEL.RPN.IN_FEATURES": ["p2", "p3", "p4", "p5", "p6"], + "MODEL.RPN.PRE_NMS_TOPK_TRAIN": 2000, + "MODEL.RPN.PRE_NMS_TOPK_TEST": 1000, + "MODEL.RPN.POST_NMS_TOPK_TRAIN": 1000, + "MODEL.POST_NMS_TOPK_TEST": 1000, + "MODEL.ROI_HEADS.NAME": "StandardROIHeads", + "MODEL.ROI_HEADS.NUM_CLASSES": 5, + "MODEL.ROI_HEADS.IN_FEATURES": ["p2", "p3", "p4", "p5"], + "MODEL.ROI_BOX_HEAD.NAME": "FastRCNNConvFCHead", + "MODEL.ROI_BOX_HEAD.NUM_FC": 2, + "MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION": 14, + "MODEL.ROI_MASK_HEAD.NAME": "MaskRCNNConvUpsampleHead", + "MODEL.ROI_MASK_HEAD.NUM_CONV": 4, + "MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION": 7, + "MODEL.RESNETS.DEPTH": 101, + "MODEL.RESNETS.SIZES": [[32], [64], [128], [256], [512]], + "MODEL.RESNETS.ASPECT_RATIOS": [[0.5, 1.0, 2.0]], + "MODEL.RESNETS.OUT_FEATURES": ["res2", "res3", "res4", "res5"], + "MODEL.RESNETS.NUM_GROUPS": 32, + "MODEL.RESNETS.WIDTH_PER_GROUP": 8, + "MODEL.RESNETS.STRIDE_IN_1X1": False, + } + + def get_detectron2_config(self): + detectron2_config = detectron2.config.get_cfg() + for k, v in self.detectron2_config_args.items(): + attributes = k.split(".") + to_set = detectron2_config + for attribute in attributes[:-1]: + to_set = getattr(to_set, attribute) + setattr(to_set, attributes[-1], v) + + return detectron2_config + + +__all__ = ["LayoutLMv2Config"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..34ed26abf1b2a708c0747cc97d28126774f92925 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_layoutlmv2.py @@ -0,0 +1,193 @@ +# 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 LayoutLMv2.""" + +import numpy as np +import torch + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images, to_pil_image +from ...image_utils import ChannelDimension, ImageInput, PILImageResampling, SizeDict +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, + is_pytesseract_available, + logging, + requires_backends, +) + + +if is_pytesseract_available(): + import pytesseract + +from torchvision.transforms.v2 import functional as tvF + + +logger = logging.get_logger(__name__) + + +def normalize_box(box, width, height): + return [ + int(1000 * (box[0] / width)), + int(1000 * (box[1] / height)), + int(1000 * (box[2] / width)), + int(1000 * (box[3] / height)), + ] + + +def apply_tesseract( + image: "np.ndarray | torch.Tensor", + lang: str | None, + tesseract_config: str | None = None, + input_data_format: str | ChannelDimension | None = None, +): + """Applies Tesseract OCR on a document image, and returns recognized words + normalized bounding boxes.""" + requires_backends(apply_tesseract, ["pytesseract"]) + + # Convert torch tensor to numpy if needed + if hasattr(image, "cpu"): + image = image.cpu().numpy() + elif not isinstance(image, np.ndarray): + image = np.array(image) + + tesseract_config = tesseract_config if tesseract_config is not None else "" + + # apply OCR + pil_image = to_pil_image(image, input_data_format=input_data_format) + image_width, image_height = pil_image.size + data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config) + words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"] + + # filter empty words and corresponding coordinates + irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()] + words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices] + left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices] + top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices] + width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices] + height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices] + + # turn coordinates into (left, top, left+width, top+height) format + actual_boxes = [] + for x, y, w, h in zip(left, top, width, height): + actual_box = [x, y, x + w, y + h] + actual_boxes.append(actual_box) + + # finally, normalize the bounding boxes + normalized_boxes = [] + for box in actual_boxes: + normalized_boxes.append(normalize_box(box, image_width, image_height)) + + assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes" + + return words, normalized_boxes + + +class LayoutLMv2ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + apply_ocr (`bool`, *optional*, defaults to `self.apply_ocr`): + Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by + the `apply_ocr` parameter in the `preprocess` method. + ocr_lang (`str`, *optional*): + The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is + used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method. + tesseract_config (`str`, *optional*): + Any additional custom configuration flags that are forwarded to the `config` parameter when calling + Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the + `preprocess` method. + """ + + apply_ocr: bool + ocr_lang: str | None + tesseract_config: str | None + + +@auto_docstring +class LayoutLMv2ImageProcessor(TorchvisionBackend): + valid_kwargs = LayoutLMv2ImageProcessorKwargs + resample = PILImageResampling.BILINEAR + size = {"height": 224, "width": 224} + rescale_factor = None + do_resize = True + apply_ocr = True + ocr_lang = None + tesseract_config = "" + + def __init__(self, **kwargs: Unpack[LayoutLMv2ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LayoutLMv2ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + apply_ocr: bool = True, + ocr_lang: str | None = None, + tesseract_config: str | None = None, + **kwargs, + ) -> BatchFeature: + # Tesseract OCR to get words + normalized bounding boxes + if apply_ocr: + requires_backends(self, "pytesseract") + words_batch = [] + boxes_batch = [] + for image in images: + if image.is_cuda: + logger.warning_once( + "apply_ocr can only be performed on cpu. Tensors will be transferred to cpu before processing." + ) + words, boxes = apply_tesseract( + image.cpu(), ocr_lang, tesseract_config, input_data_format=ChannelDimension.FIRST + ) + words_batch.append(words) + boxes_batch.append(boxes) + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, resample=resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + # flip color channels from RGB to BGR (as Detectron2 requires this) + stacked_images = stacked_images.flip(1) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + data = BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + if apply_ocr: + data["words"] = words_batch + data["boxes"] = boxes_batch + + return data + + +__all__ = ["LayoutLMv2ImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_pil_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_pil_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..5d6cad0ac1a1bcecf9598306bdb463a1752d727e --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/image_processing_pil_layoutlmv2.py @@ -0,0 +1,181 @@ +# 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 LayoutLMv2.""" + +from typing import TYPE_CHECKING + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import flip_channel_order, to_pil_image +from ...image_utils import ( + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, requires_backends + + +if TYPE_CHECKING: + pass + +try: + import pytesseract +except ImportError: + pytesseract = None + + +# Adapted from transformers.models.layoutlmv2.image_processing_layoutlmv2.LayoutLMv2ImageProcessorKwargs +class LayoutLMv2ImageProcessorKwargs(ImagesKwargs, total=False): + r""" + apply_ocr (`bool`, *optional*, defaults to `self.apply_ocr`): + Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by + the `apply_ocr` parameter in the `preprocess` method. + ocr_lang (`str`, *optional*): + The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is + used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method. + tesseract_config (`str`, *optional*): + Any additional custom configuration flags that are forwarded to the `config` parameter when calling + Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the + `preprocess` method. + """ + + apply_ocr: bool + ocr_lang: str | None + tesseract_config: str | None + + +# Adapted from transformers.models.layoutlmv2.image_processing_layoutlmv2.normalize_box +def normalize_box(box, width, height): + return [ + int(1000 * (box[0] / width)), + int(1000 * (box[1] / height)), + int(1000 * (box[2] / width)), + int(1000 * (box[3] / height)), + ] + + +# Adapted from transformers.models.layoutlmv2.image_processing_layoutlmv2.apply_tesseract +def apply_tesseract( + image: np.ndarray, + lang: str | None, + tesseract_config: str | None = None, + input_data_format: str | ChannelDimension | None = None, +): + """Applies Tesseract OCR on a document image, and returns recognized words + normalized bounding boxes.""" + requires_backends(apply_tesseract, ["pytesseract"]) + + # Convert torch tensor to numpy if needed + if hasattr(image, "cpu"): + image = image.cpu().numpy() + elif not isinstance(image, np.ndarray): + image = np.array(image) + + tesseract_config = tesseract_config if tesseract_config is not None else "" + + # apply OCR + pil_image = to_pil_image(image, input_data_format=input_data_format) + image_width, image_height = pil_image.size + data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config) + words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"] + + # filter empty words and corresponding coordinates + irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()] + words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices] + left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices] + top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices] + width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices] + height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices] + + # turn coordinates into (left, top, left+width, top+height) format + actual_boxes = [] + for x, y, w, h in zip(left, top, width, height): + actual_box = [x, y, x + w, y + h] + actual_boxes.append(actual_box) + + # finally, normalize the bounding boxes + normalized_boxes = [] + for box in actual_boxes: + normalized_boxes.append(normalize_box(box, image_width, image_height)) + + assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes" + + return words, normalized_boxes + + +@auto_docstring +class LayoutLMv2ImageProcessorPil(PilBackend): + valid_kwargs = LayoutLMv2ImageProcessorKwargs + resample = PILImageResampling.BILINEAR + size = {"height": 224, "width": 224} + rescale_factor = None + do_resize = True + apply_ocr = True + ocr_lang = None + tesseract_config = "" + + def __init__(self, **kwargs: Unpack[LayoutLMv2ImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LayoutLMv2ImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + return_tensors: str | TensorType | None, + apply_ocr: bool = True, + ocr_lang: str | None = None, + tesseract_config: str | None = None, + **kwargs, + ) -> BatchFeature: + # Tesseract OCR to get words + normalized bounding boxes + if apply_ocr: + requires_backends(self, "pytesseract") + words_batch = [] + boxes_batch = [] + for image in images: + words, boxes = apply_tesseract( + image, ocr_lang, tesseract_config, input_data_format=ChannelDimension.FIRST + ) + words_batch.append(words) + boxes_batch.append(boxes) + + processed_images = [] + for image in images: + if do_resize: + image = self.resize(image=image, size=size, resample=resample) + + # flip color channels from RGB to BGR (as Detectron2 requires this) + image = flip_channel_order(image, input_data_format=ChannelDimension.FIRST) + + processed_images.append(image) + + data = BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + if apply_ocr: + data["words"] = words_batch + data["boxes"] = boxes_batch + + return data + + +__all__ = ["LayoutLMv2ImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/modeling_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/modeling_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..5b8811d615ce6a742445bb620f8dc51866ea7afc --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/modeling_layoutlmv2.py @@ -0,0 +1,1298 @@ +# Copyright 2021 Microsoft Research 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 LayoutLMv2 model.""" + +import math + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import auto_docstring, is_detectron2_available, logging, requires_backends +from ...utils.generic import TransformersKwargs, can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_layoutlmv2 import LayoutLMv2Config + + +# soft dependency +if is_detectron2_available(): + import detectron2 + from detectron2.modeling import META_ARCH_REGISTRY + + # This is needed as otherwise their overload will break sequential loading by overwriting buffer over and over. See + # https://github.com/facebookresearch/detectron2/blob/9604f5995cc628619f0e4fd913453b4d7d61db3f/detectron2/layers/batch_norm.py#L83-L86 + detectron2.layers.batch_norm.FrozenBatchNorm2d._load_from_state_dict = torch.nn.Module._load_from_state_dict + +logger = logging.get_logger(__name__) + + +class LayoutLMv2Embeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + + self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size) + self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def _calc_spatial_position_embeddings(self, bbox): + try: + left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0]) + upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1]) + right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2]) + lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3]) + except IndexError as e: + raise IndexError("The `bbox` coordinate values should be within 0-1000 range.") from e + + h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3] - bbox[:, :, 1]) + w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2] - bbox[:, :, 0]) + + spatial_position_embeddings = torch.cat( + [ + left_position_embeddings, + upper_position_embeddings, + right_position_embeddings, + lower_position_embeddings, + h_position_embeddings, + w_position_embeddings, + ], + dim=-1, + ) + return spatial_position_embeddings + + +class LayoutLMv2SelfAttention(nn.Module): + def __init__(self, config): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + self.fast_qkv = config.fast_qkv + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + if config.fast_qkv: + self.qkv_linear = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=False) + self.q_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size)) + self.v_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size)) + else: + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def compute_qkv(self, hidden_states): + if self.fast_qkv: + qkv = self.qkv_linear(hidden_states) + q, k, v = torch.chunk(qkv, 3, dim=-1) + if q.ndimension() == self.q_bias.ndimension(): + q = q + self.q_bias + v = v + self.v_bias + else: + _sz = (1,) * (q.ndimension() - 1) + (-1,) + q = q + self.q_bias.view(*_sz) + v = v + self.v_bias.view(*_sz) + else: + q = self.query(hidden_states) + k = self.key(hidden_states) + v = self.value(hidden_states) + return q, k, v + + def forward( + self, + hidden_states, + attention_mask=None, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size = hidden_states.shape[0] + query, key, value = self.compute_qkv(hidden_states) + + # (B, L, H*D) -> (B, H, L, D) + query_layer = query.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + key_layer = key.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + value_layer = value.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + + query_layer = query_layer / math.sqrt(self.attention_head_size) + # [BSZ, NAT, L, L] + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + if self.has_relative_attention_bias: + attention_scores += rel_pos + if self.has_spatial_attention_bias: + attention_scores += rel_2d_pos + attention_scores = attention_scores.float().masked_fill_( + attention_mask.to(torch.bool), torch.finfo(attention_scores.dtype).min + ) + attention_probs = nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32).type_as(value_layer) + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer, attention_probs + + +class LayoutLMv2Attention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = LayoutLMv2SelfAttention(config) + self.output = LayoutLMv2SelfOutput(config) + + def forward( + self, + hidden_states, + attention_mask=None, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + attention_output, _ = self.self( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + attention_output = self.output(attention_output, residual) + return attention_output + + +class LayoutLMv2SelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->LayoutLMv2 +class LayoutLMv2Intermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->LayoutLM +class LayoutLMv2Output(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class LayoutLMv2Layer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = LayoutLMv2Attention(config) + self.intermediate = LayoutLMv2Intermediate(config) + self.output = LayoutLMv2Output(config) + + def forward( + self, + hidden_states, + attention_mask=None, + output_attentions=False, + rel_pos=None, + rel_2d_pos=None, + **kwargs: Unpack[TransformersKwargs], + ): + attention_output = self.attention( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + ) + + layer_output = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output + ) + + return layer_output + + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +def relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + Translate relative position to a bucket number for relative attention. The relative position is defined as + memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to + position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small + absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions + >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should + allow for more graceful generalization to longer sequences than the model has been trained on. + + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + + Returns: + a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) + """ + + ret = 0 + if bidirectional: + num_buckets //= 2 + ret += (relative_position > 0).long() * num_buckets + n = torch.abs(relative_position) + else: + n = torch.max(-relative_position, torch.zeros_like(relative_position)) + # now n is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = n < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + val_if_large = max_exact + ( + torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) + ).to(torch.long) + val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1)) + + ret += torch.where(is_small, n, val_if_large) + return ret + + +class LayoutLMv2Encoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([LayoutLMv2Layer(config) for _ in range(config.num_hidden_layers)]) + + self.has_relative_attention_bias = config.has_relative_attention_bias + self.has_spatial_attention_bias = config.has_spatial_attention_bias + + if self.has_relative_attention_bias: + self.rel_pos_bins = config.rel_pos_bins + self.max_rel_pos = config.max_rel_pos + self.rel_pos_bias = nn.Linear(self.rel_pos_bins, config.num_attention_heads, bias=False) + + if self.has_spatial_attention_bias: + self.max_rel_2d_pos = config.max_rel_2d_pos + self.rel_2d_pos_bins = config.rel_2d_pos_bins + self.rel_pos_x_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + self.rel_pos_y_bias = nn.Linear(self.rel_2d_pos_bins, config.num_attention_heads, bias=False) + + self.gradient_checkpointing = False + + def _calculate_1d_position_embeddings(self, position_ids): + rel_pos_mat = position_ids.unsqueeze(-2) - position_ids.unsqueeze(-1) + rel_pos = relative_position_bucket( + rel_pos_mat, + num_buckets=self.rel_pos_bins, + max_distance=self.max_rel_pos, + ) + # Since this is a simple indexing operation that is independent of the input, + # no need to track gradients for this operation + # + # Without this no_grad context, training speed slows down significantly + with torch.no_grad(): + rel_pos = self.rel_pos_bias.weight.t()[rel_pos].permute(0, 3, 1, 2) + rel_pos = rel_pos.contiguous() + return rel_pos + + def _calculate_2d_position_embeddings(self, bbox): + position_coord_x = bbox[:, :, 0] + position_coord_y = bbox[:, :, 3] + rel_pos_x_2d_mat = position_coord_x.unsqueeze(-2) - position_coord_x.unsqueeze(-1) + rel_pos_y_2d_mat = position_coord_y.unsqueeze(-2) - position_coord_y.unsqueeze(-1) + rel_pos_x = relative_position_bucket( + rel_pos_x_2d_mat, + num_buckets=self.rel_2d_pos_bins, + max_distance=self.max_rel_2d_pos, + ) + rel_pos_y = relative_position_bucket( + rel_pos_y_2d_mat, + num_buckets=self.rel_2d_pos_bins, + max_distance=self.max_rel_2d_pos, + ) + # Since this is a simple indexing operation that is independent of the input, + # no need to track gradients for this operation + # + # Without this no_grad context, training speed slows down significantly + with torch.no_grad(): + rel_pos_x = self.rel_pos_x_bias.weight.t()[rel_pos_x].permute(0, 3, 1, 2) + rel_pos_y = self.rel_pos_y_bias.weight.t()[rel_pos_y].permute(0, 3, 1, 2) + rel_pos_x = rel_pos_x.contiguous() + rel_pos_y = rel_pos_y.contiguous() + rel_2d_pos = rel_pos_x + rel_pos_y + return rel_2d_pos + + def forward( + self, + hidden_states, + attention_mask=None, + bbox=None, + position_ids=None, + **kwargs: Unpack[TransformersKwargs], + ): + rel_pos = self._calculate_1d_position_embeddings(position_ids) if self.has_relative_attention_bias else None + rel_2d_pos = self._calculate_2d_position_embeddings(bbox) if self.has_spatial_attention_bias else None + + for layer_module in self.layer: + hidden_states = layer_module( + hidden_states, + attention_mask, + rel_pos=rel_pos, + rel_2d_pos=rel_2d_pos, + **kwargs, + ) + return BaseModelOutput(last_hidden_state=hidden_states) + + +@auto_docstring +class LayoutLMv2PreTrainedModel(PreTrainedModel): + config: LayoutLMv2Config + base_model_prefix = "layoutlmv2" + input_modalities = ("image", "text") + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, LayoutLMv2SelfAttention): + if self.config.fast_qkv: + init.zeros_(module.q_bias) + init.zeros_(module.v_bias) + elif isinstance(module, LayoutLMv2Embeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + elif isinstance(module, LayoutLMv2VisualBackbone): + num_channels = len(module.cfg.MODEL.PIXEL_MEAN) + init.copy_(module.pixel_mean, torch.Tensor(module.cfg.MODEL.PIXEL_MEAN).view(num_channels, 1, 1)) + init.copy_(module.pixel_std, torch.Tensor(module.cfg.MODEL.PIXEL_STD).view(num_channels, 1, 1)) + elif isinstance(module, LayoutLMv2Model): + if hasattr(module, "visual_segment_embedding"): + init.normal_(module.visual_segment_embedding, mean=0.0, std=self.config.initializer_range) + # We check the existence of each one since detectron2 seems to do weird things + elif isinstance(module, detectron2.layers.FrozenBatchNorm2d): + init.ones_(module.weight) + init.zeros_(module.bias) + init.zeros_(module.running_mean) + init.constant_(module.running_var, 1.0 - module.eps) + + +def my_convert_sync_batchnorm(module, process_group=None): + # same as `nn.modules.SyncBatchNorm.convert_sync_batchnorm` but allowing converting from `detectron2.layers.FrozenBatchNorm2d` + if isinstance(module, torch.nn.modules.batchnorm._BatchNorm): + return nn.modules.SyncBatchNorm.convert_sync_batchnorm(module, process_group) + module_output = module + if isinstance(module, detectron2.layers.FrozenBatchNorm2d): + module_output = torch.nn.SyncBatchNorm( + num_features=module.num_features, + eps=module.eps, + affine=True, + track_running_stats=True, + process_group=process_group, + ) + module_output.weight = torch.nn.Parameter(module.weight) + module_output.bias = torch.nn.Parameter(module.bias) + module_output.running_mean = module.running_mean + module_output.running_var = module.running_var + module_output.num_batches_tracked = torch.tensor(0, dtype=torch.long, device=module.running_mean.device) + for name, child in module.named_children(): + module_output.add_module(name, my_convert_sync_batchnorm(child, process_group)) + del module + return module_output + + +class LayoutLMv2VisualBackbone(nn.Module): + def __init__(self, config): + super().__init__() + self.cfg = config.get_detectron2_config() + meta_arch = self.cfg.MODEL.META_ARCHITECTURE + model = META_ARCH_REGISTRY.get(meta_arch)(self.cfg) + assert isinstance(model.backbone, detectron2.modeling.backbone.FPN) + self.backbone = model.backbone + + assert len(self.cfg.MODEL.PIXEL_MEAN) == len(self.cfg.MODEL.PIXEL_STD) + num_channels = len(self.cfg.MODEL.PIXEL_MEAN) + self.register_buffer( + "pixel_mean", + torch.Tensor(self.cfg.MODEL.PIXEL_MEAN).view(num_channels, 1, 1), + persistent=False, + ) + self.register_buffer( + "pixel_std", torch.Tensor(self.cfg.MODEL.PIXEL_STD).view(num_channels, 1, 1), persistent=False + ) + self.out_feature_key = "p2" + if torch.are_deterministic_algorithms_enabled(): + logger.warning("using `AvgPool2d` instead of `AdaptiveAvgPool2d`") + input_shape = (224, 224) + backbone_stride = self.backbone.output_shape()[self.out_feature_key].stride + self.pool = nn.AvgPool2d( + ( + math.ceil(math.ceil(input_shape[0] / backbone_stride) / config.image_feature_pool_shape[0]), + math.ceil(math.ceil(input_shape[1] / backbone_stride) / config.image_feature_pool_shape[1]), + ) + ) + else: + self.pool = nn.AdaptiveAvgPool2d(config.image_feature_pool_shape[:2]) + if len(config.image_feature_pool_shape) == 2: + config.image_feature_pool_shape.append(self.backbone.output_shape()[self.out_feature_key].channels) + assert self.backbone.output_shape()[self.out_feature_key].channels == config.image_feature_pool_shape[2] + + def forward(self, images): + images_input = ((images if torch.is_tensor(images) else images.tensor) - self.pixel_mean) / self.pixel_std + features = self.backbone(images_input) + features = features[self.out_feature_key] + features = self.pool(features).flatten(start_dim=2).transpose(1, 2).contiguous() + return features + + def synchronize_batch_norm(self): + if not ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and torch.distributed.get_rank() > -1 + ): + raise RuntimeError("Make sure torch.distributed is set up properly.") + + self_rank = torch.distributed.get_rank() + node_size = torch.cuda.device_count() + world_size = torch.distributed.get_world_size() + if not (world_size % node_size == 0): + raise RuntimeError("Make sure the number of processes can be divided by the number of nodes") + + node_global_ranks = [list(range(i * node_size, (i + 1) * node_size)) for i in range(world_size // node_size)] + sync_bn_groups = [ + torch.distributed.new_group(ranks=node_global_ranks[i]) for i in range(world_size // node_size) + ] + node_rank = self_rank // node_size + + self.backbone = my_convert_sync_batchnorm(self.backbone, process_group=sync_bn_groups[node_rank]) + + +class LayoutLMv2Pooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states): + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +@auto_docstring +class LayoutLMv2Model(LayoutLMv2PreTrainedModel): + _can_record_outputs = {"hidden_states": LayoutLMv2Layer, "attentions": LayoutLMv2SelfAttention} + + def __init__(self, config): + requires_backends(self, "detectron2") + super().__init__(config) + self.config = config + self.has_visual_segment_embedding = config.has_visual_segment_embedding + self.embeddings = LayoutLMv2Embeddings(config) + + self.visual = LayoutLMv2VisualBackbone(config) + self.visual_proj = nn.Linear(config.image_feature_pool_shape[-1], config.hidden_size) + if self.has_visual_segment_embedding: + self.visual_segment_embedding = nn.Parameter(nn.Embedding(1, config.hidden_size).weight[0]) + self.visual_LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.visual_dropout = nn.Dropout(config.hidden_dropout_prob) + + self.encoder = LayoutLMv2Encoder(config) + self.pooler = LayoutLMv2Pooler(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def _calc_text_embeddings(self, input_ids, bbox, position_ids, token_type_ids, inputs_embeds=None): + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + + if position_ids is None: + position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) + position_ids = position_ids.unsqueeze(0).expand_as(input_ids) + if token_type_ids is None: + token_type_ids = torch.zeros_like(input_ids) + + if inputs_embeds is None: + inputs_embeds = self.embeddings.word_embeddings(input_ids) + position_embeddings = self.embeddings.position_embeddings(position_ids) + spatial_position_embeddings = self.embeddings._calc_spatial_position_embeddings(bbox) + token_type_embeddings = self.embeddings.token_type_embeddings(token_type_ids) + + embeddings = inputs_embeds + position_embeddings + spatial_position_embeddings + token_type_embeddings + embeddings = self.embeddings.LayerNorm(embeddings) + embeddings = self.embeddings.dropout(embeddings) + return embeddings + + def _calc_img_embeddings(self, image, bbox, position_ids): + visual_embeddings = self.visual_proj(self.visual(image)) + position_embeddings = self.embeddings.position_embeddings(position_ids) + spatial_position_embeddings = self.embeddings._calc_spatial_position_embeddings(bbox) + embeddings = visual_embeddings + position_embeddings + spatial_position_embeddings + if self.has_visual_segment_embedding: + embeddings += self.visual_segment_embedding + embeddings = self.visual_LayerNorm(embeddings) + embeddings = self.visual_dropout(embeddings) + return embeddings + + def _calc_visual_bbox(self, image_feature_pool_shape, bbox, device, final_shape): + visual_bbox_x = torch.div( + torch.arange( + 0, + 1000 * (image_feature_pool_shape[1] + 1), + 1000, + device=device, + dtype=bbox.dtype, + ), + self.config.image_feature_pool_shape[1], + rounding_mode="floor", + ) + visual_bbox_y = torch.div( + torch.arange( + 0, + 1000 * (self.config.image_feature_pool_shape[0] + 1), + 1000, + device=device, + dtype=bbox.dtype, + ), + self.config.image_feature_pool_shape[0], + rounding_mode="floor", + ) + visual_bbox = torch.stack( + [ + visual_bbox_x[:-1].repeat(image_feature_pool_shape[0], 1), + visual_bbox_y[:-1].repeat(image_feature_pool_shape[1], 1).transpose(0, 1), + visual_bbox_x[1:].repeat(image_feature_pool_shape[0], 1), + visual_bbox_y[1:].repeat(image_feature_pool_shape[1], 1).transpose(0, 1), + ], + dim=-1, + ).view(-1, bbox.size(-1)) + + visual_bbox = visual_bbox.repeat(final_shape[0], 1, 1) + + return visual_bbox + + def _get_input_shape(self, input_ids=None, inputs_embeds=None): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + return input_ids.size() + elif inputs_embeds is not None: + return inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + image: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + bbox (`torch.LongTensor` of shape `((batch_size, sequence_length), 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + image (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `detectron.structures.ImageList` whose `tensors` is of shape `(batch_size, num_channels, height, width)`): + Batch of document images. + + Examples: + + ```python + >>> from transformers import AutoProcessor, LayoutLMv2Model, set_seed + >>> from PIL import Image + >>> import torch + >>> from datasets import load_dataset + + >>> set_seed(0) + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased") + >>> model = LayoutLMv2Model.from_pretrained("microsoft/layoutlmv2-base-uncased") + + + >>> dataset = load_dataset("hf-internal-testing/fixtures_docvqa") + >>> image = dataset["test"][0]["image"] + + >>> encoding = processor(image, return_tensors="pt") + + >>> outputs = model(**encoding) + >>> last_hidden_states = outputs.last_hidden_state + + >>> last_hidden_states.shape + torch.Size([1, 342, 768]) + ``` + """ + input_shape = self._get_input_shape(input_ids, inputs_embeds) + device = input_ids.device if input_ids is not None else inputs_embeds.device + + visual_shape = list(input_shape) + visual_shape[1] = self.config.image_feature_pool_shape[0] * self.config.image_feature_pool_shape[1] + visual_shape = torch.Size(visual_shape) + # needs a new copy of input_shape for tracing. Otherwise wrong dimensions will occur + final_shape = list(self._get_input_shape(input_ids, inputs_embeds)) + final_shape[1] += visual_shape[1] + final_shape = torch.Size(final_shape) + + visual_bbox = self._calc_visual_bbox(self.config.image_feature_pool_shape, bbox, device, final_shape) + final_bbox = torch.cat([bbox, visual_bbox], dim=1) + + if attention_mask is None: + attention_mask = torch.ones(input_shape, device=device) + + visual_attention_mask = torch.ones(visual_shape, device=device) + final_attention_mask = torch.cat([attention_mask, visual_attention_mask], dim=1) + + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + if position_ids is None: + seq_length = input_shape[1] + position_ids = self.embeddings.position_ids[:, :seq_length] + position_ids = position_ids.expand(input_shape) + + visual_position_ids = torch.arange(0, visual_shape[1], dtype=torch.long, device=device).repeat( + input_shape[0], 1 + ) + final_position_ids = torch.cat([position_ids, visual_position_ids], dim=1) + + if bbox is None: + bbox = torch.zeros(tuple(list(input_shape) + [4]), dtype=torch.long, device=device) + + text_layout_emb = self._calc_text_embeddings( + input_ids=input_ids, + bbox=bbox, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + ) + + visual_emb = self._calc_img_embeddings( + image=image, + bbox=visual_bbox, + position_ids=visual_position_ids, + ) + final_emb = torch.cat([text_layout_emb, visual_emb], dim=1) + + extended_attention_mask = final_attention_mask.unsqueeze(1).unsqueeze(2) + + extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min + + encoder_outputs = self.encoder( + final_emb, + extended_attention_mask, + bbox=final_bbox, + position_ids=final_position_ids, + **kwargs, + ) + sequence_output = encoder_outputs.last_hidden_state + pooled_output = self.pooler(sequence_output) + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + ) + + +@auto_docstring( + custom_intro=""" + LayoutLMv2 Model with a sequence classification head on top (a linear layer on top of the concatenation of the + final hidden state of the [CLS] token, average-pooled initial visual embeddings and average-pooled final visual + embeddings, e.g. for document image classification tasks such as the + [RVL-CDIP](https://www.cs.cmu.edu/~aharley/rvl-cdip/) dataset. + """ +) +class LayoutLMv2ForSequenceClassification(LayoutLMv2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.layoutlmv2 = LayoutLMv2Model(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size * 3, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv2.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + image: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SequenceClassifierOutput: + r""" + input_ids (`torch.LongTensor` of shape `batch_size, sequence_length`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + image (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `detectron.structures.ImageList` whose `tensors` is of shape `(batch_size, num_channels, height, width)`): + Batch of document images. + token_type_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + Example: + + ```python + >>> from transformers import AutoProcessor, LayoutLMv2ForSequenceClassification, set_seed + >>> from PIL import Image + >>> import torch + >>> from datasets import load_dataset + + >>> set_seed(0) + + >>> dataset = load_dataset("aharley/rvl_cdip", split="train", streaming=True) + >>> data = next(iter(dataset)) + >>> image = data["image"].convert("RGB") + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased") + >>> model = LayoutLMv2ForSequenceClassification.from_pretrained( + ... "microsoft/layoutlmv2-base-uncased", num_labels=dataset.info.features["label"].num_classes + ... ) + + >>> encoding = processor(image, return_tensors="pt") + >>> sequence_label = torch.tensor([data["label"]]) + + >>> outputs = model(**encoding, labels=sequence_label) + + >>> loss, logits = outputs.loss, outputs.logits + >>> predicted_idx = logits.argmax(dim=-1).item() + >>> predicted_answer = dataset.info.features["label"].names[4] + >>> predicted_idx, predicted_answer # results are not good without further fine-tuning + (7, 'advertisement') + ``` + """ + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + visual_shape = list(input_shape) + visual_shape[1] = self.config.image_feature_pool_shape[0] * self.config.image_feature_pool_shape[1] + visual_shape = torch.Size(visual_shape) + final_shape = list(input_shape) + final_shape[1] += visual_shape[1] + final_shape = torch.Size(final_shape) + + visual_bbox = self.layoutlmv2._calc_visual_bbox( + self.config.image_feature_pool_shape, bbox, device, final_shape + ) + + visual_position_ids = torch.arange(0, visual_shape[1], dtype=torch.long, device=device).repeat( + input_shape[0], 1 + ) + + initial_image_embeddings = self.layoutlmv2._calc_img_embeddings( + image=image, + bbox=visual_bbox, + position_ids=visual_position_ids, + ) + + outputs: BaseModelOutputWithPooling = self.layoutlmv2( + input_ids=input_ids, + bbox=bbox, + image=image, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + sequence_output, final_image_embeddings = ( + outputs.last_hidden_state[:, :seq_length], + outputs.last_hidden_state[:, seq_length:], + ) + + cls_final_output = sequence_output[:, 0, :] + + # average-pool the visual embeddings + pooled_initial_image_embeddings = initial_image_embeddings.mean(dim=1) + pooled_final_image_embeddings = final_image_embeddings.mean(dim=1) + # concatenate with cls_final_output + sequence_output = torch.cat( + [cls_final_output, pooled_initial_image_embeddings, pooled_final_image_embeddings], dim=1 + ) + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + LayoutLMv2 Model with a token classification head on top (a linear layer on top of the text part of the hidden + states) e.g. for sequence labeling (information extraction) tasks such as + [FUNSD](https://guillaumejaume.github.io/FUNSD/), [SROIE](https://rrc.cvc.uab.es/?ch=13), + [CORD](https://github.com/clovaai/cord) and [Kleister-NDA](https://github.com/applicaai/kleister-nda). + """ +) +class LayoutLMv2ForTokenClassification(LayoutLMv2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.layoutlmv2 = LayoutLMv2Model(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv2.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + image: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | TokenClassifierOutput: + r""" + input_ids (`torch.LongTensor` of shape `batch_size, sequence_length`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + image (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `detectron.structures.ImageList` whose `tensors` is of shape `(batch_size, num_channels, height, width)`): + Batch of document images. + token_type_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + + Example: + + ```python + >>> from transformers import AutoProcessor, LayoutLMv2ForTokenClassification, set_seed + >>> from PIL import Image + >>> from datasets import load_dataset + + >>> set_seed(0) + + >>> datasets = load_dataset("nielsr/funsd", split="test") + >>> labels = datasets.features["ner_tags"].feature.names + >>> id2label = {v: k for v, k in enumerate(labels)} + + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased", revision="no_ocr") + >>> model = LayoutLMv2ForTokenClassification.from_pretrained( + ... "microsoft/layoutlmv2-base-uncased", num_labels=len(labels) + ... ) + + >>> data = datasets[0] + >>> image = Image.open(data["image_path"]).convert("RGB") + >>> words = data["words"] + >>> boxes = data["bboxes"] # make sure to normalize your bounding boxes + >>> word_labels = data["ner_tags"] + >>> encoding = processor( + ... image, + ... words, + ... boxes=boxes, + ... word_labels=word_labels, + ... padding="max_length", + ... truncation=True, + ... return_tensors="pt", + ... ) + + >>> outputs = model(**encoding) + >>> logits, loss = outputs.logits, outputs.loss + + >>> predicted_token_class_ids = logits.argmax(-1) + >>> predicted_tokens_classes = [id2label[t.item()] for t in predicted_token_class_ids[0]] + >>> predicted_tokens_classes[:5] # results are not good without further fine-tuning + ['I-HEADER', 'I-HEADER', 'I-QUESTION', 'I-HEADER', 'I-QUESTION'] + ``` + """ + + outputs: BaseModelOutputWithPooling = self.layoutlmv2( + input_ids=input_ids, + bbox=bbox, + image=image, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + # only take the text part of the output representations + sequence_output = outputs.last_hidden_state[:, :seq_length] + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class LayoutLMv2ForQuestionAnswering(LayoutLMv2PreTrainedModel): + def __init__(self, config, has_visual_segment_embedding=True): + r""" + has_visual_segment_embedding (`bool`, *optional*, defaults to `True`): + Whether or not to add visual segment embeddings. + """ + super().__init__(config) + self.num_labels = config.num_labels + config.has_visual_segment_embedding = has_visual_segment_embedding + self.layoutlmv2 = LayoutLMv2Model(config) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.layoutlmv2.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: torch.LongTensor | None = None, + image: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | QuestionAnsweringModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `batch_size, sequence_length`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + image (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `detectron.structures.ImageList` whose `tensors` is of shape `(batch_size, num_channels, height, width)`): + Batch of document images. + token_type_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `batch_size, sequence_length`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + + Example: + + In this example below, we give the LayoutLMv2 model an image (of texts) and ask it a question. It will give us + a prediction of what it thinks the answer is (the span of the answer within the texts parsed from the image). + + ```python + >>> from transformers import AutoProcessor, LayoutLMv2ForQuestionAnswering, set_seed + >>> import torch + >>> from PIL import Image + >>> from datasets import load_dataset + + >>> set_seed(0) + >>> processor = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased") + >>> model = LayoutLMv2ForQuestionAnswering.from_pretrained("microsoft/layoutlmv2-base-uncased") + + >>> dataset = load_dataset("hf-internal-testing/fixtures_docvqa") + >>> image = dataset["test"][0]["image"] + >>> question = "When is coffee break?" + >>> encoding = processor(image, question, return_tensors="pt") + + >>> outputs = model(**encoding) + >>> predicted_start_idx = outputs.start_logits.argmax(-1).item() + >>> predicted_end_idx = outputs.end_logits.argmax(-1).item() + >>> predicted_start_idx, predicted_end_idx + (30, 191) + + >>> predicted_answer_tokens = encoding.input_ids.squeeze()[predicted_start_idx : predicted_end_idx + 1] + >>> predicted_answer = processor.tokenizer.decode(predicted_answer_tokens) + >>> predicted_answer # results are not good without further fine-tuning + '44 a. m. to 12 : 25 p. m. 12 : 25 to 12 : 58 p. m. 12 : 58 to 4 : 00 p. m. 2 : 00 to 5 : 00 p. m. coffee break coffee will be served for men and women in the lobby adjacent to exhibit area. please move into exhibit area. ( exhibits open ) trrf general session ( part | ) presiding : lee a. waller trrf vice president " introductory remarks " lee a. waller, trrf vice presi - dent individual interviews with trrf public board members and sci - entific advisory council mem - bers conducted by trrf treasurer philip g. kuehn to get answers which the public refrigerated warehousing industry is looking for. plus questions from' + ``` + + ```python + >>> target_start_index = torch.tensor([7]) + >>> target_end_index = torch.tensor([14]) + >>> outputs = model(**encoding, start_positions=target_start_index, end_positions=target_end_index) + >>> predicted_answer_span_start = outputs.start_logits.argmax(-1).item() + >>> predicted_answer_span_end = outputs.end_logits.argmax(-1).item() + >>> predicted_answer_span_start, predicted_answer_span_end + (30, 191) + ``` + """ + + outputs: BaseModelOutputWithPooling = self.layoutlmv2( + input_ids=input_ids, + bbox=bbox, + image=image, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + # only take the text part of the output representations + sequence_output = outputs.last_hidden_state[:, :seq_length] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "LayoutLMv2ForQuestionAnswering", + "LayoutLMv2ForSequenceClassification", + "LayoutLMv2ForTokenClassification", + "LayoutLMv2Layer", + "LayoutLMv2Model", + "LayoutLMv2PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/processing_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/processing_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..84e6b8d621ef809da59dac3ad8f8f96a1ccf5c02 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/processing_layoutlmv2.py @@ -0,0 +1,124 @@ +# 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. +""" +Processor class for LayoutLMv2. +""" + +from ...processing_utils import ProcessorMixin +from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy +from ...utils import TensorType, auto_docstring + + +@auto_docstring +class LayoutLMv2Processor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, **kwargs): + super().__init__(image_processor, tokenizer) + + @auto_docstring + def __call__( + self, + images, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = False, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchEncoding: + # verify input + if self.image_processor.apply_ocr and (boxes is not None): + raise ValueError( + "You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True." + ) + + if self.image_processor.apply_ocr and (word_labels is not None): + raise ValueError( + "You cannot provide word labels if you initialized the image processor with apply_ocr set to True." + ) + + if return_overflowing_tokens is True and return_offsets_mapping is False: + raise ValueError("You cannot return overflowing tokens without returning the offsets mapping.") + + # first, apply the image processor + features = self.image_processor(images=images, return_tensors=return_tensors) + + # second, apply the tokenizer + if text is not None and self.image_processor.apply_ocr and text_pair is None: + if isinstance(text, str): + text = [text] # add batch dimension (as the image processor always adds a batch dimension) + text_pair = features["words"] + + encoded_inputs = self.tokenizer( + text=text if text is not None else features["words"], + text_pair=text_pair if text_pair is not None else None, + boxes=boxes if boxes is not None else features["boxes"], + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + return_tensors=return_tensors, + **kwargs, + ) + + # add pixel values + images = features.pop("pixel_values") + if return_overflowing_tokens is True: + images = self.get_overflowing_images(images, encoded_inputs["overflow_to_sample_mapping"]) + encoded_inputs["image"] = images + + return encoded_inputs + + def get_overflowing_images(self, images, overflow_to_sample_mapping): + # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image + images_with_overflow = [] + for sample_idx in overflow_to_sample_mapping: + images_with_overflow.append(images[sample_idx]) + + if len(images_with_overflow) != len(overflow_to_sample_mapping): + raise ValueError( + "Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got" + f" {len(images_with_overflow)} and {len(overflow_to_sample_mapping)}" + ) + + return images_with_overflow + + @property + def model_input_names(self): + return ["input_ids", "bbox", "token_type_ids", "attention_mask", "image"] + + +__all__ = ["LayoutLMv2Processor"] diff --git a/third_party/transformers/src/transformers/models/layoutlmv2/tokenization_layoutlmv2.py b/third_party/transformers/src/transformers/models/layoutlmv2/tokenization_layoutlmv2.py new file mode 100644 index 0000000000000000000000000000000000000000..d8677d06facf5c44475fe43a18eea429c869a336 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlmv2/tokenization_layoutlmv2.py @@ -0,0 +1,868 @@ +# 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 LayoutLMv2. Based on WordPiece. +""" + +from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors + +from ...tokenization_utils_base import ( + BatchEncoding, + EncodedInput, + PaddingStrategy, + PreTokenizedInput, + TensorType, + TextInput, + TextInputPair, + TruncationStrategy, +) +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import add_end_docstrings, logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"} + +# Docstring constants for encode methods +LAYOUTLMV2_ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to encode the sequences with the special tokens relative to their model. + padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + stride (`int`, *optional*, defaults to 0): + If set to a number along with `max_length`, the overflowing tokens returned when + `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence + returned to provide some overlap between truncated and overflowing sequences. The value of this + argument defines the number of overlapping tokens. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). +""" + +LAYOUTLMV2_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method +""" + + +class LayoutLMv2Tokenizer(TokenizersBackend): + r""" + Construct a "fast" LayoutLMv2 tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece. + + This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + File containing the vocabulary. + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + unk_token (`str`, *optional*, defaults to `"[UNK]"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `"[SEP]"`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + pad_token (`str`, *optional*, defaults to `"[PAD]"`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `"[CLS]"`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + mask_token (`str`, *optional*, defaults to `"[MASK]"`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + cls_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [CLS] token. + sep_token_box (`List[int]`, *optional*, defaults to `[1000, 1000, 1000, 1000]`): + The bounding box to use for the special [SEP] token. + pad_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [PAD] token. + pad_token_label (`int`, *optional*, defaults to -100): + The label to use for padding tokens. Defaults to -100, which is the `ignore_index` of PyTorch's + CrossEntropyLoss. + only_label_first_subword (`bool`, *optional*, defaults to `True`): + Whether or not to only label the first subword, in case word labels are provided. + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this + issue](https://github.com/huggingface/transformers/issues/328)). + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase` (as in the original LayoutLMv2). + """ + + vocab_files_names = VOCAB_FILES_NAMES + model = models.WordPiece + model_input_names = ["input_ids", "attention_mask", "token_type_ids"] + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + do_lower_case=True, + unk_token="[UNK]", + sep_token="[SEP]", + pad_token="[PAD]", + cls_token="[CLS]", + mask_token="[MASK]", + cls_token_box=[0, 0, 0, 0], + sep_token_box=[1000, 1000, 1000, 1000], + pad_token_box=[0, 0, 0, 0], + pad_token_label=-100, + only_label_first_subword=True, + tokenize_chinese_chars=True, + strip_accents=None, + model_max_length=512, + **kwargs, + ): + self.do_lower_case = do_lower_case + + if vocab is not None: + self._vocab = vocab + else: + self._vocab = { + str(pad_token): 0, + str(unk_token): 1, + str(cls_token): 2, + str(sep_token): 3, + str(mask_token): 4, + } + + self._tokenizer = Tokenizer(models.WordPiece(vocab=self._vocab, unk_token=str(unk_token))) + self._tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + + self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + self._tokenizer.decoder = decoders.WordPiece(prefix="##") + super().__init__( + do_lower_case=do_lower_case, + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + cls_token_box=cls_token_box, + sep_token_box=sep_token_box, + pad_token_box=pad_token_box, + pad_token_label=pad_token_label, + only_label_first_subword=only_label_first_subword, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + model_max_length=model_max_length, + **kwargs, + ) + + self.cls_token_box = cls_token_box + self.sep_token_box = sep_token_box + self.pad_token_box = pad_token_box + self.pad_token_label = pad_token_label + self.only_label_first_subword = only_label_first_subword + + # Now set post_processor with actual token IDs + cls = str(self.cls_token) + sep = str(self.sep_token) + cls_token_id = self.cls_token_id + sep_token_id = self.sep_token_id + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + + @add_end_docstrings(LAYOUTLMV2_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV2_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences with word-level normalized bounding boxes and optional labels. + + Args: + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string, a list of strings + (words of a single example or questions of a batch of examples) or a list of list of strings (batch of + words). + text_pair (`List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence should be a list of strings + (pretokenized string). + boxes (`List[List[int]]`, `List[List[List[int]]]`): + Word-level bounding boxes. Each bounding box should be normalized to be on a 0-1000 scale. + word_labels (`List[int]`, `List[List[int]]`, *optional*): + Word-level integer labels (for token classification tasks such as FUNSD, CORD). + """ + + # Input type checking for clearer error + def _is_valid_text_input(t): + if isinstance(t, str): + # Strings are fine + return True + elif isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... empty + return True + elif isinstance(t[0], str): + # ... list of strings + return True + elif isinstance(t[0], (list, tuple)): + # ... list with an empty list or with a list of strings + return len(t[0]) == 0 or isinstance(t[0][0], str) + else: + return False + else: + return False + + if text_pair is not None: + # in case text + text_pair are provided, text = questions, text_pair = words + if not _is_valid_text_input(text): + raise ValueError("text input must of type `str` (single example) or `List[str]` (batch of examples). ") + if not isinstance(text_pair, (list, tuple)): + raise ValueError( + "Words must be of type `List[str]` (single pretokenized example), " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + else: + # in case only text is provided => must be words + if not isinstance(text, (list, tuple)): + raise ValueError( + "Words must be of type `List[str]` (single pretokenized example), " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if text_pair is not None: + is_batched = isinstance(text, (list, tuple)) + else: + is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)) + + words = text if text_pair is None else text_pair + if boxes is None: + raise ValueError("You must provide corresponding bounding boxes") + if is_batched: + if len(words) != len(boxes): + raise ValueError("You must provide words and boxes for an equal amount of examples") + for words_example, boxes_example in zip(words, boxes): + if len(words_example) != len(boxes_example): + raise ValueError("You must provide as many words as there are bounding boxes") + else: + if len(words) != len(boxes): + raise ValueError("You must provide as many words as there are bounding boxes") + + if is_batched: + if text_pair is not None and len(text) != len(text_pair): + raise ValueError( + f"batch length of `text`: {len(text)} does not match batch length of `text_pair`:" + f" {len(text_pair)}." + ) + batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text + is_pair = bool(text_pair is not None) + return self.batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + else: + return self.encode_plus( + text=text, + text_pair=text_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + @add_end_docstrings(LAYOUTLMV2_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV2_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def batch_encode_plus( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def tokenize(self, text: str, pair: str | None = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + batched_input = [(text, pair)] if pair else [text] + encodings = self._tokenizer.encode_batch( + batched_input, add_special_tokens=add_special_tokens, is_pretokenized=False, **kwargs + ) + + return encodings[0].tokens if encodings else [] + + @add_end_docstrings(LAYOUTLMV2_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV2_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def encode_plus( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[int] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a sequence or a pair of sequences. .. warning:: This method is deprecated, + `__call__` should be used instead. + + Args: + text (`str`, `List[str]`, `List[List[str]]`): + The first sequence to be encoded. This can be a string, a list of strings or a list of list of strings. + text_pair (`List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a list of strings (words of a single example) or a + list of list of strings (words of a batch of examples). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._encode_plus( + text=text, + boxes=boxes, + text_pair=text_pair, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + ) -> BatchEncoding: + if not isinstance(batch_text_or_text_pairs, list): + raise TypeError(f"batch_text_or_text_pairs has to be a list (got {type(batch_text_or_text_pairs)})") + + # Set the truncation and padding strategy and restore the initial configuration + self.set_truncation_and_padding( + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + ) + + if is_pair: + batch_text_or_text_pairs = [(text.split(), text_pair) for text, text_pair in batch_text_or_text_pairs] + + encodings = self._tokenizer.encode_batch( + batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + is_pretokenized=True, # we set this to True as LayoutLMv2 always expects pretokenized inputs + ) + + # Convert encoding to dict + # `Tokens` has type: Tuple[ + # List[Dict[str, List[List[int]]]] or List[Dict[str, 2D-Tensor]], + # List[EncodingFast] + # ] + # with nested dimensions corresponding to batch, overflows, sequence length + tokens_and_encodings = [ + self._convert_encoding( + encoding=encoding, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=True + if word_labels is not None + else return_offsets_mapping, # we use offsets to create the labels + return_length=return_length, + verbose=verbose, + ) + for encoding in encodings + ] + + # Convert the output to have dict[list] from list[dict] and remove the additional overflows dimension + # From (variable) shape (batch, overflows, sequence length) to ~ (batch * overflows, sequence length) + # (we say ~ because the number of overflow varies with the example in the batch) + # + # To match each overflowing sample with the original sample in the batch + # we add an overflow_to_sample_mapping array (see below) + sanitized_tokens = {} + for key in tokens_and_encodings[0][0]: + stack = [e for item, _ in tokens_and_encodings for e in item[key]] + sanitized_tokens[key] = stack + sanitized_encodings = [e for _, item in tokens_and_encodings for e in item] + + # If returning overflowing tokens, we need to return a mapping + # from the batch idx to the original sample + if return_overflowing_tokens: + overflow_to_sample_mapping = [] + for i, (toks, _) in enumerate(tokens_and_encodings): + overflow_to_sample_mapping += [i] * len(toks["input_ids"]) + sanitized_tokens["overflow_to_sample_mapping"] = overflow_to_sample_mapping + + for input_ids in sanitized_tokens["input_ids"]: + self._eventual_warn_about_too_long_sequence(input_ids, max_length, verbose) + + # create the token boxes + token_boxes = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + token_boxes_example = [] + for id, sequence_id, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_encodings[batch_index].sequence_ids, + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if is_pair and sequence_id == 0: + token_boxes_example.append(self.pad_token_box) + else: + token_boxes_example.append(boxes[original_index][word_id]) + else: + if id == self.cls_token_id: + token_boxes_example.append(self.cls_token_box) + elif id == self.sep_token_id: + token_boxes_example.append(self.sep_token_box) + elif id == self.pad_token_id: + token_boxes_example.append(self.pad_token_box) + else: + raise ValueError("Id not recognized") + token_boxes.append(token_boxes_example) + + sanitized_tokens["bbox"] = token_boxes + + # optionally, create the labels + if word_labels is not None: + labels = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + labels_example = [] + for id, offset, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_tokens["offset_mapping"][batch_index], + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if self.only_label_first_subword: + if offset[0] == 0: + # Use the real label id for the first token of the word, and padding ids for the remaining tokens + labels_example.append(word_labels[original_index][word_id]) + else: + labels_example.append(self.pad_token_label) + else: + labels_example.append(word_labels[original_index][word_id]) + else: + labels_example.append(self.pad_token_label) + labels.append(labels_example) + + sanitized_tokens["labels"] = labels + # finally, remove offsets if the user didn't want them + if not return_offsets_mapping: + del sanitized_tokens["offset_mapping"] + + return BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors) + + def _encode_plus( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[int] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: bool | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # make it a batched input + # 2 options: + # 1) only text, in case text must be a list of str + # 2) text + text_pair, in which case text = str and text_pair a list of str + batched_input = [(text, text_pair)] if text_pair else [text] + batched_boxes = [boxes] + batched_word_labels = [word_labels] if word_labels is not None else None + batched_output = self._batch_encode_plus( + batched_input, + is_pair=bool(text_pair is not None), + boxes=batched_boxes, + word_labels=batched_word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + # Return tensor is None, then we can remove the leading batch axis + # Overflowing tokens are returned as a batch of output so we keep them in this case + if return_tensors is None and not return_overflowing_tokens: + batched_output = BatchEncoding( + { + key: value[0] if len(value) > 0 and isinstance(value[0], list) else value + for key, value in batched_output.items() + }, + batched_output.encodings, + ) + + self._eventual_warn_about_too_long_sequence(batched_output["input_ids"], max_length, verbose) + + return batched_output + + def _pad( + self, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + ) -> dict: + """ + Pad encoded inputs (on left/right and up to predefined length or max length in the batch) + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + padding_strategy: PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in self.padding_side: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + required_input = encoded_inputs[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(required_input) + + if needs_to_be_padded: + difference = max_length - len(required_input) + padding_side = padding_side if padding_side is not None else self.padding_side + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = encoded_inputs["attention_mask"] + [0] * difference + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = ( + encoded_inputs["token_type_ids"] + [self.pad_token_type_id] * difference + ) + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = encoded_inputs["bbox"] + [self.pad_token_box] * difference + if "labels" in encoded_inputs: + encoded_inputs["labels"] = encoded_inputs["labels"] + [self.pad_token_label] * difference + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference + encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[ + "token_type_ids" + ] + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = [self.pad_token_box] * difference + encoded_inputs["bbox"] + if "labels" in encoded_inputs: + encoded_inputs["labels"] = [self.pad_token_label] * difference + encoded_inputs["labels"] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input + else: + raise ValueError("Invalid padding strategy:" + str(padding_side)) + + return encoded_inputs + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. A BERT sequence has the following format: + + - single sequence: `[CLS] X [SEP]` + - pair of sequences: `[CLS] A [SEP] B [SEP]` + + Args: + token_ids_0 (`List[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + """ + output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + + if token_ids_1: + output += token_ids_1 + [self.sep_token_id] + + return output + + +__all__ = ["LayoutLMv2Tokenizer"] + +# Backward alias +LayoutLMv2TokenizerFast = LayoutLMv2Tokenizer diff --git a/third_party/transformers/src/transformers/models/lightglue/__init__.py b/third_party/transformers/src/transformers/models/lightglue/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70ffc03a6bea72ca40bc51fa7c28ba73041a98fa --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_lightglue import * + from .image_processing_lightglue import * + from .image_processing_pil_lightglue import * + from .modeling_lightglue import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/lightglue/configuration_lightglue.py b/third_party/transformers/src/transformers/models/lightglue/configuration_lightglue.py new file mode 100644 index 0000000000000000000000000000000000000000..16ef166bf0017f8cfe73d043496be717662396c0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/configuration_lightglue.py @@ -0,0 +1,99 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lightglue/modular_lightglue.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_lightglue.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig +from ..superpoint import SuperPointConfig + + +@auto_docstring(checkpoint="ETH-CVG/lightglue_superpoint") +@strict +class LightGlueConfig(PreTrainedConfig): + r""" + keypoint_detector_config (`Union[AutoConfig, dict]`, *optional*, defaults to `SuperPointConfig`): + The config object or dictionary of the keypoint detector. + descriptor_dim (`int`, *optional*, defaults to 256): + The dimension of the descriptors. + depth_confidence (`float`, *optional*, defaults to 0.95): + The confidence threshold used to perform early stopping + width_confidence (`float`, *optional*, defaults to 0.99): + The confidence threshold used to prune points + filter_threshold (`float`, *optional*, defaults to 0.1): + The confidence threshold used to filter matches + + Examples: + ```python + >>> from transformers import LightGlueConfig, LightGlueForKeypointMatching + + >>> # Initializing a LightGlue style configuration + >>> configuration = LightGlueConfig() + + >>> # Initializing a model from the LightGlue style configuration + >>> model = LightGlueForKeypointMatching(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lightglue" + sub_configs = {"keypoint_detector_config": AutoConfig} + + keypoint_detector_config: dict | SuperPointConfig | None = None + descriptor_dim: int = 256 + num_hidden_layers: int = 9 + num_attention_heads: int = 4 + num_key_value_heads: int | None = None + depth_confidence: float = 0.95 + width_confidence: float = 0.99 + filter_threshold: float = 0.1 + initializer_range: float = 0.02 + hidden_act: str = "gelu" + attention_dropout: float | int = 0.0 + attention_bias: bool = True + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + # Keypoint Detector is forced into eager attention mode because SuperPoint does not have Attention + # See https://github.com/huggingface/transformers/pull/31718#discussion_r2109733153 + if isinstance(self.keypoint_detector_config, dict): + self.keypoint_detector_config["model_type"] = self.keypoint_detector_config.get("model_type", "superpoint") + self.keypoint_detector_config = CONFIG_MAPPING[self.keypoint_detector_config["model_type"]]( + **self.keypoint_detector_config, attn_implementation="eager" + ) + elif self.keypoint_detector_config is None: + self.keypoint_detector_config = CONFIG_MAPPING["superpoint"](attn_implementation="eager") + + self.intermediate_size = self.descriptor_dim * 2 + self.hidden_size = self.descriptor_dim + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.descriptor_dim % self.num_attention_heads != 0: + raise ValueError("descriptor_dim % num_heads is different from zero") + + +__all__ = ["LightGlueConfig"] diff --git a/third_party/transformers/src/transformers/models/lightglue/convert_lightglue_to_hf.py b/third_party/transformers/src/transformers/models/lightglue/convert_lightglue_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..4a9971e91c918fcfa95f6abbb63f20eab1a87f76 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/convert_lightglue_to_hf.py @@ -0,0 +1,278 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import gc +import os +import re + +import torch +from datasets import load_dataset + +from transformers import ( + AutoModelForKeypointDetection, + LightGlueForKeypointMatching, + LightGlueImageProcessor, +) +from transformers.models.lightglue.configuration_lightglue import LightGlueConfig + + +DEFAULT_CHECKPOINT_URL = "https://github.com/cvg/LightGlue/releases/download/v0.1_arxiv/superpoint_lightglue.pth" + + +def prepare_imgs(): + dataset = load_dataset("hf-internal-testing/image-matching-test-dataset", split="train") + image0 = dataset[0]["image"] + image1 = dataset[1]["image"] + image2 = dataset[2]["image"] + # [image1, image1] on purpose to test the model early stopping + return [[image2, image0], [image1, image1]] + + +def verify_model_outputs(model, device): + images = prepare_imgs() + preprocessor = LightGlueImageProcessor() + inputs = preprocessor(images=images, return_tensors="pt").to(device) + model.to(device) + with torch.no_grad(): + outputs = model(**inputs, output_hidden_states=True, output_attentions=True) + + predicted_matches_values = outputs.matches[0, 0, 20:30] + predicted_matching_scores_values = outputs.matching_scores[0, 0, 20:30] + + predicted_number_of_matches = torch.sum(outputs.matches[0][0] != -1).item() + + expected_max_number_keypoints = 866 + expected_matches_shape = torch.Size((len(images), 2, expected_max_number_keypoints)) + expected_matching_scores_shape = torch.Size((len(images), 2, expected_max_number_keypoints)) + + expected_matches_values = torch.tensor([-1, -1, 5, -1, -1, 19, -1, 10, -1, 11], dtype=torch.int64).to(device) + expected_matching_scores_values = torch.tensor([0, 0, 0.2997, 0, 0, 0.6762, 0, 0.8826, 0, 0.5583]).to(device) + + expected_number_of_matches = 140 + + assert outputs.matches.shape == expected_matches_shape + assert outputs.matching_scores.shape == expected_matching_scores_shape + + assert torch.allclose(predicted_matches_values, expected_matches_values, atol=1e-2) + assert torch.allclose(predicted_matching_scores_values, expected_matching_scores_values, atol=1e-2) + + assert predicted_number_of_matches == expected_number_of_matches + + +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + r"posenc.Wr": r"positional_encoder.projector", + r"self_attn.(\d+).Wqkv": r"transformer_layers.\1.self_attention.Wqkv", + r"self_attn.(\d+).out_proj": r"transformer_layers.\1.self_attention.o_proj", + r"self_attn.(\d+).ffn.0": r"transformer_layers.\1.self_mlp.fc1", + r"self_attn.(\d+).ffn.1": r"transformer_layers.\1.self_mlp.layer_norm", + r"self_attn.(\d+).ffn.3": r"transformer_layers.\1.self_mlp.fc2", + r"cross_attn.(\d+).to_qk": r"transformer_layers.\1.cross_attention.to_qk", + r"cross_attn.(\d+).to_v": r"transformer_layers.\1.cross_attention.v_proj", + r"cross_attn.(\d+).to_out": r"transformer_layers.\1.cross_attention.o_proj", + r"cross_attn.(\d+).ffn.0": r"transformer_layers.\1.cross_mlp.fc1", + r"cross_attn.(\d+).ffn.1": r"transformer_layers.\1.cross_mlp.layer_norm", + r"cross_attn.(\d+).ffn.3": r"transformer_layers.\1.cross_mlp.fc2", + r"log_assignment.(\d+).matchability": r"match_assignment_layers.\1.matchability", + r"log_assignment.(\d+).final_proj": r"match_assignment_layers.\1.final_projection", + r"token_confidence.(\d+).token.0": r"token_confidence.\1.token", +} + + +def convert_old_keys_to_new_keys(state_dict_keys: list[str]): + """ + This function should be applied only once, on the concatenated keys to efficiently rename using + the key mappings. + """ + output_dict = {} + if state_dict_keys is not None: + old_text = "\n".join(state_dict_keys) + new_text = old_text + for pattern, replacement in ORIGINAL_TO_CONVERTED_KEY_MAPPING.items(): + if replacement is None: + new_text = re.sub(pattern, "", new_text) # an empty line + continue + new_text = re.sub(pattern, replacement, new_text) + output_dict = dict(zip(old_text.split("\n"), new_text.split("\n"))) + return output_dict + + +def add_keypoint_detector_state_dict(lightglue_state_dict): + keypoint_detector = AutoModelForKeypointDetection.from_pretrained("magic-leap-community/superpoint") + keypoint_detector_state_dict = keypoint_detector.state_dict() + for k, v in keypoint_detector_state_dict.items(): + lightglue_state_dict[f"keypoint_detector.{k}"] = v + return lightglue_state_dict + + +def split_weights(state_dict): + for i in range(9): + # Remove unused r values + log_assignment_r_key = f"log_assignment.{i}.r" + if state_dict.get(log_assignment_r_key, None) is not None: + state_dict.pop(log_assignment_r_key) + + Wqkv_weight = state_dict.pop(f"transformer_layers.{i}.self_attention.Wqkv.weight") + Wqkv_bias = state_dict.pop(f"transformer_layers.{i}.self_attention.Wqkv.bias") + Wqkv_weight = Wqkv_weight.reshape(256, 3, 256) + Wqkv_bias = Wqkv_bias.reshape(256, 3) + query_weight, key_weight, value_weight = Wqkv_weight[:, 0], Wqkv_weight[:, 1], Wqkv_weight[:, 2] + query_bias, key_bias, value_bias = Wqkv_bias[:, 0], Wqkv_bias[:, 1], Wqkv_bias[:, 2] + state_dict[f"transformer_layers.{i}.self_attention.q_proj.weight"] = query_weight + state_dict[f"transformer_layers.{i}.self_attention.k_proj.weight"] = key_weight + state_dict[f"transformer_layers.{i}.self_attention.v_proj.weight"] = value_weight + state_dict[f"transformer_layers.{i}.self_attention.q_proj.bias"] = query_bias + state_dict[f"transformer_layers.{i}.self_attention.k_proj.bias"] = key_bias + state_dict[f"transformer_layers.{i}.self_attention.v_proj.bias"] = value_bias + + to_qk_weight = state_dict.pop(f"transformer_layers.{i}.cross_attention.to_qk.weight") + to_qk_bias = state_dict.pop(f"transformer_layers.{i}.cross_attention.to_qk.bias") + state_dict[f"transformer_layers.{i}.cross_attention.q_proj.weight"] = to_qk_weight + state_dict[f"transformer_layers.{i}.cross_attention.q_proj.bias"] = to_qk_bias + state_dict[f"transformer_layers.{i}.cross_attention.k_proj.weight"] = to_qk_weight + state_dict[f"transformer_layers.{i}.cross_attention.k_proj.bias"] = to_qk_bias + + return state_dict + + +@torch.no_grad() +def write_model( + model_path, + checkpoint_url, + organization, + push_to_hub=False, +): + os.makedirs(model_path, exist_ok=True) + + # ------------------------------------------------------------ + # LightGlue config + # ------------------------------------------------------------ + + config = LightGlueConfig( + descriptor_dim=256, + num_hidden_layers=9, + num_attention_heads=4, + ) + config.architectures = ["LightGlueForKeypointMatching"] + config.save_pretrained(model_path) + print("Model config saved successfully...") + + # ------------------------------------------------------------ + # Convert weights + # ------------------------------------------------------------ + + print(f"Fetching all parameters from the checkpoint at {checkpoint_url}...") + original_state_dict = torch.hub.load_state_dict_from_url(checkpoint_url) + + print("Converting model...") + all_keys = list(original_state_dict.keys()) + new_keys = convert_old_keys_to_new_keys(all_keys) + + state_dict = {} + for key in all_keys: + new_key = new_keys[key] + state_dict[new_key] = original_state_dict.pop(key).contiguous().clone() + + del original_state_dict + gc.collect() + state_dict = split_weights(state_dict) + state_dict = add_keypoint_detector_state_dict(state_dict) + + print("Loading the checkpoint in a LightGlue model...") + device = "cuda" + with torch.device(device): + model = LightGlueForKeypointMatching(config) + model.load_state_dict(state_dict) + print("Checkpoint loaded successfully...") + del model.config._name_or_path + + print("Saving the model...") + model.save_pretrained(model_path) + del state_dict, model + + # Safety check: reload the converted model + gc.collect() + print("Reloading the model to check if it's saved correctly.") + model = LightGlueForKeypointMatching.from_pretrained(model_path) + print("Model reloaded successfully.") + + model_name = "lightglue" + if "superpoint" in checkpoint_url: + model_name += "_superpoint" + if checkpoint_url == DEFAULT_CHECKPOINT_URL: + print("Checking the model outputs...") + verify_model_outputs(model, device) + print("Model outputs verified successfully.") + + if push_to_hub: + print("Pushing model to the hub...") + model.push_to_hub( + repo_id=f"{organization}/{model_name}", + commit_message="Add model", + ) + config.push_to_hub(repo_id=f"{organization}/{model_name}", commit_message="Add config") + + write_image_processor(model_path, model_name, organization, push_to_hub=push_to_hub) + + +def write_image_processor(save_dir, model_name, organization, push_to_hub=False): + if "superpoint" in model_name: + image_processor = LightGlueImageProcessor(do_grayscale=True) + else: + image_processor = LightGlueImageProcessor() + image_processor.save_pretrained(save_dir) + + if push_to_hub: + print("Pushing image processor to the hub...") + image_processor.push_to_hub( + repo_id=f"{organization}/{model_name}", + commit_message="Add image processor", + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--checkpoint_url", + default=DEFAULT_CHECKPOINT_URL, + type=str, + help="URL of the original LightGlue checkpoint you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + required=True, + help="Path to the output PyTorch model directory.", + ) + parser.add_argument("--save_model", action="store_true", help="Save model to local") + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Push model and image preprocessor to the hub", + ) + parser.add_argument( + "--organization", + default="ETH-CVG", + type=str, + help="Hub organization in which you want the model to be uploaded.", + ) + + args = parser.parse_args() + write_model( + args.pytorch_dump_folder_path, + args.checkpoint_url, + args.organization, + push_to_hub=args.push_to_hub, + ) diff --git a/third_party/transformers/src/transformers/models/lightglue/image_processing_lightglue.py b/third_party/transformers/src/transformers/models/lightglue/image_processing_lightglue.py new file mode 100644 index 0000000000000000000000000000000000000000..c2c5b02922d5c90eff91fe51b914a4dc31791d41 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/image_processing_lightglue.py @@ -0,0 +1,322 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lightglue/modular_lightglue.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_lightglue.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +import torch +from PIL import Image, ImageDraw +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + ImageInput, + ImageType, + PILImageResampling, + SizeDict, + get_image_type, + is_pil_image, + is_valid_image, + to_numpy_array, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +if TYPE_CHECKING: + from .modeling_lightglue import LightGlueKeypointMatchingOutput + + +class LightGlueImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`): + Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method. + """ + + do_grayscale: bool + + +def _is_valid_image(image): + return is_pil_image(image) or ( + is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3 + ) + + +def validate_and_format_image_pairs(images: ImageInput): + error_message = ( + "Input images must be a one of the following :", + " - A pair of PIL images.", + " - A pair of 3D arrays.", + " - A list of pairs of PIL images.", + " - A list of pairs of 3D arrays.", + ) + + def _is_valid_image(image): + """images is a PIL Image or a 3D array.""" + return is_pil_image(image) or ( + is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3 + ) + + if isinstance(images, list): + if len(images) == 2 and all((_is_valid_image(image)) for image in images): + return images + if all( + isinstance(image_pair, list) + and len(image_pair) == 2 + and all(_is_valid_image(image) for image in image_pair) + for image_pair in images + ): + return [image for image_pair in images for image in image_pair] + raise ValueError(error_message) + + +def is_grayscale( + image: "torch.Tensor", +): + """Checks if an image is grayscale (all RGB channels are identical).""" + if image.ndim < 3 or image.shape[0 if image.ndim == 3 else 1] == 1: + return True + return torch.all(image[..., 0, :, :] == image[..., 1, :, :]) and torch.all( + image[..., 1, :, :] == image[..., 2, :, :] + ) + + +def convert_to_grayscale( + image: "torch.Tensor", +) -> "torch.Tensor": + """ + Converts an image to grayscale format using the NTSC formula. Only support torch.Tensor. + + This function is supposed to return a 1-channel image, but it returns a 3-channel image with the same value in each + channel, because of an issue that is discussed in : + https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446 + + Args: + image (torch.Tensor): + The image to convert. + """ + if is_grayscale(image): + return image + return tvF.rgb_to_grayscale(image, num_output_channels=3) + + +@auto_docstring +class LightGlueImageProcessor(TorchvisionBackend): + valid_kwargs = LightGlueImageProcessorKwargs + resample = PILImageResampling.BILINEAR + size = {"height": 480, "width": 640} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = None + do_grayscale = True + + def __init__(self, **kwargs: Unpack[LightGlueImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LightGlueImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _prepare_images_structure( + self, + images: ImageInput, + **kwargs, + ) -> ImageInput: + # we need to handle image pairs validation and flattening + images = self.fetch_images(images) + return validate_and_format_image_pairs(images) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + do_grayscale: bool = True, + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + processed_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size=size, resample=resample) + processed_images_grouped[shape] = stacked_images + resized_images = reorder_images(processed_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_rescale: + stacked_images = self.rescale(stacked_images, rescale_factor) + if do_grayscale: + stacked_images = convert_to_grayscale(stacked_images) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + # Convert back to pairs format + image_pairs = [processed_images[i : i + 2] for i in range(0, len(processed_images), 2)] + + # Stack each pair into a single tensor to match slow processor format + stacked_pairs = [torch.stack(pair, dim=0) for pair in image_pairs] + + # Return in same format as slow processor + + return BatchFeature(data={"pixel_values": stacked_pairs}, tensor_type=return_tensors) + + def post_process_keypoint_matching( + self, + outputs: "LightGlueKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, torch.Tensor]]: + """ + Converts the raw output of [`LightGlueKeypointMatchingOutput`] into lists of keypoints, scores and descriptors + with coordinates absolute to the original image sizes. + Args: + outputs ([`LightGlueKeypointMatchingOutput`]): + Raw outputs of the model. + target_sizes (`torch.Tensor` or `list[tuple[tuple[int, int]]]`, *optional*): + Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`tuple[int, int]`) containing the + target size `(height, width)` of each image in the batch. This must be the original image size (before + any processing). + threshold (`float`, *optional*, defaults to `0.0`): + Threshold to filter out the matches with low scores. + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image + of the pair, the matching scores and the matching indices. + """ + if outputs.mask.shape[0] != len(target_sizes): + raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask") + if not all(len(target_size) == 2 for target_size in target_sizes): + raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch") + + if isinstance(target_sizes, list): + image_pair_sizes = torch.tensor(target_sizes, device=outputs.mask.device) + else: + if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2: + raise ValueError( + "Each element of target_sizes must contain the size (h, w) of each image of the batch" + ) + image_pair_sizes = target_sizes + + keypoints = outputs.keypoints.clone() + keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2) + keypoints = keypoints.to(torch.int32) + + results = [] + for mask_pair, keypoints_pair, matches, scores in zip( + outputs.mask, keypoints, outputs.matches[:, 0], outputs.matching_scores[:, 0] + ): + mask0 = mask_pair[0] > 0 + mask1 = mask_pair[1] > 0 + keypoints0 = keypoints_pair[0][mask0] + keypoints1 = keypoints_pair[1][mask1] + matches0 = matches[mask0] + scores0 = scores[mask0] + + # Filter out matches with low scores, invalid matches, and out-of-bounds indices + valid_matches = (scores0 > threshold) & (matches0 > -1) & (matches0 < keypoints1.shape[0]) + + matched_keypoints0 = keypoints0[valid_matches] + matched_keypoints1 = keypoints1[matches0[valid_matches]] + matching_scores = scores0[valid_matches] + + results.append( + { + "keypoints0": matched_keypoints0, + "keypoints1": matched_keypoints1, + "matching_scores": matching_scores, + } + ) + + return results + + def visualize_keypoint_matching( + self, + images, + keypoint_matching_output: list[dict[str, torch.Tensor]], + ) -> list["Image.Image"]: + """ + Plots the image pairs side by side with the detected keypoints as well as the matching between them. + + Args: + images: + Image pairs to plot. Same as `EfficientLoFTRImageProcessor.preprocess`. Expects either a list of 2 + images or a list of list of 2 images list with pixel values ranging from 0 to 255. + keypoint_matching_output (List[Dict[str, torch.Tensor]]]): + A post processed keypoint matching output + + Returns: + `List[PIL.Image.Image]`: A list of PIL images, each containing the image pairs side by side with the detected + keypoints as well as the matching between them. + """ + + images = validate_and_format_image_pairs(images) + images = [to_numpy_array(image) for image in images] + image_pairs = [images[i : i + 2] for i in range(0, len(images), 2)] + + results = [] + for image_pair, pair_output in zip(image_pairs, keypoint_matching_output): + height0, width0 = image_pair[0].shape[:2] + height1, width1 = image_pair[1].shape[:2] + plot_image = torch.zeros((max(height0, height1), width0 + width1, 3), dtype=torch.uint8) + plot_image[:height0, :width0] = torch.from_numpy(image_pair[0]) + plot_image[:height1, width0:] = torch.from_numpy(image_pair[1]) + + plot_image_pil = Image.fromarray(plot_image.numpy()) + draw = ImageDraw.Draw(plot_image_pil) + + keypoints0_x, keypoints0_y = pair_output["keypoints0"].unbind(1) + keypoints1_x, keypoints1_y = pair_output["keypoints1"].unbind(1) + for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip( + keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, pair_output["matching_scores"] + ): + color = self._get_color(matching_score) + draw.line( + (keypoint0_x, keypoint0_y, keypoint1_x + width0, keypoint1_y), + fill=color, + width=3, + ) + draw.ellipse((keypoint0_x - 2, keypoint0_y - 2, keypoint0_x + 2, keypoint0_y + 2), fill="black") + draw.ellipse( + (keypoint1_x + width0 - 2, keypoint1_y - 2, keypoint1_x + width0 + 2, keypoint1_y + 2), + fill="black", + ) + + results.append(plot_image_pil) + return results + + def _get_color(self, score): + """Maps a score to a color.""" + r = int(255 * (1 - score)) + g = int(255 * score) + b = 0 + return r, g, b + + +__all__ = ["LightGlueImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/lightglue/image_processing_pil_lightglue.py b/third_party/transformers/src/transformers/models/lightglue/image_processing_pil_lightglue.py new file mode 100644 index 0000000000000000000000000000000000000000..9f43fe1bbc7a031145a54b6784dde28f640c15bb --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/image_processing_pil_lightglue.py @@ -0,0 +1,304 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lightglue/modular_lightglue.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_lightglue.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +import numpy as np +from PIL import Image, ImageDraw + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_utils import ( + ImageInput, + ImageType, + PILImageResampling, + SizeDict, + get_image_type, + is_pil_image, + is_valid_image, + to_numpy_array, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available +from ...utils.import_utils import requires + + +if TYPE_CHECKING: + from .modeling_lightglue import LightGlueKeypointMatchingOutput +if is_torch_available(): + import torch + + +class LightGlueImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`): + Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method. + """ + + do_grayscale: bool + + +def is_grayscale(image: np.ndarray): + if image.shape[0] == 1: + return True + return np.all(image[0, ...] == image[1, ...]) and np.all(image[1, ...] == image[2, ...]) + + +def convert_to_grayscale(image: ImageInput) -> ImageInput: + """ + Converts an image to grayscale format using the NTSC formula. Only support numpy and PIL Image. + + This function is supposed to return a 1-channel image, but it returns a 3-channel image with the same value in each + channel, because of an issue that is discussed in : + https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446 + + Args: + image (Image): + The image to convert. + """ + + if isinstance(image, np.ndarray): + if is_grayscale(image): + return image + + gray_image = image[0, ...] * 0.2989 + image[1, ...] * 0.5870 + image[2, ...] * 0.1140 + gray_image = np.stack([gray_image] * 3, axis=0) + return gray_image + + if not isinstance(image, Image.Image): + return image + + image = image.convert("L") + return image + + +# Adapted from transformers.models.lightglue.image_processing_lightglue.validate_and_format_image_pairs +def validate_and_format_image_pairs(images: ImageInput): + error_message = ( + "Input images must be a one of the following :", + " - A pair of PIL images.", + " - A pair of 3D arrays.", + " - A list of pairs of PIL images.", + " - A list of pairs of 3D arrays.", + ) + + def _is_valid_image(image): + """images is a PIL Image or a 3D array.""" + return is_pil_image(image) or ( + is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3 + ) + + if isinstance(images, list): + if len(images) == 2 and all((_is_valid_image(image)) for image in images): + return images + if all( + isinstance(image_pair, list) + and len(image_pair) == 2 + and all(_is_valid_image(image) for image in image_pair) + for image_pair in images + ): + return [image for image_pair in images for image in image_pair] + raise ValueError(error_message) + + +@requires(backends=("torch",)) +class LightGlueImageProcessorPil(PilBackend): + valid_kwargs = LightGlueImageProcessorKwargs + resample = PILImageResampling.BILINEAR + size = {"height": 480, "width": 640} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = None + do_grayscale = True + + def __init__(self, **kwargs: Unpack[LightGlueImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[LightGlueImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _prepare_images_structure(self, images: ImageInput, **kwargs) -> ImageInput: + # we need to handle image pairs validation and flattening + images = self.fetch_images(images) + return validate_and_format_image_pairs(images) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + return_tensors: str | TensorType | None, + do_grayscale: bool = True, + **kwargs, + ) -> BatchFeature: + all_images = [] + for image in images: + if do_resize: + image = self.resize(image=image, size=size, resample=resample) + + if do_rescale: + image = self.rescale(image=image, scale=rescale_factor) + + if do_grayscale: + image = convert_to_grayscale(image) + + all_images.append(image) + + # Convert back the flattened list of images into a list of pairs of images. + image_pairs = [all_images[i : i + 2] for i in range(0, len(all_images), 2)] + + data = {"pixel_values": image_pairs} + + return BatchFeature(data=data, tensor_type=return_tensors) + + @requires(backends=("torch",)) + def post_process_keypoint_matching( + self, + outputs: "LightGlueKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, "torch.Tensor"]]: + """ + Converts the raw output of [`LightGlueKeypointMatchingOutput`] into lists of keypoints, scores and descriptors + with coordinates absolute to the original image sizes. + Args: + outputs ([`LightGlueKeypointMatchingOutput`]): + Raw outputs of the model. + target_sizes (`torch.Tensor` or `list[tuple[tuple[int, int]]]`, *optional*): + Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`tuple[int, int]`) containing the + target size `(height, width)` of each image in the batch. This must be the original image size (before + any processing). + threshold (`float`, *optional*, defaults to `0.0`): + Threshold to filter out the matches with low scores. + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image + of the pair, the matching scores and the matching indices. + """ + import torch + + if outputs.mask.shape[0] != len(target_sizes): + raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask") + if not all(len(target_size) == 2 for target_size in target_sizes): + raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch") + + if isinstance(target_sizes, list): + image_pair_sizes = torch.tensor(target_sizes, device=outputs.mask.device) + else: + if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2: + raise ValueError( + "Each element of target_sizes must contain the size (h, w) of each image of the batch" + ) + image_pair_sizes = target_sizes + + keypoints = outputs.keypoints.clone() + keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2) + keypoints = keypoints.to(torch.int32) + + results = [] + for mask_pair, keypoints_pair, matches, scores in zip( + outputs.mask, keypoints, outputs.matches[:, 0], outputs.matching_scores[:, 0] + ): + mask0 = mask_pair[0] > 0 + mask1 = mask_pair[1] > 0 + keypoints0 = keypoints_pair[0][mask0] + keypoints1 = keypoints_pair[1][mask1] + matches0 = matches[mask0] + scores0 = scores[mask0] + + # Filter out matches with low scores, invalid matches, and out-of-bounds indices + valid_matches = (scores0 > threshold) & (matches0 > -1) & (matches0 < keypoints1.shape[0]) + + matched_keypoints0 = keypoints0[valid_matches] + matched_keypoints1 = keypoints1[matches0[valid_matches]] + matching_scores = scores0[valid_matches] + + results.append( + { + "keypoints0": matched_keypoints0, + "keypoints1": matched_keypoints1, + "matching_scores": matching_scores, + } + ) + + return results + + def visualize_keypoint_matching( + self, images: ImageInput, keypoint_matching_output: list[dict[str, "torch.Tensor"]] + ) -> list["Image.Image"]: + """ + Plots the image pairs side by side with the detected keypoints as well as the matching between them. + + Args: + images (`ImageInput`): + Image pairs to plot. Same as `LightGlueImageProcessor.preprocess`. Expects either a list of 2 + images or a list of list of 2 images list with pixel values ranging from 0 to 255. + keypoint_matching_output (List[Dict[str, torch.Tensor]]]): + A post processed keypoint matching output + + Returns: + `List[PIL.Image.Image]`: A list of PIL images, each containing the image pairs side by side with the detected + keypoints as well as the matching between them. + """ + images = validate_and_format_image_pairs(images) + images = [to_numpy_array(image) for image in images] + image_pairs = [images[i : i + 2] for i in range(0, len(images), 2)] + + results = [] + for image_pair, pair_output in zip(image_pairs, keypoint_matching_output): + height0, width0 = image_pair[0].shape[:2] + height1, width1 = image_pair[1].shape[:2] + plot_image = np.zeros((max(height0, height1), width0 + width1, 3), dtype=np.uint8) + plot_image[:height0, :width0] = image_pair[0] + plot_image[:height1, width0:] = image_pair[1] + + plot_image_pil = Image.fromarray(plot_image) + draw = ImageDraw.Draw(plot_image_pil) + + keypoints0_x, keypoints0_y = pair_output["keypoints0"].unbind(1) + keypoints1_x, keypoints1_y = pair_output["keypoints1"].unbind(1) + for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip( + keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, pair_output["matching_scores"] + ): + color = self._get_color(matching_score) + draw.line((keypoint0_x, keypoint0_y, keypoint1_x + width0, keypoint1_y), fill=color, width=3) + draw.ellipse((keypoint0_x - 2, keypoint0_y - 2, keypoint0_x + 2, keypoint0_y + 2), fill="black") + draw.ellipse( + (keypoint1_x + width0 - 2, keypoint1_y - 2, keypoint1_x + width0 + 2, keypoint1_y + 2), + fill="black", + ) + + results.append(plot_image_pil) + return results + + def _get_color(self, score): + """Maps a score to a color.""" + r = int(255 * (1 - score)) + g = int(255 * score) + b = 0 + return (r, g, b) + + +__all__ = ["LightGlueImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/lightglue/modeling_lightglue.py b/third_party/transformers/src/transformers/models/lightglue/modeling_lightglue.py new file mode 100644 index 0000000000000000000000000000000000000000..1c5b90e0f4fdb59e3ed3e2a426279486390155aa --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/modeling_lightglue.py @@ -0,0 +1,916 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lightglue/modular_lightglue.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_lightglue.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn +from torch.nn.utils.rnn import pad_sequence + +from ...activations import ACT2FN +from ...integrations import use_kernelized_func +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple +from ..auto.modeling_auto import AutoModelForKeypointDetection +from .configuration_lightglue import LightGlueConfig + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of LightGlue keypoint matching models. Due to the nature of keypoint detection and matching, + the number of keypoints is not fixed and can vary from image to image, which makes batching non-trivial. In the + batch of images, the maximum number of matches is set as the dimension of the matches and matching scores. The mask + tensor is used to indicate which values in the keypoints, matches, matching_scores and prune tensors are keypoint + matching information. + """ +) +class LightGlueKeypointMatchingOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Loss computed during training. + matches (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Index of keypoint matched in the other image. + matching_scores (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Scores of predicted matches. + keypoints (`torch.FloatTensor` of shape `(batch_size, num_keypoints, 2)`): + Absolute (x, y) coordinates of predicted keypoints in a given image. + prune (`torch.IntTensor` of shape `(batch_size, num_keypoints)`): + Pruning mask indicating which keypoints are removed and at which layer. + mask (`torch.BoolTensor` of shape `(batch_size, num_keypoints)`): + Mask indicating which values in matches, matching_scores, keypoints and prune are keypoint matching + information. + hidden_states (`Tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, 2, num_channels, + num_keypoints)` returned when `output_hidden_states=True` is passed or when + `config.output_hidden_states=True` + attentions (`Tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, 2, num_heads, num_keypoints, + num_keypoints)` returned when `output_attentions=True` is passed or when + `config.output_attentions=True` + """ + + loss: torch.FloatTensor | None = None + matches: torch.FloatTensor | None = None + matching_scores: torch.FloatTensor | None = None + keypoints: torch.FloatTensor | None = None + prune: torch.IntTensor | None = None + mask: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +class LightGluePositionalEncoder(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + self.projector = nn.Linear(2, config.descriptor_dim // config.num_attention_heads // 2, bias=False) + + def forward( + self, keypoints: torch.Tensor, output_hidden_states: bool | None = False + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor]: + projected_keypoints = self.projector(keypoints) + embeddings = projected_keypoints.repeat_interleave(2, dim=-1) + cosines = torch.cos(embeddings) + sines = torch.sin(embeddings) + embeddings = (cosines, sines) + output = (embeddings, projected_keypoints) if output_hidden_states else (embeddings,) + return output + + +def rotate_half(x): + # Split and rotate. Note that this function is different from e.g. Llama. + x1 = x[..., ::2] + x2 = x[..., 1::2] + rot_x = torch.stack([-x2, x1], dim=-1).flatten(-2) + return rot_x + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + dtype = q.dtype + q = q.float() + k = k.float() + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed.to(dtype=dtype), k_embed.to(dtype=dtype) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernelized_func(apply_rotary_pos_emb) +class LightGlueAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LightGlueConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + is_cross_attention = encoder_hidden_states is not None + current_states = encoder_hidden_states if is_cross_attention else hidden_states + current_attention_mask = encoder_attention_mask if is_cross_attention else attention_mask + + key_states = self.k_proj(current_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(current_states).view(hidden_shape).transpose(1, 2) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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, + current_attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class LightGlueMLP(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.intermediate_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + self.layer_norm = nn.LayerNorm(config.intermediate_size, elementwise_affine=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class LightGlueTransformerLayer(nn.Module): + def __init__(self, config: LightGlueConfig, layer_idx: int): + super().__init__() + self.self_attention = LightGlueAttention(config, layer_idx) + self.self_mlp = LightGlueMLP(config) + self.cross_attention = LightGlueAttention(config, layer_idx) + self.cross_mlp = LightGlueMLP(config) + + def forward( + self, + descriptors: torch.Tensor, + keypoints: torch.Tensor, + attention_mask: torch.Tensor, + output_hidden_states: bool | None = False, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, tuple[torch.Tensor] | None, tuple[torch.Tensor] | None]: + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + if output_hidden_states: + all_hidden_states = all_hidden_states + (descriptors,) + + batch_size, num_keypoints, descriptor_dim = descriptors.shape + + # Self attention block + attention_output, self_attentions = self.self_attention( + descriptors, + position_embeddings=keypoints, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + intermediate_states = torch.cat([descriptors, attention_output], dim=-1) + output_states = self.self_mlp(intermediate_states) + self_attention_descriptors = descriptors + output_states + + if output_hidden_states: + self_attention_hidden_states = (intermediate_states, output_states) + + # Reshape hidden_states to group by image_pairs : + # (batch_size, num_keypoints, descriptor_dim) -> (batch_size, 2, num_keypoints, descriptor_dim) + # Flip dimension 1 to perform cross attention : + # (image0, image1) -> (image1, image0) + # Reshape back to original shape : + # (batch_size, 2, num_keypoints, descriptor_dim) -> (batch_size, num_keypoints, descriptor_dim) + encoder_hidden_states = ( + self_attention_descriptors.reshape(-1, 2, num_keypoints, descriptor_dim) + .flip(1) + .reshape(batch_size, num_keypoints, descriptor_dim) + ) + # Same for mask + encoder_attention_mask = ( + attention_mask.reshape(-1, 2, 1, 1, num_keypoints).flip(1).reshape(batch_size, 1, 1, num_keypoints) + if attention_mask is not None + else None + ) + + # Cross attention block + cross_attention_output, cross_attentions = self.cross_attention( + self_attention_descriptors, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + ) + cross_intermediate_states = torch.cat([self_attention_descriptors, cross_attention_output], dim=-1) + cross_output_states = self.cross_mlp(cross_intermediate_states) + descriptors = self_attention_descriptors + cross_output_states + + if output_hidden_states: + cross_attention_hidden_states = (cross_intermediate_states, cross_output_states) + all_hidden_states = ( + all_hidden_states + + (self_attention_descriptors.reshape(batch_size, num_keypoints, descriptor_dim),) + + self_attention_hidden_states + + (descriptors.reshape(batch_size, num_keypoints, descriptor_dim),) + + cross_attention_hidden_states + ) + + if output_attentions: + all_attentions = all_attentions + (self_attentions,) + (cross_attentions,) + + return descriptors, all_hidden_states, all_attentions + + +def sigmoid_log_double_softmax( + similarity: torch.Tensor, matchability0: torch.Tensor, matchability1: torch.Tensor +) -> torch.Tensor: + """create the log assignment matrix from logits and similarity""" + batch_size, num_keypoints_0, num_keypoints_1 = similarity.shape + certainties = nn.functional.logsigmoid(matchability0) + nn.functional.logsigmoid(matchability1).transpose(1, 2) + scores0 = nn.functional.log_softmax(similarity, 2) + scores1 = nn.functional.log_softmax(similarity.transpose(-1, -2).contiguous(), 2).transpose(-1, -2) + scores = similarity.new_full((batch_size, num_keypoints_0 + 1, num_keypoints_1 + 1), 0) + scores[:, :num_keypoints_0, :num_keypoints_1] = scores0 + scores1 + certainties + scores[:, :-1, -1] = nn.functional.logsigmoid(-matchability0.squeeze(-1)) + scores[:, -1, :-1] = nn.functional.logsigmoid(-matchability1.squeeze(-1)) + return scores + + +class LightGlueMatchAssignmentLayer(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + + self.descriptor_dim = config.descriptor_dim + self.final_projection = nn.Linear(self.descriptor_dim, self.descriptor_dim, bias=True) + self.matchability = nn.Linear(self.descriptor_dim, 1, bias=True) + + def forward(self, descriptors: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + batch_size, num_keypoints, descriptor_dim = descriptors.shape + # Final projection and similarity computation + m_descriptors = self.final_projection(descriptors) + m_descriptors = m_descriptors / torch.tensor(self.descriptor_dim, device=m_descriptors.device) ** 0.25 + m_descriptors = m_descriptors.reshape(batch_size // 2, 2, num_keypoints, descriptor_dim) + m_descriptors0 = m_descriptors[:, 0] + m_descriptors1 = m_descriptors[:, 1] + similarity = m_descriptors0 @ m_descriptors1.transpose(-1, -2) + if mask is not None: + mask = mask.reshape(batch_size // 2, 2, num_keypoints) + mask0 = mask[:, 0].unsqueeze(-1) + mask1 = mask[:, 1].unsqueeze(-1).transpose(-1, -2) + mask = mask0 * mask1 + similarity = similarity.masked_fill(mask == 0, torch.finfo(similarity.dtype).min) + + # Compute matchability of descriptors + matchability = self.matchability(descriptors) + matchability = matchability.reshape(batch_size // 2, 2, num_keypoints, 1) + matchability_0 = matchability[:, 0] + matchability_1 = matchability[:, 1] + + # Compute scores from similarity and matchability + scores = sigmoid_log_double_softmax(similarity, matchability_0, matchability_1) + return scores + + def get_matchability(self, descriptors: torch.Tensor) -> torch.Tensor: + """Get matchability of descriptors as a probability""" + matchability = self.matchability(descriptors) + matchability = nn.functional.sigmoid(matchability).squeeze(-1) + return matchability + + +class LightGlueTokenConfidenceLayer(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + + self.token = nn.Linear(config.descriptor_dim, 1) + + def forward(self, descriptors: torch.Tensor) -> torch.Tensor: + token = self.token(descriptors.detach()) + token = nn.functional.sigmoid(token).squeeze(-1) + return token + + +@auto_docstring +class LightGluePreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: LightGlueConfig + base_model_prefix = "lightglue" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = False + _supports_flash_attn = True + _supports_sdpa = True + + +def get_matches_from_scores(scores: torch.Tensor, threshold: float) -> tuple[torch.Tensor, torch.Tensor]: + """obtain matches from a score matrix [Bx M+1 x N+1]""" + batch_size, _, _ = scores.shape + # For each keypoint, get the best match + max0 = scores[:, :-1, :-1].max(2) + max1 = scores[:, :-1, :-1].max(1) + matches0 = max0.indices + matches1 = max1.indices + + # Mutual check for matches + indices0 = torch.arange(matches0.shape[1], device=matches0.device)[None] + indices1 = torch.arange(matches1.shape[1], device=matches1.device)[None] + mutual0 = indices0 == matches1.gather(1, matches0) + mutual1 = indices1 == matches0.gather(1, matches1) + + # Get matching scores and filter based on mutual check and thresholding + max0 = max0.values.exp() + zero = max0.new_tensor(0) + matching_scores0 = torch.where(mutual0, max0, zero) + matching_scores1 = torch.where(mutual1, matching_scores0.gather(1, matches1), zero) + valid0 = mutual0 & (matching_scores0 > threshold) + valid1 = mutual1 & valid0.gather(1, matches1) + + # Filter matches based on mutual check and thresholding of scores + matches0 = torch.where(valid0, matches0, -1) + matches1 = torch.where(valid1, matches1, -1) + matches = torch.stack([matches0, matches1]).transpose(0, 1).reshape(batch_size * 2, -1) + matching_scores = torch.stack([matching_scores0, matching_scores1]).transpose(0, 1).reshape(batch_size * 2, -1) + + return matches, matching_scores + + +def normalize_keypoints(keypoints: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + Normalize keypoints locations based on image image_shape + + Args: + keypoints (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`): + Keypoints locations in (x, y) format. + height (`int`): + Image height. + width (`int`): + Image width. + + Returns: + Normalized keypoints locations of shape (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`). + """ + size = torch.tensor([width, height], device=keypoints.device, dtype=keypoints.dtype)[None] + shift = size / 2 + scale = size.max(-1).values / 2 + keypoints = (keypoints - shift[..., None, :]) / scale[..., None, None] + return keypoints + + +@auto_docstring( + custom_intro=""" + LightGlue model taking images as inputs and outputting the matching of them. + """ +) +class LightGlueForKeypointMatching(LightGluePreTrainedModel): + """ + LightGlue is a model matching keypoints in images by leveraging detections from a keypoint detector such as + SuperPoint. It is based on the SuperGlue architecture and is designed to be lightweight and efficient. + It consists of : + 1. Keypoint Encoder + 2. A Graph Neural Network with self and cross attention layers + 3. Matching Assignment layers + + The correspondence ids use -1 to indicate non-matching points. + + Philipp Lindenberger, Paul-Edouard Sarlin and Marc Pollefeys. LightGlue: Local Feature Matching at Light Speed. + In ICCV 2023. https://huggingface.co/papers/2306.13643 + """ + + def __init__(self, config: LightGlueConfig): + super().__init__(config) + self.keypoint_detector = AutoModelForKeypointDetection.from_config(config.keypoint_detector_config) + + self.keypoint_detector_descriptor_dim = config.keypoint_detector_config.descriptor_decoder_dim + self.descriptor_dim = config.descriptor_dim + self.num_layers = config.num_hidden_layers + self.filter_threshold = config.filter_threshold + self.depth_confidence = config.depth_confidence + self.width_confidence = config.width_confidence + + if self.descriptor_dim != self.keypoint_detector_descriptor_dim: + self.input_projection = nn.Linear(self.keypoint_detector_descriptor_dim, self.descriptor_dim, bias=True) + else: + self.input_projection = nn.Identity() + + self.positional_encoder = LightGluePositionalEncoder(config) + + self.transformer_layers = nn.ModuleList( + [LightGlueTransformerLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)] + ) + self.match_assignment_layers = nn.ModuleList( + [LightGlueMatchAssignmentLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.token_confidence = nn.ModuleList( + [LightGlueTokenConfidenceLayer(config) for _ in range(config.num_hidden_layers - 1)] + ) + + self.post_init() + + def _get_confidence_threshold(self, layer_index: int) -> float: + """scaled confidence threshold for a given layer""" + threshold = 0.8 + 0.1 * np.exp(-4.0 * layer_index / self.num_layers) + return np.clip(threshold, 0, 1) + + def _keypoint_processing( + self, descriptors: torch.Tensor, keypoints: torch.Tensor, output_hidden_states: bool | None = False + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + descriptors = descriptors.detach().contiguous() + projected_descriptors = self.input_projection(descriptors) + keypoint_encoding_output = self.positional_encoder(keypoints, output_hidden_states=output_hidden_states) + return projected_descriptors, keypoint_encoding_output + + def _get_early_stopped_image_pairs( + self, keypoint_confidences: torch.Tensor, layer_index: int, mask: torch.Tensor, num_points: torch.Tensor + ) -> torch.Tensor: + """evaluate whether we should stop inference based on the confidence of the keypoints""" + batch_size, _ = mask.shape + if layer_index < self.num_layers - 1: + # If the current layer is not the last layer, we compute the confidence of the keypoints and check + # if we should stop the forward pass through the transformer layers for each pair of images. + keypoint_confidences = keypoint_confidences.masked_fill(mask == 0, 1) + keypoint_confidences = keypoint_confidences.reshape(batch_size // 2, -1) + threshold = self._get_confidence_threshold(layer_index) + ratio_confident = 1.0 - (keypoint_confidences < threshold).float().sum(dim=1) / num_points + early_stopped_pairs = ratio_confident > self.depth_confidence + else: + # If the current layer is the last layer, we stop the forward pass through the transformer layers for + # all pairs of images. + early_stopped_pairs = torch.ones(batch_size, dtype=torch.bool) + return early_stopped_pairs + + def _get_keypoint_matching(self, descriptors, mask, layer_index, early_stops=None): + if early_stops is not None: + descriptors = descriptors[early_stops] + mask = mask[early_stops] + scores = self.match_assignment_layers[layer_index](descriptors, mask) + matches, matching_scores = get_matches_from_scores(scores, self.filter_threshold) + return matches, matching_scores + + def _get_pruning_mask(self, confidences: torch.Tensor, scores: torch.Tensor, layer_index: int) -> torch.Tensor: + """mask points which should be removed""" + keep = scores > (1 - self.width_confidence) + if confidences is not None: # Low-confidence points are never pruned. + keep |= confidences <= self._get_confidence_threshold(layer_index) + return keep + + def _do_layer_keypoint_pruning( + self, + descriptors: torch.Tensor, + keypoints: torch.Tensor, + mask: torch.Tensor, + indices: torch.Tensor, + prune_output: torch.Tensor, + keypoint_confidences: torch.Tensor, + layer_index: int, + ): + """ + For a given layer, prune keypoints based on the confidence of the keypoints and the matchability of the + descriptors. + """ + batch_size, _, _ = descriptors.shape + descriptors_matchability = self.match_assignment_layers[layer_index].get_matchability(descriptors) + pruned_keypoints_mask = self._get_pruning_mask(keypoint_confidences, descriptors_matchability, layer_index) + pruned_keypoints_mask = pruned_keypoints_mask.masked_fill(mask == 0, torch.tensor(False)) + + # For each image, we extract the pruned indices and the corresponding descriptors and keypoints. + pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask, pruned_indices = ( + [t[mask] for t, mask in zip(tensor, pruned_keypoints_mask)] + for tensor in [descriptors, keypoints[0], keypoints[1], pruned_keypoints_mask, indices] + ) + for i in range(batch_size): + prune_output[i, pruned_indices[i]] += 1 + + # Pad the pruned descriptors, keypoints, indices and mask to have the same shape across the batch. + pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask = ( + pad_sequence(pruned_tensor, batch_first=True) + for pruned_tensor in [pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask] + ) + pruned_keypoints = (pruned_keypoints_0, pruned_keypoints_1) + pruned_indices = pad_sequence(pruned_indices, batch_first=True, padding_value=-1) + + return pruned_descriptors, pruned_keypoints, pruned_indices, pruned_mask, prune_output + + def _concat_early_stopped_outputs( + self, + early_stops_indices, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + matches, + matching_scores, + ): + early_stops_indices = torch.stack(early_stops_indices) + # Rearrange tensors to have the same order as the input batch + ids = torch.arange(early_stops_indices.shape[0]) + order_indices = early_stops_indices[ids] + early_stops_indices = early_stops_indices[order_indices] + matches, final_pruned_keypoints_indices = ( + pad_sequence(tensor, batch_first=True, padding_value=-1) + for tensor in [matches, final_pruned_keypoints_indices] + ) + matching_scores, final_pruned_keypoints_iterations = ( + pad_sequence(tensor, batch_first=True, padding_value=0) + for tensor in [matching_scores, final_pruned_keypoints_iterations] + ) + matches, matching_scores, final_pruned_keypoints_indices, final_pruned_keypoints_iterations = ( + tensor[early_stops_indices] + for tensor in [ + matches, + matching_scores, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + ] + ) + return final_pruned_keypoints_indices, final_pruned_keypoints_iterations, matches, matching_scores + + def _do_final_keypoint_pruning( + self, + indices: torch.Tensor, + matches: torch.Tensor, + matching_scores: torch.Tensor, + num_keypoints: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # (batch_size, num_keypoints) -> (batch_size // 2, 2, num_keypoints) -> 2 * (batch_size // 2, num_keypoints) to + # have tensors from + batch_size, _ = indices.shape + indices, matches, matching_scores = ( + tensor.reshape(batch_size // 2, 2, -1) for tensor in [indices, matches, matching_scores] + ) + indices0 = indices[:, 0] + indices1 = indices[:, 1] + matches0 = matches[:, 0] + matches1 = matches[:, 1] + matching_scores0 = matching_scores[:, 0] + matching_scores1 = matching_scores[:, 1] + + # Prepare final matches and matching scores + _matches = torch.full((batch_size // 2, 2, num_keypoints), -1, device=indices.device, dtype=matches.dtype) + _matching_scores = torch.zeros( + (batch_size // 2, 2, num_keypoints), device=indices.device, dtype=matching_scores.dtype + ) + # Fill the matches and matching scores for each image pair + for i in range(batch_size // 2): + _matches[i, 0, indices0[i]] = torch.where( + matches0[i] == -1, -1, indices1[i].gather(0, matches0[i].clamp(min=0)) + ) + _matches[i, 1, indices1[i]] = torch.where( + matches1[i] == -1, -1, indices0[i].gather(0, matches1[i].clamp(min=0)) + ) + _matching_scores[i, 0, indices0[i]] = matching_scores0[i] + _matching_scores[i, 1, indices1[i]] = matching_scores1[i] + return _matches, _matching_scores + + def _match_image_pair( + self, + keypoints: torch.Tensor, + descriptors: torch.Tensor, + height: int, + width: int, + mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple, tuple]: + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + if keypoints.shape[2] == 0: # no keypoints + shape = keypoints.shape[:-1] + return ( + keypoints.new_full(shape, -1, dtype=torch.int), + keypoints.new_zeros(shape), + keypoints.new_zeros(shape), + all_hidden_states, + all_attentions, + ) + + device = keypoints.device + batch_size, _, initial_num_keypoints, _ = keypoints.shape + num_points_per_pair = torch.sum(mask.reshape(batch_size, -1), dim=1) + # (batch_size, 2, num_keypoints, 2) -> (batch_size * 2, num_keypoints, 2) + keypoints = keypoints.reshape(batch_size * 2, initial_num_keypoints, 2) + mask = mask.reshape(batch_size * 2, initial_num_keypoints) if mask is not None else None + descriptors = descriptors.reshape(batch_size * 2, initial_num_keypoints, self.keypoint_detector_descriptor_dim) + image_indices = torch.arange(batch_size * 2, device=device) + # Keypoint normalization + keypoints = normalize_keypoints(keypoints, height, width) + + descriptors, keypoint_encoding_output = self._keypoint_processing( + descriptors, keypoints, output_hidden_states=output_hidden_states + ) + + keypoints = keypoint_encoding_output[0] + + # Early stop consists of stopping the forward pass through the transformer layers when the confidence of the + # keypoints is above a certain threshold. + do_early_stop = self.depth_confidence > 0 + # Keypoint pruning consists of removing keypoints from the input of the transformer layers when the confidence of + # the keypoints is below a certain threshold. + do_keypoint_pruning = self.width_confidence > 0 + + early_stops_indices = [] + matches = [] + matching_scores = [] + final_pruned_keypoints_indices = [] + final_pruned_keypoints_iterations = [] + + pruned_keypoints_indices = torch.arange(0, initial_num_keypoints, device=device).expand(batch_size * 2, -1) + pruned_keypoints_iterations = torch.ones_like(pruned_keypoints_indices) + + for layer_index in range(self.num_layers): + input_shape = descriptors.size() + if mask is not None: + extended_attention_mask = self.get_extended_attention_mask(mask, input_shape) + else: + extended_attention_mask = torch.ones((batch_size, input_shape[-2]), device=keypoints.device) + layer_output = self.transformer_layers[layer_index]( + descriptors, + keypoints, + attention_mask=extended_attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + ) + descriptors, hidden_states, attention = layer_output + if output_hidden_states: + all_hidden_states = all_hidden_states + hidden_states + if output_attentions: + all_attentions = all_attentions + attention + + if do_early_stop: + if layer_index < self.num_layers - 1: + # Get the confidence of the keypoints for the current layer + keypoint_confidences = self.token_confidence[layer_index](descriptors) + + # Determine which pairs of images should be early stopped based on the confidence of the keypoints for + # the current layer. + early_stopped_pairs = self._get_early_stopped_image_pairs( + keypoint_confidences, layer_index, mask, num_points=num_points_per_pair + ) + else: + # Early stopping always occurs at the last layer + early_stopped_pairs = torch.ones(batch_size, dtype=torch.bool) + + if torch.any(early_stopped_pairs): + # If a pair of images is considered early stopped, we compute the matches for the remaining + # keypoints and stop the forward pass through the transformer layers for this pair of images. + early_stops = early_stopped_pairs.repeat_interleave(2) + early_stopped_image_indices = image_indices[early_stops] + early_stopped_matches, early_stopped_matching_scores = self._get_keypoint_matching( + descriptors, mask, layer_index, early_stops=early_stops + ) + early_stops_indices.extend(list(early_stopped_image_indices)) + matches.extend(list(early_stopped_matches)) + matching_scores.extend(list(early_stopped_matching_scores)) + if do_keypoint_pruning: + final_pruned_keypoints_indices.extend(list(pruned_keypoints_indices[early_stops])) + final_pruned_keypoints_iterations.extend(list(pruned_keypoints_iterations[early_stops])) + + # Remove image pairs that have been early stopped from the forward pass + num_points_per_pair = num_points_per_pair[~early_stopped_pairs] + descriptors, keypoints_0, keypoint_1, mask, image_indices = tuple( + tensor[~early_stops] + for tensor in [descriptors, keypoints[0], keypoints[1], mask, image_indices] + ) + keypoints = (keypoints_0, keypoint_1) + if do_keypoint_pruning: + pruned_keypoints_indices, pruned_keypoints_iterations, keypoint_confidences = tuple( + tensor[~early_stops] + for tensor in [ + pruned_keypoints_indices, + pruned_keypoints_iterations, + keypoint_confidences, + ] + ) + # If all pairs of images are early stopped, we stop the forward pass through the transformer + # layers for all pairs of images. + if torch.all(early_stopped_pairs): + break + + if do_keypoint_pruning: + # Prune keypoints from the input of the transformer layers for the next iterations if the confidence of + # the keypoints is below a certain threshold. + descriptors, keypoints, pruned_keypoints_indices, mask, pruned_keypoints_iterations = ( + self._do_layer_keypoint_pruning( + descriptors, + keypoints, + mask, + pruned_keypoints_indices, + pruned_keypoints_iterations, + keypoint_confidences, + layer_index, + ) + ) + + if do_early_stop and do_keypoint_pruning: + # Concatenate early stopped outputs together and perform final keypoint pruning + final_pruned_keypoints_indices, final_pruned_keypoints_iterations, matches, matching_scores = ( + self._concat_early_stopped_outputs( + early_stops_indices, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + matches, + matching_scores, + ) + ) + matches, matching_scores = self._do_final_keypoint_pruning( + final_pruned_keypoints_indices, + matches, + matching_scores, + initial_num_keypoints, + ) + else: + matches, matching_scores = self._get_keypoint_matching(descriptors, mask, self.num_layers - 1) + final_pruned_keypoints_iterations = torch.ones_like(matching_scores) * self.num_layers + + final_pruned_keypoints_iterations = final_pruned_keypoints_iterations.reshape( + batch_size, 2, initial_num_keypoints + ) + + return ( + matches, + matching_scores, + final_pruned_keypoints_iterations, + all_hidden_states, + all_attentions, + ) + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> tuple | LightGlueKeypointMatchingOutput: + loss = None + if labels is not None: + raise ValueError("LightGlue is not trainable, no labels should be provided.") + + 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 + ) + + if pixel_values.ndim != 5 or pixel_values.size(1) != 2: + raise ValueError("Input must be a 5D tensor of shape (batch_size, 2, num_channels, height, width)") + + batch_size, _, channels, height, width = pixel_values.shape + pixel_values = pixel_values.reshape(batch_size * 2, channels, height, width) + keypoint_detections = self.keypoint_detector(pixel_values) + + keypoints, _, descriptors, mask = keypoint_detections[:4] + keypoints = keypoints.reshape(batch_size, 2, -1, 2).to(pixel_values) + descriptors = descriptors.reshape(batch_size, 2, -1, self.keypoint_detector_descriptor_dim).to(pixel_values) + mask = mask.reshape(batch_size, 2, -1) + + absolute_keypoints = keypoints.clone() + absolute_keypoints[:, :, :, 0] = absolute_keypoints[:, :, :, 0] * width + absolute_keypoints[:, :, :, 1] = absolute_keypoints[:, :, :, 1] * height + + matches, matching_scores, prune, hidden_states, attentions = self._match_image_pair( + absolute_keypoints, + descriptors, + height, + width, + mask=mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + return LightGlueKeypointMatchingOutput( + loss=loss, + matches=matches, + matching_scores=matching_scores, + keypoints=keypoints, + prune=prune, + mask=mask, + hidden_states=hidden_states, + attentions=attentions, + ) + + +__all__ = ["LightGluePreTrainedModel", "LightGlueForKeypointMatching"] diff --git a/third_party/transformers/src/transformers/models/lightglue/modular_lightglue.py b/third_party/transformers/src/transformers/models/lightglue/modular_lightglue.py new file mode 100644 index 0000000000000000000000000000000000000000..62082b678b001360a4c9bdbfd49a66efb57b2f18 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lightglue/modular_lightglue.py @@ -0,0 +1,929 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +from huggingface_hub.dataclasses import strict +from torch import nn +from torch.nn.utils.rnn import pad_sequence + +from ...configuration_utils import PreTrainedConfig +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ModelOutput, TensorType, auto_docstring, can_return_tuple, logging +from ...utils.import_utils import requires +from ..auto import CONFIG_MAPPING, AutoConfig +from ..auto.modeling_auto import AutoModelForKeypointDetection +from ..clip.modeling_clip import CLIPMLP +from ..cohere.modeling_cohere import apply_rotary_pos_emb +from ..llama.modeling_llama import LlamaAttention, eager_attention_forward +from ..superglue.image_processing_pil_superglue import SuperGlueImageProcessorPil +from ..superglue.image_processing_superglue import SuperGlueImageProcessor +from ..superpoint import SuperPointConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="ETH-CVG/lightglue_superpoint") +@strict +class LightGlueConfig(PreTrainedConfig): + r""" + keypoint_detector_config (`Union[AutoConfig, dict]`, *optional*, defaults to `SuperPointConfig`): + The config object or dictionary of the keypoint detector. + descriptor_dim (`int`, *optional*, defaults to 256): + The dimension of the descriptors. + depth_confidence (`float`, *optional*, defaults to 0.95): + The confidence threshold used to perform early stopping + width_confidence (`float`, *optional*, defaults to 0.99): + The confidence threshold used to prune points + filter_threshold (`float`, *optional*, defaults to 0.1): + The confidence threshold used to filter matches + + Examples: + ```python + >>> from transformers import LightGlueConfig, LightGlueForKeypointMatching + + >>> # Initializing a LightGlue style configuration + >>> configuration = LightGlueConfig() + + >>> # Initializing a model from the LightGlue style configuration + >>> model = LightGlueForKeypointMatching(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lightglue" + sub_configs = {"keypoint_detector_config": AutoConfig} + + keypoint_detector_config: dict | SuperPointConfig | None = None + descriptor_dim: int = 256 + num_hidden_layers: int = 9 + num_attention_heads: int = 4 + num_key_value_heads: int | None = None + depth_confidence: float = 0.95 + width_confidence: float = 0.99 + filter_threshold: float = 0.1 + initializer_range: float = 0.02 + hidden_act: str = "gelu" + attention_dropout: float | int = 0.0 + attention_bias: bool = True + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + # Keypoint Detector is forced into eager attention mode because SuperPoint does not have Attention + # See https://github.com/huggingface/transformers/pull/31718#discussion_r2109733153 + if isinstance(self.keypoint_detector_config, dict): + self.keypoint_detector_config["model_type"] = self.keypoint_detector_config.get("model_type", "superpoint") + self.keypoint_detector_config = CONFIG_MAPPING[self.keypoint_detector_config["model_type"]]( + **self.keypoint_detector_config, attn_implementation="eager" + ) + elif self.keypoint_detector_config is None: + self.keypoint_detector_config = CONFIG_MAPPING["superpoint"](attn_implementation="eager") + + self.intermediate_size = self.descriptor_dim * 2 + self.hidden_size = self.descriptor_dim + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.descriptor_dim % self.num_attention_heads != 0: + raise ValueError("descriptor_dim % num_heads is different from zero") + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of LightGlue keypoint matching models. Due to the nature of keypoint detection and matching, + the number of keypoints is not fixed and can vary from image to image, which makes batching non-trivial. In the + batch of images, the maximum number of matches is set as the dimension of the matches and matching scores. The mask + tensor is used to indicate which values in the keypoints, matches, matching_scores and prune tensors are keypoint + matching information. + """ +) +class LightGlueKeypointMatchingOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Loss computed during training. + matches (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Index of keypoint matched in the other image. + matching_scores (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`): + Scores of predicted matches. + keypoints (`torch.FloatTensor` of shape `(batch_size, num_keypoints, 2)`): + Absolute (x, y) coordinates of predicted keypoints in a given image. + prune (`torch.IntTensor` of shape `(batch_size, num_keypoints)`): + Pruning mask indicating which keypoints are removed and at which layer. + mask (`torch.BoolTensor` of shape `(batch_size, num_keypoints)`): + Mask indicating which values in matches, matching_scores, keypoints and prune are keypoint matching + information. + hidden_states (`Tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, 2, num_channels, + num_keypoints)` returned when `output_hidden_states=True` is passed or when + `config.output_hidden_states=True` + attentions (`Tuple[torch.FloatTensor, ...]`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, 2, num_heads, num_keypoints, + num_keypoints)` returned when `output_attentions=True` is passed or when + `config.output_attentions=True` + """ + + loss: torch.FloatTensor | None = None + matches: torch.FloatTensor | None = None + matching_scores: torch.FloatTensor | None = None + keypoints: torch.FloatTensor | None = None + prune: torch.IntTensor | None = None + mask: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +class LightGlueImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`): + Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method. + """ + + do_grayscale: bool + + +class LightGlueImageProcessor(SuperGlueImageProcessor): + def post_process_keypoint_matching( + self, + outputs: "LightGlueKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, torch.Tensor]]: + return super().post_process_keypoint_matching(outputs, target_sizes, threshold) + + +@requires(backends=("torch",)) +class LightGlueImageProcessorPil(SuperGlueImageProcessorPil): + @requires(backends=("torch",)) + def post_process_keypoint_matching( + self, + outputs: "LightGlueKeypointMatchingOutput", + target_sizes: TensorType | list[tuple], + threshold: float = 0.0, + ) -> list[dict[str, "torch.Tensor"]]: + return super().post_process_keypoint_matching(outputs, target_sizes, threshold) + + +class LightGluePositionalEncoder(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + self.projector = nn.Linear(2, config.descriptor_dim // config.num_attention_heads // 2, bias=False) + + def forward( + self, keypoints: torch.Tensor, output_hidden_states: bool | None = False + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor]: + projected_keypoints = self.projector(keypoints) + embeddings = projected_keypoints.repeat_interleave(2, dim=-1) + cosines = torch.cos(embeddings) + sines = torch.sin(embeddings) + embeddings = (cosines, sines) + output = (embeddings, projected_keypoints) if output_hidden_states else (embeddings,) + return output + + +class LightGlueAttention(LlamaAttention): + def __init__(self, config: LightGlueConfig, layer_idx: int): + super().__init__() + del self.rotary_emb + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + is_cross_attention = encoder_hidden_states is not None + current_states = encoder_hidden_states if is_cross_attention else hidden_states + current_attention_mask = encoder_attention_mask if is_cross_attention else attention_mask + + key_states = self.k_proj(current_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(current_states).view(hidden_shape).transpose(1, 2) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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, + current_attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class LightGlueMLP(CLIPMLP): + def __init__(self, config: LightGlueConfig): + super().__init__(config) + self.fc1 = nn.Linear(config.intermediate_size, config.intermediate_size) + self.layer_norm = nn.LayerNorm(config.intermediate_size, elementwise_affine=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.layer_norm(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class LightGlueTransformerLayer(nn.Module): + def __init__(self, config: LightGlueConfig, layer_idx: int): + super().__init__() + self.self_attention = LightGlueAttention(config, layer_idx) + self.self_mlp = LightGlueMLP(config) + self.cross_attention = LightGlueAttention(config, layer_idx) + self.cross_mlp = LightGlueMLP(config) + + def forward( + self, + descriptors: torch.Tensor, + keypoints: torch.Tensor, + attention_mask: torch.Tensor, + output_hidden_states: bool | None = False, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, tuple[torch.Tensor] | None, tuple[torch.Tensor] | None]: + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + if output_hidden_states: + all_hidden_states = all_hidden_states + (descriptors,) + + batch_size, num_keypoints, descriptor_dim = descriptors.shape + + # Self attention block + attention_output, self_attentions = self.self_attention( + descriptors, + position_embeddings=keypoints, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + intermediate_states = torch.cat([descriptors, attention_output], dim=-1) + output_states = self.self_mlp(intermediate_states) + self_attention_descriptors = descriptors + output_states + + if output_hidden_states: + self_attention_hidden_states = (intermediate_states, output_states) + + # Reshape hidden_states to group by image_pairs : + # (batch_size, num_keypoints, descriptor_dim) -> (batch_size, 2, num_keypoints, descriptor_dim) + # Flip dimension 1 to perform cross attention : + # (image0, image1) -> (image1, image0) + # Reshape back to original shape : + # (batch_size, 2, num_keypoints, descriptor_dim) -> (batch_size, num_keypoints, descriptor_dim) + encoder_hidden_states = ( + self_attention_descriptors.reshape(-1, 2, num_keypoints, descriptor_dim) + .flip(1) + .reshape(batch_size, num_keypoints, descriptor_dim) + ) + # Same for mask + encoder_attention_mask = ( + attention_mask.reshape(-1, 2, 1, 1, num_keypoints).flip(1).reshape(batch_size, 1, 1, num_keypoints) + if attention_mask is not None + else None + ) + + # Cross attention block + cross_attention_output, cross_attentions = self.cross_attention( + self_attention_descriptors, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + ) + cross_intermediate_states = torch.cat([self_attention_descriptors, cross_attention_output], dim=-1) + cross_output_states = self.cross_mlp(cross_intermediate_states) + descriptors = self_attention_descriptors + cross_output_states + + if output_hidden_states: + cross_attention_hidden_states = (cross_intermediate_states, cross_output_states) + all_hidden_states = ( + all_hidden_states + + (self_attention_descriptors.reshape(batch_size, num_keypoints, descriptor_dim),) + + self_attention_hidden_states + + (descriptors.reshape(batch_size, num_keypoints, descriptor_dim),) + + cross_attention_hidden_states + ) + + if output_attentions: + all_attentions = all_attentions + (self_attentions,) + (cross_attentions,) + + return descriptors, all_hidden_states, all_attentions + + +def sigmoid_log_double_softmax( + similarity: torch.Tensor, matchability0: torch.Tensor, matchability1: torch.Tensor +) -> torch.Tensor: + """create the log assignment matrix from logits and similarity""" + batch_size, num_keypoints_0, num_keypoints_1 = similarity.shape + certainties = nn.functional.logsigmoid(matchability0) + nn.functional.logsigmoid(matchability1).transpose(1, 2) + scores0 = nn.functional.log_softmax(similarity, 2) + scores1 = nn.functional.log_softmax(similarity.transpose(-1, -2).contiguous(), 2).transpose(-1, -2) + scores = similarity.new_full((batch_size, num_keypoints_0 + 1, num_keypoints_1 + 1), 0) + scores[:, :num_keypoints_0, :num_keypoints_1] = scores0 + scores1 + certainties + scores[:, :-1, -1] = nn.functional.logsigmoid(-matchability0.squeeze(-1)) + scores[:, -1, :-1] = nn.functional.logsigmoid(-matchability1.squeeze(-1)) + return scores + + +class LightGlueMatchAssignmentLayer(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + + self.descriptor_dim = config.descriptor_dim + self.final_projection = nn.Linear(self.descriptor_dim, self.descriptor_dim, bias=True) + self.matchability = nn.Linear(self.descriptor_dim, 1, bias=True) + + def forward(self, descriptors: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + batch_size, num_keypoints, descriptor_dim = descriptors.shape + # Final projection and similarity computation + m_descriptors = self.final_projection(descriptors) + m_descriptors = m_descriptors / torch.tensor(self.descriptor_dim, device=m_descriptors.device) ** 0.25 + m_descriptors = m_descriptors.reshape(batch_size // 2, 2, num_keypoints, descriptor_dim) + m_descriptors0 = m_descriptors[:, 0] + m_descriptors1 = m_descriptors[:, 1] + similarity = m_descriptors0 @ m_descriptors1.transpose(-1, -2) + if mask is not None: + mask = mask.reshape(batch_size // 2, 2, num_keypoints) + mask0 = mask[:, 0].unsqueeze(-1) + mask1 = mask[:, 1].unsqueeze(-1).transpose(-1, -2) + mask = mask0 * mask1 + similarity = similarity.masked_fill(mask == 0, torch.finfo(similarity.dtype).min) + + # Compute matchability of descriptors + matchability = self.matchability(descriptors) + matchability = matchability.reshape(batch_size // 2, 2, num_keypoints, 1) + matchability_0 = matchability[:, 0] + matchability_1 = matchability[:, 1] + + # Compute scores from similarity and matchability + scores = sigmoid_log_double_softmax(similarity, matchability_0, matchability_1) + return scores + + def get_matchability(self, descriptors: torch.Tensor) -> torch.Tensor: + """Get matchability of descriptors as a probability""" + matchability = self.matchability(descriptors) + matchability = nn.functional.sigmoid(matchability).squeeze(-1) + return matchability + + +class LightGlueTokenConfidenceLayer(nn.Module): + def __init__(self, config: LightGlueConfig): + super().__init__() + + self.token = nn.Linear(config.descriptor_dim, 1) + + def forward(self, descriptors: torch.Tensor) -> torch.Tensor: + token = self.token(descriptors.detach()) + token = nn.functional.sigmoid(token).squeeze(-1) + return token + + +@auto_docstring +class LightGluePreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config: LightGlueConfig + base_model_prefix = "lightglue" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = False + _supports_flash_attn = True + _supports_sdpa = True + + +def get_matches_from_scores(scores: torch.Tensor, threshold: float) -> tuple[torch.Tensor, torch.Tensor]: + """obtain matches from a score matrix [Bx M+1 x N+1]""" + batch_size, _, _ = scores.shape + # For each keypoint, get the best match + max0 = scores[:, :-1, :-1].max(2) + max1 = scores[:, :-1, :-1].max(1) + matches0 = max0.indices + matches1 = max1.indices + + # Mutual check for matches + indices0 = torch.arange(matches0.shape[1], device=matches0.device)[None] + indices1 = torch.arange(matches1.shape[1], device=matches1.device)[None] + mutual0 = indices0 == matches1.gather(1, matches0) + mutual1 = indices1 == matches0.gather(1, matches1) + + # Get matching scores and filter based on mutual check and thresholding + max0 = max0.values.exp() + zero = max0.new_tensor(0) + matching_scores0 = torch.where(mutual0, max0, zero) + matching_scores1 = torch.where(mutual1, matching_scores0.gather(1, matches1), zero) + valid0 = mutual0 & (matching_scores0 > threshold) + valid1 = mutual1 & valid0.gather(1, matches1) + + # Filter matches based on mutual check and thresholding of scores + matches0 = torch.where(valid0, matches0, -1) + matches1 = torch.where(valid1, matches1, -1) + matches = torch.stack([matches0, matches1]).transpose(0, 1).reshape(batch_size * 2, -1) + matching_scores = torch.stack([matching_scores0, matching_scores1]).transpose(0, 1).reshape(batch_size * 2, -1) + + return matches, matching_scores + + +def normalize_keypoints(keypoints: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + Normalize keypoints locations based on image image_shape + + Args: + keypoints (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`): + Keypoints locations in (x, y) format. + height (`int`): + Image height. + width (`int`): + Image width. + + Returns: + Normalized keypoints locations of shape (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`). + """ + size = torch.tensor([width, height], device=keypoints.device, dtype=keypoints.dtype)[None] + shift = size / 2 + scale = size.max(-1).values / 2 + keypoints = (keypoints - shift[..., None, :]) / scale[..., None, None] + return keypoints + + +@auto_docstring( + custom_intro=""" + LightGlue model taking images as inputs and outputting the matching of them. + """ +) +class LightGlueForKeypointMatching(LightGluePreTrainedModel): + """ + LightGlue is a model matching keypoints in images by leveraging detections from a keypoint detector such as + SuperPoint. It is based on the SuperGlue architecture and is designed to be lightweight and efficient. + It consists of : + 1. Keypoint Encoder + 2. A Graph Neural Network with self and cross attention layers + 3. Matching Assignment layers + + The correspondence ids use -1 to indicate non-matching points. + + Philipp Lindenberger, Paul-Edouard Sarlin and Marc Pollefeys. LightGlue: Local Feature Matching at Light Speed. + In ICCV 2023. https://huggingface.co/papers/2306.13643 + """ + + def __init__(self, config: LightGlueConfig): + super().__init__(config) + self.keypoint_detector = AutoModelForKeypointDetection.from_config(config.keypoint_detector_config) + + self.keypoint_detector_descriptor_dim = config.keypoint_detector_config.descriptor_decoder_dim + self.descriptor_dim = config.descriptor_dim + self.num_layers = config.num_hidden_layers + self.filter_threshold = config.filter_threshold + self.depth_confidence = config.depth_confidence + self.width_confidence = config.width_confidence + + if self.descriptor_dim != self.keypoint_detector_descriptor_dim: + self.input_projection = nn.Linear(self.keypoint_detector_descriptor_dim, self.descriptor_dim, bias=True) + else: + self.input_projection = nn.Identity() + + self.positional_encoder = LightGluePositionalEncoder(config) + + self.transformer_layers = nn.ModuleList( + [LightGlueTransformerLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)] + ) + self.match_assignment_layers = nn.ModuleList( + [LightGlueMatchAssignmentLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.token_confidence = nn.ModuleList( + [LightGlueTokenConfidenceLayer(config) for _ in range(config.num_hidden_layers - 1)] + ) + + self.post_init() + + def _get_confidence_threshold(self, layer_index: int) -> float: + """scaled confidence threshold for a given layer""" + threshold = 0.8 + 0.1 * np.exp(-4.0 * layer_index / self.num_layers) + return np.clip(threshold, 0, 1) + + def _keypoint_processing( + self, descriptors: torch.Tensor, keypoints: torch.Tensor, output_hidden_states: bool | None = False + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + descriptors = descriptors.detach().contiguous() + projected_descriptors = self.input_projection(descriptors) + keypoint_encoding_output = self.positional_encoder(keypoints, output_hidden_states=output_hidden_states) + return projected_descriptors, keypoint_encoding_output + + def _get_early_stopped_image_pairs( + self, keypoint_confidences: torch.Tensor, layer_index: int, mask: torch.Tensor, num_points: torch.Tensor + ) -> torch.Tensor: + """evaluate whether we should stop inference based on the confidence of the keypoints""" + batch_size, _ = mask.shape + if layer_index < self.num_layers - 1: + # If the current layer is not the last layer, we compute the confidence of the keypoints and check + # if we should stop the forward pass through the transformer layers for each pair of images. + keypoint_confidences = keypoint_confidences.masked_fill(mask == 0, 1) + keypoint_confidences = keypoint_confidences.reshape(batch_size // 2, -1) + threshold = self._get_confidence_threshold(layer_index) + ratio_confident = 1.0 - (keypoint_confidences < threshold).float().sum(dim=1) / num_points + early_stopped_pairs = ratio_confident > self.depth_confidence + else: + # If the current layer is the last layer, we stop the forward pass through the transformer layers for + # all pairs of images. + early_stopped_pairs = torch.ones(batch_size, dtype=torch.bool) + return early_stopped_pairs + + def _get_keypoint_matching(self, descriptors, mask, layer_index, early_stops=None): + if early_stops is not None: + descriptors = descriptors[early_stops] + mask = mask[early_stops] + scores = self.match_assignment_layers[layer_index](descriptors, mask) + matches, matching_scores = get_matches_from_scores(scores, self.filter_threshold) + return matches, matching_scores + + def _get_pruning_mask(self, confidences: torch.Tensor, scores: torch.Tensor, layer_index: int) -> torch.Tensor: + """mask points which should be removed""" + keep = scores > (1 - self.width_confidence) + if confidences is not None: # Low-confidence points are never pruned. + keep |= confidences <= self._get_confidence_threshold(layer_index) + return keep + + def _do_layer_keypoint_pruning( + self, + descriptors: torch.Tensor, + keypoints: torch.Tensor, + mask: torch.Tensor, + indices: torch.Tensor, + prune_output: torch.Tensor, + keypoint_confidences: torch.Tensor, + layer_index: int, + ): + """ + For a given layer, prune keypoints based on the confidence of the keypoints and the matchability of the + descriptors. + """ + batch_size, _, _ = descriptors.shape + descriptors_matchability = self.match_assignment_layers[layer_index].get_matchability(descriptors) + pruned_keypoints_mask = self._get_pruning_mask(keypoint_confidences, descriptors_matchability, layer_index) + pruned_keypoints_mask = pruned_keypoints_mask.masked_fill(mask == 0, torch.tensor(False)) + + # For each image, we extract the pruned indices and the corresponding descriptors and keypoints. + pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask, pruned_indices = ( + [t[mask] for t, mask in zip(tensor, pruned_keypoints_mask)] + for tensor in [descriptors, keypoints[0], keypoints[1], pruned_keypoints_mask, indices] + ) + for i in range(batch_size): + prune_output[i, pruned_indices[i]] += 1 + + # Pad the pruned descriptors, keypoints, indices and mask to have the same shape across the batch. + pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask = ( + pad_sequence(pruned_tensor, batch_first=True) + for pruned_tensor in [pruned_descriptors, pruned_keypoints_0, pruned_keypoints_1, pruned_mask] + ) + pruned_keypoints = (pruned_keypoints_0, pruned_keypoints_1) + pruned_indices = pad_sequence(pruned_indices, batch_first=True, padding_value=-1) + + return pruned_descriptors, pruned_keypoints, pruned_indices, pruned_mask, prune_output + + def _concat_early_stopped_outputs( + self, + early_stops_indices, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + matches, + matching_scores, + ): + early_stops_indices = torch.stack(early_stops_indices) + # Rearrange tensors to have the same order as the input batch + ids = torch.arange(early_stops_indices.shape[0]) + order_indices = early_stops_indices[ids] + early_stops_indices = early_stops_indices[order_indices] + matches, final_pruned_keypoints_indices = ( + pad_sequence(tensor, batch_first=True, padding_value=-1) + for tensor in [matches, final_pruned_keypoints_indices] + ) + matching_scores, final_pruned_keypoints_iterations = ( + pad_sequence(tensor, batch_first=True, padding_value=0) + for tensor in [matching_scores, final_pruned_keypoints_iterations] + ) + matches, matching_scores, final_pruned_keypoints_indices, final_pruned_keypoints_iterations = ( + tensor[early_stops_indices] + for tensor in [ + matches, + matching_scores, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + ] + ) + return final_pruned_keypoints_indices, final_pruned_keypoints_iterations, matches, matching_scores + + def _do_final_keypoint_pruning( + self, + indices: torch.Tensor, + matches: torch.Tensor, + matching_scores: torch.Tensor, + num_keypoints: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # (batch_size, num_keypoints) -> (batch_size // 2, 2, num_keypoints) -> 2 * (batch_size // 2, num_keypoints) to + # have tensors from + batch_size, _ = indices.shape + indices, matches, matching_scores = ( + tensor.reshape(batch_size // 2, 2, -1) for tensor in [indices, matches, matching_scores] + ) + indices0 = indices[:, 0] + indices1 = indices[:, 1] + matches0 = matches[:, 0] + matches1 = matches[:, 1] + matching_scores0 = matching_scores[:, 0] + matching_scores1 = matching_scores[:, 1] + + # Prepare final matches and matching scores + _matches = torch.full((batch_size // 2, 2, num_keypoints), -1, device=indices.device, dtype=matches.dtype) + _matching_scores = torch.zeros( + (batch_size // 2, 2, num_keypoints), device=indices.device, dtype=matching_scores.dtype + ) + # Fill the matches and matching scores for each image pair + for i in range(batch_size // 2): + _matches[i, 0, indices0[i]] = torch.where( + matches0[i] == -1, -1, indices1[i].gather(0, matches0[i].clamp(min=0)) + ) + _matches[i, 1, indices1[i]] = torch.where( + matches1[i] == -1, -1, indices0[i].gather(0, matches1[i].clamp(min=0)) + ) + _matching_scores[i, 0, indices0[i]] = matching_scores0[i] + _matching_scores[i, 1, indices1[i]] = matching_scores1[i] + return _matches, _matching_scores + + def _match_image_pair( + self, + keypoints: torch.Tensor, + descriptors: torch.Tensor, + height: int, + width: int, + mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple, tuple]: + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + if keypoints.shape[2] == 0: # no keypoints + shape = keypoints.shape[:-1] + return ( + keypoints.new_full(shape, -1, dtype=torch.int), + keypoints.new_zeros(shape), + keypoints.new_zeros(shape), + all_hidden_states, + all_attentions, + ) + + device = keypoints.device + batch_size, _, initial_num_keypoints, _ = keypoints.shape + num_points_per_pair = torch.sum(mask.reshape(batch_size, -1), dim=1) + # (batch_size, 2, num_keypoints, 2) -> (batch_size * 2, num_keypoints, 2) + keypoints = keypoints.reshape(batch_size * 2, initial_num_keypoints, 2) + mask = mask.reshape(batch_size * 2, initial_num_keypoints) if mask is not None else None + descriptors = descriptors.reshape(batch_size * 2, initial_num_keypoints, self.keypoint_detector_descriptor_dim) + image_indices = torch.arange(batch_size * 2, device=device) + # Keypoint normalization + keypoints = normalize_keypoints(keypoints, height, width) + + descriptors, keypoint_encoding_output = self._keypoint_processing( + descriptors, keypoints, output_hidden_states=output_hidden_states + ) + + keypoints = keypoint_encoding_output[0] + + # Early stop consists of stopping the forward pass through the transformer layers when the confidence of the + # keypoints is above a certain threshold. + do_early_stop = self.depth_confidence > 0 + # Keypoint pruning consists of removing keypoints from the input of the transformer layers when the confidence of + # the keypoints is below a certain threshold. + do_keypoint_pruning = self.width_confidence > 0 + + early_stops_indices = [] + matches = [] + matching_scores = [] + final_pruned_keypoints_indices = [] + final_pruned_keypoints_iterations = [] + + pruned_keypoints_indices = torch.arange(0, initial_num_keypoints, device=device).expand(batch_size * 2, -1) + pruned_keypoints_iterations = torch.ones_like(pruned_keypoints_indices) + + for layer_index in range(self.num_layers): + input_shape = descriptors.size() + if mask is not None: + extended_attention_mask = self.get_extended_attention_mask(mask, input_shape) + else: + extended_attention_mask = torch.ones((batch_size, input_shape[-2]), device=keypoints.device) + layer_output = self.transformer_layers[layer_index]( + descriptors, + keypoints, + attention_mask=extended_attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + ) + descriptors, hidden_states, attention = layer_output + if output_hidden_states: + all_hidden_states = all_hidden_states + hidden_states + if output_attentions: + all_attentions = all_attentions + attention + + if do_early_stop: + if layer_index < self.num_layers - 1: + # Get the confidence of the keypoints for the current layer + keypoint_confidences = self.token_confidence[layer_index](descriptors) + + # Determine which pairs of images should be early stopped based on the confidence of the keypoints for + # the current layer. + early_stopped_pairs = self._get_early_stopped_image_pairs( + keypoint_confidences, layer_index, mask, num_points=num_points_per_pair + ) + else: + # Early stopping always occurs at the last layer + early_stopped_pairs = torch.ones(batch_size, dtype=torch.bool) + + if torch.any(early_stopped_pairs): + # If a pair of images is considered early stopped, we compute the matches for the remaining + # keypoints and stop the forward pass through the transformer layers for this pair of images. + early_stops = early_stopped_pairs.repeat_interleave(2) + early_stopped_image_indices = image_indices[early_stops] + early_stopped_matches, early_stopped_matching_scores = self._get_keypoint_matching( + descriptors, mask, layer_index, early_stops=early_stops + ) + early_stops_indices.extend(list(early_stopped_image_indices)) + matches.extend(list(early_stopped_matches)) + matching_scores.extend(list(early_stopped_matching_scores)) + if do_keypoint_pruning: + final_pruned_keypoints_indices.extend(list(pruned_keypoints_indices[early_stops])) + final_pruned_keypoints_iterations.extend(list(pruned_keypoints_iterations[early_stops])) + + # Remove image pairs that have been early stopped from the forward pass + num_points_per_pair = num_points_per_pair[~early_stopped_pairs] + descriptors, keypoints_0, keypoint_1, mask, image_indices = tuple( + tensor[~early_stops] + for tensor in [descriptors, keypoints[0], keypoints[1], mask, image_indices] + ) + keypoints = (keypoints_0, keypoint_1) + if do_keypoint_pruning: + pruned_keypoints_indices, pruned_keypoints_iterations, keypoint_confidences = tuple( + tensor[~early_stops] + for tensor in [ + pruned_keypoints_indices, + pruned_keypoints_iterations, + keypoint_confidences, + ] + ) + # If all pairs of images are early stopped, we stop the forward pass through the transformer + # layers for all pairs of images. + if torch.all(early_stopped_pairs): + break + + if do_keypoint_pruning: + # Prune keypoints from the input of the transformer layers for the next iterations if the confidence of + # the keypoints is below a certain threshold. + descriptors, keypoints, pruned_keypoints_indices, mask, pruned_keypoints_iterations = ( + self._do_layer_keypoint_pruning( + descriptors, + keypoints, + mask, + pruned_keypoints_indices, + pruned_keypoints_iterations, + keypoint_confidences, + layer_index, + ) + ) + + if do_early_stop and do_keypoint_pruning: + # Concatenate early stopped outputs together and perform final keypoint pruning + final_pruned_keypoints_indices, final_pruned_keypoints_iterations, matches, matching_scores = ( + self._concat_early_stopped_outputs( + early_stops_indices, + final_pruned_keypoints_indices, + final_pruned_keypoints_iterations, + matches, + matching_scores, + ) + ) + matches, matching_scores = self._do_final_keypoint_pruning( + final_pruned_keypoints_indices, + matches, + matching_scores, + initial_num_keypoints, + ) + else: + matches, matching_scores = self._get_keypoint_matching(descriptors, mask, self.num_layers - 1) + final_pruned_keypoints_iterations = torch.ones_like(matching_scores) * self.num_layers + + final_pruned_keypoints_iterations = final_pruned_keypoints_iterations.reshape( + batch_size, 2, initial_num_keypoints + ) + + return ( + matches, + matching_scores, + final_pruned_keypoints_iterations, + all_hidden_states, + all_attentions, + ) + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> tuple | LightGlueKeypointMatchingOutput: + loss = None + if labels is not None: + raise ValueError("LightGlue is not trainable, no labels should be provided.") + + 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 + ) + + if pixel_values.ndim != 5 or pixel_values.size(1) != 2: + raise ValueError("Input must be a 5D tensor of shape (batch_size, 2, num_channels, height, width)") + + batch_size, _, channels, height, width = pixel_values.shape + pixel_values = pixel_values.reshape(batch_size * 2, channels, height, width) + keypoint_detections = self.keypoint_detector(pixel_values) + + keypoints, _, descriptors, mask = keypoint_detections[:4] + keypoints = keypoints.reshape(batch_size, 2, -1, 2).to(pixel_values) + descriptors = descriptors.reshape(batch_size, 2, -1, self.keypoint_detector_descriptor_dim).to(pixel_values) + mask = mask.reshape(batch_size, 2, -1) + + absolute_keypoints = keypoints.clone() + absolute_keypoints[:, :, :, 0] = absolute_keypoints[:, :, :, 0] * width + absolute_keypoints[:, :, :, 1] = absolute_keypoints[:, :, :, 1] * height + + matches, matching_scores, prune, hidden_states, attentions = self._match_image_pair( + absolute_keypoints, + descriptors, + height, + width, + mask=mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + return LightGlueKeypointMatchingOutput( + loss=loss, + matches=matches, + matching_scores=matching_scores, + keypoints=keypoints, + prune=prune, + mask=mask, + hidden_states=hidden_states, + attentions=attentions, + ) + + +__all__ = [ + "LightGluePreTrainedModel", + "LightGlueForKeypointMatching", + "LightGlueConfig", + "LightGlueImageProcessor", + "LightGlueImageProcessorPil", +] diff --git a/third_party/transformers/src/transformers/models/longcat_flash/__init__.py b/third_party/transformers/src/transformers/models/longcat_flash/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3c6b725c76bcc5c083064e0116d3bf671cd851fd --- /dev/null +++ b/third_party/transformers/src/transformers/models/longcat_flash/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2025 Meituan and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_longcat_flash import * + from .modeling_longcat_flash import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/third_party/transformers/src/transformers/models/longcat_flash/configuration_longcat_flash.py new file mode 100644 index 0000000000000000000000000000000000000000..65755300e5e2f7fe7215672f63bc71c5cf25e594 --- /dev/null +++ b/third_party/transformers/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -0,0 +1,133 @@ +# Copyright 2025 Meituan 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. + +"""LongCat Flash model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="meituan-longcat/LongCat-Flash-Chat") +@strict +class LongcatFlashConfig(PreTrainedConfig): + r""" + ffn_hidden_size (`int`, *optional*, defaults to 12288): + Dimension of the MLP representations. + qk_head_dim (`int`, *optional*): + The total dimension of query/key heads. If not specified, set to `qk_nope_head_dim + qk_rope_head_dim`. + moe_topk (`int`, *optional*, defaults to 12): + Number of experts to route to for each token in the MoE layer. + zero_expert_num (`int`, *optional*, defaults to 256): + Number of zero experts (identity function) to add to the expert pool. + expert_ffn_hidden_size (`int`, *optional*, defaults to 2048): + Hidden size of individual expert FFN layers. + + ```python + >>> from transformers import LongcatFlashModel, LongcatFlashConfig + + >>> # Initializing a LongCat Flash style configuration + >>> configuration = LongcatFlashConfig() + + >>> # Initializing a model from the configuration + >>> model = LongcatFlashModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "longcat_flash" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 10000000.0 + base_model_tp_plan = { + "layers.*.self_attn.*.q_b_proj": "colwise", + "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj", + "layers.*.self_attn.*.kv_b_proj": "colwise", + "layers.*.self_attn.*.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts.identity_expert": "moe_identity_expert", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlps.*.gate_proj": "colwise", + "layers.*.mlps.*.up_proj": "colwise", + "layers.*.mlps.*.down_proj": "rowwise", + } + + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 131072 + hidden_size: int = 6144 + num_hidden_layers: int = 56 + num_layers: int = 28 + num_attention_heads: int = 64 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 131072 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + ffn_hidden_size: int = 12288 + q_lora_rank: int | None = 1536 + kv_lora_rank: int = 512 + qk_nope_head_dim: int = 128 + qk_rope_head_dim: int = 64 + head_dim: int = 64 + v_head_dim: int = 128 + qk_head_dim: int | None = None + moe_topk: int = 12 + n_routed_experts: int = 512 + zero_expert_num: int = 256 + expert_ffn_hidden_size: int = 2048 + routed_scaling_factor: float = 6.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.qk_head_dim is None: + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + super().__post_init__(**kwargs) + + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs + + +__all__ = ["LongcatFlashConfig"] diff --git a/third_party/transformers/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/third_party/transformers/src/transformers/models/longcat_flash/modeling_longcat_flash.py new file mode 100644 index 0000000000000000000000000000000000000000..d5ac6e237742d9d79c573fae5b393c67ae52e481 --- /dev/null +++ b/third_party/transformers/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -0,0 +1,725 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/longcat_flash/modular_longcat_flash.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_longcat_flash.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Meituan and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_longcat_flash import LongcatFlashConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class LongcatFlashRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + LongcatFlashRMSNorm 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 LongcatFlashRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: LongcatFlashConfig, 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: LongcatFlashConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class LongcatFlashMLP(nn.Module): + def __init__(self, config, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = config.ffn_hidden_size if intermediate_size is None else intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class LongcatFlashTopkRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.n_routed_experts = config.n_routed_experts + (config.zero_expert_num or 0) + self.register_buffer("e_score_correction_bias", torch.zeros(self.n_routed_experts)) + + self.top_k = config.moe_topk + self.routed_scaling_factor = config.routed_scaling_factor + self.router_bias = getattr(config, "router_bias", False) + self.classifier = nn.Linear(config.hidden_size, self.n_routed_experts, bias=self.router_bias) + + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(hidden_states.type(torch.float32), self.classifier.weight.type(torch.float32)) + scores = router_logits.softmax(dim=-1) + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + topk_weights = topk_weights * self.routed_scaling_factor + return topk_weights.to(router_logits.dtype), topk_indices + + @torch.no_grad() + def get_topk_indices(self, scores): + scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) + topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + return topk_indices + + +class LongcatFlashExperts(nn.Module): + def __init__(self, config): + super().__init__() + self.intermediate_size = config.expert_ffn_hidden_size + self.hidden_size = config.hidden_size + self.num_routed_experts = config.n_routed_experts + self.zero_expert_num = config.zero_expert_num or 0 + self.total_experts = self.num_routed_experts + self.zero_expert_num + self.act_fn = ACT2FN[config.hidden_act] + self.identity_expert = nn.Identity() + + if self.num_routed_experts > 0: + self.gate_up_proj = nn.Parameter( + torch.empty(self.total_experts, 2 * self.intermediate_size, self.hidden_size) + ) + self.down_proj = nn.Parameter( + torch.empty(self.num_routed_experts, self.hidden_size, self.intermediate_size) + ) + else: + self.register_parameter("gate_up_proj", None) + self.register_parameter("down_proj", None) + + def forward(self, hidden_states, top_k_index, top_k_weights): + final_hidden_states = torch.zeros_like(hidden_states) + if top_k_index.numel() == 0: + return final_hidden_states + + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.total_experts).permute(2, 1, 0) + + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False) + for expert_idx_tensor in expert_hit: + expert_idx = int(expert_idx_tensor.item()) + selection_idx, token_idx = torch.where(expert_mask[expert_idx].squeeze(0)) + if token_idx.numel() == 0: + continue + current_state = hidden_states[token_idx] + + if expert_idx >= self.num_routed_experts or self.gate_up_proj is None: + current_hidden_states = self.identity_expert(current_state) + else: + gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx]) + + current_hidden_states = current_hidden_states * top_k_weights[token_idx, selection_idx, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(hidden_states.dtype)) + return final_hidden_states + + +# remap config key expert_ffn_hidden_size -> moe_intermediate_size +class LongcatFlashMoE(nn.Module): + """ + A mixed expert module containing zero compute (identity) experts. + """ + + def __init__(self, config): + super().__init__() + self.intermediate_size = config.expert_ffn_hidden_size + self.config = config + self.experts = LongcatFlashExperts(config) + + self.router = LongcatFlashTopkRouter(config) + + def forward(self, hidden_states): + orig_shape = hidden_states.shape + topk_weights, topk_indices = self.router(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) + return hidden_states + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + r""" + TODO let's just use the original freqcis computation to not have the view + transpose + reshape! This is not optimized! + Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class LongcatFlashMLA(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.attention_dropout = config.attention_dropout + self.num_heads = config.num_attention_heads + + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_head_dim = config.qk_head_dim + + self.is_causal = True + if self.q_lora_rank is None: + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias) + self.q_a_layernorm = LongcatFlashRMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = LongcatFlashRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + + self.scaling = self.qk_head_dim ** (-0.5) + if self.config.rope_parameters.get("rope_type", "default") != "default": + mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_parameters["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.scaling = self.scaling * mscale * mscale + + self.mla_scale_q_lora = (config.hidden_size / self.q_lora_rank) ** 0.5 + self.mla_scale_kv_lora = (config.hidden_size / self.kv_lora_rank) ** 0.5 + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.qk_head_dim) + key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) + # we always do a lora for queries as well + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_pass = self.kv_a_layernorm(k_pass) + + # apply LoRA scaling + q_pass = q_pass * self.mla_scale_q_lora + q_rot = q_rot * self.mla_scale_q_lora + k_pass = k_pass * self.mla_scale_kv_lora + + k_pass = self.kv_b_proj(k_pass).view(key_shape).transpose(1, 2) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + cos, sin = position_embeddings + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class LongcatFlashDecoderLayer(GradientCheckpointingLayer): + """ + LongCat decoder layer with dual-sublayer + shortcut MoE architecture. + + Each logical layer contains: + - 2 attention sublayers (with layer indices: layer_idx*2, layer_idx*2+1) + - 2 MLP sublayers + - 1 shortcut MoE connection + """ + + def __init__(self, config, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + + self.mlp = LongcatFlashMoE(config) + + self.self_attn = nn.ModuleList([LongcatFlashMLA(config=config, layer_idx=layer_idx * 2 + i) for i in [0, 1]]) + self.mlps = nn.ModuleList([LongcatFlashMLP(config) for _ in [0, 1]]) + self.input_layernorm = nn.ModuleList( + [LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in [0, 1]] + ) + self.post_attention_layernorm = nn.ModuleList( + [LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in [0, 1]] + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm[0](hidden_states) + + hidden_states, _ = self.self_attn[0]( + 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 + + residual = hidden_states + hidden_states = self.post_attention_layernorm[0](hidden_states) + + shortcut_mlp_output = self.mlp(hidden_states) + hidden_states = self.mlps[0](hidden_states) + hidden_states = residual + hidden_states + + # shortcut connection after second sublayer + residual = hidden_states + hidden_states = self.input_layernorm[1](hidden_states) + + hidden_states, _ = self.self_attn[1]( + 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 + + residual = hidden_states + hidden_states = self.post_attention_layernorm[1](hidden_states) + + hidden_states = self.mlps[1](hidden_states) + hidden_states = residual + hidden_states + shortcut_mlp_output + + return hidden_states + + +@auto_docstring +class LongcatFlashPreTrainedModel(PreTrainedModel): + config: LongcatFlashConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LongcatFlashDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": LongcatFlashDecoderLayer, + "attentions": LongcatFlashMLA, + } + _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] + _keep_in_fp32_modules = [ + "classifier.weight" + ] # TODO let's make sure orignal code base has this, for now it fixes quantization + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, LongcatFlashTopkRouter): + init.normal_(module.classifier.weight, mean=0.0, std=self.config.initializer_range) + init.zeros_(module.e_score_correction_bias) + if isinstance(module, LongcatFlashExperts): + if module.gate_up_proj is not None: + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class LongcatFlashModel(LongcatFlashPreTrainedModel): + def __init__(self, config): + 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( + [LongcatFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers)] + ) + self.norm = LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LongcatFlashRotaryEmbedding(config=config) + self.gradient_checkpointing = False + # Each layer above has 2 sublayers, config hack to have a correct cache (to avoid a checkpoint change) + self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used) + + self.config.num_hidden_layers = 2 * config.num_layers + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=None, + attentions=None, + ) + + +@auto_docstring +class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] + + def __init__(self, config): + super().__init__(config) + self.model = LongcatFlashModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, LongcatFlashForCausalLM + + >>> model = LongcatFlashForCausalLM.from_pretrained("meta-longcat_flash/LongcatFlash-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-longcat_flash/LongcatFlash-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = ["LongcatFlashPreTrainedModel", "LongcatFlashModel", "LongcatFlashForCausalLM"] diff --git a/third_party/transformers/src/transformers/models/longcat_flash/modular_longcat_flash.py b/third_party/transformers/src/transformers/models/longcat_flash/modular_longcat_flash.py new file mode 100644 index 0000000000000000000000000000000000000000..205802f456909052c562a6f21e1ed2a098761c52 --- /dev/null +++ b/third_party/transformers/src/transformers/models/longcat_flash/modular_longcat_flash.py @@ -0,0 +1,434 @@ +# Copyright 2025 Meituan and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable + +import torch +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import is_flash_attention_requested +from ..deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3Attention, + DeepseekV3ForCausalLM, + DeepseekV3MLP, + DeepseekV3Model, + DeepseekV3RMSNorm, + DeepseekV3RotaryEmbedding, + DeepseekV3TopkRouter, + apply_rotary_pos_emb_interleave, + eager_attention_forward, +) +from .configuration_longcat_flash import LongcatFlashConfig + + +logger = logging.get_logger(__name__) + + +class LongcatFlashRMSNorm(DeepseekV3RMSNorm): + pass + + +class LongcatFlashRotaryEmbedding(DeepseekV3RotaryEmbedding): + pass + + +# TODO remap config key ffn_hidden_size -> intermediate_size +class LongcatFlashMLP(DeepseekV3MLP): + def __init__(self, config, hidden_size=None, intermediate_size=None): + super().__init__(config) + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = config.ffn_hidden_size if intermediate_size is None else intermediate_size + + +# TODO remap config key moe_topk -> num_experts_per_tok +class LongcatFlashTopkRouter(DeepseekV3TopkRouter): + def __init__(self, config): + super().__init__(config) + del self.n_group + del self.topk_group + del self.weight + del self.norm_topk_prob + + self.top_k = config.moe_topk + self.n_routed_experts = config.n_routed_experts + (config.zero_expert_num or 0) + self.routed_scaling_factor = config.routed_scaling_factor + self.register_buffer("e_score_correction_bias", torch.zeros(self.n_routed_experts)) + self.router_bias = getattr(config, "router_bias", False) + self.classifier = nn.Linear(config.hidden_size, self.n_routed_experts, bias=self.router_bias) + + @torch.no_grad() + def get_topk_indices(self, scores): + scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) + topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + return topk_indices + + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(hidden_states.type(torch.float32), self.classifier.weight.type(torch.float32)) + scores = router_logits.softmax(dim=-1) + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + topk_weights = topk_weights * self.routed_scaling_factor + return topk_weights.to(router_logits.dtype), topk_indices + + +class LongcatFlashExperts(nn.Module): + def __init__(self, config): + super().__init__() + self.intermediate_size = config.expert_ffn_hidden_size + self.hidden_size = config.hidden_size + self.num_routed_experts = config.n_routed_experts + self.zero_expert_num = config.zero_expert_num or 0 + self.total_experts = self.num_routed_experts + self.zero_expert_num + self.act_fn = ACT2FN[config.hidden_act] + self.identity_expert = nn.Identity() + + if self.num_routed_experts > 0: + self.gate_up_proj = nn.Parameter( + torch.empty(self.total_experts, 2 * self.intermediate_size, self.hidden_size) + ) + self.down_proj = nn.Parameter( + torch.empty(self.num_routed_experts, self.hidden_size, self.intermediate_size) + ) + else: + self.register_parameter("gate_up_proj", None) + self.register_parameter("down_proj", None) + + def forward(self, hidden_states, top_k_index, top_k_weights): + final_hidden_states = torch.zeros_like(hidden_states) + if top_k_index.numel() == 0: + return final_hidden_states + + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.total_experts).permute(2, 1, 0) + + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False) + for expert_idx_tensor in expert_hit: + expert_idx = int(expert_idx_tensor.item()) + selection_idx, token_idx = torch.where(expert_mask[expert_idx].squeeze(0)) + if token_idx.numel() == 0: + continue + current_state = hidden_states[token_idx] + + if expert_idx >= self.num_routed_experts or self.gate_up_proj is None: + current_hidden_states = self.identity_expert(current_state) + else: + gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx]) + + current_hidden_states = current_hidden_states * top_k_weights[token_idx, selection_idx, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(hidden_states.dtype)) + return final_hidden_states + + +# remap config key expert_ffn_hidden_size -> moe_intermediate_size +class LongcatFlashMoE(nn.Module): + """ + A mixed expert module containing zero compute (identity) experts. + """ + + def __init__(self, config): + super().__init__() + self.intermediate_size = config.expert_ffn_hidden_size + self.config = config + self.experts = LongcatFlashExperts(config) + + self.router = LongcatFlashTopkRouter(config) + + def forward(self, hidden_states): + orig_shape = hidden_states.shape + topk_weights, topk_indices = self.router(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) + return hidden_states + + +class LongcatFlashMLA(DeepseekV3Attention): + def __init__(self, config, layer_idx: int): + super().__init__(config, layer_idx) + + self.mla_scale_q_lora = (config.hidden_size / self.q_lora_rank) ** 0.5 + self.mla_scale_kv_lora = (config.hidden_size / self.kv_lora_rank) ** 0.5 + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.qk_head_dim) + key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) + # we always do a lora for queries as well + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_pass = self.kv_a_layernorm(k_pass) + + # apply LoRA scaling + q_pass = q_pass * self.mla_scale_q_lora + q_rot = q_rot * self.mla_scale_q_lora + k_pass = k_pass * self.mla_scale_kv_lora + + k_pass = self.kv_b_proj(k_pass).view(key_shape).transpose(1, 2) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + cos, sin = position_embeddings + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class LongcatFlashDecoderLayer(GradientCheckpointingLayer): + """ + LongCat decoder layer with dual-sublayer + shortcut MoE architecture. + + Each logical layer contains: + - 2 attention sublayers (with layer indices: layer_idx*2, layer_idx*2+1) + - 2 MLP sublayers + - 1 shortcut MoE connection + """ + + def __init__(self, config, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + + self.mlp = LongcatFlashMoE(config) + + self.self_attn = nn.ModuleList([LongcatFlashMLA(config=config, layer_idx=layer_idx * 2 + i) for i in [0, 1]]) + self.mlps = nn.ModuleList([LongcatFlashMLP(config) for _ in [0, 1]]) + self.input_layernorm = nn.ModuleList( + [LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in [0, 1]] + ) + self.post_attention_layernorm = nn.ModuleList( + [LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in [0, 1]] + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm[0](hidden_states) + + hidden_states, _ = self.self_attn[0]( + 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 + + residual = hidden_states + hidden_states = self.post_attention_layernorm[0](hidden_states) + + shortcut_mlp_output = self.mlp(hidden_states) + hidden_states = self.mlps[0](hidden_states) + hidden_states = residual + hidden_states + + # shortcut connection after second sublayer + residual = hidden_states + hidden_states = self.input_layernorm[1](hidden_states) + + hidden_states, _ = self.self_attn[1]( + 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 + + residual = hidden_states + hidden_states = self.post_attention_layernorm[1](hidden_states) + + hidden_states = self.mlps[1](hidden_states) + hidden_states = residual + hidden_states + shortcut_mlp_output + + return hidden_states + + +@auto_docstring +class LongcatFlashPreTrainedModel(PreTrainedModel): + config: LongcatFlashConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LongcatFlashDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": LongcatFlashDecoderLayer, + "attentions": LongcatFlashMLA, + } + _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] + _keep_in_fp32_modules = [ + "classifier.weight" + ] # TODO let's make sure orignal code base has this, for now it fixes quantization + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, LongcatFlashTopkRouter): + init.normal_(module.classifier.weight, mean=0.0, std=self.config.initializer_range) + init.zeros_(module.e_score_correction_bias) + if isinstance(module, LongcatFlashExperts): + if module.gate_up_proj is not None: + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + +class LongcatFlashModel(DeepseekV3Model): + def __init__(self, config): + super().__init__(config) + self.layers = nn.ModuleList( + [LongcatFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers)] + ) + # Each layer above has 2 sublayers, config hack to have a correct cache (to avoid a checkpoint change) + self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used) + + self.config.num_hidden_layers = 2 * config.num_layers + self.norm = LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LongcatFlashRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=None, + attentions=None, + ) + + +class LongcatFlashForCausalLM(DeepseekV3ForCausalLM): + _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] + + def __init__(self, config): + super().__init__(config) + self.model = LongcatFlashModel(config) + + +__all__ = ["LongcatFlashPreTrainedModel", "LongcatFlashModel", "LongcatFlashForCausalLM"] diff --git a/third_party/transformers/src/transformers/models/megatron_bert/__init__.py b/third_party/transformers/src/transformers/models/megatron_bert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..09d805dc1bf67d12cd7855d961a51d9461f03431 --- /dev/null +++ b/third_party/transformers/src/transformers/models/megatron_bert/__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_megatron_bert import * + from .modeling_megatron_bert import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/megatron_bert/configuration_megatron_bert.py b/third_party/transformers/src/transformers/models/megatron_bert/configuration_megatron_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..042ed3092215bac4538f6b57dd1508e10b3556a2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/megatron_bert/configuration_megatron_bert.py @@ -0,0 +1,64 @@ +# Copyright 2021- NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MEGATRON_BERT model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="nvidia/megatron-bert-uncased-345m") +@strict +class MegatronBertConfig(PreTrainedConfig): + r""" + Examples: + + ```python + >>> from transformers import MegatronBertConfig, MegatronBertModel + + >>> # Initializing a MEGATRON_BERT google-bert/bert-base-uncased style configuration + >>> configuration = MegatronBertConfig() + + >>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration + >>> model = MegatronBertModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "megatron-bert" + + vocab_size: int = 29056 + hidden_size: int = 1024 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + intermediate_size: int = 4096 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.1 + attention_probs_dropout_prob: float | int = 0.1 + max_position_embeddings: int = 512 + type_vocab_size: int = 2 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + pad_token_id: int | None = 0 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + use_cache: bool = True + is_decoder: bool = False + add_cross_attention: bool = False + tie_word_embeddings: bool = True + + +__all__ = ["MegatronBertConfig"] diff --git a/third_party/transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py b/third_party/transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..69ebed5aa8e0f1a5d908654938fc85623e27cace --- /dev/null +++ b/third_party/transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py @@ -0,0 +1,334 @@ +#################################################################################################### + +# Copyright (c) 2021-, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#################################################################################################### + +# +# Note: If when running this conversion script you're getting an exception: +# ModuleNotFoundError: No module named 'megatron.model.enums' +# you need to tell python where to find the clone of Megatron-LM, e.g.: +# +# cd /tmp +# git clone https://github.com/NVIDIA/Megatron-LM +# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py ... +# +# if you already have it cloned elsewhere, simply adjust the path to the existing path +# +# If the training was done using a Megatron-LM fork, e.g., +# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one +# in your path, i.e., /path/to/Megatron-DeepSpeed/ +# + +import argparse +import os +import re +import zipfile + +import torch + +from transformers import MegatronBertConfig + + +#################################################################################################### + + +def recursive_print(name, val, spaces=0): + # Format the message. + if name is None: + msg = None + else: + fmt = "." * max(0, spaces - 2) + "# {:" + str(50 - spaces) + "s}" + msg = fmt.format(name) + + # Print and recurse (if needed). + if isinstance(val, dict): + if msg is not None: + print(msg) + for k in val: + recursive_print(k, val[k], spaces + 2) + elif isinstance(val, torch.Tensor): + print(msg, ":", val.size()) + else: + print(msg, ":", val) + + +def fix_query_key_value_ordering(param, checkpoint_version, num_splits, num_heads, hidden_size): + # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :] + # for compatibility with later versions of NVIDIA Megatron-LM. + # The inverse operation is performed inside Megatron-LM to read checkpoints: + # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209 + # If param is the weight tensor of the self-attention block, the returned tensor + # will have to be transposed one more time to be read by HuggingFace BERT. + input_shape = param.size() + if checkpoint_version == 1.0: + # version 1.0 stores [num_heads * hidden_size * num_splits, :] + saved_shape = (num_heads, hidden_size, num_splits) + input_shape[1:] + param = param.view(*saved_shape) + param = param.transpose(0, 2) + param = param.transpose(1, 2).contiguous() + elif checkpoint_version >= 2.0: + # other versions store [num_heads * num_splits * hidden_size, :] + saved_shape = (num_heads, num_splits, hidden_size) + input_shape[1:] + param = param.view(*saved_shape) + param = param.transpose(0, 1).contiguous() + param = param.view(*input_shape) + return param + + +#################################################################################################### + + +def convert_megatron_checkpoint(args, input_state_dict, config): + # The converted output model. + output_state_dict = {} + + # old versions did not store training args + ds_args = input_state_dict.get("args", None) + if ds_args is not None: + # do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint + # from pprint import pprint + # pprint(vars(ds_args)) + + config.tokenizer_type = ds_args.tokenizer_type + config.vocab_size = ds_args.padded_vocab_size + config.max_position_embeddings = ds_args.max_position_embeddings + config.hidden_size = ds_args.hidden_size + config.num_hidden_layers = ds_args.num_layers + config.num_attention_heads = ds_args.num_attention_heads + config.intermediate_size = ds_args.ffn_hidden_size if "ffn_hidden_size" in ds_args else 4 * ds_args.hidden_size + # pprint(config) + + # The number of heads. + heads = config.num_attention_heads + # The hidden_size per head. + hidden_size_per_head = config.hidden_size // heads + # Megatron-LM checkpoint version + if "checkpoint_version" in input_state_dict: + checkpoint_version = input_state_dict["checkpoint_version"] + else: + checkpoint_version = 0.0 + + # The model. + model = input_state_dict["model"] + # The language model. + lm = model["language_model"] + # The embeddings. + embeddings = lm["embedding"] + + # The word embeddings. + word_embeddings = embeddings["word_embeddings"]["weight"] + # Truncate the embedding table to vocab_size rows. + word_embeddings = word_embeddings[: config.vocab_size, :] + # Store the word embeddings. + output_state_dict["bert.embeddings.word_embeddings.weight"] = word_embeddings + + # The position embeddings. + pos_embeddings = embeddings["position_embeddings"]["weight"] + assert pos_embeddings.size(0) == config.max_position_embeddings and pos_embeddings.size(1) == config.hidden_size + # Store the position embeddings. + output_state_dict["bert.embeddings.position_embeddings.weight"] = pos_embeddings + + # The token-type embeddings. + tokentype_embeddings = embeddings["tokentype_embeddings"]["weight"] + # Store the position embeddings. + output_state_dict["bert.embeddings.token_type_embeddings.weight"] = tokentype_embeddings + + # The transformer. + transformer = lm["transformer"] if "transformer" in lm else lm["encoder"] + + # The regex to extract layer names. + layer_re = re.compile(r"layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)") + + # The simple map of names for "automated" rules. + megatron_to_transformers = { + "attention.dense": ".attention.output.dense.", + "self_attention.dense": ".attention.output.dense.", + "mlp.dense_h_to_4h": ".intermediate.dense.", + "mlp.dense_4h_to_h": ".output.dense.", + } + + # Keep track of the attention/query/value tensor. + attention_qkv_weight = None + + # Extract the layers. + for key, val in transformer.items(): + # Match the name. + m = layer_re.match(key) + + # Stop if that's not a layer + if m is None: + break + + # The index of the layer. + layer_idx = int(m.group(1)) + # The name of the operation. + op_name = m.group(2) + # Is it a weight or a bias? + weight_or_bias = m.group(3) + + # The name of the layer. + layer_name = f"bert.encoder.layer.{layer_idx}" + + # For layernorm(s), simply store the layer norm. + if op_name.endswith("layernorm"): + ln_name = "attention.ln" if op_name.startswith("input") else "ln" + output_state_dict[layer_name + "." + ln_name + "." + weight_or_bias] = val + + # Transpose the QKV matrix. + elif ( + op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value" + ) and weight_or_bias == "weight": + # Make sure the QKV pointer is nil. + assert attention_qkv_weight is None, "" + + out_val = fix_query_key_value_ordering(val, checkpoint_version, 3, heads, hidden_size_per_head) + # Store the tensor as we need the bias as well to interleave QKV and biases. + attention_qkv_weight = out_val + + # Transpose the bias. + elif ( + op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value" + ) and weight_or_bias == "bias": + # Make sure we read the weight tensor. + assert attention_qkv_weight is not None, "" + + # Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved. + q = attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :] + k = attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :] + v = attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :] + + out_val = fix_query_key_value_ordering(val, checkpoint_version, 3, heads, hidden_size_per_head) + # Split the bias. + q_bias = out_val[0 * config.hidden_size : 1 * config.hidden_size] + k_bias = out_val[1 * config.hidden_size : 2 * config.hidden_size] + v_bias = out_val[2 * config.hidden_size : 3 * config.hidden_size] + + # Store. + output_state_dict[f"{layer_name}.attention.self.query.weight"] = q + output_state_dict[f"{layer_name}.attention.self.query.bias"] = q_bias + output_state_dict[f"{layer_name}.attention.self.key.weight"] = k + output_state_dict[f"{layer_name}.attention.self.key.bias"] = k_bias + output_state_dict[f"{layer_name}.attention.self.value.weight"] = v + output_state_dict[f"{layer_name}.attention.self.value.bias"] = v_bias + + # Clear the stored tensor. + attention_qkv_weight = None + + # Copy weights and biases as is. + elif weight_or_bias in ["weight", "bias"]: + out_name = megatron_to_transformers[op_name] + output_state_dict[layer_name + out_name + weight_or_bias] = val + + # The final layernorm. + output_state_dict["bert.encoder.ln.weight"] = transformer["final_layernorm.weight"] + output_state_dict["bert.encoder.ln.bias"] = transformer["final_layernorm.bias"] + + # The pooler. + pooler = lm["pooler"] + + # Store the matrix and the bias. + output_state_dict["bert.pooler.dense.weight"] = pooler["dense.weight"] + output_state_dict["bert.pooler.dense.bias"] = pooler["dense.bias"] + + # The LM head from Megatron (for RACE). + lm_head = model["lm_head"] + + # The transform matrix. + output_state_dict["cls.predictions.transform.dense.weight"] = lm_head["dense.weight"] + output_state_dict["cls.predictions.transform.dense.bias"] = lm_head["dense.bias"] + + # The transform LN. + output_state_dict["cls.predictions.transform.LayerNorm.weight"] = lm_head["layernorm.weight"] + output_state_dict["cls.predictions.transform.LayerNorm.bias"] = lm_head["layernorm.bias"] + + # For the decoder, we replicate the weights. + output_state_dict["cls.predictions.decoder.weight"] = word_embeddings + output_state_dict["cls.predictions.bias"] = lm_head["bias"] + + # The classifier from Megatron (for MLNI). + binary_head = model["binary_head"] + + # Store the classifier. + output_state_dict["cls.seq_relationship.weight"] = binary_head["weight"] + output_state_dict["cls.seq_relationship.bias"] = binary_head["bias"] + + # It should be done! + return output_state_dict + + +#################################################################################################### + + +def main(): + # Create the argument parser. + parser = argparse.ArgumentParser() + parser.add_argument("--print-checkpoint-structure", action="store_true") + parser.add_argument("path_to_checkpoint", type=str, help="Path to the ZIP file containing the checkpoint") + parser.add_argument( + "--config_file", + default="", + type=str, + help="An optional config json file describing the pre-trained model.", + ) + args = parser.parse_args() + + # Extract the basename. + basename = os.path.dirname(args.path_to_checkpoint) + + # Load the model. + # the .zip is very optional, let's keep it for backward compatibility + print(f'Extracting PyTorch state dictionary from "{args.path_to_checkpoint}"') + if args.path_to_checkpoint.endswith(".zip"): + with zipfile.ZipFile(args.path_to_checkpoint, "r") as checkpoint: + with checkpoint.open("release/mp_rank_00/model_optim_rng.pt") as pytorch_dict: + input_state_dict = torch.load(pytorch_dict, map_location="cpu", weights_only=True) + else: + input_state_dict = torch.load(args.path_to_checkpoint, map_location="cpu", weights_only=True) + + if args.config_file == "": + # Default config of megatron-bert 345m + config = MegatronBertConfig() + + # different megatron-bert-*-345m models have different vocab sizes, so override the default + # config (which is for megatron-bert-cased-345m) with the actual vocab dimension + config.vocab_size = input_state_dict["model"]["lm_head"]["bias"].numel() + else: + config = MegatronBertConfig.from_json_file(args.config_file) + + # Convert. + print("Converting") + output_state_dict = convert_megatron_checkpoint(args, input_state_dict, config) + + # Print the structure of converted state dict. + if args.print_checkpoint_structure: + recursive_print(None, output_state_dict) + + # Store the config to file. + print("Saving config") + config.save_pretrained(basename) + + # Store the state_dict to file. + output_checkpoint_file = os.path.join(basename, "pytorch_model.bin") + print(f'Saving checkpoint to "{output_checkpoint_file}"') + torch.save(output_state_dict, output_checkpoint_file) + + +#################################################################################################### + +if __name__ == "__main__": + main() + +#################################################################################################### diff --git a/third_party/transformers/src/transformers/models/megatron_bert/modeling_megatron_bert.py b/third_party/transformers/src/transformers/models/megatron_bert/modeling_megatron_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..cb8500a04c63deaca324bed58b057fa4f6e85588 --- /dev/null +++ b/third_party/transformers/src/transformers/models/megatron_bert/modeling_megatron_bert.py @@ -0,0 +1,1428 @@ +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch MegatronBERT model.""" + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + NextSentencePredictorOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_megatron_bert import MegatronBertConfig + + +logger = logging.get_logger(__name__) + + +class MegatronBertEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + # In Megatron, layer-norm is applied after the 1st dropout. + # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.LongTensor | None = None, + past_key_values_length: int = 0, + ) -> torch.Tensor: + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + seq_length = input_shape[1] + + if position_ids is None: + position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length] + + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = inputs_embeds + token_type_embeddings + + position_embeddings = self.position_embeddings(position_ids) + embeddings += position_embeddings + + # Megatron BERT moves that layer norm after the drop-out (and to each layer). + # embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +# copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->MegatronBert +class MegatronBertSelfAttention(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.is_decoder = config.is_decoder + self.layer_idx = layer_idx + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor]: + batch_size, seq_length, _ = hidden_states.shape + query_layer = self.query(hidden_states) + query_layer = query_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose( + 1, 2 + ) + + is_updated = False + is_cross_attention = encoder_hidden_states is not None + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_layer from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + + current_states = encoder_hidden_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_layer = curr_past_key_values.layers[self.layer_idx].keys + value_layer = curr_past_key_values.layers[self.layer_idx].values + else: + key_layer = self.key(current_states) + key_layer = key_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose( + 1, 2 + ) + value_layer = self.value(current_states) + value_layer = value_layer.view( + batch_size, -1, self.num_attention_heads, self.attention_head_size + ).transpose(1, 2) + + if past_key_values is not None: + # save all key/value_layer to cache to be re-used for fast auto-regressive generation + key_layer, value_layer = curr_past_key_values.update(key_layer, value_layer, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in MegatronBertModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + return context_layer, attention_probs + + +# Based transformers.models.bert.modeling_bert.BertSelfOutput. Moved LayerNorm to MegatronBertAttention below. +class MegatronBertSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return residual + hidden_states + + +# Based transformers.models.bert.modeling_bert.BertAttention. Added LayerNorm. +class MegatronBertAttention(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + self.ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.self = MegatronBertSelfAttention(config, layer_idx=layer_idx) + self.output = MegatronBertSelfOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor]: + ln_outputs = self.ln(hidden_states) + self_outputs = self.self( + ln_outputs, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + attention_output = self.output(self_outputs[0], hidden_states) + outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->MegatronBert +class MegatronBertIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +# Based on transformers.models.bert.modeling_bert.BertOutput. Moved LayerNorm to MegatronBertLayer below. +class MegatronBertOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return input_tensor + hidden_states + + +# Based on transformers.models.bert.modeling_bert.BertLayer. Added LayerNorm. +class MegatronBertLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = MegatronBertAttention(config, layer_idx=layer_idx) + self.is_decoder = config.is_decoder + self.add_cross_attention = config.add_cross_attention + if self.add_cross_attention: + if not self.is_decoder: + raise TypeError(f"{self} should be used as a decoder model if cross attention is added") + self.crossattention = MegatronBertAttention(config, layer_idx=layer_idx) + self.ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.intermediate = MegatronBertIntermediate(config) + self.output = MegatronBertOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor]: + # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 + self_attention_outputs = self.attention( + hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + past_key_values=past_key_values, + ) + attention_output = self_attention_outputs[0] + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + if self.is_decoder and encoder_hidden_states is not None: + if not hasattr(self, "crossattention"): + raise AttributeError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers" + " by setting `config.add_cross_attention=True`" + ) + + cross_attention_outputs = self.crossattention( + attention_output, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + attention_output = cross_attention_outputs[0] + outputs = outputs + cross_attention_outputs[1:] # add cross 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 + ) + return (layer_output,) + outputs + + def feed_forward_chunk(self, attention_output): + ln_output = self.ln(attention_output) + intermediate_output = self.intermediate(ln_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class MegatronBertEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([MegatronBertLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]) + + # The final layer norm. We removed the 1st LN, moved LN to each hidden layer and this one + # is simply the final LN (Transformer's BERT has it attached to each hidden layer). + self.ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = False, + output_hidden_states: bool | None = False, + return_dict: bool | None = True, + **kwargs, + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_outputs = layer_module( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_values, + output_attentions, + ) + + # Because we moved the layer-norm at the end of the hidden layer, we have non-normali- + # zed data here. If that's really needed, we must apply LN to match Transformer's BERT. + + hidden_states = layer_outputs[0] + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + if self.config.add_cross_attention: + all_cross_attentions = all_cross_attentions + (layer_outputs[2],) + + # Finalize the hidden states. + hidden_states = self.ln(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + past_key_values, + all_hidden_states, + all_self_attentions, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + +# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->MegatronBert +class MegatronBertPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->MegatronBert +class MegatronBertPredictionHeadTransform(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + if isinstance(config.hidden_act, str): + self.transform_act_fn = ACT2FN[config.hidden_act] + else: + self.transform_act_fn = config.hidden_act + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->MegatronBert +class MegatronBertLMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = MegatronBertPredictionHeadTransform(config) + + # The output weights are the same as the input embeddings, but there is + # an output-only bias for each token. + self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True) + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + def forward(self, hidden_states): + hidden_states = self.transform(hidden_states) + hidden_states = self.decoder(hidden_states) + return hidden_states + + +# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->MegatronBert +class MegatronBertOnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = MegatronBertLMPredictionHead(config) + + def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +# Copied from transformers.models.bert.modeling_bert.BertOnlyNSPHead with Bert->MegatronBert +class MegatronBertOnlyNSPHead(nn.Module): + def __init__(self, config): + super().__init__() + self.seq_relationship = nn.Linear(config.hidden_size, 2) + + def forward(self, pooled_output): + seq_relationship_score = self.seq_relationship(pooled_output) + return seq_relationship_score + + +# Copied from transformers.models.bert.modeling_bert.BertPreTrainingHeads with Bert->MegatronBert +class MegatronBertPreTrainingHeads(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = MegatronBertLMPredictionHead(config) + self.seq_relationship = nn.Linear(config.hidden_size, 2) + + def forward(self, sequence_output, pooled_output): + prediction_scores = self.predictions(sequence_output) + seq_relationship_score = self.seq_relationship(pooled_output) + return prediction_scores, seq_relationship_score + + +@auto_docstring +class MegatronBertPreTrainedModel(PreTrainedModel): + config: MegatronBertConfig + base_model_prefix = "bert" + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, MegatronBertLMPredictionHead): + init.zeros_(module.bias) + elif isinstance(module, MegatronBertEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`MegatronBertForPreTraining`]. + """ +) +# Copied from transformers.models.bert.modeling_bert.BertForPreTrainingOutput with Bert->MegatronBert +class MegatronBertForPreTrainingOutput(ModelOutput): + r""" + loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): + Total loss as the sum of the masked language modeling loss and the next sequence prediction + (classification) loss. + prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`): + Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation + before SoftMax). + """ + + loss: torch.FloatTensor | None = None + prediction_logits: torch.FloatTensor | None = None + seq_relationship_logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@auto_docstring +class MegatronBertModel(MegatronBertPreTrainedModel): + """ + + The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of + cross-attention is added between the self-attention layers, following the architecture described in [Attention is + all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, + Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. + + To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set + to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and + `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. + """ + + def __init__(self, config, add_pooling_layer=True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + self.embeddings = MegatronBertEmbeddings(config) + self.encoder = MegatronBertEncoder(config) + + self.pooler = MegatronBertPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPoolingAndCrossAttentions: + 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.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + + past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() + + if attention_mask is None: + attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +@auto_docstring( + custom_intro=""" + MegatronBert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a + `next sentence prediction (classification)` head. + """ +) +class MegatronBertForPreTraining(MegatronBertPreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight", + "cls.predictions.decoder.bias": "cls.predictions.bias", + } + + def __init__(self, config, add_binary_head=True): + r""" + add_binary_head (`bool`, *optional*, defaults to `True`): + Whether or not to add a binary head. + """ + super().__init__(config) + + self.bert = MegatronBertModel(config) + self.cls = MegatronBertPreTrainingHeads(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + next_sentence_label: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MegatronBertForPreTrainingOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair + (see `input_ids` docstring) Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + + Example: + + ```python + >>> from transformers import AutoTokenizer, MegatronBertForPreTraining + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m") + >>> model = MegatronBertForPreTraining.from_pretrained("nvidia/megatron-bert-cased-345m") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> prediction_logits = outputs.prediction_logits + >>> seq_relationship_logits = outputs.seq_relationship_logits + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output, pooled_output = outputs[:2] + prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) + + total_loss = None + if labels is not None and next_sentence_label is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)) + total_loss = masked_lm_loss + next_sentence_loss + + if not return_dict: + output = (prediction_scores, seq_relationship_score) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return MegatronBertForPreTrainingOutput( + loss=total_loss, + prediction_logits=prediction_scores, + seq_relationship_logits=seq_relationship_score, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + MegatronBert Model with a `language modeling` head on top for CLM fine-tuning. + """ +) +class MegatronBertForCausalLM(MegatronBertPreTrainedModel, GenerationMixin): + _tied_weights_keys = { + "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight", + "cls.predictions.decoder.bias": "cls.predictions.bias", + } + + def __init__(self, config): + super().__init__(config) + + if not config.is_decoder: + logger.warning("If you want to use `MegatronBertForCausalLM` as a standalone, add `is_decoder=True.`") + + self.bert = MegatronBertModel(config, add_pooling_layer=False) + self.cls = MegatronBertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithCrossAttentions: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` + + Example: + + ```python + >>> from transformers import AutoTokenizer, MegatronBertForCausalLM, MegatronBertConfig + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m") + >>> model = MegatronBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True) + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> prediction_logits = outputs.logits + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + if labels is not None: + use_cache = False + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.cls(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + + +@auto_docstring +class MegatronBertForMaskedLM(MegatronBertPreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight", + "cls.predictions.decoder.bias": "cls.predictions.bias", + } + + def __init__(self, config): + super().__init__(config) + + if config.is_decoder: + logger.warning( + "If you want to use `MegatronBertForMaskedLM` make sure `config.is_decoder=False` for " + "bi-directional self-attention." + ) + + self.bert = MegatronBertModel(config, add_pooling_layer=False) + self.cls = MegatronBertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MaskedLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + MegatronBert Model with a `next sentence prediction (classification)` head on top. + """ +) +class MegatronBertForNextSentencePrediction(MegatronBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.bert = MegatronBertModel(config) + self.cls = MegatronBertOnlyNSPHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | NextSentencePredictorOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair + (see `input_ids` docstring). Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + + Example: + + ```python + >>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m") + >>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." + >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") + + >>> outputs = model(**encoding, labels=torch.LongTensor([1])) + >>> logits = outputs.logits + >>> assert logits[0, 0] < logits[0, 1] # next sentence was random + ```""" + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + seq_relationship_scores = self.cls(pooled_output) + + next_sentence_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1)) + + if not return_dict: + output = (seq_relationship_scores,) + outputs[2:] + return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output + + return NextSentencePredictorOutput( + loss=next_sentence_loss, + logits=seq_relationship_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + MegatronBert Model transformer with a sequence classification/regression head on top (a linear layer on top of the + pooled output) e.g. for GLUE tasks. + """ +) +class MegatronBertForSequenceClassification(MegatronBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.bert = MegatronBertModel(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SequenceClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class MegatronBertForMultipleChoice(MegatronBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.bert = MegatronBertModel(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, 1) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MultipleChoiceModelOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None + inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + output = (reshaped_logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return MultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class MegatronBertForTokenClassification(MegatronBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.bert = MegatronBertModel(config, add_pooling_layer=False) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class MegatronBertForQuestionAnswering(MegatronBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.bert = MegatronBertModel(config, add_pooling_layer=False) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | QuestionAnsweringModelOutput: + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "MegatronBertForCausalLM", + "MegatronBertForMaskedLM", + "MegatronBertForMultipleChoice", + "MegatronBertForNextSentencePrediction", + "MegatronBertForPreTraining", + "MegatronBertForQuestionAnswering", + "MegatronBertForSequenceClassification", + "MegatronBertForTokenClassification", + "MegatronBertModel", + "MegatronBertPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/mobilevit/__init__.py b/third_party/transformers/src/transformers/models/mobilevit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ee9c65f364990069cdab3a66ef273b832822a5d --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_mobilevit import * + from .image_processing_mobilevit import * + from .image_processing_pil_mobilevit import * + from .modeling_mobilevit import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/mobilevit/configuration_mobilevit.py b/third_party/transformers/src/transformers/models/mobilevit/configuration_mobilevit.py new file mode 100644 index 0000000000000000000000000000000000000000..6d811285eca664ec8750ed56453e7a3a149aef48 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/configuration_mobilevit.py @@ -0,0 +1,75 @@ +# 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. +"""MobileViT model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/mobilenet_v2_1.0_224") +@strict +class MobileViTConfig(PreTrainedConfig): + r""" + neck_hidden_sizes (`list[int]`, *optional*, defaults to `[16, 32, 64, 96, 128, 160, 640]`): + The number of channels for the feature maps of the backbone. + aspp_out_channels (`int`, *optional*, defaults to 256): + Number of output channels used in the ASPP layer for semantic segmentation. + atrous_rates (`list[int]`, *optional*, defaults to `[6, 12, 18]`): + Dilation (atrous) factors used in the ASPP layer for semantic segmentation. + aspp_dropout_prob (`float`, *optional*, defaults to 0.1): + The dropout ratio for the ASPP layer for semantic segmentation. + + Example: + + ```python + >>> from transformers import MobileViTConfig, MobileViTModel + + >>> # Initializing a mobilevit-small style configuration + >>> configuration = MobileViTConfig() + + >>> # Initializing a model from the mobilevit-small style configuration + >>> model = MobileViTModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mobilevit" + + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 256 + patch_size: int | list[int] | tuple[int, int] = 2 + hidden_sizes: list[int] | tuple[int, ...] = (144, 192, 240) + neck_hidden_sizes: list[int] | tuple[int, ...] = (16, 32, 64, 96, 128, 160, 640) + num_attention_heads: int = 4 + mlp_ratio: float = 2.0 + expand_ratio: float = 4.0 + hidden_act: str = "silu" + conv_kernel_size: int = 3 + output_stride: int = 32 + hidden_dropout_prob: float | int = 0.1 + attention_probs_dropout_prob: float | int = 0.0 + classifier_dropout_prob: float | int = 0.1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + qkv_bias: bool = True + aspp_out_channels: int = 256 + atrous_rates: list[int] | tuple[int, ...] = (6, 12, 18) + aspp_dropout_prob: float | int = 0.1 + semantic_loss_ignore_index: int = 255 + + +__all__ = ["MobileViTConfig"] diff --git a/third_party/transformers/src/transformers/models/mobilevit/convert_mlcvnets_to_pytorch.py b/third_party/transformers/src/transformers/models/mobilevit/convert_mlcvnets_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..7a3ce6c1433eb09c30049bc58084d54d67230cca --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/convert_mlcvnets_to_pytorch.py @@ -0,0 +1,314 @@ +# Copyright 2022 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert MobileViT checkpoints from the ml-cvnets library.""" + +import argparse +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import ( + MobileViTConfig, + MobileViTForImageClassification, + MobileViTForSemanticSegmentation, + MobileViTImageProcessor, +) +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_mobilevit_config(mobilevit_name): + config = MobileViTConfig() + + # size of the architecture + if "mobilevit_s" in mobilevit_name: + config.hidden_sizes = [144, 192, 240] + config.neck_hidden_sizes = [16, 32, 64, 96, 128, 160, 640] + elif "mobilevit_xs" in mobilevit_name: + config.hidden_sizes = [96, 120, 144] + config.neck_hidden_sizes = [16, 32, 48, 64, 80, 96, 384] + elif "mobilevit_xxs" in mobilevit_name: + config.hidden_sizes = [64, 80, 96] + config.neck_hidden_sizes = [16, 16, 24, 48, 64, 80, 320] + config.hidden_dropout_prob = 0.05 + config.expand_ratio = 2.0 + + if mobilevit_name.startswith("deeplabv3_"): + config.image_size = 512 + config.output_stride = 16 + config.num_labels = 21 + filename = "pascal-voc-id2label.json" + else: + config.num_labels = 1000 + filename = "imagenet-1k-id2label.json" + + repo_id = "huggingface/label-files" + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + + return config + + +def rename_key(name, base_model=False): + for i in range(1, 6): + if f"layer_{i}." in name: + name = name.replace(f"layer_{i}.", f"encoder.layer.{i - 1}.") + + if "conv_1." in name: + name = name.replace("conv_1.", "conv_stem.") + if ".block." in name: + name = name.replace(".block.", ".") + + if "exp_1x1" in name: + name = name.replace("exp_1x1", "expand_1x1") + if "red_1x1" in name: + name = name.replace("red_1x1", "reduce_1x1") + if ".local_rep.conv_3x3." in name: + name = name.replace(".local_rep.conv_3x3.", ".conv_kxk.") + if ".local_rep.conv_1x1." in name: + name = name.replace(".local_rep.conv_1x1.", ".conv_1x1.") + if ".norm." in name: + name = name.replace(".norm.", ".normalization.") + if ".conv." in name: + name = name.replace(".conv.", ".convolution.") + if ".conv_proj." in name: + name = name.replace(".conv_proj.", ".conv_projection.") + + for i in range(0, 2): + for j in range(0, 4): + if f".{i}.{j}." in name: + name = name.replace(f".{i}.{j}.", f".{i}.layer.{j}.") + + for i in range(2, 6): + for j in range(0, 4): + if f".{i}.{j}." in name: + name = name.replace(f".{i}.{j}.", f".{i}.") + if "expand_1x1" in name: + name = name.replace("expand_1x1", "downsampling_layer.expand_1x1") + if "conv_3x3" in name: + name = name.replace("conv_3x3", "downsampling_layer.conv_3x3") + if "reduce_1x1" in name: + name = name.replace("reduce_1x1", "downsampling_layer.reduce_1x1") + + for i in range(2, 5): + if f".global_rep.{i}.weight" in name: + name = name.replace(f".global_rep.{i}.weight", ".layernorm.weight") + if f".global_rep.{i}.bias" in name: + name = name.replace(f".global_rep.{i}.bias", ".layernorm.bias") + + if ".global_rep." in name: + name = name.replace(".global_rep.", ".transformer.") + if ".pre_norm_mha.0." in name: + name = name.replace(".pre_norm_mha.0.", ".layernorm_before.") + if ".pre_norm_mha.1.out_proj." in name: + name = name.replace(".pre_norm_mha.1.out_proj.", ".attention.output.dense.") + if ".pre_norm_ffn.0." in name: + name = name.replace(".pre_norm_ffn.0.", ".layernorm_after.") + if ".pre_norm_ffn.1." in name: + name = name.replace(".pre_norm_ffn.1.", ".intermediate.dense.") + if ".pre_norm_ffn.4." in name: + name = name.replace(".pre_norm_ffn.4.", ".output.dense.") + if ".transformer." in name: + name = name.replace(".transformer.", ".transformer.layer.") + + if ".aspp_layer." in name: + name = name.replace(".aspp_layer.", ".") + if ".aspp_pool." in name: + name = name.replace(".aspp_pool.", ".") + if "seg_head." in name: + name = name.replace("seg_head.", "segmentation_head.") + if "segmentation_head.classifier.classifier." in name: + name = name.replace("segmentation_head.classifier.classifier.", "segmentation_head.classifier.") + + if "classifier.fc." in name: + name = name.replace("classifier.fc.", "classifier.") + elif (not base_model) and ("segmentation_head." not in name): + name = "mobilevit." + name + + return name + + +def convert_state_dict(orig_state_dict, model, base_model=False): + if base_model: + model_prefix = "" + else: + model_prefix = "mobilevit." + + for key in orig_state_dict.copy(): + val = orig_state_dict.pop(key) + + if key[:8] == "encoder.": + key = key[8:] + + if "qkv" in key: + key_split = key.split(".") + layer_num = int(key_split[0][6:]) - 1 + transformer_num = int(key_split[3]) + layer = model.get_submodule(f"{model_prefix}encoder.layer.{layer_num}") + dim = layer.transformer.layer[transformer_num].attention.attention.all_head_size + prefix = ( + f"{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention." + ) + if "weight" in key: + orig_state_dict[prefix + "query.weight"] = val[:dim, :] + orig_state_dict[prefix + "key.weight"] = val[dim : dim * 2, :] + orig_state_dict[prefix + "value.weight"] = val[-dim:, :] + else: + orig_state_dict[prefix + "query.bias"] = val[:dim] + orig_state_dict[prefix + "key.bias"] = val[dim : dim * 2] + orig_state_dict[prefix + "value.bias"] = val[-dim:] + else: + orig_state_dict[rename_key(key, base_model)] = val + + return orig_state_dict + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_movilevit_checkpoint(mobilevit_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False): + """ + Copy/paste/tweak model's weights to our MobileViT structure. + """ + config = get_mobilevit_config(mobilevit_name) + + # load original state_dict + state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + # load 🤗 model + if mobilevit_name.startswith("deeplabv3_"): + model = MobileViTForSemanticSegmentation(config).eval() + else: + model = MobileViTForImageClassification(config).eval() + + new_state_dict = convert_state_dict(state_dict, model) + model.load_state_dict(new_state_dict) + + # Check outputs on an image, prepared by MobileViTImageProcessor + image_processor = MobileViTImageProcessor(crop_size=config.image_size, size=config.image_size + 32) + encoding = image_processor(images=prepare_img(), return_tensors="pt") + outputs = model(**encoding) + logits = outputs.logits + + if mobilevit_name.startswith("deeplabv3_"): + assert logits.shape == (1, 21, 32, 32) + + if mobilevit_name == "deeplabv3_mobilevit_s": + expected_logits = torch.tensor( + [ + [[6.2065, 6.1292, 6.2070], [6.1079, 6.1254, 6.1747], [6.0042, 6.1071, 6.1034]], + [[-6.9253, -6.8653, -7.0398], [-7.3218, -7.3983, -7.3670], [-7.1961, -7.2482, -7.1569]], + [[-4.4723, -4.4348, -4.3769], [-5.3629, -5.4632, -5.4598], [-5.1587, -5.3402, -5.5059]], + ] + ) + elif mobilevit_name == "deeplabv3_mobilevit_xs": + expected_logits = torch.tensor( + [ + [[5.4449, 5.5733, 5.6314], [5.1815, 5.3930, 5.5963], [5.1656, 5.4333, 5.4853]], + [[-9.4423, -9.7766, -9.6714], [-9.1581, -9.5720, -9.5519], [-9.1006, -9.6458, -9.5703]], + [[-7.7721, -7.3716, -7.1583], [-8.4599, -8.0624, -7.7944], [-8.4172, -7.8366, -7.5025]], + ] + ) + elif mobilevit_name == "deeplabv3_mobilevit_xxs": + expected_logits = torch.tensor( + [ + [[6.9811, 6.9743, 7.3123], [7.1777, 7.1931, 7.3938], [7.5633, 7.8050, 7.8901]], + [[-10.5536, -10.2332, -10.2924], [-10.2336, -9.8624, -9.5964], [-10.8840, -10.8158, -10.6659]], + [[-3.4938, -3.0631, -2.8620], [-3.4205, -2.8135, -2.6875], [-3.4179, -2.7945, -2.8750]], + ] + ) + else: + raise ValueError(f"Unknown mobilevit_name: {mobilevit_name}") + + assert torch.allclose(logits[0, :3, :3, :3], expected_logits, atol=1e-4) + else: + assert logits.shape == (1, 1000) + + if mobilevit_name == "mobilevit_s": + expected_logits = torch.tensor([-0.9866, 0.2392, -1.1241]) + elif mobilevit_name == "mobilevit_xs": + expected_logits = torch.tensor([-2.4761, -0.9399, -1.9587]) + elif mobilevit_name == "mobilevit_xxs": + expected_logits = torch.tensor([-1.9364, -1.2327, -0.4653]) + else: + raise ValueError(f"Unknown mobilevit_name: {mobilevit_name}") + + assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4) + + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model {mobilevit_name} to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + model_mapping = { + "mobilevit_s": "mobilevit-small", + "mobilevit_xs": "mobilevit-x-small", + "mobilevit_xxs": "mobilevit-xx-small", + "deeplabv3_mobilevit_s": "deeplabv3-mobilevit-small", + "deeplabv3_mobilevit_xs": "deeplabv3-mobilevit-x-small", + "deeplabv3_mobilevit_xxs": "deeplabv3-mobilevit-xx-small", + } + + print("Pushing to the hub...") + model_name = model_mapping[mobilevit_name] + image_processor.push_to_hub(model_name, organization="apple") + model.push_to_hub(model_name, organization="apple") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--mobilevit_name", + default="mobilevit_s", + type=str, + help=( + "Name of the MobileViT model you'd like to convert. Should be one of 'mobilevit_s', 'mobilevit_xs'," + " 'mobilevit_xxs', 'deeplabv3_mobilevit_s', 'deeplabv3_mobilevit_xs', 'deeplabv3_mobilevit_xxs'." + ), + ) + parser.add_argument( + "--checkpoint_path", required=True, type=str, help="Path to the original state dict (.pt file)." + ) + parser.add_argument( + "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory." + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether or not to push the converted model to the Hugging Face hub.", + ) + + args = parser.parse_args() + convert_movilevit_checkpoint( + args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub + ) diff --git a/third_party/transformers/src/transformers/models/mobilevit/image_processing_mobilevit.py b/third_party/transformers/src/transformers/models/mobilevit/image_processing_mobilevit.py new file mode 100644 index 0000000000000000000000000000000000000000..d94c1912fbd9fb7e29d11de407ca3f45d0940a5f --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/image_processing_mobilevit.py @@ -0,0 +1,231 @@ +# 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 MobileViT.""" + +from typing import Union + +import torch +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, + logging, + requires_backends, +) + + +logger = logging.get_logger(__name__) + + +class MobileVitImageProcessorKwargs(ImagesKwargs, total=False): + """ + do_flip_channel_order (`bool`, *optional*, defaults to `self.do_flip_channel_order`): + Whether to flip the color channels from RGB to BGR or vice versa. + do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`): + Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. + ADE20k). The background label will be replaced by 255. + """ + + do_flip_channel_order: bool + do_reduce_labels: bool + + +@auto_docstring +class MobileViTImageProcessor(TorchvisionBackend): + """Torchvision backend for MobileViT with flip_channel_order and reduce_label support.""" + + valid_kwargs = MobileVitImageProcessorKwargs + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"shortest_edge": 224} + default_to_square = False + crop_size = {"height": 256, "width": 256} + do_resize = True + do_center_crop = True + do_rescale = True + do_normalize = None + do_convert_rgb = None + do_flip_channel_order = True + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[MobileVitImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[MobileVitImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to preprocess. + """ + return super().preprocess(images, segmentation_maps, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + return_tensors: str | TensorType | None, + device: Union[str, "torch.device"] | None = None, + **kwargs, + ) -> BatchFeature: + """Handle extra inputs beyond images.""" + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update( + { + "do_rescale": False, + "do_flip_channel_order": False, + # Nearest interpolation is used for segmentation maps instead of BICUBIC. + "resample": PILImageResampling.NEAREST, + } + ) + + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + processed_segmentation_maps = [ + processed_segmentation_map.squeeze(0).to(torch.int64) + for processed_segmentation_map in processed_segmentation_maps + ] + data["labels"] = processed_segmentation_maps + + return BatchFeature(data=data, tensor_type=return_tensors) + + def reduce_label(self, labels: list["torch.Tensor"]) -> list["torch.Tensor"]: + """Reduce label values by 1, replacing 0 with 255.""" + for idx in range(len(labels)): + label = labels[idx] + label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype, device=label.device), label) + label = label - 1 + label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype, device=label.device), label) + labels[idx] = label + return labels + + def flip_channel_order(self, images: "torch.Tensor") -> "torch.Tensor": + """Flip RGB to BGR or vice versa.""" + if images.ndim == 3: + # Single image: (C, H, W) + flipped = images.clone() + flipped[0:3] = images[[2, 1, 0]] + return flipped + elif images.ndim == 4: + # Batched images: (B, C, H, W) + flipped = images.clone() + flipped[:, 0:3] = images[:, [2, 1, 0]] + return flipped + return images + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + disable_grouping: bool | None, + do_reduce_labels: bool = False, + do_flip_channel_order: bool = True, + **kwargs, + ) -> list["torch.Tensor"]: + """Custom preprocessing for MobileViT.""" + if do_reduce_labels: + images = self.reduce_label(images) + + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size, resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + if do_rescale: + stacked_images = self.rescale(stacked_images, rescale_factor) + if do_flip_channel_order: + stacked_images = self.flip_channel_order(stacked_images) + processed_images_grouped[shape] = stacked_images + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + return processed_images + + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """Converts the output of [`MobileViTForSemanticSegmentation`] into semantic segmentation maps.""" + requires_backends(self, "torch") + logits = outputs.logits + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + semantic_segmentation = [] + for idx in range(len(logits)): + resized_logits = torch.nn.functional.interpolate( + logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + return semantic_segmentation + + +__all__ = ["MobileViTImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/mobilevit/image_processing_pil_mobilevit.py b/third_party/transformers/src/transformers/models/mobilevit/image_processing_pil_mobilevit.py new file mode 100644 index 0000000000000000000000000000000000000000..893e27fe4ccff64513a4a708ebae87d777375fd5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/image_processing_pil_mobilevit.py @@ -0,0 +1,212 @@ +# 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 MobileViT.""" + +from typing import TYPE_CHECKING + +import numpy as np + + +if TYPE_CHECKING: + pass + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import flip_channel_order as np_flip_channel_order +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, logging +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + + +# Adapted from transformers.models.mobilevit.image_processing_mobilevit.MobileVitImageProcessorKwargs +class MobileVitImageProcessorKwargs(ImagesKwargs, total=False): + """ + do_flip_channel_order (`bool`, *optional*, defaults to `self.do_flip_channel_order`): + Whether to flip the color channels from RGB to BGR or vice versa. + do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`): + Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. + ADE20k). The background label will be replaced by 255. + """ + + do_flip_channel_order: bool + do_reduce_labels: bool + + +@auto_docstring +class MobileViTImageProcessorPil(PilBackend): + """PIL backend for MobileViT with flip_channel_order and reduce_label support.""" + + valid_kwargs = MobileVitImageProcessorKwargs + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"shortest_edge": 224} + default_to_square = False + crop_size = {"height": 256, "width": 256} + do_resize = True + do_center_crop = True + do_rescale = True + do_normalize = None + do_convert_rgb = None + do_flip_channel_order = True + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[MobileVitImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[MobileVitImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to preprocess. + """ + return super().preprocess(images, segmentation_maps, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + return_tensors: str | TensorType | None, + device: str | None = None, + **kwargs, + ) -> BatchFeature: + """Handle extra inputs beyond images.""" + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update( + { + "do_rescale": False, + "do_flip_channel_order": False, + # Nearest interpolation is used for segmentation maps instead of BICUBIC. + "resample": PILImageResampling.NEAREST, + } + ) + + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + processed_segmentation_maps = [ + processed_segmentation_map.squeeze(0).astype(np.int64) + for processed_segmentation_map in processed_segmentation_maps + ] + data["labels"] = processed_segmentation_maps + + return BatchFeature(data=data, tensor_type=return_tensors) + + def reduce_label(self, image: np.ndarray) -> np.ndarray: + """Reduce label values by 1, replacing 0 with 255.""" + image[image == 0] = 255 + image = image - 1 + image[image == 254] = 255 + return image + + def flip_channel_order(self, image: np.ndarray) -> np.ndarray: + """Flip RGB to BGR or vice versa.""" + return np_flip_channel_order( + image, data_format=ChannelDimension.FIRST, input_data_format=ChannelDimension.FIRST + ) + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_reduce_labels: bool = False, + do_flip_channel_order: bool = True, + **kwargs, + ) -> list[np.ndarray]: + """Custom preprocessing for MobileViT.""" + processed_images = [] + for image in images: + if do_reduce_labels: + image = self.reduce_label(image) + if do_resize: + image = self.resize(image, size, resample) + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_center_crop: + image = self.center_crop(image, crop_size) + if do_flip_channel_order: + image = self.flip_channel_order(image) + processed_images.append(image) + return processed_images + + @requires(backends=("torch",)) + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """Converts the output of [`MobileViTForSemanticSegmentation`] into semantic segmentation maps.""" + import torch + + logits = outputs.logits + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + semantic_segmentation = [] + for idx in range(len(logits)): + resized_logits = torch.nn.functional.interpolate( + logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + return semantic_segmentation + + +__all__ = ["MobileViTImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/mobilevit/modeling_mobilevit.py b/third_party/transformers/src/transformers/models/mobilevit/modeling_mobilevit.py new file mode 100644 index 0000000000000000000000000000000000000000..aff43a5daa7a9bba95e0e5e41908159d82e77918 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevit/modeling_mobilevit.py @@ -0,0 +1,974 @@ +# Copyright 2022 Apple Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Original license: https://github.com/apple/ml-cvnets/blob/main/LICENSE +"""PyTorch MobileViT model.""" + +import math + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithNoAttention, + BaseModelOutputWithPoolingAndNoAttention, + ImageClassifierOutputWithNoAttention, + SemanticSegmenterOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging, torch_int +from .configuration_mobilevit import MobileViTConfig + + +logger = logging.get_logger(__name__) + + +def make_divisible(value: int, divisor: int = 8, min_value: int | None = None) -> int: + """ + Ensure that all layers have a channel count that is divisible by `divisor`. + """ + if min_value is None: + min_value = divisor + new_value = max(min_value, int(value + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_value < 0.9 * value: + new_value += divisor + return int(new_value) + + +class MobileViTConvLayer(nn.Module): + def __init__( + self, + config: MobileViTConfig, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + groups: int = 1, + bias: bool = False, + dilation: int = 1, + use_normalization: bool = True, + use_activation: bool | str = True, + ) -> None: + super().__init__() + padding = int((kernel_size - 1) / 2) * dilation + + if in_channels % groups != 0: + raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups.") + if out_channels % groups != 0: + raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups.") + + self.convolution = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode="zeros", + ) + + if use_normalization: + self.normalization = nn.BatchNorm2d( + num_features=out_channels, + eps=1e-5, + momentum=0.1, + affine=True, + track_running_stats=True, + ) + else: + self.normalization = None + + if use_activation: + if isinstance(use_activation, str): + self.activation = ACT2FN[use_activation] + elif isinstance(config.hidden_act, str): + self.activation = ACT2FN[config.hidden_act] + else: + self.activation = config.hidden_act + else: + self.activation = None + + def forward(self, features: torch.Tensor) -> torch.Tensor: + features = self.convolution(features) + if self.normalization is not None: + features = self.normalization(features) + if self.activation is not None: + features = self.activation(features) + return features + + +class MobileViTInvertedResidual(nn.Module): + """ + Inverted residual block (MobileNetv2): https://huggingface.co/papers/1801.04381 + """ + + def __init__( + self, config: MobileViTConfig, in_channels: int, out_channels: int, stride: int, dilation: int = 1 + ) -> None: + super().__init__() + expanded_channels = make_divisible(int(round(in_channels * config.expand_ratio)), 8) + + if stride not in [1, 2]: + raise ValueError(f"Invalid stride {stride}.") + + self.use_residual = (stride == 1) and (in_channels == out_channels) + + self.expand_1x1 = MobileViTConvLayer( + config, in_channels=in_channels, out_channels=expanded_channels, kernel_size=1 + ) + + self.conv_3x3 = MobileViTConvLayer( + config, + in_channels=expanded_channels, + out_channels=expanded_channels, + kernel_size=3, + stride=stride, + groups=expanded_channels, + dilation=dilation, + ) + + self.reduce_1x1 = MobileViTConvLayer( + config, + in_channels=expanded_channels, + out_channels=out_channels, + kernel_size=1, + use_activation=False, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + residual = features + + features = self.expand_1x1(features) + features = self.conv_3x3(features) + features = self.reduce_1x1(features) + + return residual + features if self.use_residual else features + + +class MobileViTMobileNetLayer(nn.Module): + def __init__( + self, config: MobileViTConfig, in_channels: int, out_channels: int, stride: int = 1, num_stages: int = 1 + ) -> None: + super().__init__() + + self.layer = nn.ModuleList() + for i in range(num_stages): + layer = MobileViTInvertedResidual( + config, + in_channels=in_channels, + out_channels=out_channels, + stride=stride if i == 0 else 1, + ) + self.layer.append(layer) + in_channels = out_channels + + def forward(self, features: torch.Tensor) -> torch.Tensor: + for layer_module in self.layer: + features = layer_module(features) + return features + + +class MobileViTSelfAttention(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int) -> None: + super().__init__() + + if hidden_size % config.num_attention_heads != 0: + raise ValueError( + f"The hidden size {hidden_size} is not a multiple of the number of attention " + f"heads {config.num_attention_heads}." + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias) + self.key = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias) + self.value = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, seq_length, _ = hidden_states.shape + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + return context_layer + + +class MobileViTSelfOutput(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int) -> None: + super().__init__() + self.dense = nn.Linear(hidden_size, hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class MobileViTAttention(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int) -> None: + super().__init__() + self.attention = MobileViTSelfAttention(config, hidden_size) + self.output = MobileViTSelfOutput(config, hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self_outputs = self.attention(hidden_states) + attention_output = self.output(self_outputs) + return attention_output + + +class MobileViTIntermediate(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.dense = nn.Linear(hidden_size, intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +class MobileViTOutput(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.dense = nn.Linear(intermediate_size, hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +class MobileViTTransformerLayer(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.attention = MobileViTAttention(config, hidden_size) + self.intermediate = MobileViTIntermediate(config, hidden_size, intermediate_size) + self.output = MobileViTOutput(config, hidden_size, intermediate_size) + self.layernorm_before = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + self.layernorm_after = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + attention_output = self.attention(self.layernorm_before(hidden_states)) + hidden_states = attention_output + hidden_states + + layer_output = self.layernorm_after(hidden_states) + layer_output = self.intermediate(layer_output) + layer_output = self.output(layer_output, hidden_states) + return layer_output + + +class MobileViTTransformer(nn.Module): + def __init__(self, config: MobileViTConfig, hidden_size: int, num_stages: int) -> None: + super().__init__() + + self.layer = nn.ModuleList() + for _ in range(num_stages): + transformer_layer = MobileViTTransformerLayer( + config, + hidden_size=hidden_size, + intermediate_size=int(hidden_size * config.mlp_ratio), + ) + self.layer.append(transformer_layer) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for layer_module in self.layer: + hidden_states = layer_module(hidden_states) + return hidden_states + + +class MobileViTLayer(GradientCheckpointingLayer): + """ + MobileViT block: https://huggingface.co/papers/2110.02178 + """ + + def __init__( + self, + config: MobileViTConfig, + in_channels: int, + out_channels: int, + stride: int, + hidden_size: int, + num_stages: int, + dilation: int = 1, + ) -> None: + super().__init__() + self.patch_width = config.patch_size + self.patch_height = config.patch_size + + if stride == 2: + self.downsampling_layer = MobileViTInvertedResidual( + config, + in_channels=in_channels, + out_channels=out_channels, + stride=stride if dilation == 1 else 1, + dilation=dilation // 2 if dilation > 1 else 1, + ) + in_channels = out_channels + else: + self.downsampling_layer = None + + self.conv_kxk = MobileViTConvLayer( + config, + in_channels=in_channels, + out_channels=in_channels, + kernel_size=config.conv_kernel_size, + ) + + self.conv_1x1 = MobileViTConvLayer( + config, + in_channels=in_channels, + out_channels=hidden_size, + kernel_size=1, + use_normalization=False, + use_activation=False, + ) + + self.transformer = MobileViTTransformer( + config, + hidden_size=hidden_size, + num_stages=num_stages, + ) + + self.layernorm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + + self.conv_projection = MobileViTConvLayer( + config, in_channels=hidden_size, out_channels=in_channels, kernel_size=1 + ) + + self.fusion = MobileViTConvLayer( + config, in_channels=2 * in_channels, out_channels=in_channels, kernel_size=config.conv_kernel_size + ) + + def unfolding(self, features: torch.Tensor) -> tuple[torch.Tensor, dict]: + patch_width, patch_height = self.patch_width, self.patch_height + patch_area = int(patch_width * patch_height) + + batch_size, channels, orig_height, orig_width = features.shape + + new_height = ( + torch_int(torch.ceil(orig_height / patch_height) * patch_height) + if torch.jit.is_tracing() + else int(math.ceil(orig_height / patch_height) * patch_height) + ) + new_width = ( + torch_int(torch.ceil(orig_width / patch_width) * patch_width) + if torch.jit.is_tracing() + else int(math.ceil(orig_width / patch_width) * patch_width) + ) + + interpolate = False + if new_width != orig_width or new_height != orig_height: + # Note: Padding can be done, but then it needs to be handled in attention function. + features = nn.functional.interpolate( + features, size=(new_height, new_width), mode="bilinear", align_corners=False + ) + interpolate = True + + # number of patches along width and height + num_patch_width = new_width // patch_width + num_patch_height = new_height // patch_height + num_patches = num_patch_height * num_patch_width + + # convert from shape (batch_size, channels, orig_height, orig_width) + # to the shape (batch_size * patch_area, num_patches, channels) + patches = features.reshape( + batch_size * channels * num_patch_height, patch_height, num_patch_width, patch_width + ) + patches = patches.transpose(1, 2) + patches = patches.reshape(batch_size, channels, num_patches, patch_area) + patches = patches.transpose(1, 3) + patches = patches.reshape(batch_size * patch_area, num_patches, -1) + + info_dict = { + "orig_size": (orig_height, orig_width), + "batch_size": batch_size, + "channels": channels, + "interpolate": interpolate, + "num_patches": num_patches, + "num_patches_width": num_patch_width, + "num_patches_height": num_patch_height, + } + return patches, info_dict + + def folding(self, patches: torch.Tensor, info_dict: dict) -> torch.Tensor: + patch_width, patch_height = self.patch_width, self.patch_height + patch_area = int(patch_width * patch_height) + + batch_size = info_dict["batch_size"] + channels = info_dict["channels"] + num_patches = info_dict["num_patches"] + num_patch_height = info_dict["num_patches_height"] + num_patch_width = info_dict["num_patches_width"] + + # convert from shape (batch_size * patch_area, num_patches, channels) + # back to shape (batch_size, channels, orig_height, orig_width) + features = patches.contiguous().view(batch_size, patch_area, num_patches, -1) + features = features.transpose(1, 3) + features = features.reshape( + batch_size * channels * num_patch_height, num_patch_width, patch_height, patch_width + ) + features = features.transpose(1, 2) + features = features.reshape( + batch_size, channels, num_patch_height * patch_height, num_patch_width * patch_width + ) + + if info_dict["interpolate"]: + features = nn.functional.interpolate( + features, size=info_dict["orig_size"], mode="bilinear", align_corners=False + ) + + return features + + def forward(self, features: torch.Tensor) -> torch.Tensor: + # reduce spatial dimensions if needed + if self.downsampling_layer: + features = self.downsampling_layer(features) + + residual = features + + # local representation + features = self.conv_kxk(features) + features = self.conv_1x1(features) + + # convert feature map to patches + patches, info_dict = self.unfolding(features) + + # learn global representations + patches = self.transformer(patches) + patches = self.layernorm(patches) + + # convert patches back to feature maps + features = self.folding(patches, info_dict) + + features = self.conv_projection(features) + features = self.fusion(torch.cat((residual, features), dim=1)) + return features + + +class MobileViTEncoder(nn.Module): + def __init__(self, config: MobileViTConfig) -> None: + super().__init__() + self.config = config + + self.layer = nn.ModuleList() + self.gradient_checkpointing = False + + # segmentation architectures like DeepLab and PSPNet modify the strides + # of the classification backbones + dilate_layer_4 = dilate_layer_5 = False + if config.output_stride == 8: + dilate_layer_4 = True + dilate_layer_5 = True + elif config.output_stride == 16: + dilate_layer_5 = True + + dilation = 1 + + layer_1 = MobileViTMobileNetLayer( + config, + in_channels=config.neck_hidden_sizes[0], + out_channels=config.neck_hidden_sizes[1], + stride=1, + num_stages=1, + ) + self.layer.append(layer_1) + + layer_2 = MobileViTMobileNetLayer( + config, + in_channels=config.neck_hidden_sizes[1], + out_channels=config.neck_hidden_sizes[2], + stride=2, + num_stages=3, + ) + self.layer.append(layer_2) + + layer_3 = MobileViTLayer( + config, + in_channels=config.neck_hidden_sizes[2], + out_channels=config.neck_hidden_sizes[3], + stride=2, + hidden_size=config.hidden_sizes[0], + num_stages=2, + ) + self.layer.append(layer_3) + + if dilate_layer_4: + dilation *= 2 + + layer_4 = MobileViTLayer( + config, + in_channels=config.neck_hidden_sizes[3], + out_channels=config.neck_hidden_sizes[4], + stride=2, + hidden_size=config.hidden_sizes[1], + num_stages=4, + dilation=dilation, + ) + self.layer.append(layer_4) + + if dilate_layer_5: + dilation *= 2 + + layer_5 = MobileViTLayer( + config, + in_channels=config.neck_hidden_sizes[4], + out_channels=config.neck_hidden_sizes[5], + stride=2, + hidden_size=config.hidden_sizes[2], + num_stages=3, + dilation=dilation, + ) + self.layer.append(layer_5) + + def forward( + self, + hidden_states: torch.Tensor, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> tuple | BaseModelOutputWithNoAttention: + all_hidden_states = () if output_hidden_states else None + + for i, layer_module in enumerate(self.layer): + hidden_states = layer_module(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] if v is not None) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=all_hidden_states) + + +@auto_docstring +class MobileViTPreTrainedModel(PreTrainedModel): + config: MobileViTConfig + base_model_prefix = "mobilevit" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MobileViTLayer"] + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + """Initialize the weights""" + if isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + if getattr(module, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + elif isinstance(module, nn.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + + +@auto_docstring +class MobileViTModel(MobileViTPreTrainedModel): + def __init__(self, config: MobileViTConfig, expand_output: bool = True): + r""" + expand_output (`bool`, *optional*, defaults to `True`): + Whether to expand the output of the model using a 1x1 convolution. If `True`, the model will apply an additional + 1x1 convolution to expand the output channels from `config.neck_hidden_sizes[5]` to `config.neck_hidden_sizes[6]`. + """ + super().__init__(config) + self.config = config + self.expand_output = expand_output + + self.conv_stem = MobileViTConvLayer( + config, + in_channels=config.num_channels, + out_channels=config.neck_hidden_sizes[0], + kernel_size=3, + stride=2, + ) + + self.encoder = MobileViTEncoder(config) + + if self.expand_output: + self.conv_1x1_exp = MobileViTConvLayer( + config, + in_channels=config.neck_hidden_sizes[5], + out_channels=config.neck_hidden_sizes[6], + kernel_size=1, + ) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPoolingAndNoAttention: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + embedding_output = self.conv_stem(pixel_values) + + encoder_outputs = self.encoder( + embedding_output, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.expand_output: + last_hidden_state = self.conv_1x1_exp(encoder_outputs[0]) + + # global average pooling: (batch_size, channels, height, width) -> (batch_size, channels) + pooled_output = torch.mean(last_hidden_state, dim=[-2, -1], keepdim=False) + else: + last_hidden_state = encoder_outputs[0] + pooled_output = None + + if not return_dict: + output = (last_hidden_state, pooled_output) if pooled_output is not None else (last_hidden_state,) + return output + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndNoAttention( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + MobileViT model with an image classification head on top (a linear layer on top of the pooled features), e.g. for + ImageNet. + """ +) +class MobileViTForImageClassification(MobileViTPreTrainedModel): + def __init__(self, config: MobileViTConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.mobilevit = MobileViTModel(config) + + # Classifier head + self.dropout = nn.Dropout(config.classifier_dropout_prob, inplace=True) + self.classifier = ( + nn.Linear(config.neck_hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + labels: torch.Tensor | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ImageClassifierOutputWithNoAttention: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss). If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.mobilevit(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict) + + pooled_output = outputs.pooler_output if return_dict else outputs[1] + + logits = self.classifier(self.dropout(pooled_output)) + + loss = None + if labels is not None: + loss = self.loss_function(labels, logits, self.config) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutputWithNoAttention( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + ) + + +class MobileViTASPPPooling(nn.Module): + def __init__(self, config: MobileViTConfig, in_channels: int, out_channels: int) -> None: + super().__init__() + + self.global_pool = nn.AdaptiveAvgPool2d(output_size=1) + + self.conv_1x1 = MobileViTConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + use_normalization=True, + use_activation="relu", + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + spatial_size = features.shape[-2:] + features = self.global_pool(features) + features = self.conv_1x1(features) + features = nn.functional.interpolate(features, size=spatial_size, mode="bilinear", align_corners=False) + return features + + +class MobileViTASPP(nn.Module): + """ + ASPP module defined in DeepLab papers: https://huggingface.co/papers/1606.00915, https://huggingface.co/papers/1706.05587 + """ + + def __init__(self, config: MobileViTConfig) -> None: + super().__init__() + + in_channels = config.neck_hidden_sizes[-2] + out_channels = config.aspp_out_channels + + if len(config.atrous_rates) != 3: + raise ValueError("Expected 3 values for atrous_rates") + + self.convs = nn.ModuleList() + + in_projection = MobileViTConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + use_activation="relu", + ) + self.convs.append(in_projection) + + self.convs.extend( + [ + MobileViTConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + dilation=rate, + use_activation="relu", + ) + for rate in config.atrous_rates + ] + ) + + pool_layer = MobileViTASPPPooling(config, in_channels, out_channels) + self.convs.append(pool_layer) + + self.project = MobileViTConvLayer( + config, in_channels=5 * out_channels, out_channels=out_channels, kernel_size=1, use_activation="relu" + ) + + self.dropout = nn.Dropout(p=config.aspp_dropout_prob) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + pyramid = [] + for conv in self.convs: + pyramid.append(conv(features)) + pyramid = torch.cat(pyramid, dim=1) + + pooled_features = self.project(pyramid) + pooled_features = self.dropout(pooled_features) + return pooled_features + + +class MobileViTDeepLabV3(nn.Module): + """ + DeepLabv3 architecture: https://huggingface.co/papers/1706.05587 + """ + + def __init__(self, config: MobileViTConfig) -> None: + super().__init__() + self.aspp = MobileViTASPP(config) + + self.dropout = nn.Dropout2d(config.classifier_dropout_prob) + + self.classifier = MobileViTConvLayer( + config, + in_channels=config.aspp_out_channels, + out_channels=config.num_labels, + kernel_size=1, + use_normalization=False, + use_activation=False, + bias=True, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + features = self.aspp(hidden_states[-1]) + features = self.dropout(features) + features = self.classifier(features) + return features + + +@auto_docstring( + custom_intro=""" + MobileViT model with a semantic segmentation head on top, e.g. for Pascal VOC. + """ +) +class MobileViTForSemanticSegmentation(MobileViTPreTrainedModel): + def __init__(self, config: MobileViTConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.mobilevit = MobileViTModel(config, expand_output=False) + self.segmentation_head = MobileViTDeepLabV3(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SemanticSegmenterOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> import httpx + >>> from io import BytesIO + >>> import torch + >>> from PIL import Image + >>> from transformers import AutoImageProcessor, MobileViTForSemanticSegmentation + + >>> 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("apple/deeplabv3-mobilevit-small") + >>> model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-small") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # logits are of shape (batch_size, num_labels, height, width) + >>> logits = outputs.logits + ```""" + 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 labels is not None and self.config.num_labels == 1: + raise ValueError("The number of labels should be greater than one") + + outputs = self.mobilevit( + pixel_values, + output_hidden_states=True, # we need the intermediate hidden states + return_dict=return_dict, + ) + + encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1] + + logits = self.segmentation_head(encoder_hidden_states) + + loss = None + if labels is not None: + # upsample logits to the images' original size + upsampled_logits = nn.functional.interpolate( + logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index) + loss = loss_fct(upsampled_logits, labels) + + if not return_dict: + if output_hidden_states: + output = (logits,) + outputs[1:] + else: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SemanticSegmenterOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=None, + ) + + +__all__ = [ + "MobileViTForImageClassification", + "MobileViTForSemanticSegmentation", + "MobileViTModel", + "MobileViTPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/mobilevitv2/__init__.py b/third_party/transformers/src/transformers/models/mobilevitv2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..15e10a1d9815935629fadf9405effdae5a16fbb2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevitv2/__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_mobilevitv2 import * + from .modeling_mobilevitv2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/mobilevitv2/configuration_mobilevitv2.py b/third_party/transformers/src/transformers/models/mobilevitv2/configuration_mobilevitv2.py new file mode 100644 index 0000000000000000000000000000000000000000..df8e48d9a45e5802278e9f7af8922e0f521769f9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevitv2/configuration_mobilevitv2.py @@ -0,0 +1,82 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MobileViTV2 model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="apple/mobilevitv2-1.0") +@strict +class MobileViTV2Config(PreTrainedConfig): + r""" + aspp_out_channels (`int`, *optional*, defaults to 512): + Number of output channels used in the ASPP layer for semantic segmentation. + atrous_rates (`list[int]`, *optional*, defaults to `[6, 12, 18]`): + Dilation (atrous) factors used in the ASPP layer for semantic segmentation. + aspp_dropout_prob (`float`, *optional*, defaults to 0.1): + The dropout ratio for the ASPP layer for semantic segmentation. + n_attn_blocks (`list[int]`, *optional*, defaults to `[2, 4, 3]`): + The number of attention blocks in each MobileViTV2Layer + base_attn_unit_dims (`list[int]`, *optional*, defaults to `[128, 192, 256]`): + The base multiplier for dimensions of attention blocks in each MobileViTV2Layer + width_multiplier (`float`, *optional*, defaults to 1.0): + The width multiplier for MobileViTV2. + ffn_multiplier (`int`, *optional*, defaults to 2): + The FFN multiplier for MobileViTV2. + ffn_dropout (`float`, *optional*, defaults to 0.0): + The dropout between FFN layers. + + Example: + + ```python + >>> from transformers import MobileViTV2Config, MobileViTV2Model + + >>> # Initializing a mobilevitv2-small style configuration + >>> configuration = MobileViTV2Config() + + >>> # Initializing a model from the mobilevitv2-small style configuration + >>> model = MobileViTV2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mobilevitv2" + + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 256 + patch_size: int | list[int] | tuple[int, int] = 2 + expand_ratio: float = 2.0 + hidden_act: str = "swish" + conv_kernel_size: int = 3 + output_stride: int = 32 + classifier_dropout_prob: float | int = 0.1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + aspp_out_channels: int = 512 + atrous_rates: list[int] | tuple[int, ...] = (6, 12, 18) + aspp_dropout_prob: float | int = 0.1 + semantic_loss_ignore_index: int = 255 + n_attn_blocks: list[int] | tuple[int, ...] = (2, 4, 3) + base_attn_unit_dims: list[int] | tuple[int, ...] = (128, 192, 256) + width_multiplier: float | int = 1.0 + ffn_multiplier: int = 2 + attn_dropout: float | int = 0.0 + ffn_dropout: float | int = 0.0 + + +__all__ = ["MobileViTV2Config"] diff --git a/third_party/transformers/src/transformers/models/mobilevitv2/convert_mlcvnets_to_pytorch.py b/third_party/transformers/src/transformers/models/mobilevitv2/convert_mlcvnets_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..e5feb223bf4442632838300c8c9c3756773368c7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevitv2/convert_mlcvnets_to_pytorch.py @@ -0,0 +1,333 @@ +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert MobileViTV2 checkpoints from the ml-cvnets library.""" + +import argparse +import collections +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +import yaml +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import ( + MobileViTImageProcessor, + MobileViTV2Config, + MobileViTV2ForImageClassification, + MobileViTV2ForSemanticSegmentation, +) +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def load_orig_config_file(orig_cfg_file): + print("Loading config file...") + + def flatten_yaml_as_dict(d, parent_key="", sep="."): + items = [] + for k, v in d.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, collections.abc.MutableMapping): + items.extend(flatten_yaml_as_dict(v, new_key, sep=sep).items()) + else: + items.append((new_key, v)) + return dict(items) + + config = argparse.Namespace() + with open(orig_cfg_file, "r") as yaml_file: + try: + cfg = yaml.load(yaml_file, Loader=yaml.FullLoader) + + flat_cfg = flatten_yaml_as_dict(cfg) + for k, v in flat_cfg.items(): + setattr(config, k, v) + except yaml.YAMLError as exc: + logger.error(f"Error while loading config file: {orig_cfg_file}. Error message: {str(exc)}") + return config + + +def get_mobilevitv2_config(task_name, orig_cfg_file): + config = MobileViTV2Config() + + is_segmentation_model = False + + # dataset + if task_name.startswith("imagenet1k_"): + config.num_labels = 1000 + if int(task_name.strip().split("_")[-1]) == 384: + config.image_size = 384 + else: + config.image_size = 256 + filename = "imagenet-1k-id2label.json" + elif task_name.startswith("imagenet21k_to_1k_"): + config.num_labels = 21000 + if int(task_name.strip().split("_")[-1]) == 384: + config.image_size = 384 + else: + config.image_size = 256 + filename = "imagenet-22k-id2label.json" + elif task_name.startswith("ade20k_"): + config.num_labels = 151 + config.image_size = 512 + filename = "ade20k-id2label.json" + is_segmentation_model = True + elif task_name.startswith("voc_"): + config.num_labels = 21 + config.image_size = 512 + filename = "pascal-voc-id2label.json" + is_segmentation_model = True + + # orig_config + orig_config = load_orig_config_file(orig_cfg_file) + assert getattr(orig_config, "model.classification.name", -1) == "mobilevit_v2", "Invalid model" + config.width_multiplier = getattr(orig_config, "model.classification.mitv2.width_multiplier", 1.0) + assert getattr(orig_config, "model.classification.mitv2.attn_norm_layer", -1) == "layer_norm_2d", ( + "Norm layers other than layer_norm_2d is not supported" + ) + config.hidden_act = getattr(orig_config, "model.classification.activation.name", "swish") + # config.image_size == getattr(orig_config, 'sampler.bs.crop_size_width', 256) + + if is_segmentation_model: + config.output_stride = getattr(orig_config, "model.segmentation.output_stride", 16) + if "_deeplabv3" in task_name: + config.atrous_rates = getattr(orig_config, "model.segmentation.deeplabv3.aspp_rates", [12, 24, 36]) + config.aspp_out_channels = getattr(orig_config, "model.segmentation.deeplabv3.aspp_out_channels", 512) + config.aspp_dropout_prob = getattr(orig_config, "model.segmentation.deeplabv3.aspp_dropout", 0.1) + + # id2label + repo_id = "huggingface/label-files" + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + + return config + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +def create_rename_keys(state_dict, base_model=False): + if base_model: + model_prefix = "" + else: + model_prefix = "mobilevitv2." + + rename_keys = [] + for k in state_dict: + if k[:8] == "encoder.": + k_new = k[8:] + else: + k_new = k + + if ".block." in k: + k_new = k_new.replace(".block.", ".") + if ".conv." in k: + k_new = k_new.replace(".conv.", ".convolution.") + if ".norm." in k: + k_new = k_new.replace(".norm.", ".normalization.") + + if "conv_1." in k: + k_new = k_new.replace("conv_1.", f"{model_prefix}conv_stem.") + for i in [1, 2]: + if f"layer_{i}." in k: + k_new = k_new.replace(f"layer_{i}.", f"{model_prefix}encoder.layer.{i - 1}.layer.") + if ".exp_1x1." in k: + k_new = k_new.replace(".exp_1x1.", ".expand_1x1.") + if ".red_1x1." in k: + k_new = k_new.replace(".red_1x1.", ".reduce_1x1.") + + for i in [3, 4, 5]: + if f"layer_{i}.0." in k: + k_new = k_new.replace(f"layer_{i}.0.", f"{model_prefix}encoder.layer.{i - 1}.downsampling_layer.") + if f"layer_{i}.1.local_rep.0." in k: + k_new = k_new.replace(f"layer_{i}.1.local_rep.0.", f"{model_prefix}encoder.layer.{i - 1}.conv_kxk.") + if f"layer_{i}.1.local_rep.1." in k: + k_new = k_new.replace(f"layer_{i}.1.local_rep.1.", f"{model_prefix}encoder.layer.{i - 1}.conv_1x1.") + + for i in [3, 4, 5]: + if i == 3: + j_in = [0, 1] + elif i == 4: + j_in = [0, 1, 2, 3] + elif i == 5: + j_in = [0, 1, 2] + + for j in j_in: + if f"layer_{i}.1.global_rep.{j}." in k: + k_new = k_new.replace( + f"layer_{i}.1.global_rep.{j}.", f"{model_prefix}encoder.layer.{i - 1}.transformer.layer.{j}." + ) + if f"layer_{i}.1.global_rep.{j + 1}." in k: + k_new = k_new.replace( + f"layer_{i}.1.global_rep.{j + 1}.", f"{model_prefix}encoder.layer.{i - 1}.layernorm." + ) + + if f"layer_{i}.1.conv_proj." in k: + k_new = k_new.replace( + f"layer_{i}.1.conv_proj.", f"{model_prefix}encoder.layer.{i - 1}.conv_projection." + ) + + if "pre_norm_attn.0." in k: + k_new = k_new.replace("pre_norm_attn.0.", "layernorm_before.") + if "pre_norm_attn.1." in k: + k_new = k_new.replace("pre_norm_attn.1.", "attention.") + if "pre_norm_ffn.0." in k: + k_new = k_new.replace("pre_norm_ffn.0.", "layernorm_after.") + if "pre_norm_ffn.1." in k: + k_new = k_new.replace("pre_norm_ffn.1.", "ffn.conv1.") + if "pre_norm_ffn.3." in k: + k_new = k_new.replace("pre_norm_ffn.3.", "ffn.conv2.") + + if "classifier.1." in k: + k_new = k_new.replace("classifier.1.", "classifier.") + + if "seg_head." in k: + k_new = k_new.replace("seg_head.", "segmentation_head.") + if ".aspp_layer." in k: + k_new = k_new.replace(".aspp_layer.", ".") + if ".aspp_pool." in k: + k_new = k_new.replace(".aspp_pool.", ".") + + rename_keys.append((k, k_new)) + return rename_keys + + +def remove_unused_keys(state_dict): + """remove unused keys (e.g.: seg_head.aux_head)""" + keys_to_ignore = [] + for k in state_dict: + if k.startswith("seg_head.aux_head."): + keys_to_ignore.append(k) + for k in keys_to_ignore: + state_dict.pop(k, None) + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + # url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_mobilevitv2_checkpoint(task_name, checkpoint_path, orig_config_path, pytorch_dump_folder_path): + """ + Copy/paste/tweak model's weights to our MobileViTV2 structure. + """ + config = get_mobilevitv2_config(task_name, orig_config_path) + + # load original state_dict + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + # load huggingface model + if task_name.startswith("ade20k_") or task_name.startswith("voc_"): + model = MobileViTV2ForSemanticSegmentation(config).eval() + base_model = False + else: + model = MobileViTV2ForImageClassification(config).eval() + base_model = False + + # remove and rename some keys of load the original model + state_dict = checkpoint + remove_unused_keys(state_dict) + rename_keys = create_rename_keys(state_dict, base_model=base_model) + for rename_key_src, rename_key_dest in rename_keys: + rename_key(state_dict, rename_key_src, rename_key_dest) + + # load modified state_dict + model.load_state_dict(state_dict) + + # Check outputs on an image, prepared by MobileViTImageProcessor + image_processor = MobileViTImageProcessor(crop_size=config.image_size, size=config.image_size + 32) + encoding = image_processor(images=prepare_img(), return_tensors="pt") + outputs = model(**encoding) + + # verify classification model + if task_name.startswith("imagenet"): + logits = outputs.logits + predicted_class_idx = logits.argmax(-1).item() + print("Predicted class:", model.config.id2label[predicted_class_idx]) + if task_name.startswith("imagenet1k_256") and config.width_multiplier == 1.0: + # expected_logits for base variant + expected_logits = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01]) + assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4) + + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model {task_name} to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--task", + default="imagenet1k_256", + type=str, + help=( + "Name of the task for which the MobileViTV2 model you'd like to convert is trained on . " + """ + Classification (ImageNet-1k) + - MobileViTV2 (256x256) : imagenet1k_256 + - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384 + - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) : + imagenet21k_to_1k_256 + - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on + ImageNet-1k 384x384) : imagenet21k_to_1k_384 + Segmentation + - ADE20K Dataset : ade20k_deeplabv3 + - Pascal VOC 2012 Dataset: voc_deeplabv3 + """ + ), + choices=[ + "imagenet1k_256", + "imagenet1k_384", + "imagenet21k_to_1k_256", + "imagenet21k_to_1k_384", + "ade20k_deeplabv3", + "voc_deeplabv3", + ], + ) + + parser.add_argument( + "--orig_checkpoint_path", required=True, type=str, help="Path to the original state dict (.pt file)." + ) + parser.add_argument( + "--orig_config_path", + required=True, + type=str, + help="Path to the original config file. yaml.load will be used to load the file, please be wary of which file you're loading.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory." + ) + + args = parser.parse_args() + convert_mobilevitv2_checkpoint( + args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path + ) diff --git a/third_party/transformers/src/transformers/models/mobilevitv2/modeling_mobilevitv2.py b/third_party/transformers/src/transformers/models/mobilevitv2/modeling_mobilevitv2.py new file mode 100644 index 0000000000000000000000000000000000000000..2e28695ae1f2d1521a5b606903335bf3985cadb6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mobilevitv2/modeling_mobilevitv2.py @@ -0,0 +1,942 @@ +# Copyright 2023 Apple Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Original license: https://github.com/apple/ml-cvnets/blob/main/LICENSE +"""PyTorch MobileViTV2 model.""" + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithNoAttention, + BaseModelOutputWithPoolingAndNoAttention, + ImageClassifierOutputWithNoAttention, + SemanticSegmenterOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_mobilevitv2 import MobileViTV2Config + + +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.make_divisible +def make_divisible(value: int, divisor: int = 8, min_value: int | None = None) -> int: + """ + Ensure that all layers have a channel count that is divisible by `divisor`. + """ + if min_value is None: + min_value = divisor + new_value = max(min_value, int(value + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_value < 0.9 * value: + new_value += divisor + return int(new_value) + + +def clip(value: float, min_val: float = float("-inf"), max_val: float = float("inf")) -> float: + return max(min_val, min(max_val, value)) + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.MobileViTConvLayer with MobileViT->MobileViTV2 +class MobileViTV2ConvLayer(nn.Module): + def __init__( + self, + config: MobileViTV2Config, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + groups: int = 1, + bias: bool = False, + dilation: int = 1, + use_normalization: bool = True, + use_activation: bool | str = True, + ) -> None: + super().__init__() + padding = int((kernel_size - 1) / 2) * dilation + + if in_channels % groups != 0: + raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups.") + if out_channels % groups != 0: + raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups.") + + self.convolution = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode="zeros", + ) + + if use_normalization: + self.normalization = nn.BatchNorm2d( + num_features=out_channels, + eps=1e-5, + momentum=0.1, + affine=True, + track_running_stats=True, + ) + else: + self.normalization = None + + if use_activation: + if isinstance(use_activation, str): + self.activation = ACT2FN[use_activation] + elif isinstance(config.hidden_act, str): + self.activation = ACT2FN[config.hidden_act] + else: + self.activation = config.hidden_act + else: + self.activation = None + + def forward(self, features: torch.Tensor) -> torch.Tensor: + features = self.convolution(features) + if self.normalization is not None: + features = self.normalization(features) + if self.activation is not None: + features = self.activation(features) + return features + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.MobileViTInvertedResidual with MobileViT->MobileViTV2 +class MobileViTV2InvertedResidual(nn.Module): + """ + Inverted residual block (MobileNetv2): https://huggingface.co/papers/1801.04381 + """ + + def __init__( + self, config: MobileViTV2Config, in_channels: int, out_channels: int, stride: int, dilation: int = 1 + ) -> None: + super().__init__() + expanded_channels = make_divisible(int(round(in_channels * config.expand_ratio)), 8) + + if stride not in [1, 2]: + raise ValueError(f"Invalid stride {stride}.") + + self.use_residual = (stride == 1) and (in_channels == out_channels) + + self.expand_1x1 = MobileViTV2ConvLayer( + config, in_channels=in_channels, out_channels=expanded_channels, kernel_size=1 + ) + + self.conv_3x3 = MobileViTV2ConvLayer( + config, + in_channels=expanded_channels, + out_channels=expanded_channels, + kernel_size=3, + stride=stride, + groups=expanded_channels, + dilation=dilation, + ) + + self.reduce_1x1 = MobileViTV2ConvLayer( + config, + in_channels=expanded_channels, + out_channels=out_channels, + kernel_size=1, + use_activation=False, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + residual = features + + features = self.expand_1x1(features) + features = self.conv_3x3(features) + features = self.reduce_1x1(features) + + return residual + features if self.use_residual else features + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.MobileViTMobileNetLayer with MobileViT->MobileViTV2 +class MobileViTV2MobileNetLayer(nn.Module): + def __init__( + self, config: MobileViTV2Config, in_channels: int, out_channels: int, stride: int = 1, num_stages: int = 1 + ) -> None: + super().__init__() + + self.layer = nn.ModuleList() + for i in range(num_stages): + layer = MobileViTV2InvertedResidual( + config, + in_channels=in_channels, + out_channels=out_channels, + stride=stride if i == 0 else 1, + ) + self.layer.append(layer) + in_channels = out_channels + + def forward(self, features: torch.Tensor) -> torch.Tensor: + for layer_module in self.layer: + features = layer_module(features) + return features + + +class MobileViTV2LinearSelfAttention(nn.Module): + """ + This layer applies a self-attention with linear complexity, as described in MobileViTV2 paper: + https://huggingface.co/papers/2206.02680 + + Args: + config (`MobileVitv2Config`): + Model configuration object + embed_dim (`int`): + `input_channels` from an expected input of size :math:`(batch_size, input_channels, height, width)` + """ + + def __init__(self, config: MobileViTV2Config, embed_dim: int) -> None: + super().__init__() + + self.qkv_proj = MobileViTV2ConvLayer( + config=config, + in_channels=embed_dim, + out_channels=1 + (2 * embed_dim), + bias=True, + kernel_size=1, + use_normalization=False, + use_activation=False, + ) + + self.attn_dropout = nn.Dropout(p=config.attn_dropout) + self.out_proj = MobileViTV2ConvLayer( + config=config, + in_channels=embed_dim, + out_channels=embed_dim, + bias=True, + kernel_size=1, + use_normalization=False, + use_activation=False, + ) + self.embed_dim = embed_dim + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # (batch_size, embed_dim, num_pixels_in_patch, num_patches) --> (batch_size, 1+2*embed_dim, num_pixels_in_patch, num_patches) + qkv = self.qkv_proj(hidden_states) + + # Project hidden_states into query, key and value + # Query --> [batch_size, 1, num_pixels_in_patch, num_patches] + # value, key --> [batch_size, embed_dim, num_pixels_in_patch, num_patches] + query, key, value = torch.split(qkv, split_size_or_sections=[1, self.embed_dim, self.embed_dim], dim=1) + + # apply softmax along num_patches dimension + context_scores = torch.nn.functional.softmax(query, dim=-1) + context_scores = self.attn_dropout(context_scores) + + # Compute context vector + # [batch_size, embed_dim, num_pixels_in_patch, num_patches] x [batch_size, 1, num_pixels_in_patch, num_patches] -> [batch_size, embed_dim, num_pixels_in_patch, num_patches] + context_vector = key * context_scores + # [batch_size, embed_dim, num_pixels_in_patch, num_patches] --> [batch_size, embed_dim, num_pixels_in_patch, 1] + context_vector = torch.sum(context_vector, dim=-1, keepdim=True) + + # combine context vector with values + # [batch_size, embed_dim, num_pixels_in_patch, num_patches] * [batch_size, embed_dim, num_pixels_in_patch, 1] --> [batch_size, embed_dim, num_pixels_in_patch, num_patches] + out = torch.nn.functional.relu(value) * context_vector.expand_as(value) + out = self.out_proj(out) + return out + + +class MobileViTV2FFN(nn.Module): + def __init__( + self, + config: MobileViTV2Config, + embed_dim: int, + ffn_latent_dim: int, + ffn_dropout: float = 0.0, + ) -> None: + super().__init__() + self.conv1 = MobileViTV2ConvLayer( + config=config, + in_channels=embed_dim, + out_channels=ffn_latent_dim, + kernel_size=1, + stride=1, + bias=True, + use_normalization=False, + use_activation=True, + ) + self.dropout1 = nn.Dropout(ffn_dropout) + + self.conv2 = MobileViTV2ConvLayer( + config=config, + in_channels=ffn_latent_dim, + out_channels=embed_dim, + kernel_size=1, + stride=1, + bias=True, + use_normalization=False, + use_activation=False, + ) + self.dropout2 = nn.Dropout(ffn_dropout) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv1(hidden_states) + hidden_states = self.dropout1(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.dropout2(hidden_states) + return hidden_states + + +class MobileViTV2TransformerLayer(nn.Module): + def __init__( + self, + config: MobileViTV2Config, + embed_dim: int, + ffn_latent_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.layernorm_before = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=config.layer_norm_eps) + self.attention = MobileViTV2LinearSelfAttention(config, embed_dim) + self.dropout1 = nn.Dropout(p=dropout) + self.layernorm_after = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=config.layer_norm_eps) + self.ffn = MobileViTV2FFN(config, embed_dim, ffn_latent_dim, config.ffn_dropout) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + layernorm_1_out = self.layernorm_before(hidden_states) + attention_output = self.attention(layernorm_1_out) + hidden_states = attention_output + hidden_states + + layer_output = self.layernorm_after(hidden_states) + layer_output = self.ffn(layer_output) + + layer_output = layer_output + hidden_states + return layer_output + + +class MobileViTV2Transformer(nn.Module): + def __init__(self, config: MobileViTV2Config, n_layers: int, d_model: int) -> None: + super().__init__() + + ffn_multiplier = config.ffn_multiplier + + ffn_dims = [ffn_multiplier * d_model] * n_layers + + # ensure that dims are multiple of 16 + ffn_dims = [int((d // 16) * 16) for d in ffn_dims] + + self.layer = nn.ModuleList() + for block_idx in range(n_layers): + transformer_layer = MobileViTV2TransformerLayer( + config, embed_dim=d_model, ffn_latent_dim=ffn_dims[block_idx] + ) + self.layer.append(transformer_layer) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for layer_module in self.layer: + hidden_states = layer_module(hidden_states) + return hidden_states + + +class MobileViTV2Layer(GradientCheckpointingLayer): + """ + MobileViTV2 layer: https://huggingface.co/papers/2206.02680 + """ + + def __init__( + self, + config: MobileViTV2Config, + in_channels: int, + out_channels: int, + attn_unit_dim: int, + n_attn_blocks: int = 2, + dilation: int = 1, + stride: int = 2, + ) -> None: + super().__init__() + self.patch_width = config.patch_size + self.patch_height = config.patch_size + + cnn_out_dim = attn_unit_dim + + if stride == 2: + self.downsampling_layer = MobileViTV2InvertedResidual( + config, + in_channels=in_channels, + out_channels=out_channels, + stride=stride if dilation == 1 else 1, + dilation=dilation // 2 if dilation > 1 else 1, + ) + in_channels = out_channels + else: + self.downsampling_layer = None + + # Local representations + self.conv_kxk = MobileViTV2ConvLayer( + config, + in_channels=in_channels, + out_channels=in_channels, + kernel_size=config.conv_kernel_size, + groups=in_channels, + ) + self.conv_1x1 = MobileViTV2ConvLayer( + config, + in_channels=in_channels, + out_channels=cnn_out_dim, + kernel_size=1, + use_normalization=False, + use_activation=False, + ) + + # Global representations + self.transformer = MobileViTV2Transformer(config, d_model=attn_unit_dim, n_layers=n_attn_blocks) + + # self.layernorm = MobileViTV2LayerNorm2D(attn_unit_dim, eps=config.layer_norm_eps) + self.layernorm = nn.GroupNorm(num_groups=1, num_channels=attn_unit_dim, eps=config.layer_norm_eps) + + # Fusion + self.conv_projection = MobileViTV2ConvLayer( + config, + in_channels=cnn_out_dim, + out_channels=in_channels, + kernel_size=1, + use_normalization=True, + use_activation=False, + ) + + def unfolding(self, feature_map: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]: + batch_size, in_channels, img_height, img_width = feature_map.shape + patches = nn.functional.unfold( + feature_map, + kernel_size=(self.patch_height, self.patch_width), + stride=(self.patch_height, self.patch_width), + ) + patches = patches.reshape(batch_size, in_channels, self.patch_height * self.patch_width, -1) + + return patches, (img_height, img_width) + + def folding(self, patches: torch.Tensor, output_size: tuple[int, int]) -> torch.Tensor: + batch_size, in_dim, patch_size, n_patches = patches.shape + patches = patches.reshape(batch_size, in_dim * patch_size, n_patches) + + feature_map = nn.functional.fold( + patches, + output_size=output_size, + kernel_size=(self.patch_height, self.patch_width), + stride=(self.patch_height, self.patch_width), + ) + + return feature_map + + def forward(self, features: torch.Tensor) -> torch.Tensor: + # reduce spatial dimensions if needed + if self.downsampling_layer: + features = self.downsampling_layer(features) + + # local representation + features = self.conv_kxk(features) + features = self.conv_1x1(features) + + # convert feature map to patches + patches, output_size = self.unfolding(features) + + # learn global representations + patches = self.transformer(patches) + patches = self.layernorm(patches) + + # convert patches back to feature maps + # [batch_size, patch_height, patch_width, input_dim] --> [batch_size, input_dim, patch_height, patch_width] + features = self.folding(patches, output_size) + + features = self.conv_projection(features) + return features + + +class MobileViTV2Encoder(nn.Module): + def __init__(self, config: MobileViTV2Config) -> None: + super().__init__() + self.config = config + + self.layer = nn.ModuleList() + self.gradient_checkpointing = False + + # segmentation architectures like DeepLab and PSPNet modify the strides + # of the classification backbones + dilate_layer_4 = dilate_layer_5 = False + if config.output_stride == 8: + dilate_layer_4 = True + dilate_layer_5 = True + elif config.output_stride == 16: + dilate_layer_5 = True + + dilation = 1 + + layer_0_dim = make_divisible( + clip(value=32 * config.width_multiplier, min_val=16, max_val=64), divisor=8, min_value=16 + ) + + layer_1_dim = make_divisible(64 * config.width_multiplier, divisor=16) + layer_2_dim = make_divisible(128 * config.width_multiplier, divisor=8) + layer_3_dim = make_divisible(256 * config.width_multiplier, divisor=8) + layer_4_dim = make_divisible(384 * config.width_multiplier, divisor=8) + layer_5_dim = make_divisible(512 * config.width_multiplier, divisor=8) + + layer_1 = MobileViTV2MobileNetLayer( + config, + in_channels=layer_0_dim, + out_channels=layer_1_dim, + stride=1, + num_stages=1, + ) + self.layer.append(layer_1) + + layer_2 = MobileViTV2MobileNetLayer( + config, + in_channels=layer_1_dim, + out_channels=layer_2_dim, + stride=2, + num_stages=2, + ) + self.layer.append(layer_2) + + layer_3 = MobileViTV2Layer( + config, + in_channels=layer_2_dim, + out_channels=layer_3_dim, + attn_unit_dim=make_divisible(config.base_attn_unit_dims[0] * config.width_multiplier, divisor=8), + n_attn_blocks=config.n_attn_blocks[0], + ) + self.layer.append(layer_3) + + if dilate_layer_4: + dilation *= 2 + + layer_4 = MobileViTV2Layer( + config, + in_channels=layer_3_dim, + out_channels=layer_4_dim, + attn_unit_dim=make_divisible(config.base_attn_unit_dims[1] * config.width_multiplier, divisor=8), + n_attn_blocks=config.n_attn_blocks[1], + dilation=dilation, + ) + self.layer.append(layer_4) + + if dilate_layer_5: + dilation *= 2 + + layer_5 = MobileViTV2Layer( + config, + in_channels=layer_4_dim, + out_channels=layer_5_dim, + attn_unit_dim=make_divisible(config.base_attn_unit_dims[2] * config.width_multiplier, divisor=8), + n_attn_blocks=config.n_attn_blocks[2], + dilation=dilation, + ) + self.layer.append(layer_5) + + def forward( + self, + hidden_states: torch.Tensor, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> tuple | BaseModelOutputWithNoAttention: + all_hidden_states = () if output_hidden_states else None + + for i, layer_module in enumerate(self.layer): + hidden_states = layer_module(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] if v is not None) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=all_hidden_states) + + +@auto_docstring +class MobileViTV2PreTrainedModel(PreTrainedModel): + config: MobileViTV2Config + base_model_prefix = "mobilevitv2" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MobileViTV2Layer"] + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + """Initialize the weights""" + if isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + if getattr(module, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + elif isinstance(module, nn.GroupNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + + +@auto_docstring +class MobileViTV2Model(MobileViTV2PreTrainedModel): + def __init__(self, config: MobileViTV2Config, expand_output: bool = True): + r""" + expand_output (`bool`, *optional*, defaults to `True`): + Whether to expand the output of the model. If `True`, the model will output pooled features in addition to + hidden states. If `False`, only the hidden states will be returned. + """ + super().__init__(config) + self.config = config + self.expand_output = expand_output + + layer_0_dim = make_divisible( + clip(value=32 * config.width_multiplier, min_val=16, max_val=64), divisor=8, min_value=16 + ) + + self.conv_stem = MobileViTV2ConvLayer( + config, + in_channels=config.num_channels, + out_channels=layer_0_dim, + kernel_size=3, + stride=2, + use_normalization=True, + use_activation=True, + ) + self.encoder = MobileViTV2Encoder(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPoolingAndNoAttention: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + embedding_output = self.conv_stem(pixel_values) + + encoder_outputs = self.encoder( + embedding_output, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if self.expand_output: + last_hidden_state = encoder_outputs[0] + + # global average pooling: (batch_size, channels, height, width) -> (batch_size, channels) + pooled_output = torch.mean(last_hidden_state, dim=[-2, -1], keepdim=False) + else: + last_hidden_state = encoder_outputs[0] + pooled_output = None + + if not return_dict: + output = (last_hidden_state, pooled_output) if pooled_output is not None else (last_hidden_state,) + return output + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndNoAttention( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + MobileViTV2 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for + ImageNet. + """ +) +class MobileViTV2ForImageClassification(MobileViTV2PreTrainedModel): + def __init__(self, config: MobileViTV2Config) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.mobilevitv2 = MobileViTV2Model(config) + + out_channels = make_divisible(512 * config.width_multiplier, divisor=8) # layer 5 output dimension + # Classifier head + self.classifier = ( + nn.Linear(in_features=out_channels, out_features=config.num_labels) + if config.num_labels > 0 + else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + labels: torch.Tensor | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ImageClassifierOutputWithNoAttention: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss). If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.mobilevitv2(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict) + + pooled_output = outputs.pooler_output if return_dict else outputs[1] + + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + loss = self.loss_function(labels, logits, self.config) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutputWithNoAttention( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + ) + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.MobileViTASPPPooling with MobileViT->MobileViTV2 +class MobileViTV2ASPPPooling(nn.Module): + def __init__(self, config: MobileViTV2Config, in_channels: int, out_channels: int) -> None: + super().__init__() + + self.global_pool = nn.AdaptiveAvgPool2d(output_size=1) + + self.conv_1x1 = MobileViTV2ConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + use_normalization=True, + use_activation="relu", + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + spatial_size = features.shape[-2:] + features = self.global_pool(features) + features = self.conv_1x1(features) + features = nn.functional.interpolate(features, size=spatial_size, mode="bilinear", align_corners=False) + return features + + +class MobileViTV2ASPP(nn.Module): + """ + ASPP module defined in DeepLab papers: https://huggingface.co/papers/1606.00915, https://huggingface.co/papers/1706.05587 + """ + + def __init__(self, config: MobileViTV2Config) -> None: + super().__init__() + + encoder_out_channels = make_divisible(512 * config.width_multiplier, divisor=8) # layer 5 output dimension + in_channels = encoder_out_channels + out_channels = config.aspp_out_channels + + if len(config.atrous_rates) != 3: + raise ValueError("Expected 3 values for atrous_rates") + + self.convs = nn.ModuleList() + + in_projection = MobileViTV2ConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + use_activation="relu", + ) + self.convs.append(in_projection) + + self.convs.extend( + [ + MobileViTV2ConvLayer( + config, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + dilation=rate, + use_activation="relu", + ) + for rate in config.atrous_rates + ] + ) + + pool_layer = MobileViTV2ASPPPooling(config, in_channels, out_channels) + self.convs.append(pool_layer) + + self.project = MobileViTV2ConvLayer( + config, in_channels=5 * out_channels, out_channels=out_channels, kernel_size=1, use_activation="relu" + ) + + self.dropout = nn.Dropout(p=config.aspp_dropout_prob) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + pyramid = [] + for conv in self.convs: + pyramid.append(conv(features)) + pyramid = torch.cat(pyramid, dim=1) + + pooled_features = self.project(pyramid) + pooled_features = self.dropout(pooled_features) + return pooled_features + + +# Copied from transformers.models.mobilevit.modeling_mobilevit.MobileViTDeepLabV3 with MobileViT->MobileViTV2 +class MobileViTV2DeepLabV3(nn.Module): + """ + DeepLabv3 architecture: https://huggingface.co/papers/1706.05587 + """ + + def __init__(self, config: MobileViTV2Config) -> None: + super().__init__() + self.aspp = MobileViTV2ASPP(config) + + self.dropout = nn.Dropout2d(config.classifier_dropout_prob) + + self.classifier = MobileViTV2ConvLayer( + config, + in_channels=config.aspp_out_channels, + out_channels=config.num_labels, + kernel_size=1, + use_normalization=False, + use_activation=False, + bias=True, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + features = self.aspp(hidden_states[-1]) + features = self.dropout(features) + features = self.classifier(features) + return features + + +@auto_docstring( + custom_intro=""" + MobileViTV2 model with a semantic segmentation head on top, e.g. for Pascal VOC. + """ +) +class MobileViTV2ForSemanticSegmentation(MobileViTV2PreTrainedModel): + def __init__(self, config: MobileViTV2Config) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.mobilevitv2 = MobileViTV2Model(config, expand_output=False) + self.segmentation_head = MobileViTV2DeepLabV3(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SemanticSegmenterOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> import httpx + >>> from io import BytesIO + >>> import torch + >>> from PIL import Image + >>> from transformers import AutoImageProcessor, MobileViTV2ForSemanticSegmentation + + >>> 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("apple/mobilevitv2-1.0-imagenet1k-256") + >>> model = MobileViTV2ForSemanticSegmentation.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # logits are of shape (batch_size, num_labels, height, width) + >>> logits = outputs.logits + ```""" + 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 labels is not None and self.config.num_labels == 1: + raise ValueError("The number of labels should be greater than one") + + outputs = self.mobilevitv2( + pixel_values, + output_hidden_states=True, # we need the intermediate hidden states + return_dict=return_dict, + ) + + encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1] + + logits = self.segmentation_head(encoder_hidden_states) + + loss = None + if labels is not None: + # upsample logits to the images' original size + upsampled_logits = nn.functional.interpolate( + logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index) + loss = loss_fct(upsampled_logits, labels) + + if not return_dict: + if output_hidden_states: + output = (logits,) + outputs[1:] + else: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SemanticSegmenterOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=None, + ) + + +__all__ = [ + "MobileViTV2ForImageClassification", + "MobileViTV2ForSemanticSegmentation", + "MobileViTV2Model", + "MobileViTV2PreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/modernvbert/__init__.py b/third_party/transformers/src/transformers/models/modernvbert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d74edb24382cc318866799a4e2817b5f10d7654e --- /dev/null +++ b/third_party/transformers/src/transformers/models/modernvbert/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2026 Illuin Technology and contributors, and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_modernvbert import * + from .image_processing_modernvbert import * + from .image_processing_modernvbert_fast import * + from .modeling_modernvbert import * + from .processing_modernvbert import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/modernvbert/configuration_modernvbert.py b/third_party/transformers/src/transformers/models/modernvbert/configuration_modernvbert.py new file mode 100644 index 0000000000000000000000000000000000000000..7b134f7b38698f9857b0aba2485823da4f192abe --- /dev/null +++ b/third_party/transformers/src/transformers/models/modernvbert/configuration_modernvbert.py @@ -0,0 +1,88 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/modernvbert/modular_modernvbert.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_modernvbert.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 Illuin Technology and contributors, and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Literal + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig + + +@auto_docstring(checkpoint="ModernVBERT/modernvbert") +@strict +class ModernVBertConfig(PreTrainedConfig): + r""" + pixel_shuffle_factor (`int | None`, *optional*, defaults to 4): + Scale factor used by any pixel-shuffle / upsampling operations in the vision head. + initializer_cutoff_factor (`float | None`, *optional*, defaults to 2.0): + The cutoff factor for the truncated_normal_initializer for initializing all weight matrices. + classifier_pooling (`Literal["cls", "mean"]`, *optional*, defaults to `"cls"`): + The pooling strategy to use for classification tasks. + classifier_bias (`bool | None`, *optional*, defaults to `False`): + Whether to add a bias term to the classification head + + Example: + ```python + >>> from transformers import ModernVBertConfig + + >>> # Initializing configuration + >>> configuration = ModernVBertConfig() + + >>> # Initializing a model from the configuration (model class is implemented in + >>> # `modernvbert.modeling_modernvbert`) + + >>> from transformers import ModernVBertModel + >>> model = ModernVBertModel(configuration) + + >>> # Accessing the model configuration + >>> cfg = model.config + ```""" + + model_type = "modernvbert" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + text_config: PreTrainedConfig | dict | None = None + vision_config: PreTrainedConfig | dict | None = None + image_token_id: int = 50407 + pixel_shuffle_factor: int = 4 + initializer_range: float = 0.02 + initializer_cutoff_factor: float = 2.0 + classifier_pooling: Literal["cls", "mean"] = "cls" + classifier_dropout: float | int = 0.0 + classifier_bias: bool = False + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if self.text_config is None: + self.text_config = CONFIG_MAPPING["modernbert"]() + elif isinstance(self.text_config, dict): + self.text_config = CONFIG_MAPPING["modernbert"](**self.text_config) + + if self.vision_config is None: + self.vision_config = CONFIG_MAPPING["siglip_vision_model"]() + elif isinstance(self.vision_config, dict): + self.vision_config = CONFIG_MAPPING["siglip_vision_model"](**self.vision_config) + + super().__post_init__(**kwargs) + + +__all__ = ["ModernVBertConfig"] diff --git a/third_party/transformers/src/transformers/models/modernvbert/modeling_modernvbert.py b/third_party/transformers/src/transformers/models/modernvbert/modeling_modernvbert.py new file mode 100644 index 0000000000000000000000000000000000000000..b12a6a57f5f6fdc3a5dcb42927524bba6c99c889 --- /dev/null +++ b/third_party/transformers/src/transformers/models/modernvbert/modeling_modernvbert.py @@ -0,0 +1,696 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/modernvbert/modular_modernvbert.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_modernvbert.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 Illuin Technology and contributors, and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + MaskedLMOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils.generic import can_return_tuple +from ..auto import AutoModel +from .configuration_modernvbert import ModernVBertConfig + + +@dataclass +class ModernVBertBaseModelOutput(BaseModelOutput): + """ + Base class for ModernVBERT model's outputs. + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder + """ + + last_hidden_state: torch.FloatTensor = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +class ModernVBertMaskedLMOutput(MaskedLMOutput): + """ + Base class for ModernVBERT model's outputs with masked language modeling loss. + Args: + loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + Masked language modeling (MLM) loss. + logits (`torch.FloatTensor`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + image_hidden_states: torch.FloatTensor | None = None + + +class ModernVBertConnector(nn.Module): + """ + Connector module for ModernVBERT. It performs a pixel shuffle operation followed by a linear projection to match the text model's hidden size. + Based on https://pytorch.org/docs/stable/generated/torch.nn.PixelShuffle.html + """ + + def __init__(self, config): + super().__init__() + self.pixel_shuffle_factor = config.pixel_shuffle_factor + self.modality_projection = nn.Linear( + config.vision_config.hidden_size * (config.pixel_shuffle_factor**2), + config.text_config.hidden_size, + bias=False, + ) + + def pixel_shuffle(self, image_hidden_states, pixel_shuffle_factor): + batch_size, seq_length, embed_dim = image_hidden_states.size() + height = width = int(seq_length**0.5) + image_hidden_states = image_hidden_states.view(batch_size, height, width, embed_dim) + image_hidden_states = image_hidden_states.view( + batch_size, height, int(width / pixel_shuffle_factor), embed_dim * pixel_shuffle_factor + ) + image_hidden_states = image_hidden_states.permute(0, 2, 1, 3) + image_hidden_states = image_hidden_states.reshape( + batch_size, + int(width / pixel_shuffle_factor), + int(height / pixel_shuffle_factor), + embed_dim * (pixel_shuffle_factor**2), + ) + image_hidden_states = image_hidden_states.permute(0, 2, 1, 3) + return image_hidden_states.reshape( + batch_size, int(seq_length / (pixel_shuffle_factor**2)), embed_dim * (pixel_shuffle_factor**2) + ) + + def forward(self, image_hidden_states): + image_hidden_states = self.pixel_shuffle(image_hidden_states, self.pixel_shuffle_factor) + return self.modality_projection(image_hidden_states) + + +@auto_docstring +class ModernVBertPreTrainedModel(PreTrainedModel): + config: ModernVBertConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = [] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + config_class = ModernVBertConfig + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + + def init_weight(module: nn.Module, std: float): + cutoff_factor = getattr(self.config, "initializer_cutoff_factor", 2.0) + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-cutoff_factor * std, + b=cutoff_factor * std, + ) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + if module.bias is not None: + init.zeros_(module.bias) + + if isinstance(module, ModernVBertConnector): + out_std = self.config.initializer_range / math.sqrt(2.0 * self.config.text_config.num_hidden_layers) + init_weight(module.modality_projection, out_std) + elif isinstance(module, ModernVBertForMaskedLM): + out_std = self.config.initializer_range / math.sqrt(2.0 * self.config.text_config.num_hidden_layers) + init_weight(module.lm_head, out_std) + elif isinstance( + module, + ( + ModernVBertForSequenceClassification, + ModernVBertForTokenClassification, + ), + ): + final_out_std = self.config.initializer_range / math.sqrt(self.config.text_config.hidden_size) + init_weight(module.classifier, final_out_std) + + +@auto_docstring( + custom_intro=""" + ModernVBertModel is a model that combines a vision encoder (SigLIP) and a text encoder (ModernBert). + + ModernVBert is the base model of the visual retriver ColModernVBert, and was introduced in the following paper: + [*ModernVBERT: Towards Smaller Visual Document Retrievers*](https://arxiv.org/abs/2510.01149). + """ +) +class ModernVBertModel(ModernVBertPreTrainedModel): + """ + A subclass of Idefics3Model. We do *not* remove or block the call to inputs_merger + in forward. Instead, we override inputs_merger here with custom logic. + """ + + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + self.padding_idx = self.config.text_config.pad_token_id + self.vocab_size = self.config.text_config.vocab_size + self.vision_model = AutoModel.from_config(config.vision_config) + + # init components + self.connector = ModernVBertConnector(config) + self.text_model = AutoModel.from_config(config.text_config) + + self.image_seq_len = int( + ((config.vision_config.image_size // config.vision_config.patch_size) ** 2) + / (config.pixel_shuffle_factor**2) + ) + self.image_token_id = self.config.image_token_id + + self.post_init() + + def get_input_embeddings(self): + return self.text_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.text_model.set_input_embeddings(value) + + def inputs_merger( + self, input_ids: torch.LongTensor, inputs_embeds: torch.Tensor, image_hidden_states: torch.Tensor + ): + """ + This method aims at merging the token embeddings with the image hidden states into one single sequence of vectors that are fed to the transformer LM. + The merging happens as follows: + - The text token sequence is: `tok_1 tok_2 tok_3 ... tok_4`. + - We get the image hidden states for the image through the vision encoder and that hidden state, after a pixel shuffle operation, is then projected into the text embedding space. + We thus have a sequence of image hidden states of size (1, image_seq_len, hidden_dim), where 1 is for batch_size of 1 image and hidden_dim is the hidden_dim of the LM transformer. + - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM. + - To fit the format of that sequence, `input_ids`, `inputs_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states. + """ + _, patch_size, _ = image_hidden_states.shape + + if input_ids is None: + image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + image_mask = image_mask[..., 0] # slice off the hidden dim + else: + image_mask = input_ids == self.config.image_token_id + + num_image_tokens = image_mask.sum(dim=1) + torch_compilable_check( + torch.all(num_image_tokens % patch_size == 0), + "At least one sample has tokens not divisible by patch_size.", + ) + blocks_per_sample = num_image_tokens // patch_size + + offsets = torch.nn.functional.pad(blocks_per_sample.cumsum(dim=0), (1, 0), value=0) + block_offset = offsets[:-1] + row_cum = image_mask.cumsum(dim=-1) + chunk_idx = (row_cum - 1) // patch_size + local_idx = (row_cum - 1) % patch_size + block_idx = block_offset.unsqueeze(1) + chunk_idx + + image_embeds = torch.zeros_like(inputs_embeds) + image_embeds[image_mask] = image_hidden_states[block_idx[image_mask], local_idx[image_mask], :] + + merged_embeds = torch.where(image_mask.unsqueeze(-1), image_embeds, inputs_embeds) + return merged_embeds + + @can_return_tuple + @auto_docstring( + custom_intro="Encodes images into continuous embeddings that can be forwarded to the language model." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + batch_size, num_images, num_channels, height, width = pixel_values.shape + pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility + pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:]) + + # Remove padding images - padding images are full 0. + nb_values_per_image = pixel_values.shape[1:].numel() + real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image + + # If no images, leave one empty image. + real_images_inds[0] |= ~torch.any(real_images_inds) + + pixel_values = pixel_values[real_images_inds].contiguous() + # Handle the vision attention mask + if pixel_attention_mask is None: + pixel_attention_mask = torch.ones( + size=[pixel_values.shape[i] for i in (0, 2, 3)], + dtype=torch.bool, + device=pixel_values.device, + ) + else: + # Remove padding images from the mask + pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:]) + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + patch_size = self.config.vision_config.patch_size + patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size) + patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size) + patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() + + # Get sequence from the vision encoder + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state + + # Modality projection & resampling + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | ModernVBertBaseModelOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + """ + + if inputs_embeds is None: + inputs_embeds = self.text_model.get_input_embeddings()(input_ids).to(input_ids.device) + + # Images processing + if pixel_values is not None: + image_hidden_states = self.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask + ).pooler_output + + # Merge image and text embeddings + if image_hidden_states is not None: + image_hidden_states = image_hidden_states.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device) + inputs_embeds = self.inputs_merger( + input_ids=input_ids, inputs_embeds=inputs_embeds, image_hidden_states=image_hidden_states + ) + + # Language model pass + outputs = self.text_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + + return ModernVBertBaseModelOutput( + last_hidden_state=outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_hidden_states, + ) + + +class ModernVBertPredictionHead(nn.Module): + def __init__(self, config: ModernVBertConfig): + super().__init__() + self.config = config + self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.classifier_bias) + self.act = ACT2FN[config.classifier_activation] + self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(self.act(self.dense(hidden_states))) + + +@auto_docstring +class ModernVBertForMaskedLM(ModernVBertPreTrainedModel): + _tied_weights_keys = {"lm_head.weight": "model.text_model.embeddings.tok_embeddings.weight"} + + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + + self.vocab_size = config.text_config.vocab_size + + self.model = ModernVBertModel(config) + self.projection_head = ModernVBertPredictionHead(config.text_config) + self.lm_head = nn.Linear(config.text_config.hidden_size, self.vocab_size, bias=config.text_config.decoder_bias) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | ModernVBertMaskedLMOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + hidden_states = outputs[0] + + logits = self.lm_head(self.projection_head(hidden_states)) + + loss = None + if labels is not None: + criterion = CrossEntropyLoss() + loss = criterion(logits.view(-1, self.vocab_size), labels.view(-1)) + + return ModernVBertMaskedLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + The ModernVBert Model with a sequence classification head on top that performs pooling. + """ +) +class ModernVBertForSequenceClassification(ModernVBertPreTrainedModel): + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + + self.model = ModernVBertModel(config) + self.head = ModernVBertPredictionHead(config.text_config) + self.drop = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SequenceClassifierOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + last_hidden_state = outputs[0] + + if self.config.classifier_pooling == "cls": + last_hidden_state = last_hidden_state[:, 0] + elif self.config.classifier_pooling == "mean": + if inputs_embeds is not None: + batch_size, seq_len = inputs_embeds.shape[:2] + else: + batch_size, seq_len = input_ids.shape[:2] + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) + last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum( + dim=1, keepdim=True + ) + + pooled_output = self.head(last_hidden_state) + pooled_output = self.drop(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The ModernVBert Model with a token classification head on top, e.g. for Named Entity Recognition (NER) tasks. + """ +) +class ModernVBertForTokenClassification(ModernVBertPreTrainedModel): + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + self.num_labels = config.num_labels + + self.model = ModernVBertModel(config) + self.head = ModernVBertPredictionHead(config.text_config) + self.drop = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | TokenClassifierOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + last_hidden_state = outputs[0] + + last_hidden_state = self.head(last_hidden_state) + last_hidden_state = self.drop(last_hidden_state) + logits = self.classifier(last_hidden_state) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "ModernVBertPreTrainedModel", + "ModernVBertModel", + "ModernVBertForMaskedLM", + "ModernVBertForSequenceClassification", + "ModernVBertForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/modernvbert/modular_modernvbert.py b/third_party/transformers/src/transformers/models/modernvbert/modular_modernvbert.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec870a5a4cc543ac643b4e7e9c09f81fa825b90 --- /dev/null +++ b/third_party/transformers/src/transformers/models/modernvbert/modular_modernvbert.py @@ -0,0 +1,626 @@ +# Copyright 2026 Illuin Technology and contributors, and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import Literal + +import torch +import torch.nn as nn +from huggingface_hub.dataclasses import strict +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...configuration_utils import PreTrainedConfig +from ...modeling_outputs import ( + BaseModelOutput, + MaskedLMOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import can_return_tuple +from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel +from ..modernbert.modeling_modernbert import ModernBertPredictionHead +from ..smolvlm.modeling_smolvlm import SmolVLMModel, SmolVLMPreTrainedModel + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="ModernVBERT/modernvbert") +@strict +class ModernVBertConfig(PreTrainedConfig): + r""" + pixel_shuffle_factor (`int | None`, *optional*, defaults to 4): + Scale factor used by any pixel-shuffle / upsampling operations in the vision head. + initializer_cutoff_factor (`float | None`, *optional*, defaults to 2.0): + The cutoff factor for the truncated_normal_initializer for initializing all weight matrices. + classifier_pooling (`Literal["cls", "mean"]`, *optional*, defaults to `"cls"`): + The pooling strategy to use for classification tasks. + classifier_bias (`bool | None`, *optional*, defaults to `False`): + Whether to add a bias term to the classification head + + Example: + ```python + >>> from transformers import ModernVBertConfig + + >>> # Initializing configuration + >>> configuration = ModernVBertConfig() + + >>> # Initializing a model from the configuration (model class is implemented in + >>> # `modernvbert.modeling_modernvbert`) + + >>> from transformers import ModernVBertModel + >>> model = ModernVBertModel(configuration) + + >>> # Accessing the model configuration + >>> cfg = model.config + ```""" + + model_type = "modernvbert" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + text_config: PreTrainedConfig | dict | None = None + vision_config: PreTrainedConfig | dict | None = None + image_token_id: int = 50407 + pixel_shuffle_factor: int = 4 + initializer_range: float = 0.02 + initializer_cutoff_factor: float = 2.0 + classifier_pooling: Literal["cls", "mean"] = "cls" + classifier_dropout: float | int = 0.0 + classifier_bias: bool = False + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if self.text_config is None: + self.text_config = CONFIG_MAPPING["modernbert"]() + elif isinstance(self.text_config, dict): + self.text_config = CONFIG_MAPPING["modernbert"](**self.text_config) + + if self.vision_config is None: + self.vision_config = CONFIG_MAPPING["siglip_vision_model"]() + elif isinstance(self.vision_config, dict): + self.vision_config = CONFIG_MAPPING["siglip_vision_model"](**self.vision_config) + + super().__post_init__(**kwargs) + + +@dataclass +class ModernVBertBaseModelOutput(BaseModelOutput): + """ + Base class for ModernVBERT model's outputs. + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder + """ + + last_hidden_state: torch.FloatTensor = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +class ModernVBertMaskedLMOutput(MaskedLMOutput): + """ + Base class for ModernVBERT model's outputs with masked language modeling loss. + Args: + loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + Masked language modeling (MLM) loss. + logits (`torch.FloatTensor`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, + sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + image_hidden_states: torch.FloatTensor | None = None + + +class ModernVBertConnector(nn.Module): + """ + Connector module for ModernVBERT. It performs a pixel shuffle operation followed by a linear projection to match the text model's hidden size. + Based on https://pytorch.org/docs/stable/generated/torch.nn.PixelShuffle.html + """ + + def __init__(self, config): + super().__init__() + self.pixel_shuffle_factor = config.pixel_shuffle_factor + self.modality_projection = nn.Linear( + config.vision_config.hidden_size * (config.pixel_shuffle_factor**2), + config.text_config.hidden_size, + bias=False, + ) + + def pixel_shuffle(self, image_hidden_states, pixel_shuffle_factor): + batch_size, seq_length, embed_dim = image_hidden_states.size() + height = width = int(seq_length**0.5) + image_hidden_states = image_hidden_states.view(batch_size, height, width, embed_dim) + image_hidden_states = image_hidden_states.view( + batch_size, height, int(width / pixel_shuffle_factor), embed_dim * pixel_shuffle_factor + ) + image_hidden_states = image_hidden_states.permute(0, 2, 1, 3) + image_hidden_states = image_hidden_states.reshape( + batch_size, + int(width / pixel_shuffle_factor), + int(height / pixel_shuffle_factor), + embed_dim * (pixel_shuffle_factor**2), + ) + image_hidden_states = image_hidden_states.permute(0, 2, 1, 3) + return image_hidden_states.reshape( + batch_size, int(seq_length / (pixel_shuffle_factor**2)), embed_dim * (pixel_shuffle_factor**2) + ) + + def forward(self, image_hidden_states): + image_hidden_states = self.pixel_shuffle(image_hidden_states, self.pixel_shuffle_factor) + return self.modality_projection(image_hidden_states) + + +@auto_docstring +class ModernVBertPreTrainedModel(SmolVLMPreTrainedModel): + config_class = ModernVBertConfig + _no_split_modules = [] + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + + def init_weight(module: nn.Module, std: float): + cutoff_factor = getattr(self.config, "initializer_cutoff_factor", 2.0) + init.trunc_normal_( + module.weight, + mean=0.0, + std=std, + a=-cutoff_factor * std, + b=cutoff_factor * std, + ) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + if module.bias is not None: + init.zeros_(module.bias) + + if isinstance(module, ModernVBertConnector): + out_std = self.config.initializer_range / math.sqrt(2.0 * self.config.text_config.num_hidden_layers) + init_weight(module.modality_projection, out_std) + elif isinstance(module, ModernVBertForMaskedLM): + out_std = self.config.initializer_range / math.sqrt(2.0 * self.config.text_config.num_hidden_layers) + init_weight(module.lm_head, out_std) + elif isinstance( + module, + ( + ModernVBertForSequenceClassification, + ModernVBertForTokenClassification, + ), + ): + final_out_std = self.config.initializer_range / math.sqrt(self.config.text_config.hidden_size) + init_weight(module.classifier, final_out_std) + + +@auto_docstring( + custom_intro=""" + ModernVBertModel is a model that combines a vision encoder (SigLIP) and a text encoder (ModernBert). + + ModernVBert is the base model of the visual retriver ColModernVBert, and was introduced in the following paper: + [*ModernVBERT: Towards Smaller Visual Document Retrievers*](https://arxiv.org/abs/2510.01149). + """ +) +class ModernVBertModel(SmolVLMModel): + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + + # init components + self.connector = ModernVBertConnector(config) + self.text_model = AutoModel.from_config(config.text_config) + self.vision_model = AutoModel.from_config(config.vision_config) + + self.image_seq_len = int( + ((config.vision_config.image_size // config.vision_config.patch_size) ** 2) + / (config.pixel_shuffle_factor**2) + ) + + # initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | ModernVBertBaseModelOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + """ + + if inputs_embeds is None: + inputs_embeds = self.text_model.get_input_embeddings()(input_ids).to(input_ids.device) + + # Images processing + if pixel_values is not None: + image_hidden_states = self.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask + ).pooler_output + + # Merge image and text embeddings + if image_hidden_states is not None: + image_hidden_states = image_hidden_states.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device) + inputs_embeds = self.inputs_merger( + input_ids=input_ids, inputs_embeds=inputs_embeds, image_hidden_states=image_hidden_states + ) + + # Language model pass + outputs = self.text_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + + return ModernVBertBaseModelOutput( + last_hidden_state=outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_hidden_states, + ) + + +class ModernVBertPredictionHead(ModernBertPredictionHead): + pass + + +@auto_docstring +class ModernVBertForMaskedLM(ModernVBertPreTrainedModel): + _tied_weights_keys = {"lm_head.weight": "model.text_model.embeddings.tok_embeddings.weight"} + + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + + self.vocab_size = config.text_config.vocab_size + + self.model = ModernVBertModel(config) + self.projection_head = ModernVBertPredictionHead(config.text_config) + self.lm_head = nn.Linear(config.text_config.hidden_size, self.vocab_size, bias=config.text_config.decoder_bias) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | ModernVBertMaskedLMOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + hidden_states = outputs[0] + + logits = self.lm_head(self.projection_head(hidden_states)) + + loss = None + if labels is not None: + criterion = CrossEntropyLoss() + loss = criterion(logits.view(-1, self.vocab_size), labels.view(-1)) + + return ModernVBertMaskedLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + The ModernVBert Model with a sequence classification head on top that performs pooling. + """ +) +class ModernVBertForSequenceClassification(ModernVBertPreTrainedModel): + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + + self.model = ModernVBertModel(config) + self.head = ModernVBertPredictionHead(config.text_config) + self.drop = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SequenceClassifierOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + last_hidden_state = outputs[0] + + if self.config.classifier_pooling == "cls": + last_hidden_state = last_hidden_state[:, 0] + elif self.config.classifier_pooling == "mean": + if inputs_embeds is not None: + batch_size, seq_len = inputs_embeds.shape[:2] + else: + batch_size, seq_len = input_ids.shape[:2] + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) + last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum( + dim=1, keepdim=True + ) + + pooled_output = self.head(last_hidden_state) + pooled_output = self.drop(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The ModernVBert Model with a token classification head on top, e.g. for Named Entity Recognition (NER) tasks. + """ +) +class ModernVBertForTokenClassification(ModernVBertPreTrainedModel): + def __init__(self, config: ModernVBertConfig): + super().__init__(config) + self.num_labels = config.num_labels + + self.model = ModernVBertModel(config) + self.head = ModernVBertPredictionHead(config.text_config) + self.drop = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """, + checkpoint="ModernVBERT/modernvbert", + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | TokenClassifierOutput: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + text_config.]` or `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.]`. + """ + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + **kwargs, + ) + last_hidden_state = outputs[0] + + last_hidden_state = self.head(last_hidden_state) + last_hidden_state = self.drop(last_hidden_state) + logits = self.classifier(last_hidden_state) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "ModernVBertConfig", + "ModernVBertPreTrainedModel", + "ModernVBertModel", + "ModernVBertForMaskedLM", + "ModernVBertForSequenceClassification", + "ModernVBertForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/oneformer/__init__.py b/third_party/transformers/src/transformers/models/oneformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b8dcfac01991d238e12406df77467e8d0827cfae --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/__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_oneformer import * + from .image_processing_oneformer import * + from .image_processing_pil_oneformer import * + from .modeling_oneformer import * + from .processing_oneformer import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/oneformer/configuration_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/configuration_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f3ac0e9bae888c917994fdd80d60054edabcbd41 --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/configuration_oneformer.py @@ -0,0 +1,158 @@ +# Copyright 2022 SHI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OneFormer model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import AutoConfig + + +@auto_docstring(checkpoint="shi-labs/oneformer_ade20k_swin_tiny") +@strict +class OneFormerConfig(PreTrainedConfig): + r""" + ignore_value (`int`, *optional*, defaults to 255): + Values to be ignored in GT label while calculating loss. + num_queries (`int`, *optional*, defaults to 150): + Number of object queries. + contrastive_weight (`float`, *optional*, defaults to 0.5): + Weight for contrastive loss. + contrastive_temperature (`float`, *optional*, defaults to 0.07): + Initial value for scaling the contrastive logits. + train_num_points (`int`, *optional*, defaults to 12544): + Number of points to sample while calculating losses on mask predictions. + oversample_ratio (`float`, *optional*, defaults to 3.0): + Ratio to decide how many points to oversample. + importance_sample_ratio (`float`, *optional*, defaults to 0.75): + Ratio of points that are sampled via importance sampling.. + is_training (`bool`, *optional*, defaults to `False`): + Whether to run in training or inference mode. + output_auxiliary_logits (`bool`, *optional*, defaults to `True`): + Whether to return intermediate predictions from transformer decoder. + strides (`list`, *optional*, defaults to `[4, 8, 16, 32]`): + List containing the strides for feature maps in the encoder. + task_seq_len (`int`, *optional*, defaults to 77): + Sequence length for tokenizing text list input. + text_encoder_width (`int`, *optional*, defaults to 256): + Hidden size for text encoder. + text_encoder_context_length (`int`, *optional*, defaults to 77): + Input sequence length for text encoder. + text_encoder_num_layers (`int`, *optional*, defaults to 6): + Number of layers for transformer in text encoder. + text_encoder_vocab_size (`int`, *optional*, defaults to 49408): + Vocabulary size for tokenizer. + text_encoder_proj_layers (`int`, *optional*, defaults to 2): + Number of layers in MLP for project text queries. + text_encoder_n_ctx (`int`, *optional*, defaults to 16): + Number of learnable text context queries. + conv_dim (`int`, *optional*, defaults to 256): + Feature map dimension to map outputs from the backbone. + mask_dim (`int`, *optional*, defaults to 256): + Dimension for feature maps in pixel decoder. + hidden_dim (`int`, *optional*, defaults to 256): + Dimension for hidden states in transformer decoder. + encoder_feedforward_dim (`int`, *optional*, defaults to 1024): + Dimension for FFN layer in pixel decoder. + norm (`str`, *optional*, defaults to `"GN"`): + Type of normalization. + use_task_norm (`bool`, *optional*, defaults to `True`): + Whether to normalize the task token. + dim_feedforward (`int`, *optional*, defaults to 2048): + Dimension for FFN layer in transformer decoder. + pre_norm (`bool`, *optional*, defaults to `False`): + Whether to normalize hidden states before attention layers in transformer decoder. + enforce_input_proj (`bool`, *optional*, defaults to `False`): + Whether to project hidden states in transformer decoder. + query_dec_layers (`int`, *optional*, defaults to 2): + Number of layers in query transformer. + common_stride (`int`, *optional*, defaults to 4): + Common stride used for features in pixel decoder. + + Examples: + ```python + >>> from transformers import OneFormerConfig, OneFormerModel + + >>> # Initializing a OneFormer shi-labs/oneformer_ade20k_swin_tiny configuration + >>> configuration = OneFormerConfig() + >>> # Initializing a model (with random weights) from the shi-labs/oneformer_ade20k_swin_tiny style configuration + >>> model = OneFormerModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "oneformer" + sub_configs = {"backbone_config": AutoConfig} + attribute_map = {"hidden_size": "hidden_dim", "num_hidden_layers": "decoder_layers"} + + backbone_config: dict | PreTrainedConfig | None = None + ignore_value: int = 255 + num_queries: int = 150 + no_object_weight: float = 0.1 + class_weight: float = 2.0 + mask_weight: float = 5.0 + dice_weight: float = 5.0 + contrastive_weight: float = 0.5 + contrastive_temperature: float = 0.07 + train_num_points: int = 12544 + oversample_ratio: float = 3.0 + importance_sample_ratio: float = 0.75 + init_std: float = 0.02 + init_xavier_std: float = 1.0 + layer_norm_eps: float = 1e-05 + is_training: bool = False + use_auxiliary_loss: bool = True + output_auxiliary_logits: bool = True + strides: list[int] | tuple[int, ...] = (4, 8, 16, 32) + task_seq_len: int = 77 + text_encoder_width: int = 256 + text_encoder_context_length: int = 77 + text_encoder_num_layers: int = 6 + text_encoder_vocab_size: int = 49408 + text_encoder_proj_layers: int = 2 + text_encoder_n_ctx: int = 16 + conv_dim: int = 256 + mask_dim: int = 256 + hidden_dim: int = 256 + encoder_feedforward_dim: int = 1024 + norm: str = "GN" + encoder_layers: int = 6 + decoder_layers: int = 10 + use_task_norm: bool = True + num_attention_heads: int = 8 + dropout: float | int = 0.1 + dim_feedforward: int = 2048 + pre_norm: bool = False + enforce_input_proj: bool = False + query_dec_layers: int = 2 + common_stride: int = 4 + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="swin", + default_config_kwargs={ + "drop_path_rate": 0.3, + "out_features": ["stage1", "stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + super().__post_init__(**kwargs) + + +__all__ = ["OneFormerConfig"] diff --git a/third_party/transformers/src/transformers/models/oneformer/convert_to_hf_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/convert_to_hf_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f2752741d4c4e6744f6db9967356576114307f00 --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/convert_to_hf_oneformer.py @@ -0,0 +1,1185 @@ +# Copyright 2022 SHI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert OneFormer checkpoints from the original repository. URL: https://github.com/SHI-Labs/OneFormer""" + +import os +import sys +from argparse import ArgumentParser +from collections.abc import Iterator +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from pprint import pformat +from typing import Any + +import httpx +import torch +import torchvision.transforms as T +from PIL import Image +from torch import Tensor, nn + + +try: + from detectron2.checkpoint import DetectionCheckpointer + from detectron2.config import get_cfg + from detectron2.data import MetadataCatalog + from detectron2.projects.deeplab import add_deeplab_config +except ImportError: + pass +from transformers import CLIPTokenizer, DinatConfig, SwinConfig +from transformers.models.oneformer.image_processing_oneformer import OneFormerImageProcessor +from transformers.models.oneformer.modeling_oneformer import ( + OneFormerConfig, + OneFormerForUniversalSegmentation, + OneFormerForUniversalSegmentationOutput, + OneFormerModel, + OneFormerModelOutput, +) +from transformers.models.oneformer.processing_oneformer import OneFormerProcessor +from transformers.utils import logging + + +StateDict = dict[str, Tensor] + +logging.set_verbosity_info() +logger = logging.get_logger() + +torch.manual_seed(0) + + +class TrackedStateDict: + def __init__(self, to_track: dict): + """This class "tracks" a python dictionary by keeping track of which item is accessed. + + Args: + to_track (Dict): The dictionary we wish to track + """ + self.to_track = to_track + self._seen: set[str] = set() + + def __getitem__(self, key: str) -> Any: + return self.to_track[key] + + def __setitem__(self, key: str, item: Any): + self._seen.add(key) + self.to_track[key] = item + + def diff(self) -> list[str]: + """This method returns a set difference between the keys in the tracked state dict and the one we have access so far. + This is an effective method to check if we have update all the keys + + Returns: + list[str]: List of keys not yet updated + """ + return set(self.to_track.keys()) - self._seen + + def copy(self) -> dict: + # proxy the call to the internal dictionary + return self.to_track.copy() + + +# Image to verify the result +def prepare_img(): + url = "https://praeclarumjj3.github.io/files/coco.jpeg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@dataclass +class Args: + """Fake command line arguments needed by oneformer/detectron2 implementation""" + + config_file: str + + +def setup_cfg(args: Args): + # load config from file and command-line arguments + cfg = get_cfg() + add_deeplab_config(cfg) + add_common_config(cfg) + add_oneformer_config(cfg) + add_swin_config(cfg) + add_dinat_config(cfg) + cfg.merge_from_file(args.config_file) + cfg.freeze() + return cfg + + +class OriginalOneFormerConfigToOursConverter: + def __call__(self, original_config: object, is_swin: bool) -> OneFormerConfig: + model = original_config.MODEL + + dataset_catalog = MetadataCatalog.get(original_config.DATASETS.TEST_PANOPTIC[0]) + id2label = dict(enumerate(dataset_catalog.stuff_classes)) + label2id = {label: idx for idx, label in id2label.items()} + + if is_swin: + if model.SWIN.EMBED_DIM == 96: + backbone_config = SwinConfig.from_pretrained( + "microsoft/swin-tiny-patch4-window7-224", + drop_path_rate=model.SWIN.DROP_PATH_RATE, + out_features=["stage1", "stage2", "stage3", "stage4"], + ) + elif model.SWIN.EMBED_DIM == 192: + backbone_config = SwinConfig.from_pretrained( + "microsoft/swin-large-patch4-window12-384", + drop_path_rate=model.SWIN.DROP_PATH_RATE, + out_features=["stage1", "stage2", "stage3", "stage4"], + ) + else: + raise ValueError(f"embed dim {model.SWIN.EMBED_DIM} not supported for Swin!") + else: + backbone_config = DinatConfig.from_pretrained( + "shi-labs/dinat-large-11x11-in22k-in1k-384", + dilations=model.DiNAT.DILATIONS, + kernel_size=model.DiNAT.KERNEL_SIZE, + out_features=["stage1", "stage2", "stage3", "stage4"], + ) + + config: OneFormerConfig = OneFormerConfig( + backbone_config=backbone_config, + output_attentions=True, + output_hidden_states=True, + return_dict=True, + ignore_value=model.SEM_SEG_HEAD.IGNORE_VALUE, + num_classes=model.SEM_SEG_HEAD.NUM_CLASSES, + num_queries=model.ONE_FORMER.NUM_OBJECT_QUERIES, + no_object_weight=model.ONE_FORMER.NO_OBJECT_WEIGHT, + class_weight=model.ONE_FORMER.CLASS_WEIGHT, + mask_weight=model.ONE_FORMER.MASK_WEIGHT, + dice_weight=model.ONE_FORMER.DICE_WEIGHT, + contrastive_weight=model.ONE_FORMER.CONTRASTIVE_WEIGHT, + contrastive_temperature=model.ONE_FORMER.CONTRASTIVE_TEMPERATURE, + train_num_points=model.ONE_FORMER.TRAIN_NUM_POINTS, + oversample_ratio=model.ONE_FORMER.OVERSAMPLE_RATIO, + importance_sample_ratio=model.ONE_FORMER.IMPORTANCE_SAMPLE_RATIO, + init_std=0.02, + init_xavier_std=1.0, + layer_norm_eps=1e-05, + is_training=False, + use_auxiliary_loss=model.ONE_FORMER.DEEP_SUPERVISION, + output_auxiliary_logits=True, + strides=[4, 8, 16, 32], + task_seq_len=original_config.INPUT.TASK_SEQ_LEN, + max_seq_len=original_config.INPUT.MAX_SEQ_LEN, + text_encoder_width=model.TEXT_ENCODER.WIDTH, + text_encoder_context_length=model.TEXT_ENCODER.CONTEXT_LENGTH, + text_encoder_num_layers=model.TEXT_ENCODER.NUM_LAYERS, + text_encoder_vocab_size=model.TEXT_ENCODER.VOCAB_SIZE, + text_encoder_proj_layers=model.TEXT_ENCODER.PROJ_NUM_LAYERS, + text_encoder_n_ctx=model.TEXT_ENCODER.N_CTX, + conv_dim=model.SEM_SEG_HEAD.CONVS_DIM, + mask_dim=model.SEM_SEG_HEAD.MASK_DIM, + hidden_dim=model.ONE_FORMER.HIDDEN_DIM, + norm=model.SEM_SEG_HEAD.NORM, + encoder_layers=model.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS, + encoder_feedforward_dim=1024, + decoder_layers=model.ONE_FORMER.DEC_LAYERS, + use_task_norm=model.ONE_FORMER.USE_TASK_NORM, + num_attention_heads=model.ONE_FORMER.NHEADS, + dropout=model.ONE_FORMER.DROPOUT, + dim_feedforward=model.ONE_FORMER.DIM_FEEDFORWARD, + pre_norm=model.ONE_FORMER.PRE_NORM, + enforce_input_proj=model.ONE_FORMER.ENFORCE_INPUT_PROJ, + query_dec_layers=model.ONE_FORMER.CLASS_DEC_LAYERS, + common_stride=model.SEM_SEG_HEAD.COMMON_STRIDE, + id2label=id2label, + label2id=label2id, + ) + + return config + + +class OriginalOneFormerConfigToProcessorConverter: + def __call__(self, original_config: object, model_repo: str) -> OneFormerProcessor: + model = original_config.MODEL + model_input = original_config.INPUT + dataset_catalog = MetadataCatalog.get(original_config.DATASETS.TEST_PANOPTIC[0]) + + if "ade20k" in model_repo: + class_info_file = "ade20k_panoptic.json" + elif "coco" in model_repo: + class_info_file = "coco_panoptic.json" + elif "cityscapes" in model_repo: + class_info_file = "cityscapes_panoptic.json" + else: + raise ValueError("Invalid Dataset!") + + image_processor = OneFormerImageProcessor( + image_mean=(torch.tensor(model.PIXEL_MEAN) / 255).tolist(), + image_std=(torch.tensor(model.PIXEL_STD) / 255).tolist(), + size=model_input.MIN_SIZE_TEST, + max_size=model_input.MAX_SIZE_TEST, + num_labels=model.SEM_SEG_HEAD.NUM_CLASSES, + ignore_index=dataset_catalog.ignore_label, + class_info_file=class_info_file, + ) + + tokenizer = CLIPTokenizer.from_pretrained(model_repo) + + return OneFormerProcessor( + image_processor=image_processor, + tokenizer=tokenizer, + task_seq_length=original_config.INPUT.TASK_SEQ_LEN, + max_seq_length=original_config.INPUT.MAX_SEQ_LEN, + ) + + +class OriginalOneFormerCheckpointToOursConverter: + def __init__(self, original_model: nn.Module, config: OneFormerConfig): + self.original_model = original_model + self.config = config + + def pop_all(self, renamed_keys: list[tuple[str, str]], dst_state_dict: StateDict, src_state_dict: StateDict): + for src_key, dst_key in renamed_keys: + dst_state_dict[dst_key] = src_state_dict.pop(src_key) + + # Swin Backbone + def replace_swin_backbone(self, dst_state_dict: StateDict, src_state_dict: StateDict, config: OneFormerConfig): + dst_prefix: str = "pixel_level_module.encoder" + src_prefix: str = "backbone" + + renamed_keys = [ + ( + f"{src_prefix}.patch_embed.proj.weight", + f"{dst_prefix}.embeddings.patch_embeddings.projection.weight", + ), + (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.embeddings.patch_embeddings.projection.bias"), + (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.embeddings.norm.weight"), + (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.embeddings.norm.bias"), + ] + num_layers = len(config.backbone_config.depths) + for layer_idx in range(num_layers): + for block_idx in range(config.backbone_config.depths[layer_idx]): + renamed_keys.extend( + [ # src, dst + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table", + ), + ] + ) + # now we need to handle the attentions + # read in weights + bias of input projection layer of cross-attention + + src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"] + src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"] + + size = src_att_weight.shape[0] + offset = size // 3 + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight" + ] = src_att_weight[:offset, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias" + ] = src_att_bias[:offset] + + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight" + ] = src_att_weight[offset : offset * 2, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias" + ] = src_att_bias[offset : offset * 2] + + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight" + ] = src_att_weight[-offset:, :] + dst_state_dict[ + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias" + ] = src_att_bias[-offset:] + + # let's pop them + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight") + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias") + # proj + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias", + ), + ] + ) + + # second norm + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias", + ), + ] + ) + + # mlp + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias", + ), + ] + ) + + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index", + f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index", + ) + ] + ) + + if layer_idx < num_layers - 1: + # patch merging + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.reduction.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias", + f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.bias", + ), + ] + ) + + # hidden states norms + renamed_keys.extend( + [ + ( + f"{src_prefix}.norm{layer_idx}.weight", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.weight", + ), + ( + f"{src_prefix}.norm{layer_idx}.bias", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.bias", + ), + ] + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + # Dinat Backbone + def replace_dinat_backbone(self, dst_state_dict: StateDict, src_state_dict: StateDict, config: OneFormerConfig): + dst_prefix: str = "pixel_level_module.encoder" + src_prefix: str = "backbone" + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + renamed_keys = rename_keys_for_weight_bias(f"{src_prefix}.patch_embed.norm", f"{dst_prefix}.embeddings.norm") + + for i in range(2): + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.patch_embed.proj.{i}", + f"{dst_prefix}.embeddings.patch_embeddings.projection.{i}", + ) + ) + + num_layers = len(config.backbone_config.depths) + for layer_idx in range(num_layers): + for block_idx in range(config.backbone_config.depths[layer_idx]): + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.norm1", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.layernorm_before", + ) + ) + + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.norm2", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.layernorm_after", + ) + ) + + renamed_keys.extend( + [ # src, dst + ( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.rpb", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.rpb", + ), + ] + ) + # now we need to handle the attentions + # read in weights + bias of input projection layer of cross-attention + + src_att_weight = src_state_dict[f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"] + src_att_bias = src_state_dict[f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"] + + size = src_att_weight.shape[0] + offset = size // 3 + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.query.weight" + ] = src_att_weight[:offset, :] + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.query.bias" + ] = src_att_bias[:offset] + + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.key.weight" + ] = src_att_weight[offset : offset * 2, :] + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.key.bias" + ] = src_att_bias[offset : offset * 2] + + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.value.weight" + ] = src_att_weight[-offset:, :] + dst_state_dict[ + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.self.value.bias" + ] = src_att_bias[-offset:] + + # let's pop them + src_state_dict.pop(f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.qkv.weight") + src_state_dict.pop(f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.qkv.bias") + # proj + + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.attn.proj", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.attention.output.dense", + ) + ) + + # mlp + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.mlp.fc1", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.intermediate.dense", + ) + ) + + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.levels.{layer_idx}.blocks.{block_idx}.mlp.fc2", + f"{dst_prefix}.encoder.levels.{layer_idx}.layers.{block_idx}.output.dense", + ) + ) + + if layer_idx < num_layers - 1: + # patch merging + renamed_keys.extend( + [ + ( + f"{src_prefix}.levels.{layer_idx}.downsample.reduction.weight", + f"{dst_prefix}.encoder.levels.{layer_idx}.downsample.reduction.weight", + ), + ( + f"{src_prefix}.levels.{layer_idx}.downsample.norm.weight", + f"{dst_prefix}.encoder.levels.{layer_idx}.downsample.norm.weight", + ), + ( + f"{src_prefix}.levels.{layer_idx}.downsample.norm.bias", + f"{dst_prefix}.encoder.levels.{layer_idx}.downsample.norm.bias", + ), + ] + ) + + # hidden states norms + renamed_keys.extend( + [ + ( + f"{src_prefix}.norm{layer_idx}.weight", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.weight", + ), + ( + f"{src_prefix}.norm{layer_idx}.bias", + f"{dst_prefix}.hidden_states_norms.stage{layer_idx + 1}.bias", + ), + ] + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + # Backbone + Pixel Decoder + def replace_pixel_module(self, dst_state_dict: StateDict, src_state_dict: StateDict, is_swin: bool): + dst_prefix: str = "pixel_level_module.decoder" + src_prefix: str = "sem_seg_head.pixel_decoder" + + if is_swin: + self.replace_swin_backbone(dst_state_dict, src_state_dict, self.config) + else: + self.replace_dinat_backbone(dst_state_dict, src_state_dict, self.config) + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + def rename_keys_for_self_attn(src_prefix: str, dst_prefix: str): + self_attn_keys = [] + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.attention_weights", f"{dst_prefix}.attention_weights") + ) + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.output_proj", f"{dst_prefix}.output_proj") + ) + self_attn_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.sampling_offsets", f"{dst_prefix}.sampling_offsets") + ) + self_attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.value_proj", f"{dst_prefix}.value_proj")) + + return self_attn_keys + + def rename_keys_for_encoder_layer(src_prefix: str, dst_prefix: str): + encoder_keys = [] + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear1", f"{dst_prefix}.fc1")) + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear2", f"{dst_prefix}.fc2")) + encoder_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.norm1", f"{dst_prefix}.self_attn_layer_norm") + ) + encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm2", f"{dst_prefix}.final_layer_norm")) + encoder_keys.extend(rename_keys_for_self_attn(f"{src_prefix}.self_attn", f"{dst_prefix}.self_attn")) + + return encoder_keys + + # convolution layer for final features + renamed_keys = [ + (f"{src_prefix}.adapter_1.weight", f"{dst_prefix}.adapter_1.0.weight"), + (f"{src_prefix}.adapter_1.norm.weight", f"{dst_prefix}.adapter_1.1.weight"), + (f"{src_prefix}.adapter_1.norm.bias", f"{dst_prefix}.adapter_1.1.bias"), + ] + + renamed_keys.extend( + [ + (f"{src_prefix}.layer_1.weight", f"{dst_prefix}.layer_1.0.weight"), + (f"{src_prefix}.layer_1.norm.weight", f"{dst_prefix}.layer_1.1.weight"), + (f"{src_prefix}.layer_1.norm.bias", f"{dst_prefix}.layer_1.1.bias"), + ] + ) + + # proj layers + for i in range(3): + for j in range(2): + renamed_keys.extend( + [ + (f"{src_prefix}.input_proj.{i}.{j}.weight", f"{dst_prefix}.input_projections.{i}.{j}.weight"), + (f"{src_prefix}.input_proj.{i}.{j}.bias", f"{dst_prefix}.input_projections.{i}.{j}.bias"), + ] + ) + + renamed_keys.extend([(f"{src_prefix}.transformer.level_embed", f"{dst_prefix}.level_embed")]) + + # layers + for layer_idx in range(self.config.encoder_layers): + renamed_keys.extend( + rename_keys_for_encoder_layer( + f"{src_prefix}.transformer.encoder.layers.{layer_idx}", f"{dst_prefix}.encoder.layers.{layer_idx}" + ) + ) + + # proj + renamed_keys.extend( + [ + (f"{src_prefix}.mask_features.weight", f"{dst_prefix}.mask_projection.weight"), + (f"{src_prefix}.mask_features.bias", f"{dst_prefix}.mask_projection.bias"), + ] + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + # Transformer Decoder + def replace_keys_qkv_transformer_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder.layers" + src_prefix: str = "sem_seg_head.predictor" + for i in range(self.config.decoder_layers - 1): + # read in weights + bias of input projection layer of self-attention + in_proj_weight = src_state_dict.pop( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_weight" + ) + in_proj_bias = src_state_dict.pop( + f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_bias" + ) + # next, add query, keys and values (in that order) to the state dict + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.q_proj.weight"] = in_proj_weight[:256, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.q_proj.bias"] = in_proj_bias[:256] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.k_proj.weight"] = in_proj_weight[256:512, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.k_proj.bias"] = in_proj_bias[256:512] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.v_proj.weight"] = in_proj_weight[-256:, :] + dst_state_dict[f"{dst_prefix}.{i}.self_attn.self_attn.v_proj.bias"] = in_proj_bias[-256:] + + def replace_transformer_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module" + src_prefix: str = "sem_seg_head.predictor" + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + def rename_keys_for_attn(src_prefix: str, dst_prefix: str): + attn_keys = [ + (f"{src_prefix}.in_proj_bias", f"{dst_prefix}.in_proj_bias"), + (f"{src_prefix}.in_proj_weight", f"{dst_prefix}.in_proj_weight"), + ] + attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.out_proj", f"{dst_prefix}.out_proj")) + + return attn_keys + + def rename_keys_for_self_attn(src_prefix: str, dst_prefix: str): + attn_keys = [] + attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.out_proj", f"{dst_prefix}.out_proj")) + + return attn_keys + + def rename_keys_for_query_transformer_layer(src_prefix: str, dst_prefix: str): + query_transformer_layer_keys = [] + + query_transformer_layer_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.linear1", f"{dst_prefix}.linear1") + ) + query_transformer_layer_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.linear2", f"{dst_prefix}.linear2") + ) + query_transformer_layer_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.norm1", f"{dst_prefix}.norm1") + ) + query_transformer_layer_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.norm2", f"{dst_prefix}.norm2") + ) + query_transformer_layer_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.norm3", f"{dst_prefix}.norm3") + ) + + query_transformer_layer_keys.extend( + rename_keys_for_attn(f"{src_prefix}.self_attn", f"{dst_prefix}.self_attn") + ) + + query_transformer_layer_keys.extend( + rename_keys_for_attn(f"{src_prefix}.multihead_attn", f"{dst_prefix}.multihead_attn") + ) + + return query_transformer_layer_keys + + def rename_keys_for_cross_attn_layer(src_prefix: str, dst_prefix: str): + cross_attn_layer_keys = [] + + cross_attn_layer_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm", f"{dst_prefix}.norm")) + cross_attn_layer_keys.extend( + rename_keys_for_attn(f"{src_prefix}.multihead_attn", f"{dst_prefix}.multihead_attn") + ) + + return cross_attn_layer_keys + + def rename_keys_for_self_attn_layer(src_prefix: str, dst_prefix: str): + self_attn_layer_keys = [] + + self_attn_layer_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm", f"{dst_prefix}.norm")) + self_attn_layer_keys.extend( + rename_keys_for_self_attn(f"{src_prefix}.self_attn", f"{dst_prefix}.self_attn") + ) + + return self_attn_layer_keys + + def rename_keys_for_ffn_layer(src_prefix: str, dst_prefix: str): + ffn_layer_keys = [] + + ffn_layer_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear1", f"{dst_prefix}.linear1")) + ffn_layer_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear2", f"{dst_prefix}.linear2")) + ffn_layer_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm", f"{dst_prefix}.norm")) + + return ffn_layer_keys + + def rename_keys_for_transformer_decoder_layer(src_prefix: str, dst_prefix: str, idx: int): + transformer_decoder_layer_keys = [] + + transformer_decoder_layer_keys.extend( + rename_keys_for_cross_attn_layer( + f"{src_prefix}.transformer_cross_attention_layers.{idx}", f"{dst_prefix}.{idx}.cross_attn" + ) + ) + + transformer_decoder_layer_keys.extend( + rename_keys_for_self_attn_layer( + f"{src_prefix}.transformer_self_attention_layers.{idx}", f"{dst_prefix}.{idx}.self_attn" + ) + ) + + transformer_decoder_layer_keys.extend( + rename_keys_for_ffn_layer(f"{src_prefix}.transformer_ffn_layers.{idx}", f"{dst_prefix}.{idx}.ffn") + ) + + return transformer_decoder_layer_keys + + # positional embedding for object queries + renamed_keys = [ + (f"{src_prefix}.query_embed.weight", f"{dst_prefix}.queries_embedder.weight"), + (f"{src_prefix}.level_embed.weight", f"{dst_prefix}.level_embed.weight"), + ] + + # norm + renamed_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.decoder_norm", f"{dst_prefix}.decoder.decoder_norm") + ) + + # proj + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.class_input_proj", f"{dst_prefix}.decoder.query_input_projection" + ) + ) + + renamed_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.class_embed", f"{dst_prefix}.decoder.class_embed") + ) + + for i in range(3): + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.mask_embed.layers.{i}", f"{dst_prefix}.decoder.mask_embed.layers.{i}.0" + ) + ) + + # norm + renamed_keys.extend( + rename_keys_for_weight_bias( + f"{src_prefix}.class_transformer.decoder.norm", f"{dst_prefix}.decoder.query_transformer.decoder.norm" + ) + ) + + # transformer to update queries with task tokens + for i in range(self.config.query_dec_layers): + renamed_keys.extend( + rename_keys_for_query_transformer_layer( + f"{src_prefix}.class_transformer.decoder.layers.{i}", + f"{dst_prefix}.decoder.query_transformer.decoder.layers.{i}", + ) + ) + + # decoder layers + for i in range(self.config.decoder_layers - 1): + renamed_keys.extend( + rename_keys_for_transformer_decoder_layer( + f"{src_prefix}", + f"{dst_prefix}.decoder.layers", + i, + ) + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + self.replace_keys_qkv_transformer_decoder(dst_state_dict, src_state_dict) + + def replace_task_mlp(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "task_encoder" + src_prefix: str = "task_mlp" + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + renamed_keys = [] + + for i in range(2): + renamed_keys.extend( + rename_keys_for_weight_bias(f"{src_prefix}.layers.{i}", f"{dst_prefix}.task_mlp.layers.{i}.0") + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_text_projector(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "text_mapper.text_projector" + src_prefix: str = "text_projector" + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + renamed_keys = [] + + for i in range(self.config.text_encoder_config["text_encoder_proj_layers"]): + renamed_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.layers.{i}", f"{dst_prefix}.{i}.0")) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_text_mapper(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "text_mapper.text_encoder" + src_prefix: str = "text_encoder" + + self.replace_text_projector(dst_state_dict, src_state_dict) + + def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str): + return [ + (f"{src_prefix}.weight", f"{dst_prefix}.weight"), + (f"{src_prefix}.bias", f"{dst_prefix}.bias"), + ] + + def rename_keys_for_attn(src_prefix: str, dst_prefix: str): + attn_keys = [ + (f"{src_prefix}.in_proj_bias", f"{dst_prefix}.in_proj_bias"), + (f"{src_prefix}.in_proj_weight", f"{dst_prefix}.in_proj_weight"), + ] + attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.out_proj", f"{dst_prefix}.out_proj")) + + return attn_keys + + def rename_keys_for_layer(src_prefix: str, dst_prefix: str): + resblock_keys = [] + + resblock_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.mlp.c_fc", f"{dst_prefix}.mlp.fc1")) + resblock_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.mlp.c_proj", f"{dst_prefix}.mlp.fc2")) + resblock_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.ln_1", f"{dst_prefix}.layer_norm1")) + resblock_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.ln_2", f"{dst_prefix}.layer_norm2")) + resblock_keys.extend(rename_keys_for_attn(f"{src_prefix}.attn", f"{dst_prefix}.self_attn")) + + return resblock_keys + + renamed_keys = [ + ("prompt_ctx.weight", "text_mapper.prompt_ctx.weight"), + ] + + renamed_keys.extend( + [ + (f"{src_prefix}.positional_embedding", f"{dst_prefix}.positional_embedding"), + (f"{src_prefix}.token_embedding.weight", f"{dst_prefix}.token_embedding.weight"), + ] + ) + + renamed_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.ln_final", f"{dst_prefix}.ln_final")) + + for i in range(self.config.text_encoder_config["text_encoder_num_layers"]): + renamed_keys.extend( + rename_keys_for_layer( + f"{src_prefix}.transformer.resblocks.{i}", f"{dst_prefix}.transformer.layers.{i}" + ) + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def convert(self, oneformer: OneFormerModel, is_swin: bool) -> OneFormerModel: + dst_state_dict = TrackedStateDict(oneformer.state_dict()) + src_state_dict = self.original_model.state_dict() + + self.replace_pixel_module(dst_state_dict, src_state_dict, is_swin) + self.replace_transformer_module(dst_state_dict, src_state_dict) + self.replace_task_mlp(dst_state_dict, src_state_dict) + if self.config.is_training: + self.replace_text_mapper(dst_state_dict, src_state_dict) + + logger.info(f"Missed keys are {pformat(dst_state_dict.diff())}") + logger.info(f"Not copied keys are {pformat(src_state_dict.keys())}") + logger.info("🙌 Done") + + oneformer.load_state_dict(dst_state_dict) + + return oneformer + + @staticmethod + def using_dirs(checkpoints_dir: Path, config_dir: Path) -> Iterator[tuple[object, Path, Path]]: + checkpoints: list[Path] = checkpoints_dir.glob("**/*.pth") + + for checkpoint in checkpoints: + logger.info(f"Converting {checkpoint.stem}") + # find associated config file + config: Path = config_dir / f"{checkpoint.stem}.yaml" + + yield config, checkpoint + + +def post_process_sem_seg_output(outputs: OneFormerForUniversalSegmentationOutput, target_size: tuple[int, int]): + # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1] + class_queries_logits = outputs.class_queries_logits + # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH] + masks_queries_logits = outputs.masks_queries_logits + if target_size is not None: + masks_queries_logits = torch.nn.functional.interpolate( + masks_queries_logits, + size=target_size, + mode="bilinear", + align_corners=False, + ) + # remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH] + masks_probs = masks_queries_logits.sigmoid() + # now we want to sum over the queries, + # $ out_{c,h,w} = \sum_q p_{q,c} * m_{q,h,w} $ + # where $ softmax(p) \in R^{q, c} $ is the mask classes + # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities + # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + + return segmentation + + +def test( + original_model, + our_model: OneFormerForUniversalSegmentation, + processor: OneFormerProcessor, + model_repo: str, +): + def _preprocess_text(text_list=None, max_length=77): + if text_list is None: + raise ValueError("tokens cannot be None.") + + tokens = tokenizer(text_list, padding="max_length", max_length=max_length, truncation=True) + + attention_masks, input_ids = tokens["attention_mask"], tokens["input_ids"] + + token_inputs = [] + for attn_mask, input_id in zip(attention_masks, input_ids): + token = torch.tensor(attn_mask) * torch.tensor(input_id) + token_inputs.append(token.unsqueeze(0)) + + token_inputs = torch.cat(token_inputs, dim=0) + return token_inputs + + with torch.no_grad(): + tokenizer = CLIPTokenizer.from_pretrained(model_repo) + original_model = original_model.eval() + our_model = our_model.eval() + + im = prepare_img() + + tr = T.Compose( + [ + T.Resize((640, 640)), + T.ToTensor(), + T.Normalize( + mean=torch.tensor([123.675, 116.280, 103.530]) / 255.0, + std=torch.tensor([58.395, 57.120, 57.375]) / 255.0, + ), + ], + ) + + x = tr(im).unsqueeze(0) + + task_input = ["the task is semantic"] + task_token = _preprocess_text(task_input, max_length=processor.task_seq_length) + + original_model_backbone_features = original_model.backbone(x.clone()) + + our_model_output: OneFormerModelOutput = our_model.model(x.clone(), task_token, output_hidden_states=True) + + for original_model_feature, our_model_feature in zip( + original_model_backbone_features.values(), our_model_output.encoder_hidden_states + ): + assert torch.allclose(original_model_feature, our_model_feature, atol=3e-3), ( + "The backbone features are not the same." + ) + mask_features, _, multi_scale_features, _, _ = original_model.sem_seg_head.pixel_decoder.forward_features( + original_model_backbone_features + ) + + original_pixel_decoder_features = [] + original_pixel_decoder_features.append(mask_features) + for i in range(len(multi_scale_features)): + original_pixel_decoder_features.append(multi_scale_features[i]) + + for original_model_feature, our_model_feature in zip( + original_pixel_decoder_features, our_model_output.pixel_decoder_hidden_states + ): + assert torch.allclose(original_model_feature, our_model_feature, atol=3e-4), ( + "The pixel decoder feature are not the same" + ) + + tr_complete = T.Compose( + [ + T.Resize((640, 640)), + T.ToTensor(), + ], + ) + + y = (tr_complete(im) * 255.0).to(torch.int).float() + + # let's test the full model + original_model_out = original_model([{"image": y.clone(), "task": "The task is semantic"}]) + + original_segmentation = original_model_out[0]["sem_seg"] + + our_model_out: OneFormerForUniversalSegmentationOutput = our_model( + x.clone(), task_token, output_hidden_states=True + ) + + our_segmentation = post_process_sem_seg_output(our_model_out, target_size=(640, 640))[0] + + assert torch.allclose(original_segmentation, our_segmentation, atol=1e-3), ( + "The segmentation image is not the same." + ) + + logger.info("Test passed!") + + +def get_name(checkpoint_file: Path): + model_name_raw: str = checkpoint_file.stem + + backbone = "swin" if "swin" in model_name_raw else "dinat" + dataset = "" + if "coco" in model_name_raw: + dataset = "coco" + elif "ade20k" in model_name_raw: + dataset = "ade20k" + elif "cityscapes" in model_name_raw: + dataset = "cityscapes" + else: + raise ValueError( + f"{model_name_raw} must be wrong since we didn't find 'coco' or 'ade20k' or 'cityscapes' in it " + ) + + backbone_types = ["tiny", "large"] + + backbone_type = list(filter(lambda x: x in model_name_raw, backbone_types))[0] + + model_name = f"oneformer_{dataset}_{backbone}_{backbone_type}" + + return model_name + + +if __name__ == "__main__": + parser = ArgumentParser( + description=( + "Command line to convert the original oneformer models (with swin backbone) to transformers" + " implementation." + ) + ) + + parser.add_argument( + "--checkpoints_dir", + type=Path, + help=( + "A directory containing the model's checkpoints. The directory has to have the following structure:" + " structure: //.pth; where name must follow the" + " following nomenclature nomenclature: oneformer___" + ), + ) + parser.add_argument( + "--configs_dir", + type=Path, + help=( + "A directory containing the model's configs, see detectron2 doc. The directory has to have the following" + " structure: //.yaml; where name must follow the" + " following nomenclature nomenclature: oneformer___" + ), + ) + parser.add_argument( + "--pytorch_dump_folder_path", + required=True, + type=Path, + help="Path to the folder to output PyTorch models.", + ) + parser.add_argument( + "--oneformer_dir", + required=True, + type=Path, + help=( + "A path to OneFormer's original implementation directory. You can download from here: " + "https://github.com/SHI-Labs/OneFormer" + ), + ) + + args = parser.parse_args() + + checkpoints_dir: Path = args.checkpoints_dir + config_dir: Path = args.configs_dir + save_directory: Path = args.pytorch_dump_folder_path + oneformer_dir: Path = args.oneformer_dir + # append the path to the parents to oneformer dir + sys.path.append(str(oneformer_dir.parent)) + # and import what's needed + from OneFormer.oneformer import add_common_config, add_dinat_config, add_oneformer_config, add_swin_config + from OneFormer.oneformer.oneformer_model import OneFormer as OriginalOneFormer + + if not save_directory.exists(): + save_directory.mkdir(parents=True) + + for config_file, checkpoint_file in OriginalOneFormerCheckpointToOursConverter.using_dirs( + checkpoints_dir, config_dir + ): + processor = OriginalOneFormerConfigToProcessorConverter()( + setup_cfg(Args(config_file=config_file)), os.path.join("shi-labs", config_file.stem) + ) + + original_config = setup_cfg(Args(config_file=config_file)) + oneformer_kwargs = OriginalOneFormer.from_config(original_config) + + original_model = OriginalOneFormer(**oneformer_kwargs).eval() + + DetectionCheckpointer(original_model).load(str(checkpoint_file)) + + is_swin = "swin" in config_file.stem + + config: OneFormerConfig = OriginalOneFormerConfigToOursConverter()(original_config, is_swin) + + oneformer = OneFormerModel(config=config).eval() + + converter = OriginalOneFormerCheckpointToOursConverter(original_model, config) + + oneformer = converter.convert(oneformer, is_swin) + + oneformer_for_universal_segmentation = OneFormerForUniversalSegmentation(config=config).eval() + + oneformer_for_universal_segmentation.model = oneformer + + test( + original_model, + oneformer_for_universal_segmentation, + processor, + os.path.join("shi-labs", config_file.stem), + ) + + model_name = get_name(checkpoint_file) + logger.info(f"Saving {model_name}") + + processor.save_pretrained(save_directory / model_name) + oneformer_for_universal_segmentation.save_pretrained(save_directory / model_name) + + model_id = f"shi-labs/{model_name}" + processor.push_to_hub(repo_id=model_id) + oneformer_for_universal_segmentation.push_to_hub(repo_id=model_id) diff --git a/third_party/transformers/src/transformers/models/oneformer/image_processing_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/image_processing_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..ccee798bc39ec853f11f6e6fc4d7c5413cc952da --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/image_processing_oneformer.py @@ -0,0 +1,969 @@ +# Copyright 2025 SHI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for OneFormer.""" + +import json +import os +from typing import Union + +import torch +from torch import nn +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, + logging, +) + + +try: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import RepositoryNotFoundError +except ImportError: + hf_hub_download = None + RepositoryNotFoundError = None + + +logger = logging.get_logger(__name__) + + +class OneFormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + repo_path (`str`, *optional*, defaults to `shi-labs/oneformer_demo`): + Path to a local directory or HuggingFace Hub repository containing model metadata. + class_info_file (`str`, *optional*): + Path to the JSON file within the repository that contains class metadata. + num_text (`int`, *optional*): + Number of text queries for the text encoder, used as task-guiding prompts. + num_labels (`int`, *optional*): + Number of semantic classes for segmentation, determining the output layer's size. + ignore_index (`int`, *optional*): + Label to ignore in segmentation maps, often used for padding. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether to decrement all label values by 1, mapping the background class to `ignore_index`. + """ + + repo_path: str | None + class_info_file: str | None + num_text: int | None + num_labels: int | None + ignore_index: int | None + do_reduce_labels: bool + + +def prepare_metadata(class_info): + metadata = {} + class_names = [] + thing_ids = [] + for key, info in class_info.items(): + metadata[key] = info["name"] + class_names.append(info["name"]) + if info["isthing"]: + thing_ids.append(int(key)) + metadata["thing_ids"] = thing_ids + metadata["class_names"] = class_names + return metadata + + +def load_metadata(repo_id, class_info_file): + fname = os.path.join("" if repo_id is None else repo_id, class_info_file) + + if not os.path.exists(fname) or not os.path.isfile(fname): + if repo_id is None: + raise ValueError(f"Could not file {fname} locally. repo_id must be defined if loading from the hub") + if hf_hub_download is None: + raise ImportError( + "huggingface_hub is required to download metadata files. Install it with `pip install huggingface_hub`" + ) + # We try downloading from a dataset by default for backward compatibility + try: + fname = hf_hub_download(repo_id, class_info_file, repo_type="dataset") + except RepositoryNotFoundError: + fname = hf_hub_download(repo_id, class_info_file) + + with open(fname, "r") as f: + class_info = json.load(f) + + return class_info + + +def make_pixel_mask(image: "torch.Tensor", output_size: tuple[int, int]) -> "torch.Tensor": + """ + Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + + Args: + image (`torch.Tensor`): + Image to make the pixel mask for. + output_size (`Tuple[int, int]`): + Output size of the mask. + """ + + input_height, input_width = image.shape[-2], image.shape[-1] + mask = torch.zeros(output_size, dtype=torch.int64) + mask[:input_height, :input_width] = 1 + return mask + + +def binary_mask_to_rle(mask): + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `numpy.array`): + A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target + segment_id or class_id. + Returns: + `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE + format. + """ + import numpy as np + + # Convert to torch tensor if numpy array + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask) + + pixels = mask.flatten() + pixels = torch.concat( + [ + torch.tensor([0], device=pixels.device, dtype=pixels.dtype), + pixels, + torch.tensor([0], device=pixels.device, dtype=pixels.dtype), + ] + ) + runs = torch.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return list(runs) + + +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor` or `numpy.array`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels): + """ + Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and + `labels`. + + Args: + masks (`torch.Tensor`): + A tensor of shape `(num_queries, height, width)`. + scores (`torch.Tensor`): + A tensor of shape `(num_queries)`. + labels (`torch.Tensor`): + A tensor of shape `(num_queries)`. + object_mask_threshold (`float`): + A number between 0 and 1 used to binarize the masks. + Raises: + `ValueError`: Raised when the first dimension doesn't match in all input tensors. + Returns: + `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region + < `object_mask_threshold`. + """ + if not (masks.shape[0] == scores.shape[0] == labels.shape[0]): + raise ValueError("mask, scores and labels must have the same shape!") + + to_keep = labels.ne(num_labels) & (scores > object_mask_threshold) + + return masks[to_keep], scores[to_keep], labels[to_keep] + + +def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8): + # Get the mask associated with the k class + mask_k = mask_labels == k + mask_k_area = mask_k.sum() + + # Compute the area of all the stuff in query k + original_area = (mask_probs[k] >= mask_threshold).sum() + mask_exists = mask_k_area > 0 and original_area > 0 + + # Eliminate disconnected tiny segments + if mask_exists: + area_ratio = mask_k_area / original_area + if not area_ratio.item() > overlap_mask_area_threshold: + mask_exists = False + + return mask_exists, mask_k + + +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_size: tuple[int, int] | None = None, +): + height = mask_probs.shape[1] if target_size is None else target_size[0] + width = mask_probs.shape[2] if target_size is None else target_size[1] + + segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = tvF.resize( + mask_probs.unsqueeze(0), + size=target_size, + interpolation=tvF.InterpolationMode.BILINEAR, + )[0] + + current_segment_id = 0 + + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +@auto_docstring +class OneFormerImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + crop_size = None + do_resize = True + do_rescale = True + do_normalize = True + default_to_square = False + do_center_crop = False + do_convert_rgb = True + rescale_factor = 1 / 255 + ignore_index = None + do_reduce_labels = False + repo_path = "shi-labs/oneformer_demo" + class_info_file = None + num_text = None + num_labels = None + valid_kwargs = OneFormerImageProcessorKwargs + model_input_names = ["pixel_values", "pixel_mask", "task_inputs"] + + def __init__(self, **kwargs: Unpack[OneFormerImageProcessorKwargs]): + super().__init__(**kwargs) + if self.class_info_file: + self.metadata = prepare_metadata(load_metadata(self.repo_path, self.class_info_file)) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + task_inputs: list[str] | None = None, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[OneFormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + task_inputs (`list[str]`, *optional*): + List of tasks (`"panoptic"`, `"instance"`, `"semantic"`) for each image in the batch. + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, task_inputs, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + task_inputs: list[str] | None, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + device: Union[str, "torch.device"] | None = None, + **kwargs: Unpack[OneFormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, task_inputs, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + task_inputs: list[str] | None, + segmentation_maps: list["torch.Tensor"], + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + + processed_images_grouped = {} + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, resample=resample) + 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) + + processed_segmentation_maps = None + if segmentation_maps is not None: + grouped_segmentation_maps, grouped_segmentation_maps_index = group_images_by_shape( + segmentation_maps, disable_grouping=disable_grouping + ) + processed_segmentation_maps_grouped = {} + for shape, stacked_segmentation_maps in grouped_segmentation_maps.items(): + if do_resize: + stacked_segmentation_maps = self.resize( + stacked_segmentation_maps, size=size, resample=tvF.InterpolationMode.NEAREST_EXACT + ) + processed_segmentation_maps_grouped[shape] = stacked_segmentation_maps + processed_segmentation_maps = reorder_images( + processed_segmentation_maps_grouped, grouped_segmentation_maps_index + ) + + encoded_inputs = self.encode_inputs( + processed_images, + task_inputs, + segmentation_maps=processed_segmentation_maps, + instance_id_to_semantic_id=instance_id_to_semantic_id, + ignore_index=ignore_index, + do_reduce_labels=do_reduce_labels, + return_tensors=return_tensors, + ) + + return encoded_inputs + + def _pad_image( + self, + image: "torch.Tensor", + output_size: tuple[int, int], + constant_values: float = 0, + ) -> "torch.Tensor": + """ + Pad an image with zeros to the given size using torch operations. + + Args: + image (`torch.Tensor`): + Image tensor in channel-first format (C, H, W). + output_size (`tuple[int, int]`): + Target output size (height, width). + constant_values (`float`, *optional*, defaults to 0): + The value to use for padding. + + Returns: + `torch.Tensor`: The padded image. + """ + input_height, input_width = image.shape[1], image.shape[2] + output_height, output_width = output_size + + pad_bottom = output_height - input_height + pad_right = output_width - input_width + + padded_image = tvF.pad(image, padding=[0, 0, pad_right, pad_bottom], fill=constant_values) + + return padded_image + + def pad( + self, + images: list["torch.Tensor"], + return_pixel_mask: bool = True, + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + """ + Pad a batch of images to the same size using torch operations. + + Args: + images (`List[torch.Tensor]`): + List of image tensors in channel-first format. + return_pixel_mask (`bool`, *optional*, defaults to `True`): + Whether to return pixel masks. + return_tensors (`str` or `TensorType`, *optional*): + The type of tensors to return. + + Returns: + `BatchFeature`: Padded images and optional pixel masks. + """ + outputs = super().pad(images, return_mask=return_pixel_mask) + padded_images = outputs[0] if return_pixel_mask else outputs + pixel_masks = outputs[1] if return_pixel_mask else None + + if return_tensors: + padded_images = torch.stack(padded_images, dim=0) + if return_pixel_mask: + pixel_masks = torch.stack(pixel_masks, dim=0) + + data = {"pixel_values": padded_images} + if return_pixel_mask: + data["pixel_mask"] = pixel_masks + + return BatchFeature(data=data, tensor_type=return_tensors) + + def convert_segmentation_map_to_binary_masks( + self, + segmentation_map: "torch.Tensor", + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, + ): + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = torch.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = torch.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] + + binary_masks = [(segmentation_map == i) for i in all_labels] + + if binary_masks: + binary_masks = torch.stack(binary_masks, dim=0) + else: + binary_masks = torch.zeros((0, *segmentation_map.shape), device=segmentation_map.device) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = torch.zeros(all_labels.shape[0], device=segmentation_map.device) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(label.item() + 1 if do_reduce_labels else label.item())] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels + + return ( + binary_masks.float(), + labels.long(), + ) + + def get_semantic_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["a semantic photo"] * self.num_text + classes = [] + masks = [] + + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] + if not torch.all(mask == 0): + if class_id not in classes: + cls_name = self.metadata[str(class_id.cpu().item())] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + else: + idx = classes.index(class_id) + masks[idx] += mask + masks[idx] = torch.clamp(masks[idx], 0, 1) + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = torch.stack(classes) + masks = torch.stack(masks) + return classes, masks, texts + + def get_instance_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["an instance photo"] * self.num_text + classes = [] + masks = [] + + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] + + if class_id in self.metadata["thing_ids"]: + if not torch.all(mask == 0): + cls_name = self.metadata[str(class_id.cpu().item())] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = torch.stack(classes) + masks = torch.stack(masks) + return classes, masks, texts + + def get_panoptic_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["an panoptic photo"] * self.num_text + classes = [] + masks = [] + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] if hasattr(annotation_masks[idx], "data") else annotation_masks[idx] + if not torch.all(mask == 0): + cls_name = self.metadata[str(class_id.cpu().item())] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = torch.stack(classes) + masks = torch.stack(masks) + return classes, masks, texts + + def encode_inputs( + self, + pixel_values_list: list["ImageInput"], + task_inputs: list[str] | None = None, + segmentation_maps: list["ImageInput"] | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + ignore_index = self.ignore_index if ignore_index is None else ignore_index + do_reduce_labels = self.do_reduce_labels if do_reduce_labels is None else do_reduce_labels + if task_inputs is None: + task_inputs = ["panoptic"] + pixel_values_list = self._prepare_image_like_inputs( + pixel_values_list, input_data_format=ChannelDimension.FIRST + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + pad_size = get_max_height_width(pixel_values_list) + encoded_inputs = self.pad(pixel_values_list, return_tensors=return_tensors) + + annotations = None + if segmentation_maps is not None: + annotations = [] + for idx, segmentation_map in enumerate(segmentation_maps): + # Use instance2class_id mapping per image + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + + # Convert segmentation map to binary masks using torch operations + masks, classes = self.convert_segmentation_map_to_binary_masks( + segmentation_map, + instance_id, + ignore_index=ignore_index, + do_reduce_labels=do_reduce_labels, + ) + + annotations.append({"masks": masks, "classes": classes}) + + if annotations is not None: + mask_labels = [] + class_labels = [] + text_inputs = [] + num_class_obj = dict.fromkeys(self.metadata["class_names"], 0) + + for i, label in enumerate(annotations): + task = task_inputs[i] + + if task == "semantic": + classes, masks, texts = self.get_semantic_annotations(label, num_class_obj) + elif task == "instance": + classes, masks, texts = self.get_instance_annotations(label, num_class_obj) + elif task == "panoptic": + classes, masks, texts = self.get_panoptic_annotations(label, num_class_obj) + else: + raise ValueError(f"{task} was not expected, expected `semantic`, `instance` or `panoptic`") + # Pad masks to max size using torch operations + padded_masks = [ + self._pad_image(image=mask, output_size=pad_size, constant_values=ignore_index) for mask in masks + ] + padded_masks = torch.cat(padded_masks, dim=0) + mask_labels.append(padded_masks) + class_labels.append(classes) + text_inputs.append(texts) + + encoded_inputs["mask_labels"] = mask_labels + encoded_inputs["class_labels"] = class_labels + encoded_inputs["text_inputs"] = text_inputs + + encoded_inputs["task_inputs"] = [f"the task is {task_input}" for task_input in task_inputs] + return encoded_inputs + + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`MaskFormerForInstanceSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`MaskFormerForInstanceSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *optional*): + List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If left to None, predictions will not be resized. + Returns: + `List[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = tvF.resize( + segmentation[idx].unsqueeze(dim=0), + size=target_sizes[idx], + interpolation=tvF.InterpolationMode.BILINEAR, + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + def post_process_instance_segmentation( + self, + outputs, + task_type: str = "instance", + is_demo: bool = True, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + ): + """ + Converts the output of [`OneFormerForUniversalSegmentationOutput`] into image instance segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`OneFormerForUniversalSegmentationOutput`]): + The outputs from [`OneFormerForUniversalSegmentationOutput`]. + task_type (`str`, *optional*, defaults to "instance"): + The post processing depends on the task token input. If the `task_type` is "panoptic", we need to + ignore the stuff predictions. + is_demo (`bool`, *optional)*, defaults to `True`): + Whether the model is in demo mode. If true, use threshold to predict final masks. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`List[Tuple]`, *optional*): + List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If left to None, predictions will not be + resized. + return_coco_annotation (`bool`, *optional)*, defaults to `False`): + Whether to return predictions in COCO format. + + Returns: + `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set + to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized + to the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + device = masks_queries_logits.device + batch_size = class_queries_logits.shape[0] + num_queries = class_queries_logits.shape[1] + num_classes = class_queries_logits.shape[-1] - 1 + + # Loop over items in batch size + results: list[dict[str, torch.Tensor]] = [] + + for i in range(batch_size): + # [Q, K] + scores = nn.functional.softmax(class_queries_logits[i], dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False) + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1) + mask_pred = masks_queries_logits[i][topk_indices] + + # Only consider scores with confidence over [threshold] for demo + if is_demo: + keep = scores_per_image > threshold + scores_per_image = scores_per_image[keep] + labels_per_image = labels_per_image[keep] + mask_pred = mask_pred[keep] + + # if this is panoptic segmentation, we only keep the "thing" classes + if task_type == "panoptic": + keep = torch.zeros_like(scores_per_image).bool() + for j, lab in enumerate(labels_per_image): + keep[j] = lab in self.metadata["thing_ids"] + + scores_per_image = scores_per_image[keep] + labels_per_image = labels_per_image[keep] + mask_pred = mask_pred[keep] + + if mask_pred.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_pred.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + if "ade20k" in self.class_info_file and not is_demo and "instance" in task_type: + for j in range(labels_per_image.shape[0]): + labels_per_image[j] = self.metadata["thing_ids"].index(labels_per_image[j].item()) + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_pred, + scores_per_image, + labels_per_image, + mask_threshold, + overlap_mask_area_threshold, + set(), + target_size, + ) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + # Copied from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.post_process_panoptic_segmentation + def post_process_panoptic_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_sizes: list[tuple[int, int]] | None = None, + ) -> list[dict]: + """ + Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`MaskFormerForInstanceSegmentationOutput`]): + The outputs from [`MaskFormerForInstanceSegmentation`]. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + label_ids_to_fuse (`Set[int]`, *optional*): + The labels in this state will have all their instances be fused together. For instance we could say + there can only be one sky in an image, but several persons, so the label ID for sky would be in that + set, but not the one for person. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If left to None, predictions will not be + resized. + + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set + to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized + to the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + + if label_ids_to_fuse is None: + logger.warning("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=label_ids_to_fuse, + target_size=target_size, + ) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + +__all__ = ["OneFormerImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/oneformer/image_processing_pil_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/image_processing_pil_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..bed17ab5f4eb0fbf2f15483efb226c6bc0d4c8b7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/image_processing_pil_oneformer.py @@ -0,0 +1,1008 @@ +# Copyright 2025 SHI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for OneFormer.""" + +import json +import os + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import PaddingMode +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available, is_torchvision_available, logging +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + +if is_torch_available(): + import torch + from torch import nn + +if is_torchvision_available(): + import torchvision.transforms.v2.functional as tvF + +try: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import RepositoryNotFoundError +except ImportError: + hf_hub_download = None + RepositoryNotFoundError = None + + +def make_pixel_mask(image: np.ndarray, output_size: tuple[int, int]) -> np.ndarray: + """ + Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + + Args: + image (`np.ndarray`): + Image to make the pixel mask for. + output_size (`Tuple[int, int]`): + Output size of the mask. + """ + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + mask = np.zeros(output_size, dtype=np.int64) + mask[:input_height, :input_width] = 1 + return mask + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.OneFormerImageProcessorKwargs +class OneFormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + repo_path (`str`, *optional*, defaults to `shi-labs/oneformer_demo`): + Path to a local directory or HuggingFace Hub repository containing model metadata. + class_info_file (`str`, *optional*): + Path to the JSON file within the repository that contains class metadata. + num_text (`int`, *optional*): + Number of text queries for the text encoder, used as task-guiding prompts. + num_labels (`int`, *optional*): + Number of semantic classes for segmentation, determining the output layer's size. + ignore_index (`int`, *optional*): + Label to ignore in segmentation maps, often used for padding. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether to decrement all label values by 1, mapping the background class to `ignore_index`. + """ + + repo_path: str | None + class_info_file: str | None + num_text: int | None + num_labels: int | None + ignore_index: int | None + do_reduce_labels: bool + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.binary_mask_to_rle +def binary_mask_to_rle(mask): + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `numpy.array`): + A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target + segment_id or class_id. + Returns: + `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE + format. + """ + from ...utils import is_torch_tensor + + if is_torch_tensor(mask): + mask = mask.numpy() + + pixels = mask.flatten() + pixels = np.concatenate([[0], pixels, [0]]) + runs = np.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return list(runs) + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.check_segment_validity +def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8): + # Get the mask associated with the k class + mask_k = mask_labels == k + mask_k_area = mask_k.sum() + + # Compute the area of all the stuff in query k + original_area = (mask_probs[k] >= mask_threshold).sum() + mask_exists = mask_k_area > 0 and original_area > 0 + + # Eliminate disconnected tiny segments + if mask_exists: + area_ratio = mask_k_area / original_area + if not area_ratio.item() > overlap_mask_area_threshold: + mask_exists = False + + return mask_exists, mask_k + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.compute_segments +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_size: tuple[int, int] | None = None, +): + height = mask_probs.shape[1] if target_size is None else target_size[0] + width = mask_probs.shape[2] if target_size is None else target_size[1] + + segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = tvF.resize( + mask_probs.unsqueeze(0), + size=target_size, + interpolation=tvF.InterpolationMode.BILINEAR, + )[0] + + current_segment_id = 0 + + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.convert_segmentation_to_rle +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor` or `numpy.array`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.load_metadata +def load_metadata(repo_id, class_info_file): + fname = os.path.join("" if repo_id is None else repo_id, class_info_file) + + if not os.path.exists(fname) or not os.path.isfile(fname): + if repo_id is None: + raise ValueError(f"Could not file {fname} locally. repo_id must be defined if loading from the hub") + if hf_hub_download is None: + raise ImportError( + "huggingface_hub is required to download metadata files. Install it with `pip install huggingface_hub`" + ) + # We try downloading from a dataset by default for backward compatibility + try: + fname = hf_hub_download(repo_id, class_info_file, repo_type="dataset") + except RepositoryNotFoundError: + fname = hf_hub_download(repo_id, class_info_file) + + with open(fname, "r") as f: + class_info = json.load(f) + + return class_info + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.prepare_metadata +def prepare_metadata(class_info): + metadata = {} + class_names = [] + thing_ids = [] + for key, info in class_info.items(): + metadata[key] = info["name"] + class_names.append(info["name"]) + if info["isthing"]: + thing_ids.append(int(key)) + metadata["thing_ids"] = thing_ids + metadata["class_names"] = class_names + return metadata + + +# Adapted from transformers.models.oneformer.image_processing_oneformer.remove_low_and_no_objects +def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels): + """ + Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and + `labels`. + + Args: + masks (`torch.Tensor`): + A tensor of shape `(num_queries, height, width)`. + scores (`torch.Tensor`): + A tensor of shape `(num_queries)`. + labels (`torch.Tensor`): + A tensor of shape `(num_queries)`. + object_mask_threshold (`float`): + A number between 0 and 1 used to binarize the masks. + Raises: + `ValueError`: Raised when the first dimension doesn't match in all input tensors. + Returns: + `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region + < `object_mask_threshold`. + """ + if not (masks.shape[0] == scores.shape[0] == labels.shape[0]): + raise ValueError("mask, scores and labels must have the same shape!") + + to_keep = labels.ne(num_labels) & (scores > object_mask_threshold) + + return masks[to_keep], scores[to_keep], labels[to_keep] + + +@auto_docstring +@requires(backends=("torch",)) +class OneFormerImageProcessorPil(PilBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + crop_size = None + do_resize = True + do_rescale = True + do_normalize = True + default_to_square = False + do_center_crop = False + do_convert_rgb = True + rescale_factor = 1 / 255 + ignore_index = None + do_reduce_labels = False + repo_path = "shi-labs/oneformer_demo" + class_info_file = None + num_text = None + num_labels = None + valid_kwargs = OneFormerImageProcessorKwargs + model_input_names = ["pixel_values", "pixel_mask", "task_inputs"] + + def __init__(self, **kwargs: Unpack[OneFormerImageProcessorKwargs]): + super().__init__(**kwargs) + if self.class_info_file: + self.metadata = prepare_metadata(load_metadata(self.repo_path, self.class_info_file)) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + task_inputs: list[str] | None = None, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[OneFormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + task_inputs (`list[str]`, *optional*): + List of tasks (`"panoptic"`, `"instance"`, `"semantic"`) for each image in the batch. + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, task_inputs, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + task_inputs: list[str] | None, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + **kwargs: Unpack[OneFormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, task_inputs, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + task_inputs: list[str] | None, + segmentation_maps: list[np.ndarray], + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + # Process images one by one (no batching in PIL backend) + processed_images = [] + processed_segmentation_maps = None + if segmentation_maps is not None: + processed_segmentation_maps = [] + + for idx, image in enumerate(images): + 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) + + if segmentation_maps is not None: + seg_map = segmentation_maps[idx] + if do_resize: + seg_map = self.resize(image=seg_map, size=size, resample=PILImageResampling.NEAREST) + processed_segmentation_maps.append(seg_map) + + encoded_inputs = self.encode_inputs( + processed_images, + task_inputs, + segmentation_maps=processed_segmentation_maps, + instance_id_to_semantic_id=instance_id_to_semantic_id, + ignore_index=ignore_index, + do_reduce_labels=do_reduce_labels, + return_tensors=return_tensors, + ) + + return encoded_inputs + + def _pad_image(self, image: np.ndarray, output_size: tuple[int, int], constant_values: float = 0) -> np.ndarray: + """ + Pad an image with zeros to the given size using numpy operations. + + Args: + image (`np.ndarray`): + Image array in channel-first format (C, H, W) or (H, W). + output_size (`tuple[int, int]`): + Target output size (height, width). + constant_values (`float`, *optional*, defaults to 0): + The value to use for padding. + + Returns: + `np.ndarray`: The padded image. + """ + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = output_size + + pad_bottom = output_height - input_height + pad_right = output_width - input_width + + # For 2D arrays (masks), use np.pad directly + # For 3D arrays (images), use np_pad which handles channel dimension + if image.ndim == 2: + padding = ((0, pad_bottom), (0, pad_right)) + padded_image = np.pad(image, padding, mode="constant", constant_values=constant_values) + else: + padding = ((0, pad_bottom), (0, pad_right)) + padded_image = np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=constant_values, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + + return padded_image + + def pad( + self, images: list[np.ndarray], return_pixel_mask: bool = True, return_tensors: str | TensorType | None = None + ) -> BatchFeature: + """ + Pad a batch of images to the same size using numpy operations. + + Args: + images (`List[np.ndarray]`): + List of image arrays in channel-first format. + return_pixel_mask (`bool`, *optional*, defaults to `True`): + Whether to return pixel masks. + return_tensors (`str` or `TensorType`, *optional*): + The type of tensors to return. + + Returns: + `BatchFeature`: Padded images and optional pixel masks. + """ + pad_size = get_max_height_width(images, input_data_format=ChannelDimension.FIRST) + + padded_images = [] + pixel_masks = [] + for image in images: + padded_image = self._pad_image(image, pad_size, constant_values=0) + padded_images.append(padded_image) + if return_pixel_mask: + pixel_mask = make_pixel_mask(image, pad_size) + pixel_masks.append(pixel_mask) + + if return_tensors == "pt": + padded_images = [torch.from_numpy(img) for img in padded_images] + padded_images = torch.stack(padded_images, dim=0) + if return_pixel_mask: + pixel_masks = [torch.from_numpy(mask) for mask in pixel_masks] + pixel_masks = torch.stack(pixel_masks, dim=0) + + data = {"pixel_values": padded_images} + if return_pixel_mask: + data["pixel_mask"] = pixel_masks + + return BatchFeature(data=data, tensor_type=return_tensors) + + def convert_segmentation_map_to_binary_masks( + self, + segmentation_map: np.ndarray, + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, + ): + """Convert segmentation map to binary masks using NumPy operations.""" + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = np.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] + + binary_masks = [(segmentation_map == i) for i in all_labels] + if binary_masks: + binary_masks = np.stack(binary_masks, axis=0) + else: + binary_masks = np.zeros((0, *segmentation_map.shape), dtype=np.float32) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = np.zeros(all_labels.shape[0], dtype=np.int64) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(int(label) + 1 if do_reduce_labels else int(label))] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels.astype(np.int64) + return binary_masks.astype(np.float32), labels + + def get_semantic_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["a semantic photo"] * self.num_text + classes = [] + masks = [] + + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] + if not np.all(mask == 0): + if class_id not in classes: + cls_name = self.metadata[str(class_id)] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + else: + idx = classes.index(class_id) + masks[idx] += mask + masks[idx] = np.clip(masks[idx], 0, 1) + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = np.array(classes) if classes else np.array([], dtype=np.int64) + # Stack masks into a 3D array (num_masks, H, W) to match torchvision version + if masks: + masks = np.stack(masks, axis=0) + else: + # Empty masks - use shape from first annotation mask if available + if annotation_masks and len(annotation_masks) > 0: + mask_shape = annotation_masks[0].shape[-2:] if hasattr(annotation_masks[0], "shape") else (0, 0) + else: + mask_shape = (0, 0) + masks = np.zeros((0, *mask_shape), dtype=np.float32) + return classes, masks, texts + + def get_instance_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["an instance photo"] * self.num_text + classes = [] + masks = [] + + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] + + if class_id in self.metadata["thing_ids"]: + if not np.all(mask == 0): + cls_name = self.metadata[str(class_id)] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = np.array(classes) if classes else np.array([], dtype=np.int64) + # Stack masks into a 3D array (num_masks, H, W) to match torchvision version + if masks: + masks = np.stack(masks, axis=0) + else: + # Empty masks - use shape from first annotation mask if available + if annotation_masks and len(annotation_masks) > 0: + mask_shape = annotation_masks[0].shape[-2:] if hasattr(annotation_masks[0], "shape") else (0, 0) + else: + mask_shape = (0, 0) + masks = np.zeros((0, *mask_shape), dtype=np.float32) + return classes, masks, texts + + def get_panoptic_annotations(self, label, num_class_obj): + annotation_classes = label["classes"] + annotation_masks = label["masks"] + + texts = ["an panoptic photo"] * self.num_text + classes = [] + masks = [] + for idx in range(len(annotation_classes)): + class_id = annotation_classes[idx] + mask = annotation_masks[idx] if hasattr(annotation_masks[idx], "data") else annotation_masks[idx] + if not np.all(mask == 0): + cls_name = self.metadata[str(class_id)] + classes.append(class_id) + masks.append(mask) + num_class_obj[cls_name] += 1 + + num = 0 + for i, cls_name in enumerate(self.metadata["class_names"]): + if num_class_obj[cls_name] > 0: + for _ in range(num_class_obj[cls_name]): + if num >= len(texts): + break + texts[num] = f"a photo with a {cls_name}" + num += 1 + + classes = np.array(classes) if classes else np.array([], dtype=np.int64) + # Stack masks into a 3D array (num_masks, H, W) to match torchvision version + if masks: + masks = np.stack(masks, axis=0) + else: + # Empty masks - use shape from first annotation mask if available + if annotation_masks and len(annotation_masks) > 0: + mask_shape = annotation_masks[0].shape[-2:] if hasattr(annotation_masks[0], "shape") else (0, 0) + else: + mask_shape = (0, 0) + masks = np.zeros((0, *mask_shape), dtype=np.float32) + return classes, masks, texts + + def encode_inputs( + self, + pixel_values_list: list[np.ndarray], + task_inputs: list[str] | None = None, + segmentation_maps: list[np.ndarray] | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + ignore_index = self.ignore_index if ignore_index is None else ignore_index + do_reduce_labels = self.do_reduce_labels if do_reduce_labels is None else do_reduce_labels + if task_inputs is None: + task_inputs = ["panoptic"] + pixel_values_list = self._prepare_image_like_inputs( + pixel_values_list, input_data_format=ChannelDimension.FIRST + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + pad_size = get_max_height_width(pixel_values_list, input_data_format=ChannelDimension.FIRST) + encoded_inputs = self.pad(pixel_values_list, return_tensors=return_tensors) + + annotations = None + if segmentation_maps is not None: + annotations = [] + for idx, segmentation_map in enumerate(segmentation_maps): + # Use instance2class_id mapping per image + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + + # Squeeze channel dimension if present + if segmentation_map.ndim == 3 and segmentation_map.shape[0] == 1: + segmentation_map = segmentation_map.squeeze(0) + + # Convert segmentation map to binary masks using numpy operations + masks, classes = self.convert_segmentation_map_to_binary_masks( + segmentation_map, instance_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels + ) + + annotations.append({"masks": masks, "classes": classes}) + + if annotations is not None: + mask_labels = [] + class_labels = [] + text_inputs = [] + num_class_obj = dict.fromkeys(self.metadata["class_names"], 0) + + for i, label in enumerate(annotations): + task = task_inputs[i] + + if task == "semantic": + classes, masks, texts = self.get_semantic_annotations(label, num_class_obj) + elif task == "instance": + classes, masks, texts = self.get_instance_annotations(label, num_class_obj) + elif task == "panoptic": + classes, masks, texts = self.get_panoptic_annotations(label, num_class_obj) + else: + raise ValueError(f"{task} was not expected, expected `semantic`, `instance` or `panoptic`") + # Pad masks to max size using numpy operations + # masks is a 3D array (num_masks, H, W), iterate to get 2D slices + padded_masks = [ + self._pad_image(image=mask, output_size=pad_size, constant_values=ignore_index) for mask in masks + ] + # Stack padded masks back into 3D array (num_masks, padded_H, padded_W) + padded_masks = ( + np.stack(padded_masks, axis=0) if padded_masks else np.zeros((0, *pad_size), dtype=np.float32) + ) + mask_labels.append(padded_masks) + class_labels.append(classes) + text_inputs.append(texts) + + encoded_inputs["mask_labels"] = [ + torch.from_numpy(mask_label) if return_tensors == "pt" else mask_label for mask_label in mask_labels + ] + encoded_inputs["class_labels"] = [ + torch.from_numpy(class_label) if return_tensors == "pt" else class_label + for class_label in class_labels + ] + encoded_inputs["text_inputs"] = text_inputs + + encoded_inputs["task_inputs"] = [f"the task is {task_input}" for task_input in task_inputs] + return encoded_inputs + + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`MaskFormerForInstanceSegmentation`] into semantic segmentation maps. Only supports + PyTorch. + + Args: + outputs ([`MaskFormerForInstanceSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple[int, int]]`, *optional*): + List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If left to None, predictions will not be resized. + Returns: + `List[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = torch.nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + def post_process_instance_segmentation( + self, + outputs, + task_type: str = "instance", + is_demo: bool = True, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + ): + """ + Converts the output of [`OneFormerForUniversalSegmentationOutput`] into image instance segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`OneFormerForUniversalSegmentationOutput`]): + The outputs from [`OneFormerForUniversalSegmentationOutput`]. + task_type (`str`, *optional*, defaults to "instance"): + The post processing depends on the task token input. If the `task_type` is "panoptic", we need to + ignore the stuff predictions. + is_demo (`bool`, *optional)*, defaults to `True`): + Whether the model is in demo mode. If true, use threshold to predict final masks. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`List[Tuple]`, *optional*): + List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If left to None, predictions will not be + resized. + return_coco_annotation (`bool`, *optional)*, defaults to `False`): + Whether to return predictions in COCO format. + + Returns: + `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set + to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized + to the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + device = masks_queries_logits.device + batch_size = class_queries_logits.shape[0] + num_queries = class_queries_logits.shape[1] + num_classes = class_queries_logits.shape[-1] - 1 + + # Loop over items in batch size + results: list[dict[str, torch.Tensor]] = [] + + for i in range(batch_size): + # [Q, K] + scores = nn.functional.softmax(class_queries_logits[i], dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False) + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1) + mask_pred = masks_queries_logits[i][topk_indices] + + # Only consider scores with confidence over [threshold] for demo + if is_demo: + keep = scores_per_image > threshold + scores_per_image = scores_per_image[keep] + labels_per_image = labels_per_image[keep] + mask_pred = mask_pred[keep] + + # if this is panoptic segmentation, we only keep the "thing" classes + if task_type == "panoptic": + keep = torch.zeros_like(scores_per_image).bool() + for j, lab in enumerate(labels_per_image): + keep[j] = lab in self.metadata["thing_ids"] + + scores_per_image = scores_per_image[keep] + labels_per_image = labels_per_image[keep] + mask_pred = mask_pred[keep] + + if mask_pred.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_pred.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + if "ade20k" in self.class_info_file and not is_demo and "instance" in task_type: + for j in range(labels_per_image.shape[0]): + labels_per_image[j] = self.metadata["thing_ids"].index(labels_per_image[j].item()) + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_pred, + scores_per_image, + labels_per_image, + mask_threshold, + overlap_mask_area_threshold, + set(), + target_size, + ) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + # Adapted from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.post_process_panoptic_segmentation + def post_process_panoptic_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_sizes: list[tuple[int, int]] | None = None, + ) -> list[dict]: + """ + Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image panoptic segmentation + predictions. Only supports PyTorch. + + Args: + outputs ([`MaskFormerForInstanceSegmentationOutput`]): + The outputs from [`MaskFormerForInstanceSegmentation`]. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + label_ids_to_fuse (`Set[int]`, *optional*): + The labels in this state will have all their instances be fused together. For instance we could say + there can only be one sky in an image, but several persons, so the label ID for sky would be in that + set, but not the one for person. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If left to None, predictions will not be + resized. + + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set + to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized + to the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + + if label_ids_to_fuse is None: + logger.warning("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=label_ids_to_fuse, + target_size=target_size, + ) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + +__all__ = ["OneFormerImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/oneformer/modeling_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/modeling_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..57214df16d8250b03b8a9f184e992acb163e5fb4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/modeling_oneformer.py @@ -0,0 +1,3182 @@ +# Copyright 2022 SHI Labs and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch OneFormer model.""" + +import copy +import math +import warnings +from dataclasses import dataclass + +import numpy as np +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import load_backbone +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import ( + ModelOutput, + auto_docstring, + is_accelerate_available, + is_scipy_available, + logging, + requires_backends, + torch_compilable_check, +) +from ...utils.generic import maybe_autocast +from .configuration_oneformer import OneFormerConfig + + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import reduce + +logger = logging.get_logger(__name__) + + +if is_scipy_available(): + from scipy.optimize import linear_sum_assignment + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def multi_scale_deformable_attention( + value: Tensor, + value_spatial_shapes: Tensor | list[tuple], + sampling_locations: Tensor, + attention_weights: Tensor, +) -> Tensor: + batch_size, _, num_heads, hidden_dim = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([height * width for height, width in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level_id, (height, width) in enumerate(value_spatial_shapes): + # batch_size, height*width, num_heads, hidden_dim + # -> batch_size, height*width, num_heads*hidden_dim + # -> batch_size, num_heads*hidden_dim, height*width + # -> batch_size*num_heads, hidden_dim, height, width + value_l_ = ( + value_list[level_id].flatten(2).transpose(1, 2).reshape(batch_size * num_heads, hidden_dim, height, width) + ) + # batch_size, num_queries, num_heads, num_points, 2 + # -> batch_size, num_heads, num_queries, num_points, 2 + # -> batch_size*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1) + # batch_size*num_heads, hidden_dim, num_queries, num_points + sampling_value_l_ = nn.functional.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (batch_size, num_queries, num_heads, num_levels, num_points) + # -> (batch_size, num_heads, num_queries, num_levels, num_points) + # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + batch_size * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(batch_size, num_heads * hidden_dim, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +# Copied from transformers.models.maskformer.modeling_maskformer.dice_loss +def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor: + r""" + Compute the DICE loss, similar to generalized IOU for masks as follows: + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$ + + In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow + + $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$ + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + num_masks (`int`): + The number of masks present in the current batch, used for normalization. + + Returns: + `torch.Tensor`: The computed loss. + """ + probs = inputs.sigmoid().flatten(1) + numerator = 2 * (probs * labels).sum(-1) + denominator = probs.sum(-1) + labels.sum(-1) + loss = 1 - (numerator + 1) / (denominator + 1) + loss = loss.sum() / num_masks + return loss + + +# Copied from transformers.models.mask2former.modeling_mask2former.sigmoid_cross_entropy_loss +def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor: + r""" + Args: + inputs (`torch.Tensor`): + A float tensor of arbitrary shape. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss. + """ + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss = criterion(inputs, labels) + + loss = cross_entropy_loss.mean(1).sum() / num_masks + return loss + + +# Copied from transformers.models.maskformer.modeling_maskformer.pair_wise_dice_loss +def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor: + """ + A pair wise version of the dice loss, see `dice_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + `torch.Tensor`: The computed loss between each pairs. + """ + inputs = inputs.sigmoid().flatten(1) + numerator = 2 * torch.matmul(inputs, labels.T) + # using broadcasting to get a [num_queries, NUM_CLASSES] matrix + denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :] + loss = 1 - (numerator + 1) / (denominator + 1) + return loss + + +# Copied from transformers.models.mask2former.modeling_mask2former.pair_wise_sigmoid_cross_entropy_loss +def pair_wise_sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + r""" + A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage. + + Args: + inputs (`torch.Tensor`): + A tensor representing a mask. + labels (`torch.Tensor`): + A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs + (0 for the negative class and 1 for the positive class). + + Returns: + loss (`torch.Tensor`): The computed loss between each pairs. + """ + + height_and_width = inputs.shape[1] + + criterion = nn.BCEWithLogitsLoss(reduction="none") + cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) + cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) + + loss_pos = torch.matmul(cross_entropy_loss_pos / height_and_width, labels.T) + loss_neg = torch.matmul(cross_entropy_loss_neg / height_and_width, (1 - labels).T) + loss = loss_pos + loss_neg + return loss + + +# Copied from transformers.models.mask2former.modeling_mask2former.sample_point +def sample_point( + input_features: torch.Tensor, point_coordinates: torch.Tensor, add_dim=False, **kwargs +) -> torch.Tensor: + """ + A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors. + + Args: + input_features (`torch.Tensor` of shape (batch_size, channels, height, width)): + A tensor that contains features map on a height * width grid + point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,: + 2)): + A tensor that contains [0, 1] * [0, 1] normalized point coordinates + add_dim (`bool`): + boolean value to keep track of added dimension + + Returns: + point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels, + height_grid, width_grid): + A tensor that contains features for points in `point_coordinates`. + """ + if point_coordinates.dim() == 3: + add_dim = True + point_coordinates = point_coordinates.unsqueeze(2) + + # use nn.function.grid_sample to get features for points in `point_coordinates` via bilinear interpolation + point_features = torch.nn.functional.grid_sample(input_features, 2.0 * point_coordinates - 1.0, **kwargs) + if add_dim: + point_features = point_features.squeeze(3) + + return point_features + + +# Refactored from https://github.com/SHI-Labs/OneFormer/blob/33ebb56ed34f970a30ae103e786c0cb64c653d9a/oneformer/modeling/matcher.py#L93 +class OneFormerHungarianMatcher(nn.Module): + def __init__( + self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0, num_points: int = 12544 + ): + """This class computes an assignment between the labels and the predictions of the network. + + For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more + predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are + un-matched (and thus treated as non-objects). + + Params: + cost_class (float, *optional*, defaults to 1.0): + This is the relative weight of the classification error in the matching cost. + cost_mask (float, *optional*, defaults to 1.0): + This is the relative weight of the sigmoid ce loss of the binary mask in the matching cost. + cost_dice (float, *optional*, defaults to 1.0): + This is the relative weight of the dice loss of the binary mask in the matching cost + num_points (int, *optional*, defaults to 12544): + Number of points to be sampled for dice and mask loss matching cost. + """ + super().__init__() + if cost_class == 0 and cost_mask == 0 and cost_dice == 0: + raise ValueError("All costs can't be 0") + self.cost_class = cost_class + self.cost_mask = cost_mask + self.cost_dice = cost_dice + self.num_points = num_points + + @torch.no_grad() + def forward(self, masks_queries_logits, class_queries_logits, mask_labels, class_labels) -> list[tuple[Tensor]]: + """Performs the matching + + Params: + masks_queries_logits (`torch.Tensor`): + A tensor` of dim `batch_size, num_queries, num_labels` with the + classification logits. + class_queries_logits (`torch.Tensor`): + A tensor` of dim `batch_size, num_queries, height, width` with the + predicted masks. + + class_labels (`torch.Tensor`): + A tensor` of dim `num_target_boxes` (where num_target_boxes is the number + of ground-truth objects in the target) containing the class labels. + mask_labels (`torch.Tensor`): + A tensor` of dim `num_target_boxes, height, width` containing the target + masks. + + Returns: + `list[tuple[Tensor]]`: A list of size batch_size, containing tuples of (index_i, index_j) where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected labels (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_targets). + """ + indices: list[tuple[np.array]] = [] + + num_queries = class_queries_logits.shape[1] + + preds_masks = masks_queries_logits + preds_probs = class_queries_logits + # iterate through batch size + for pred_probs, pred_mask, target_mask, labels in zip(preds_probs, preds_masks, mask_labels, class_labels): + pred_probs = pred_probs.softmax(-1) + # Compute the classification cost. Contrary to the loss, we don't use the NLL, + # but approximate it in 1 - proba[target class]. + # The 1 is a constant that doesn't change the matching, it can be omitted. + cost_class = -pred_probs[:, labels] + + pred_mask = pred_mask[:, None] + target_mask = target_mask[:, None].to(pred_mask.device) + + # all masks share the same set of points for efficient matching! + point_coords = torch.rand(1, self.num_points, 2, device=pred_mask.device) + + # get ground truth labels + target_mask = sample_point( + target_mask, + point_coords.repeat(target_mask.shape[0], 1, 1), + align_corners=False, + ).squeeze(1) + + pred_mask = sample_point( + pred_mask, + point_coords.repeat(pred_mask.shape[0], 1, 1), + align_corners=False, + ).squeeze(1) + + with maybe_autocast(device_type="cuda", enabled=False): + pred_mask = pred_mask.float() + target_mask = target_mask.float() + + # compute the sigmoid ce loss + cost_mask = pair_wise_sigmoid_cross_entropy_loss(pred_mask, target_mask) + # Compute the dice loss + cost_dice = pair_wise_dice_loss(pred_mask, target_mask) + # final cost matrix + cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice + cost_matrix = cost_matrix.reshape(num_queries, -1).cpu() + # do the assigmented using the hungarian algorithm in scipy + assigned_indices: tuple[np.array] = linear_sum_assignment(cost_matrix.cpu()) + indices.append(assigned_indices) + + # It could be stacked in one tensor + matched_indices = [ + (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices + ] + return matched_indices + + +class OneFormerLoss(nn.Module): + def __init__( + self, + num_classes: int, + matcher: OneFormerHungarianMatcher, + weight_dict: dict[str, float], + eos_coef: float, + num_points: int, + oversample_ratio: float, + importance_sample_ratio: float, + contrastive_temperature: float | None = None, + ): + """ + This class computes the losses using the class predictions, mask predictions and the contrastive queries. + + Oneformer calculates the classification CE loss on the class predictions. Mask predictions are used for + calculating the binary CE loss and dice loss. The contrastive queries are used for calculating the contrastive + loss. + + Args: + num_labels (`int`): + The number of classes. + matcher (`OneFormerHungarianMatcher`): + A torch module that computes the assignments between the predictions and labels. + weight_dict (`dict[str, float]`): + A dictionary of weights to be applied to the different losses. + eos_coef (`float`): + Weight to apply to the null class. + num_points (`int`): + Number of points to be sampled for dice and mask loss calculations. + oversample_ratio (`float`): + Required for pointwise loss calculation. + importance_sample_ratio (`float`): + Required for pointwise loss calculation. + contrastive_temperature (`float`): + Temperature for scaling the contrastive logits. + """ + requires_backends(self, ["scipy"]) + super().__init__() + self.num_classes = num_classes + self.matcher = matcher + self.weight_dict = weight_dict + self.eos_coef = eos_coef + empty_weight = torch.ones(self.num_classes + 1) + empty_weight[-1] = self.eos_coef + self.register_buffer("empty_weight", empty_weight) + + # pointwise mask loss parameters + self.num_points = num_points + self.oversample_ratio = oversample_ratio + self.importance_sample_ratio = importance_sample_ratio + self.contrastive_temperature = contrastive_temperature + if self.contrastive_temperature is not None: + self.logit_scale = nn.Parameter(torch.tensor(np.log(1 / contrastive_temperature))) + + def _max_by_axis(self, the_list: list[list[int]]) -> list[int]: + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + def _pad_images_to_max_in_batch(self, tensors: list[Tensor]) -> tuple[Tensor, Tensor]: + # get the maximum size in the batch + max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors]) + batch_size = len(tensors) + # compute finel size + batch_shape = [batch_size] + max_size + b, _, h, w = batch_shape + # get metadata + dtype = tensors[0].dtype + device = tensors[0].device + padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device) + padding_masks = torch.ones((b, h, w), dtype=torch.bool, device=device) + # pad the tensors to the size of the biggest one + for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks): + padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor) + padding_mask[: tensor.shape[1], : tensor.shape[2]] = False + + return padded_tensors, padding_masks + + def loss_contrastive(self, contrastive_queries_logits: Tensor, text_queries: Tensor): + """Compute the query-text contrastive loss. + + Args: + contrastive_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, hidden_dim` + text_queries (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, hidden_dim` + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: + - **loss_contrastive** -- The query-text contrastive loss computed using task-guided queries + and text queries derived from input text list. + """ + + image_queries = contrastive_queries_logits.float() + + # [batch_size, hidden_dim] + image_queries = nn.functional.normalize(image_queries.flatten(1), dim=-1) + text_queries = nn.functional.normalize(text_queries.flatten(1), dim=-1) + + logit_scale = torch.clamp(self.logit_scale.exp(), max=100) + + logits_per_text = torch.matmul(text_queries, image_queries.t()) * logit_scale + logits_per_img = logits_per_text.t() + + loss_img = nn.functional.cross_entropy( + logits_per_img, torch.arange(len(logits_per_img), device=logits_per_text.device) + ) + loss_text = nn.functional.cross_entropy( + logits_per_text, torch.arange(len(logits_per_text), device=logits_per_text.device) + ) + + loss_contrastive = loss_img + loss_text + + losses = {"loss_contrastive": loss_contrastive} + return losses + + def loss_labels( + self, class_queries_logits: Tensor, class_labels: list[Tensor], indices: tuple[np.array] + ) -> dict[str, Tensor]: + """Compute the losses related to the labels using cross entropy. + + Args: + class_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, num_labels` + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + """ + pred_logits = class_queries_logits + batch_size, num_queries, _ = pred_logits.shape + criterion = nn.CrossEntropyLoss(weight=self.empty_weight) + idx = self._get_predictions_permutation_indices(indices) + + # shape = (batch_size, num_queries) + target_classes_o = torch.cat([target[j] for target, (_, j) in zip(class_labels, indices)]) + # shape = (batch_size, num_queries) + target_classes = torch.full( + (batch_size, num_queries), fill_value=self.num_classes, dtype=torch.int64, device=pred_logits.device + ) + target_classes[idx] = target_classes_o + # permute pred_logits (batch_size, num_queries, num_labels) -> (batch_size, num_labels, num_queries) + pred_logits_transposed = pred_logits.transpose(1, 2) + loss_ce = criterion(pred_logits_transposed, target_classes) + losses = {"loss_cross_entropy": loss_ce} + return losses + + def loss_masks( + self, masks_queries_logits: Tensor, mask_labels: list[Tensor], indices: tuple[np.array], num_masks: int + ) -> dict[str, Tensor]: + """Compute the losses related to the masks using focal and dice loss. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, height, width` + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + indices (`tuple[np.array])`: + The indices computed by the Hungarian matcher. + num_masks (`int)`: + The number of masks, used for normalization. + + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys: + - **loss_mask** -- The loss computed using sigmoid ce loss on the predicted and ground truth masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth + masks. + """ + src_idx = self._get_predictions_permutation_indices(indices) + tgt_idx = self._get_targets_permutation_indices(indices) + # shape (batch_size * num_queries, height, width) + pred_masks = masks_queries_logits[src_idx] + # shape (batch_size, num_queries, height, width) + # pad all and stack the targets to the num_labels dimension + # upsample predictions to the target size, we have to add one dim to use interpolate + target_masks, _ = self._pad_images_to_max_in_batch(mask_labels) + target_masks = target_masks[tgt_idx] + + pred_masks = pred_masks[:, None] + target_masks = target_masks[:, None] + + with torch.no_grad(): + # sample point_coords + point_coords = self.sample_points_using_uncertainty( + pred_masks, + self.calculate_uncertainty, + self.num_points, + self.oversample_ratio, + self.importance_sample_ratio, + ) + # get ground-truth labels + point_labels = sample_point(target_masks, point_coords, align_corners=False).squeeze(1) + + point_logits = sample_point(pred_masks, point_coords, align_corners=False).squeeze(1) + + losses = { + "loss_mask": sigmoid_cross_entropy_loss(point_logits, point_labels, num_masks), + "loss_dice": dice_loss(point_logits, point_labels, num_masks), + } + + del pred_masks + del target_masks + return losses + + # Copied from transformers.models.mask2former.modeling_mask2former.Mask2FormerLoss.calculate_uncertainty + def calculate_uncertainty(self, logits: torch.Tensor) -> torch.Tensor: + """ + In Mask2Former paper, uncertainty is estimated as L1 distance between 0.0 and the logit prediction in 'logits' + for the foreground class in `classes`. + + Args: + logits (`torch.Tensor`): + A tensor of shape (R, 1, ...) for class-specific or class-agnostic, where R is the total number of predicted masks in all images and C is: + the number of foreground classes. The values are logits. + + Returns: + scores (`torch.Tensor`): A tensor of shape (R, 1, ...) that contains uncertainty scores with the most + uncertain locations having the highest uncertainty score. + """ + uncertainty_scores = -(torch.abs(logits)) + return uncertainty_scores + + # Copied from transformers.models.mask2former.modeling_mask2former.Mask2FormerLoss.sample_points_using_uncertainty + def sample_points_using_uncertainty( + self, + logits: torch.Tensor, + uncertainty_function, + num_points: int, + oversample_ratio: int, + importance_sample_ratio: float, + ) -> torch.Tensor: + """ + This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The + uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit + prediction as input. + + Args: + logits (`float`): + Logit predictions for P points. + uncertainty_function: + A function that takes logit predictions for P points and returns their uncertainties. + num_points (`int`): + The number of points P to sample. + oversample_ratio (`int`): + Oversampling parameter. + importance_sample_ratio (`float`): + Ratio of points that are sampled via importance sampling. + + Returns: + point_coordinates (`torch.Tensor`): + Coordinates for P sampled points. + """ + + num_boxes = logits.shape[0] + num_points_sampled = int(num_points * oversample_ratio) + + # Get random point coordinates + point_coordinates = torch.rand(num_boxes, num_points_sampled, 2, device=logits.device) + # Get sampled prediction value for the point coordinates + point_logits = sample_point(logits, point_coordinates, align_corners=False) + # Calculate the uncertainties based on the sampled prediction values of the points + point_uncertainties = uncertainty_function(point_logits) + + num_uncertain_points = int(importance_sample_ratio * num_points) + num_random_points = num_points - num_uncertain_points + + idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] + shift = num_points_sampled * torch.arange(num_boxes, dtype=torch.long, device=logits.device) + idx += shift[:, None] + point_coordinates = point_coordinates.view(-1, 2)[idx.view(-1), :].view(num_boxes, num_uncertain_points, 2) + + if num_random_points > 0: + point_coordinates = torch.cat( + [point_coordinates, torch.rand(num_boxes, num_random_points, 2, device=logits.device)], + dim=1, + ) + return point_coordinates + + def _get_predictions_permutation_indices(self, indices): + # permute predictions following indices + batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) + predictions_indices = torch.cat([src for (src, _) in indices]) + return batch_indices, predictions_indices + + def _get_targets_permutation_indices(self, indices): + # permute labels following indices + batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) + target_indices = torch.cat([tgt for (_, tgt) in indices]) + return batch_indices, target_indices + + def forward( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + contrastive_queries_logits: Tensor, + mask_labels: list[Tensor], + class_labels: list[Tensor], + text_queries: Tensor, + auxiliary_predictions: dict[str, Tensor] | None = None, + calculate_contrastive_loss: bool = True, + ) -> dict[str, Tensor]: + """ + This performs the loss computation. + + Args: + masks_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, height, width` + class_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, num_labels` + contrastive_queries_logits (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, hidden_dim` + mask_labels (`torch.Tensor`): + List of mask labels of shape `(labels, height, width)`. + class_labels (`list[torch.Tensor]`): + List of class labels of shape `(labels)`. + text_queries (`torch.Tensor`): + A tensor of shape `batch_size, num_queries, hidden_dim` + auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*): + if `use_auxiliary_loss` was set to `true` in [`OneFormerConfig`], then it contains the logits from the + inner layers of the Detr's Decoder. + calculate_contrastive_loss (`bool`, *optional*, defaults to `True`): + Whether or not to calculate the contrastive loss. + + Returns: + `dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys: + - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. + - **loss_mask** -- The loss computed using sigmoid ce loss on the predicted and ground truth masks. + - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth + masks. + - **loss_contrastive** -- The query-text contrstive loss computed using object and text queries. + if `use_auxiliary_loss` was set to `true` in [`OneFormerConfig`], the dictionary contains additional losses + for each auxiliary predictions. + """ + + # retrieve the matching between the outputs of the last layer and the labels + indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + # compute the average number of target masks for normalization purposes + num_masks = self.get_num_masks(class_labels, device=class_labels[0].device) + # get all the losses + losses: dict[str, Tensor] = { + **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), + **self.loss_labels(class_queries_logits, class_labels, indices), + } + if calculate_contrastive_loss: + losses = {**losses, **self.loss_contrastive(contrastive_queries_logits, text_queries)} + + # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if auxiliary_predictions is not None: + for idx, aux_outputs in enumerate(auxiliary_predictions): + masks_queries_logits = aux_outputs["masks_queries_logits"] + class_queries_logits = aux_outputs["class_queries_logits"] + loss_dict = self.forward( + masks_queries_logits, + class_queries_logits, + None, + mask_labels, + class_labels, + None, + calculate_contrastive_loss=False, + ) + loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} + losses.update(loss_dict) + + return losses + + def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Computes the average number of target masks across the batch, for normalization purposes. + """ + num_masks = sum(len(classes) for classes in class_labels) + num_masks = torch.as_tensor([num_masks], dtype=torch.float, device=device) + world_size = 1 + if is_accelerate_available(): + if PartialState._shared_state != {}: + num_masks = reduce(num_masks) + world_size = PartialState().num_processes + + num_masks = torch.clamp(num_masks / world_size, min=1) + return num_masks + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the Transformer decoder. This class adds attributes for class predictions, mask + predictions and contrastive logits to BaseModelOutputWithCrossAttentions. + """ +) +class OneFormerTransformerDecoderOutput(BaseModelOutput): + r""" + object_queries (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Queries representation for the region proposals. + contrastive_logits (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Queries representation for the contrastive loss. + prediction_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height, width)`): + Mask predictions from last layer of the transformer decoder. + prediction_class (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes+1)`): + Class predictions from last layer of the transformer decoder. + auxiliary_predictions (Tuple of Dict of `str, torch.FloatTensor`, *optional*): + Tuple of class and mask predictions from each layer of the transformer decoder. + """ + + object_queries: torch.FloatTensor | None = None + contrastive_logits: torch.FloatTensor | None = None + prediction_masks: torch.FloatTensor | None = None + prediction_class: torch.FloatTensor | None = None + auxiliary_predictions: tuple[dict[str, torch.FloatTensor]] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + OneFormer's pixel decoder module output, practically a Multi-Scale Deformable Attention based decoder. It returns + the mask features and the multiscale features. + """ +) +# Copied from transformers.models.mask2former.modeling_mask2former.Mask2FormerPixelDecoderOutput with Mask2->One +class OneFormerPixelDecoderOutput(ModelOutput): + r""" + multi_scale_features (`tuple(torch.FloatTensor)`): + Tuple of multi-scale features of scales [1/8, 1/16, 1/32] and shape `(batch_size, num_channels, height, + width)`from the Multi-Scale Deformable Attenntion based Pixel Decoder. + mask_features (`torch.FloatTensor`): + Tensor of shape `(batch_size, num_channels, height, width)`, 1/4 scale features from the last Pixel Decoder + Layer. + attentions (`tuple(torch.FloatTensor)`, *optional*): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights from pixel decoder. Returned when `output_attentions=True` is passed + or when `config.output_attentions=True` + """ + + multi_scale_features: tuple[torch.FloatTensor] | None = None + mask_features: torch.FloatTensor | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + OneFormer's pixel level module output. It returns both the last and (optionally) the hidden states from the + `encoder` and `decoder`. By default, the `encoder` is a Swin/Dinat Backbone and the `decoder` is a Multi-Scale + Deformable Attention based decoder. + """ +) +class OneFormerPixelLevelModuleOutput(ModelOutput): + r""" + encoder_features (List of `(torch.FloatTensor)`): + List of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden-states (also + called feature maps) of the model at the output of each stage. + decoder_features (List of `(torch.FloatTensor)`): + List of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden-states (also + called feature maps) of the model at the output of each stage. + decoder_last_feature (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)): + 1/4 scale features from the last Pixel Decoder Layer. + """ + + encoder_features: list[torch.FloatTensor] | None = None + decoder_features: list[torch.FloatTensor] | None = None + decoder_last_feature: torch.FloatTensor | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`OneFormerModel`]. This class returns all the needed hidden states to compute the logits. + """ +) +class OneFormerModelOutput(ModelOutput): + r""" + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder + model at the output of each stage. + pixel_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel + decoder model at the output of each stage. + transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the + transformer decoder at the output of each stage. + transformer_decoder_object_queries (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Output object queries from the last layer in the transformer decoder. + transformer_decoder_contrastive_queries (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Contrastive queries from the transformer decoder. + transformer_decoder_mask_predictions (`torch.FloatTensor` of shape `(batch_size, num_queries, height, width)`): + Mask Predictions from the last layer in the transformer decoder. + transformer_decoder_class_predictions (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes+1)`): + Class Predictions from the last layer in the transformer decoder. + transformer_decoder_auxiliary_predictions (Tuple of Dict of `str, torch.FloatTensor`, *optional*): + Tuple of class and mask predictions from each layer of the transformer decoder. + text_queries (`torch.FloatTensor`, *optional* of shape `(batch_size, num_queries, hidden_dim)`): + Text queries derived from the input text list used for calculating contrastive loss during training. + task_token (`torch.FloatTensor` of shape `(batch_size, hidden_dim)`): + 1D task token to condition the queries. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Self and Cross Attentions weights from transformer decoder. + """ + + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: torch.FloatTensor | None = None + transformer_decoder_object_queries: torch.FloatTensor | None = None + transformer_decoder_contrastive_queries: torch.FloatTensor | None = None + transformer_decoder_mask_predictions: torch.FloatTensor | None = None + transformer_decoder_class_predictions: torch.FloatTensor | None = None + transformer_decoder_auxiliary_predictions: tuple[dict[str, torch.FloatTensor]] | None = None + text_queries: torch.FloatTensor | None = None + task_token: torch.FloatTensor | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`OneFormerForUniversalSegmentationOutput`]. + + This output can be directly passed to [`~OneFormerImageProcessor.post_process_semantic_segmentation`] or + [`~OneFormerImageProcessor.post_process_instance_segmentation`] or + [`~OneFormerImageProcessor.post_process_panoptic_segmentation`] depending on the task. Please, see + [`~OneFormerImageProcessor] for details regarding usage. + """ +) +class OneFormerForUniversalSegmentationOutput(ModelOutput): + r""" + loss (`torch.Tensor`, *optional*): + The computed loss, returned when labels are present. + class_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each + query. Note the `+ 1` is needed because we incorporate the null class. + masks_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each + query. + auxiliary_predictions (List of Dict of `str, torch.FloatTensor`, *optional*): + List of class and mask predictions from each layer of the transformer decoder. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder + model at the output of each stage. + pixel_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel + decoder model at the output of each stage. + transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of + shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the + transformer decoder at the output of each stage. + transformer_decoder_object_queries (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Output object queries from the last layer in the transformer decoder. + transformer_decoder_contrastive_queries (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_dim)`): + Contrastive queries from the transformer decoder. + transformer_decoder_mask_predictions (`torch.FloatTensor` of shape `(batch_size, num_queries, height, width)`): + Mask Predictions from the last layer in the transformer decoder. + transformer_decoder_class_predictions (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes+1)`): + Class Predictions from the last layer in the transformer decoder. + transformer_decoder_auxiliary_predictions (List of Dict of `str, torch.FloatTensor`, *optional*): + List of class and mask predictions from each layer of the transformer decoder. + text_queries (`torch.FloatTensor`, *optional* of shape `(batch_size, num_queries, hidden_dim)`): + Text queries derived from the input text list used for calculating contrastive loss during training. + task_token (`torch.FloatTensor` of shape `(batch_size, hidden_dim)`): + 1D task token to condition the queries. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Self and Cross Attentions weights from transformer decoder. + """ + + loss: torch.FloatTensor | None = None + class_queries_logits: torch.FloatTensor | None = None + masks_queries_logits: torch.FloatTensor | None = None + auxiliary_predictions: list[dict[str, torch.FloatTensor]] = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: list[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: torch.FloatTensor | None = None + transformer_decoder_object_queries: torch.FloatTensor | None = None + transformer_decoder_contrastive_queries: torch.FloatTensor | None = None + transformer_decoder_mask_predictions: torch.FloatTensor | None = None + transformer_decoder_class_predictions: torch.FloatTensor | None = None + transformer_decoder_auxiliary_predictions: list[dict[str, torch.FloatTensor]] | None = None + text_queries: torch.FloatTensor | None = None + task_token: torch.FloatTensor | None = None + attentions: tuple[tuple[torch.FloatTensor]] | None = None + + +# Modified from transformers.models.detr.modeling_deformable_detr.DeformableDetrMultiscaleDeformableAttention with DeformableDetr->OneFormerPixelDecoderEncoder +class OneFormerPixelDecoderEncoderMultiscaleDeformableAttention(nn.Module): + """ + Multiscale deformable attention as proposed in Deformable DETR. + """ + + def __init__(self, embed_dim: int, num_heads: int, n_levels: int, n_points: int): + super().__init__() + if embed_dim % num_heads != 0: + raise ValueError( + f"embed_dim (d_model) must be divisible by num_heads, but got {embed_dim} and {num_heads}" + ) + dim_per_head = embed_dim // num_heads + # check if dim_per_head is power of 2 + if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0): + warnings.warn( + "You'd better set embed_dim (d_model) in DeformableDetrMultiscaleDeformableAttention to make the" + " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA" + " implementation." + ) + + self.im2col_step = 128 + + self.d_model = embed_dim + self.n_levels = n_levels + self.n_heads = num_heads + self.n_points = n_points + + self.sampling_offsets = nn.Linear(embed_dim, num_heads * n_levels * n_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_heads * n_levels * n_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): + return tensor if position_embeddings is None else tensor + position_embeddings + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states=None, + encoder_attention_mask=None, + position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes=None, + level_start_index=None, + output_attentions: bool = False, + ): + # add position embeddings to the hidden states before projecting to queries and keys + if position_embeddings is not None: + hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + + batch_size, num_queries, _ = hidden_states.shape + batch_size, sequence_length, _ = encoder_hidden_states.shape + torch_compilable_check( + (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == sequence_length, + "Make sure to align the spatial shapes with the sequence length of the encoder hidden states", + ) + + value = self.value_proj(encoder_hidden_states) + if attention_mask is not None: + # we invert the attention_mask + value = value.masked_fill(attention_mask[..., None], float(0)) + value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads) + sampling_offsets = self.sampling_offsets(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2 + ) + attention_weights = self.attention_weights(hidden_states).view( + batch_size, num_queries, self.n_heads, self.n_levels * self.n_points + ) + attention_weights = nn.functional.softmax(attention_weights, -1).view( + batch_size, num_queries, self.n_heads, self.n_levels, self.n_points + ) + # batch_size, num_queries, n_heads, n_levels, n_points, 2 + if reference_points.shape[-1] == 2: + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) + else: + raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}") + # PyTorch implementation + output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights) + output = self.output_proj(output) + + return output, attention_weights + + +class OneFormerPixelDecoderEncoderLayer(nn.Module): + def __init__(self, config: OneFormerConfig): + super().__init__() + self.embed_dim = config.conv_dim + self.self_attn = OneFormerPixelDecoderEncoderMultiscaleDeformableAttention( + embed_dim=self.embed_dim, + num_heads=config.num_attention_heads, + n_levels=3, + n_points=4, + ) + + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.dropout = config.dropout + self.activation_fn = nn.functional.relu + self.activation_dropout = config.dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_feedforward_dim) + self.fc2 = nn.Linear(config.encoder_feedforward_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + self.is_training = config.is_training + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_embeddings: torch.Tensor | None = None, + reference_points=None, + spatial_shapes=None, + level_start_index=None, + output_attentions: bool = False, + ): + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Input to the layer. + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Attention mask. + position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings, to be added to `hidden_states`. + reference_points (`torch.FloatTensor`, *optional*): + Reference points. + spatial_shapes (`torch.LongTensor`, *optional*): + Spatial shapes of the backbone feature maps. + level_start_index (`torch.LongTensor`, *optional*): + Level start index. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + # Apply Multi-scale Deformable Attention Module on the multi-scale feature maps. + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + position_embeddings=position_embeddings, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + output_attentions=output_attentions, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.is_training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + residual = hidden_states + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.is_training) + + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.is_training) + + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + if self.is_training: + if not torch.isfinite(hidden_states).all(): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +# Modified from from transformers.models.detr.modeling_deformable_detr.DeformableDetrEncoder with DeformableDetrEncoder->OneFormerPixelDecoderEncoderOnly +class OneFormerPixelDecoderEncoderOnly(nn.Module): + """ + Transformer encoder consisting of *config.encoder_layers* deformable attention layers. Each layer is a + [`OneFormerPixelDecoderEncoderLayer`]. + + The encoder updates the flattened multi-scale feature maps through multiple deformable attention layers. + + Args: + config: OneFormerConfig + """ + + def __init__(self, config: OneFormerConfig): + super().__init__() + + self.config = config + self.dropout = config.dropout + self.layers = nn.ModuleList([OneFormerPixelDecoderEncoderLayer(config) for _ in range(config.encoder_layers)]) + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + """ + Get reference points for each feature map. Used in decoder. + + Args: + spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`): + Spatial shapes of each feature map. + valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): + Valid ratios of each feature map. + device (`torch.device`): + Device on which to create the tensors. + Returns: + `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)` + """ + reference_points_list = [] + for lvl, (height, width) in enumerate(spatial_shapes): + ref_y, ref_x = torch.meshgrid( + torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device), + torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device), + indexing="ij", + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * height) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * width) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + return reference_points + + def forward( + self, + inputs_embeds=None, + attention_mask=None, + position_embeddings=None, + spatial_shapes=None, + level_start_index=None, + valid_ratios=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Flattened feature map (output of the backbone + projection layer) that is passed to the encoder. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`: + - 1 for pixel features that are real (i.e. **not masked**), + - 0 for pixel features that are padding (i.e. **masked**). + [What are attention masks?](../glossary#attention-mask) + position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Position embeddings that are added to the queries and keys in each self-attention layer. + spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`): + Spatial shapes of each feature map. + level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`): + Starting index of each feature map. + valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): + Ratio of valid area in each feature level. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + hidden_states = inputs_embeds + reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=inputs_embeds.device) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + for i, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + position_embeddings=position_embeddings, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + +# Modified from from transformers.models.mask2former.modeling_mask2former.Mask2FormerPixelDecoder with Mask2->One +class OneFormerPixelDecoder(nn.Module): + def __init__(self, config: OneFormerConfig, feature_channels): + super().__init__() + + self.config = config + + # positional encoding + self.position_embedding = OneFormerSinePositionEmbedding(num_pos_feats=config.conv_dim // 2, normalize=True) + self.num_feature_levels = 3 + transformer_in_channels = feature_channels[-self.num_feature_levels :] + self.transformer_feature_strides = config.strides[-self.num_feature_levels :] + self.feature_channels = feature_channels + self.level_embed = nn.Parameter(torch.Tensor(self.num_feature_levels, config.conv_dim)) + + # Create input projection layers + if self.num_feature_levels > 1: + input_projections_list = [] + for in_channels in transformer_in_channels[::-1]: + input_projections_list.append( + nn.Sequential( + nn.Conv2d(in_channels, config.conv_dim, kernel_size=1), + nn.GroupNorm(32, config.conv_dim), + ) + ) + self.input_projections = nn.ModuleList(input_projections_list) + else: + self.input_projections = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(transformer_in_channels[-1], config.conv_dim, kernel_size=1), + nn.GroupNorm(32, config.conv_dim), + ) + ] + ) + + self.encoder = OneFormerPixelDecoderEncoderOnly(config) + + self.mask_projection = nn.Conv2d( + config.conv_dim, + config.mask_dim, + kernel_size=1, + stride=1, + padding=0, + ) + + self.common_stride = config.common_stride + + # extra fpn levels + stride = min(self.transformer_feature_strides) + self.num_fpn_levels = int(np.log2(stride) - np.log2(self.common_stride)) + + lateral_convs = [] + output_convs = [] + + for idx, in_channels in enumerate(self.feature_channels[: self.num_fpn_levels]): + lateral_conv = nn.Sequential( + nn.Conv2d( + in_channels, + config.conv_dim, + kernel_size=1, + bias=False, + ), + nn.GroupNorm(32, config.conv_dim), + ) + output_conv = nn.Sequential( + nn.Conv2d( + config.conv_dim, + config.conv_dim, + kernel_size=3, + stride=1, + padding=1, + bias=False, + ), + nn.GroupNorm(32, config.conv_dim), + nn.ReLU(), + ) + self.add_module(f"adapter_{idx + 1}", lateral_conv) + self.add_module(f"layer_{idx + 1}", output_conv) + + lateral_convs.append(lateral_conv) + output_convs.append(output_conv) + # Place convs into top-down order (from low to high resolution) + # to make the top-down computation in forward clearer. + self.lateral_convs = lateral_convs[::-1] + self.output_convs = output_convs[::-1] + + def get_valid_ratio(self, mask, dtype=torch.float32): + """Get the valid ratio of all feature maps.""" + + _, height, width = mask.shape + valid_height = torch.sum(~mask[:, :, 0], 1) + valid_width = torch.sum(~mask[:, 0, :], 1) + valid_ratio_height = valid_height.to(dtype) / height + valid_ratio_width = valid_width.to(dtype) / width + valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1) + return valid_ratio + + def forward( + self, + features, + encoder_outputs=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) + sources = [] + position_embeddings_list = [] + for level, source in enumerate(features[::-1][: self.num_feature_levels]): + sources.append(self.input_projections[level](source)) + position_embeddings_list.append(self.position_embedding(source.shape, source.device, source.dtype)) + + masks = [torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in sources] + + # Prepare encoder inputs (by flattening) + source_flatten = [] + mask_flatten = [] + lvl_pos_embed_flatten = [] + spatial_shapes = [] + for level, (source, mask, pos_embed) in enumerate(zip(sources, masks, position_embeddings_list)): + batch_size, num_channels, height, width = source.shape + spatial_shape = (height, width) + spatial_shapes.append(spatial_shape) + source = source.flatten(2).transpose(1, 2) + mask = mask.flatten(1) + pos_embed = pos_embed.flatten(2).transpose(1, 2) + lvl_pos_embed = pos_embed + self.level_embed[level].view(1, 1, -1) + lvl_pos_embed_flatten.append(lvl_pos_embed) + source_flatten.append(source) + mask_flatten.append(mask) + source_flatten = torch.cat(source_flatten, 1) + mask_flatten = torch.cat(mask_flatten, 1) + lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) + spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=source_flatten.device) + level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) + valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1) + + # Fourth, sent source_flatten + mask_flatten + lvl_pos_embed_flatten (backbone + proj layer output) through encoder + # Also provide spatial_shapes, level_start_index and valid_ratios + if encoder_outputs is None: + encoder_outputs = self.encoder( + inputs_embeds=source_flatten, + attention_mask=mask_flatten, + position_embeddings=lvl_pos_embed_flatten, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + valid_ratios=valid_ratios, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + y = encoder_outputs.last_hidden_state + bs = y.shape[0] + + split_size_or_sections = [None] * self.num_feature_levels + for i in range(self.num_feature_levels): + if i < self.num_feature_levels - 1: + split_size_or_sections[i] = level_start_index[i + 1] - level_start_index[i] + else: + split_size_or_sections[i] = y.shape[1] - level_start_index[i] + y = torch.split(y, split_size_or_sections, dim=1) + + out = [] + multi_scale_features = [] + num_cur_levels = 0 + for i, z in enumerate(y): + out.append(z.transpose(1, 2).view(bs, -1, spatial_shapes[i][0], spatial_shapes[i][1])) + + # append `out` with extra FPN levels + # Reverse feature maps into top-down order (from low to high resolution) + for idx, feats in enumerate(features[: self.num_fpn_levels][::-1]): + lateral_conv = self.lateral_convs[idx] + output_conv = self.output_convs[idx] + cur_fpn = lateral_conv(feats) + # Following FPN implementation, we use nearest upsampling here + y = cur_fpn + nn.functional.interpolate( + out[-1], size=cur_fpn.shape[-2:], mode="bilinear", align_corners=False + ) + y = output_conv(y) + out.append(y) + + for o in out: + if num_cur_levels < self.num_feature_levels: + multi_scale_features.append(o) + num_cur_levels += 1 + + return OneFormerPixelDecoderOutput( + mask_features=self.mask_projection(out[-1]), + multi_scale_features=multi_scale_features, + attentions=encoder_outputs.attentions, + ) + + +# Modified from from transformers.models.mask2former.modeling_mask2former.Mask2FormerPixelLevelModule with Mask2->One +class OneFormerPixelLevelModule(nn.Module): + def __init__(self, config: OneFormerConfig): + """ + Pixel Level Module proposed in [Masked-attention Mask Transformer for Universal Image + Segmentation](https://huggingface.co/papers/2112.01527). It runs the input image through a backbone and a pixel + decoder, generating multi-scale feature maps and pixel embeddings. + + Args: + config ([`OneFormerConfig`]): + The configuration used to instantiate this model. + """ + super().__init__() + self.encoder = load_backbone(config) + self.decoder = OneFormerPixelDecoder(config, feature_channels=self.encoder.channels) + + def forward(self, pixel_values: Tensor, output_hidden_states: bool = False) -> OneFormerPixelLevelModuleOutput: + features: list[Tensor] = self.encoder(pixel_values).feature_maps + decoder_output: OneFormerPixelDecoderOutput = self.decoder(features, output_hidden_states=output_hidden_states) + return OneFormerPixelLevelModuleOutput( + encoder_features=tuple(features), + decoder_features=decoder_output.multi_scale_features, + decoder_last_feature=decoder_output.mask_features, + ) + + +# Modified from transformers.models.detr.modeling_detr.DetrAttention with Detr->OneFormer +class OneFormerAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Here, we add position embeddings to the queries and + keys (as explained in the DETR paper). + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + if self.head_dim * num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int): + return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): + return tensor if position_embeddings is None else tensor + position_embeddings + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + key_value_states: torch.Tensor | None = None, + key_value_position_embeddings: torch.Tensor | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + hidden_states = hidden_states.permute(1, 0, 2) if hidden_states is not None else None + position_embeddings = position_embeddings.permute(1, 0, 2) if position_embeddings is not None else None + key_value_states = key_value_states.permute(1, 0, 2) if key_value_states is not None else None + key_value_position_embeddings = ( + key_value_position_embeddings.permute(1, 0, 2) if key_value_position_embeddings is not None else None + ) + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + batch_size, target_len, embed_dim = hidden_states.size() + + # add position embeddings to the hidden states before projecting to queries and keys + if position_embeddings is not None: + hidden_states_original = hidden_states + hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + + # add key-value position embeddings to the key value states + if key_value_position_embeddings is not None: + key_value_states_original = key_value_states + key_value_states = self.with_pos_embed(key_value_states, key_value_position_embeddings) + + # get query proj + query_states = self.q_proj(hidden_states) * self.scaling + # get key, value proj + if is_cross_attention: + # cross_attentions + key_states = self._shape(self.k_proj(key_value_states), -1, batch_size) + value_states = self._shape(self.v_proj(key_value_states_original), -1, batch_size) + else: + # self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, batch_size) + value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size) + + proj_shape = (batch_size * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_states = value_states.view(*proj_shape) + + source_len = key_states.size(1) + + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) + + if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len): + raise ValueError( + f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (batch_size * self.num_heads, target_len, source_len): + raise ValueError( + f"Attention mask should be of size {(target_len, batch_size * self.num_heads, source_len)}, but is" + f" {attention_mask.size()}" + ) + attn_weights += attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + if output_attentions: + # this operation is a bit awkward, but it's required to + # make sure that attn_weights keeps its gradient. + # In order to do so, attn_weights have to reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len) + else: + attn_weights_reshaped = None + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = torch.bmm(attn_probs, value_states) + + if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(batch_size, target_len, embed_dim) + + attn_output = self.out_proj(attn_output).permute(1, 0, 2) + + return attn_output, attn_weights_reshaped + + +class OneFormerTransformerDecoderSelfAttentionLayer(nn.Module): + def __init__( + self, embed_dim, num_heads, dropout=0.0, activation="relu", normalize_before=False, layer_norm_eps=1e-05 + ): + super().__init__() + self.self_attn = OneFormerAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, is_decoder=True) + + self.norm = nn.LayerNorm(embed_dim, eps=layer_norm_eps) + self.dropout = nn.Dropout(dropout) + + self.activation = ACT2FN[activation] + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Tensor | None): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + output, + output_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + query_pos: Tensor | None = None, + ): + output2, attention_weights = self.self_attn( + hidden_states=output, position_embeddings=query_pos, attention_mask=output_mask, output_attentions=True + ) + output = output + self.dropout(output2) + output = self.norm(output) + + return output, attention_weights + + def forward_pre( + self, + output, + output_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + query_pos: Tensor | None = None, + ): + output2 = self.norm(output) + output2, attention_weights = self.self_attn( + hidden_states=output2, position_embeddings=query_pos, attention_mask=output_mask, output_attentions=True + ) + output = output + self.dropout(output2) + + return output, attention_weights + + def forward( + self, + output, + output_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + query_pos: Tensor | None = None, + ): + if self.normalize_before: + return self.forward_pre(output, output_mask, output_key_padding_mask, query_pos) + return self.forward_post(output, output_mask, output_key_padding_mask, query_pos) + + +class OneFormerTransformerDecoderCrossAttentionLayer(nn.Module): + def __init__( + self, embed_dim, num_heads, dropout=0.0, activation="relu", normalize_before=False, layer_norm_eps=1e-05 + ): + super().__init__() + self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout) + + self.norm = nn.LayerNorm(embed_dim, eps=layer_norm_eps) + self.dropout = nn.Dropout(dropout) + + self.activation = ACT2FN[activation] + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Tensor | None): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + output, + memory, + memory_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + output2, attention_weights = self.multihead_attn( + query=self.with_pos_embed(output, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + ) + output = output + self.dropout(output2) + output = self.norm(output) + + return output, attention_weights + + def forward_pre( + self, + output, + memory, + memory_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + output2 = self.norm(output) + output2, attention_weights = self.multihead_attn( + query=self.with_pos_embed(output2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + ) + output = output + self.dropout(output2) + + return output, attention_weights + + def forward( + self, + output, + memory, + memory_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + if self.normalize_before: + return self.forward_pre(output, memory, memory_mask, memory_key_padding_mask, pos, query_pos) + return self.forward_post(output, memory, memory_mask, memory_key_padding_mask, pos, query_pos) + + +class OneFormerTransformerDecoderFFNLayer(nn.Module): + def __init__( + self, + d_model, + dim_feedforward=2048, + dropout=0.0, + activation="relu", + normalize_before=False, + layer_norm_eps=1e-05, + ): + super().__init__() + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm = nn.LayerNorm(d_model, eps=layer_norm_eps) + + self.activation = ACT2FN[activation] + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Tensor | None): + return tensor if pos is None else tensor + pos + + def forward_post(self, output): + output2 = self.linear2(self.dropout(self.activation(self.linear1(output)))) + output = output + self.dropout(output2) + output = self.norm(output) + return output + + def forward_pre(self, output): + output2 = self.norm(output) + output2 = self.linear2(self.dropout(self.activation(self.linear1(output2)))) + output = output + self.dropout(output2) + return output + + def forward(self, output): + if self.normalize_before: + return self.forward_pre(output) + return self.forward_post(output) + + +class OneFormerMLPPredictionHead(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3): + """ + A classic Multi Layer Perceptron (MLP). + + Args: + input_dim (`int`): + The input dimensions. + hidden_dim (`int`): + The hidden dimensions. + output_dim (`int`): + The output dimensions. + num_layers (int, *optional*, defaults to 3): + The number of layers. + """ + super().__init__() + in_dims = [input_dim] + [hidden_dim] * (num_layers - 1) + out_dims = [hidden_dim] * (num_layers - 1) + [output_dim] + + layers = [] + for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)): + layers.append( + PredictionBlock(in_dim, out_dim, activation=nn.ReLU() if i < num_layers - 1 else nn.Identity()) + ) + + self.layers = nn.Sequential(*layers) + + def forward(self, input: Tensor) -> Tensor: + return self.layers(input) + + +# refactored from original implementation +class OneFormerTransformerDecoderLayer(nn.Module): + def __init__(self, config: OneFormerConfig): + super().__init__() + self.embed_dim = config.hidden_dim + self.num_feature_levels = 3 + + self.cross_attn = OneFormerTransformerDecoderCrossAttentionLayer( + embed_dim=self.embed_dim, + num_heads=config.num_attention_heads, + dropout=0.0, + normalize_before=config.pre_norm, + layer_norm_eps=config.layer_norm_eps, + ) + + self.self_attn = OneFormerTransformerDecoderSelfAttentionLayer( + embed_dim=self.embed_dim, + num_heads=config.num_attention_heads, + dropout=0.0, + normalize_before=config.pre_norm, + layer_norm_eps=config.layer_norm_eps, + ) + + self.ffn = OneFormerTransformerDecoderFFNLayer( + d_model=self.embed_dim, + dim_feedforward=config.dim_feedforward, + dropout=0.0, + normalize_before=config.pre_norm, + layer_norm_eps=config.layer_norm_eps, + ) + + def forward( + self, + index: int, + output: torch.Tensor, + multi_stage_features: list[torch.Tensor], + multi_stage_positional_embeddings: list[torch.Tensor], + attention_mask: torch.Tensor | None = None, + query_embeddings: torch.Tensor | None = None, + output_attentions: bool | None = False, + ): + """ + Args: + index (`int`): index of the layer in the Transformer decoder. + output (`torch.FloatTensor`): the object queries of shape `(N, batch, hidden_dim)` + multi_stage_features (`list[torch.Tensor]`): the multi-scale features from the pixel decoder. + multi_stage_positional_embeddings (`list[torch.Tensor]`): + positional embeddings for the multi_stage_features + attention_mask (`torch.FloatTensor`): attention mask for the masked cross attention layer + query_embeddings (`torch.FloatTensor`, *optional*): + position embeddings that are added to the queries and keys in the self-attention layer. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + + level_index = index % self.num_feature_levels + attention_mask[torch.where(attention_mask.sum(-1) == attention_mask.shape[-1])] = False + + # Masked Cross Attention + output, cross_attn_weights = self.cross_attn( + output, + multi_stage_features[level_index], + memory_mask=attention_mask, + memory_key_padding_mask=None, # here we do not apply masking on padded region + pos=multi_stage_positional_embeddings[level_index], + query_pos=query_embeddings, + ) + + # Self Attention + output, self_attn_weights = self.self_attn( + output, + output_mask=None, + output_key_padding_mask=None, + query_pos=query_embeddings, + ) + + # Fully Connected + output = self.ffn(output) + + outputs = (output,) + + if output_attentions: + outputs += (self_attn_weights, cross_attn_weights) + + return outputs + + +class OneFormerTransformerDecoderQueryTransformerDecoder(nn.Module): + def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): + super().__init__() + self.layers = _get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + + def forward( + self, + output, + memory, + output_mask: Tensor | None = None, + memory_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + intermediate = [] + + for layer in self.layers: + output = layer( + output, + memory, + output_mask=output_mask, + memory_mask=memory_mask, + output_key_padding_mask=output_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, + query_pos=query_pos, + ) + if self.return_intermediate: + intermediate.append(self.norm(output)) + + if self.norm is not None: + output = self.norm(output) + if self.return_intermediate: + intermediate.pop() + intermediate.append(output) + + if self.return_intermediate: + return torch.stack(intermediate) + + return output.unsqueeze(0) + + +class OneFormerTransformerDecoderQueryTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + layer_norm_eps=1e-05, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = ACT2FN[activation] + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Tensor | None): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + output, + memory, + output_mask: Tensor | None = None, + memory_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + q = k = self.with_pos_embed(output, query_pos) + output2 = self.self_attn(q, k, value=output, attn_mask=output_mask, key_padding_mask=output_key_padding_mask) + output2 = output2[0] + output = output + self.dropout1(output2) + output = self.norm1(output) + output2 = self.multihead_attn( + query=self.with_pos_embed(output, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + ) + output2 = output2[0] + output = output + self.dropout2(output2) + output = self.norm2(output) + output2 = self.linear2(self.dropout(self.activation(self.linear1(output)))) + output = output + self.dropout3(output2) + output = self.norm3(output) + return output + + def forward_pre( + self, + output, + memory, + output_mask: Tensor | None = None, + memory_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + output2 = self.norm1(output) + q = k = self.with_pos_embed(output2, query_pos) + output2 = self.self_attn(q, k, value=output2, attn_mask=output_mask, key_padding_mask=output_key_padding_mask) + output2 = output2[0] + output = output + self.dropout1(output2) + output2 = self.norm2(output) + output2 = self.multihead_attn( + query=self.with_pos_embed(output2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + ) + output2 = output2[0] + output = output + self.dropout2(output2) + output2 = self.norm3(output) + output2 = self.linear2(self.dropout(self.activation(self.linear1(output2)))) + output = output + self.dropout3(output2) + return output + + def forward( + self, + output, + memory, + output_mask: Tensor | None = None, + memory_mask: Tensor | None = None, + output_key_padding_mask: Tensor | None = None, + memory_key_padding_mask: Tensor | None = None, + pos: Tensor | None = None, + query_pos: Tensor | None = None, + ): + if self.normalize_before: + return self.forward_pre( + output, + memory, + output_mask, + memory_mask, + output_key_padding_mask, + memory_key_padding_mask, + pos, + query_pos, + ) + return self.forward_post( + output, + memory, + output_mask, + memory_mask, + output_key_padding_mask, + memory_key_padding_mask, + pos, + query_pos, + ) + + +class OneFormerTransformerDecoderQueryTransformer(nn.Module): + def __init__( + self, + d_model=512, + nhead=8, + num_decoder_layers=6, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + return_intermediate_dec=False, + layer_norm_eps=1e-05, + ): + super().__init__() + + decoder_layer = OneFormerTransformerDecoderQueryTransformerDecoderLayer( + d_model, nhead, dim_feedforward, dropout, activation, normalize_before, layer_norm_eps + ) + decoder_norm = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.decoder = OneFormerTransformerDecoderQueryTransformerDecoder( + decoder_layer, + num_decoder_layers, + decoder_norm, + return_intermediate=return_intermediate_dec, + ) + + self.d_model = d_model + self.nhead = nhead + + def forward(self, src, mask, query_embed, pos_embed, task_token=None): + batch_size = src.shape[0] + src = src.flatten(2).permute(2, 0, 1) + pos_embed = pos_embed.flatten(2).permute(2, 0, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, batch_size, 1) + if mask is not None: + mask = mask.flatten(1) + + if task_token is None: + queries = torch.zeros_like(query_embed) + else: + queries = task_token.repeat(query_embed.shape[0], 1, 1) + + queries = self.decoder(queries, src, memory_key_padding_mask=mask, pos=pos_embed, query_pos=query_embed) + return queries.transpose(1, 2) + + +class OneFormerTransformerDecoder(nn.Module): + """ + Transformer decoder + """ + + def __init__(self, in_channels: int, config: OneFormerConfig): + super().__init__() + self.config = config + + self.dropout = config.dropout + self.num_heads = config.num_attention_heads + self.is_training = config.is_training + self.use_task_norm = config.use_task_norm + self.use_auxiliary_loss = config.use_auxiliary_loss + + self.query_transformer = OneFormerTransformerDecoderQueryTransformer( + d_model=config.hidden_dim, + dropout=config.dropout, + nhead=config.num_attention_heads, + dim_feedforward=config.dim_feedforward, + num_decoder_layers=config.query_dec_layers, + normalize_before=config.pre_norm, + return_intermediate_dec=False, + layer_norm_eps=config.layer_norm_eps, + ) + + self.decoder_norm = nn.LayerNorm(config.hidden_dim, eps=config.layer_norm_eps) + + self.num_feature_levels = 3 + + self.layers = nn.ModuleList( + [OneFormerTransformerDecoderLayer(config) for _ in range(config.decoder_layers - 1)] + ) + + self.query_input_projection = nn.Conv2d(in_channels, config.hidden_dim, kernel_size=1) + + self.class_embed = nn.Linear(config.hidden_dim, config.num_labels + 1) + self.mask_embed = OneFormerMLPPredictionHead( + config.hidden_dim, + config.hidden_dim, + config.mask_dim, + 3, + ) + + def forward( + self, + task_token=None, + multi_stage_features=None, + multi_stage_positional_embeddings=None, + mask_features=None, + query_features=None, + query_embeddings=None, + query_embedder=None, + size_list=None, + output_attentions=None, + ): + if self.use_task_norm: + task_token = self.decoder_norm(task_token) + + object_queries = self.query_transformer( + query_features, + None, + query_embedder.weight[:-1], + self.query_input_projection(mask_features), + task_token if self.use_task_norm else None, + ) + + object_queries = object_queries[0].permute(1, 0, 2) + + queries = torch.cat([object_queries, task_token], dim=0) + + output = queries.clone() + + intermediate_class_predictions = [] + intermediate_mask_predictions = [] + + # prediction heads on learnable query features + outputs_class, outputs_mask, attention_mask = self.forward_prediction_heads( + output, mask_features, attention_mask_target_size=size_list[0] + ) + intermediate_class_predictions.append(outputs_class) + intermediate_mask_predictions.append(outputs_mask) + + attentions = () + + for index, layer in enumerate(self.layers): + layer_outputs = layer( + index=index, + output=output, + multi_stage_features=multi_stage_features, + multi_stage_positional_embeddings=multi_stage_positional_embeddings, + attention_mask=attention_mask, + query_embeddings=query_embeddings, + output_attentions=output_attentions, + ) + + output = layer_outputs[0] + attentions += (layer_outputs[1:],) + + outputs_class, outputs_mask, attention_mask = self.forward_prediction_heads( + output, mask_features, attention_mask_target_size=size_list[(index + 1) % self.num_feature_levels] + ) + intermediate_class_predictions.append(outputs_class) + intermediate_mask_predictions.append(outputs_mask) + + if not len(intermediate_mask_predictions) == len(self.layers) + 1: + raise ValueError( + "Intermediate predictions in the transformer decoder must have the same number of elements as number" + " of layers" + ) + + object_queries = layer_outputs[0].permute(1, 0, 2) + + contrastive_logits = queries.permute(1, 0, 2) + + return OneFormerTransformerDecoderOutput( + object_queries=object_queries, + contrastive_logits=contrastive_logits, + prediction_masks=intermediate_mask_predictions[-1], + prediction_class=intermediate_class_predictions[-1], + auxiliary_predictions=self._get_aux_predictions( + intermediate_class_predictions, intermediate_mask_predictions + ) + if self.use_auxiliary_loss + else None, + attentions=attentions, + ) + + def forward_prediction_heads(self, output, mask_features, attention_mask_target_size): + decoder_output = self.decoder_norm(output) + decoder_output = decoder_output.transpose(0, 1) + outputs_class = self.class_embed(decoder_output) + mask_embed = self.mask_embed(decoder_output) + outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features) + + attention_mask = nn.functional.interpolate( + outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False + ) + + # must use bool type + # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged. + attention_mask = ( + attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5 + ).bool() + attention_mask = attention_mask.detach() + + return outputs_class, outputs_mask, attention_mask + + def _get_aux_predictions(self, outputs_class, outputs_seg_masks): + aux_list = [ + {"class_queries_logits": a, "masks_queries_logits": b} + for a, b in zip(outputs_class[:-1], outputs_seg_masks[:-1]) + ] + return tuple(aux_list) + + +class OneFormerTransformerModule(nn.Module): + """ + The OneFormer's transformer module. + """ + + def __init__(self, in_features: int, config: OneFormerConfig): + super().__init__() + hidden_dim = config.hidden_dim + self.num_feature_levels = 3 + self.position_embedder = OneFormerSinePositionEmbedding(num_pos_feats=hidden_dim // 2, normalize=True) + self.queries_embedder = nn.Embedding(config.num_queries, hidden_dim) + self.input_projections = [] + + for _ in range(self.num_feature_levels): + if in_features != hidden_dim or config.enforce_input_proj: + self.input_projections.append(nn.Conv2d(in_features, hidden_dim, kernel_size=1)) + else: + self.input_projections.append(nn.Sequential()) + + self.decoder = OneFormerTransformerDecoder(in_channels=in_features, config=config) + self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim) + + def forward( + self, + multi_scale_features: list[Tensor], + mask_features: Tensor, + task_token: Tensor, + output_attentions: bool = False, + ) -> OneFormerTransformerDecoderOutput: + if not len(multi_scale_features) == self.num_feature_levels: + raise ValueError( + f"Number of elements in multi_scale_features ({len(multi_scale_features)}) and num_feature_levels" + f" ({self.num_feature_levels}) do not match!" + ) + multi_stage_features = [] + multi_stage_positional_embeddings = [] + size_list = [] + + for i in range(self.num_feature_levels): + size_list.append(multi_scale_features[i].shape[-2:]) + multi_stage_positional_embeddings.append( + self.position_embedder( + multi_scale_features[i].shape, multi_scale_features[i].device, multi_scale_features[i].dtype, None + ).flatten(2) + ) + multi_stage_features.append( + self.input_projections[i](multi_scale_features[i]).flatten(2) + + self.level_embed.weight[i][None, :, None] + ) + + # flatten NxCxHxW to HWxNxC + multi_stage_positional_embeddings[-1] = multi_stage_positional_embeddings[-1].permute(2, 0, 1) + multi_stage_features[-1] = multi_stage_features[-1].permute(2, 0, 1) + + _, batch_size, _ = multi_stage_features[0].shape + + # QxNxC + query_embeddings = self.queries_embedder.weight.unsqueeze(1).repeat(1, batch_size, 1) + task_token = task_token.unsqueeze(0) + + query_features = self.position_embedder(mask_features.shape, mask_features.device, mask_features.dtype, None) + + return self.decoder( + task_token=task_token, + multi_stage_features=multi_stage_features, + multi_stage_positional_embeddings=multi_stage_positional_embeddings, + mask_features=mask_features, + query_features=query_features, + query_embeddings=query_embeddings, + query_embedder=self.queries_embedder, + size_list=size_list, + output_attentions=output_attentions, + ) + + +# Copied from transformers.models.maskformer.modeling_maskformer.MaskFormerSinePositionEmbedding with Mask->One +class OneFormerSinePositionEmbedding(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one used by the Attention is all you + need paper, generalized to work on images. + """ + + def __init__( + self, num_pos_feats: int = 64, temperature: int = 10000, normalize: bool = False, scale: float | None = None + ): + super().__init__() + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + self.scale = 2 * math.pi if scale is None else scale + + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: Tensor | None = None, + ) -> Tensor: + if mask is None: + mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool) + not_mask = (~mask).to(dtype) + y_embed = not_mask.cumsum(1) + x_embed = not_mask.cumsum(2) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.int64, device=device).to(dtype) + dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos + + +# Copied from transformers.models.maskformer.modeling_maskformer.PredictionBlock +class PredictionBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None: + super().__init__() + self.layers = [nn.Linear(in_dim, out_dim), activation] + # Maintain submodule indexing as if part of a Sequential block + for i, layer in enumerate(self.layers): + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class OneFormerTextMapperAttention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights + self.scale = qk_scale or head_dim**-0.5 + + self.q_proj = nn.Linear(dim, dim, bias=qkv_bias) + self.k_proj = nn.Linear(dim, dim, bias=qkv_bias) + self.v_proj = nn.Linear(dim, dim, bias=qkv_bias) + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, q, k, v): + batch_size, q_sequence_length, num_channels = q.shape + if not k.shape == v.shape: + raise ValueError(f"keys ({list(k.shape)}) and values ({list(v.shape)}) have different shapes!") + batch_size, k_sequence_length, num_channels = k.shape + q = self.q_proj(q).reshape(batch_size, q_sequence_length, self.num_heads, num_channels // self.num_heads) + k = self.k_proj(k).reshape(batch_size, k_sequence_length, self.num_heads, num_channels // self.num_heads) + v = self.v_proj(v).reshape(batch_size, k_sequence_length, self.num_heads, num_channels // self.num_heads) + + attn = torch.einsum("bnkc,bmkc->bknm", q, k) * self.scale + + attn = attn.softmax(dim=-1) + + output = torch.einsum("bknm,bmkc->bnkc", attn, v).reshape(batch_size, q_sequence_length, num_channels) + + output = self.proj(output) + output = self.proj_drop(output) + return output + + +class OneFormerTextTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dropout=0.1, + layer_norm_eps=1e-05, + ): + super().__init__() + self.self_attn = OneFormerTextMapperAttention(d_model, nhead, proj_drop=dropout) + self.cross_attn = OneFormerTextMapperAttention(d_model, nhead, proj_drop=dropout) + + self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.dropout = nn.Dropout(dropout) + + self.mlp = nn.Sequential( + nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model * 4, d_model) + ) + + def forward(self, hidden_state, mem): + q = k = v = self.norm1(hidden_state) + hidden_state = hidden_state + self.self_attn(q, k, v) + q = self.norm2(hidden_state) + hidden_state = hidden_state + self.cross_attn(q, mem, mem) + hidden_state = hidden_state + self.dropout(self.mlp(self.norm3(hidden_state))) + return hidden_state + + +class OneFormerTextContextDecoder(nn.Module): + def __init__( + self, + transformer_width=256, + transformer_heads=4, + transformer_layers=6, + visual_dim=1024, + dropout=0.1, + layer_norm_eps=1e-05, + **kwargs, + ): + super().__init__() + + self.memory_proj = nn.Sequential( + nn.LayerNorm(visual_dim, eps=layer_norm_eps), + nn.Linear(visual_dim, transformer_width), + nn.LayerNorm(transformer_width, eps=layer_norm_eps), + ) + + self.text_proj = nn.Sequential( + nn.LayerNorm(visual_dim, eps=layer_norm_eps), + nn.Linear(visual_dim, transformer_width), + ) + + self.decoder = nn.ModuleList( + [ + OneFormerTextTransformerDecoderLayer(transformer_width, transformer_heads, dropout, layer_norm_eps) + for _ in range(transformer_layers) + ] + ) + + self.out_proj = nn.Sequential( + nn.LayerNorm(transformer_width, eps=layer_norm_eps), nn.Linear(transformer_width, visual_dim) + ) + + def forward(self, text, visual): + visual = self.memory_proj(visual) + hidden_state = self.text_proj(text) + + for layer in self.decoder: + hidden_state = layer(hidden_state, visual) + + return self.out_proj(hidden_state) + + +class OneFormerTextMLP(nn.Module): + def __init__( + self, + hidden_size: int | None = None, + intermediate_size: int | None = None, + output_size: int | None = None, + ): + super().__init__() + self.activation_fn = ACT2FN["quick_gelu"] + self.fc1 = nn.Linear(hidden_size, intermediate_size) + self.fc2 = nn.Linear(intermediate_size, output_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class OneFormerTextTransformerLayer(GradientCheckpointingLayer): + def __init__(self, width: int, heads: int, attn_mask: torch.Tensor, layer_norm_eps=1e-05): + super().__init__() + self.self_attn = nn.MultiheadAttention(width, heads) + self.layer_norm1 = nn.LayerNorm(width, eps=layer_norm_eps) + self.mlp = OneFormerTextMLP(width, width * 4, width) + self.layer_norm2 = nn.LayerNorm(width, eps=layer_norm_eps) + self.attn_mask = attn_mask + + def forward( + self, + hidden_states: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.self_attn( + hidden_states, + hidden_states, + hidden_states, + need_weights=False, + key_padding_mask=key_padding_mask, + )[0] + 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 OneFormerTextTransformer(nn.Module): + def __init__( + self, + width: int, + layers: int, + heads: int, + attn_mask: torch.Tensor | None = None, + use_checkpoint=False, + layer_norm_eps=1e-05, + ): + super().__init__() + self.width = width + self.num_layers = layers + self.layers = nn.Sequential( + *[OneFormerTextTransformerLayer(width, heads, attn_mask, layer_norm_eps) for _ in range(layers)] + ) + self.use_checkpoint = use_checkpoint + + def forward(self, hidden_states: torch.Tensor): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class OneFormerTextEncoder(nn.Module): + def __init__( + self, + context_length: int, + width: int, + layers: int, + vocab_size, + use_checkpoint=False, + layer_norm_eps=1e-05, + ): + super().__init__() + heads = width // 64 + self.context_length = context_length + self.width = width + self.transformer = OneFormerTextTransformer( + width=width, + layers=layers, + heads=heads, + attn_mask=self.build_attention_mask(), + use_checkpoint=use_checkpoint, + layer_norm_eps=layer_norm_eps, + ) + + self.positional_embedding = nn.Parameter(torch.empty(self.context_length, width)) + self.ln_final = nn.LayerNorm(width, eps=layer_norm_eps) + self.token_embedding = nn.Embedding(vocab_size, width) + + def build_attention_mask(self): + # lazily create causal attention mask, with full attention between the vision tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(self.context_length, self.context_length) + mask.fill_(float("-inf")) + mask.triu_(1) # zero out the lower diagonal + return mask + + def forward(self, text): + hidden_state = self.token_embedding(text) + hidden_state = hidden_state + self.positional_embedding + hidden_state = hidden_state.permute(1, 0, 2) + hidden_state = self.transformer(hidden_state) + hidden_state = hidden_state.permute(1, 0, 2) + hidden_state = self.ln_final(hidden_state) + hidden_state = hidden_state[torch.arange(hidden_state.shape[0]), text.argmax(dim=-1)] + + return hidden_state + + +class OneFormerTextMapper(nn.Module): + def __init__(self, config: OneFormerConfig): + super().__init__() + self.text_encoder = OneFormerTextEncoder( + context_length=config.text_encoder_context_length, + width=config.text_encoder_width, + layers=config.text_encoder_num_layers, + vocab_size=config.text_encoder_vocab_size, + layer_norm_eps=config.layer_norm_eps, + ) + + self.text_projector = OneFormerMLPPredictionHead( + config.text_encoder_width, + config.hidden_dim, + config.hidden_dim, + config.text_encoder_proj_layers, + ) + if config.text_encoder_n_ctx > 0: + self.prompt_ctx = nn.Embedding( + config.text_encoder_n_ctx, + config.text_encoder_width, + ) + else: + self.prompt_ctx = None + + def forward( + self, + inputs: Tensor, + ) -> Tensor: + text_queries = self.encode_text(inputs) + + return text_queries + + def encode_text(self, text): + if text.ndim is None: + raise ValueError("text must not be NoneType") + if text.ndim not in [2, 3]: + raise ValueError("Number of dimensions in text must be 2 or 3") + squeeze_dim = False + num_text = 1 + if text.ndim == 3: + num_text = text.shape[1] + batch_size, num_text, hidden_dim = text.shape + text = text.reshape(batch_size * num_text, hidden_dim) + squeeze_dim = True + + # [batch_size, num_channels] + encoded_text = self.text_encoder(text) + + text_queries = self.text_projector(encoded_text) + + if squeeze_dim: + _, hidden_dim = text_queries.shape + text_queries = text_queries.reshape(batch_size, num_text, hidden_dim) + if self.prompt_ctx is not None: + text_queries_ctx = self.prompt_ctx.weight.unsqueeze(0).repeat(text_queries.shape[0], 1, 1) + text_queries = torch.cat([text_queries, text_queries_ctx], dim=1) + + return text_queries + + +class OneFormerTaskModel(nn.Module): + def __init__(self, config: OneFormerConfig): + super().__init__() + self.task_mlp = OneFormerMLPPredictionHead( + config.task_seq_len, + config.hidden_dim, + config.hidden_dim, + 2, + ) + + def forward(self, inputs: Tensor) -> Tensor: + task_tokens = self.task_mlp(inputs) + return task_tokens + + +@auto_docstring +class OneFormerPreTrainedModel(PreTrainedModel): + config: OneFormerConfig + base_model_prefix = "model" + main_input_name = "pixel_values" + input_modalities = ("image",) + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + xavier_std = self.config.init_xavier_std + std = self.config.init_std + if isinstance(module, OneFormerTransformerModule): + if module.input_projections is not None: + for input_projection in module.input_projections: + if not isinstance(input_projection, nn.Sequential): + init.xavier_uniform_(input_projection.weight, gain=xavier_std) + init.constant_(input_projection.bias, 0) + elif isinstance(module, OneFormerTransformerDecoder): + init.xavier_uniform_(module.query_input_projection.weight, gain=xavier_std) + init.constant_(module.query_input_projection.bias, 0) + elif isinstance(module, OneFormerPixelDecoderEncoderMultiscaleDeformableAttention): + init.constant_(module.sampling_offsets.weight, 0.0) + thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (2.0 * math.pi / module.n_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(module.n_heads, 1, 1, 2) + .repeat(1, module.n_levels, module.n_points, 1) + ) + for i in range(module.n_points): + grid_init[:, :, i, :] *= i + 1 + + init.copy_(module.sampling_offsets.bias, grid_init.view(-1)) + init.constant_(module.attention_weights.weight, 0.0) + init.constant_(module.attention_weights.bias, 0.0) + init.xavier_uniform_(module.value_proj.weight) + init.constant_(module.value_proj.bias, 0.0) + init.xavier_uniform_(module.output_proj.weight) + init.constant_(module.output_proj.bias, 0.0) + elif isinstance(module, OneFormerPixelDecoder): + init.normal_(module.level_embed, std=0) + elif isinstance(module, (OneFormerTransformerDecoderLayer, OneFormerTransformerDecoderQueryTransformer)): + for p in module.parameters(): + if p.dim() > 1: + init.xavier_uniform_(p, gain=xavier_std) + elif isinstance(module, OneFormerTextTransformer): + proj_std = (module.width**-0.5) * ((2 * module.num_layers) ** -0.5) + attn_std = module.width**-0.5 + fc_std = (2 * module.width) ** -0.5 + for layer in module.layers: + init.normal_(layer.self_attn.in_proj_weight, std=attn_std) + init.normal_(layer.self_attn.out_proj.weight, std=proj_std) + init.normal_(layer.mlp.fc1.weight, std=fc_std) + init.normal_(layer.mlp.fc2.weight, std=proj_std) + elif isinstance(module, OneFormerTextEncoder): + init.normal_(module.token_embedding.weight, std=0.02) + init.normal_(module.positional_embedding, std=0.01) + if hasattr(module, "reference_points"): + init.xavier_uniform_(module.reference_points.weight, gain=1.0) + init.constant_(module.reference_points.bias, 0.0) + elif isinstance(module, OneFormerMLPPredictionHead): + for submodule in module.modules(): + if isinstance(submodule, nn.Linear): + init.xavier_uniform_(submodule.weight, gain=xavier_std) + init.constant_(submodule.bias, 0) + elif isinstance(module, nn.MultiheadAttention): + init.normal_(module.in_proj_weight, mean=0.0, std=std) + init.zeros_(module.in_proj_bias) + elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + if getattr(module, "running_mean", None) is not None: + init.zeros_(module.running_mean) + init.ones_(module.running_var) + init.zeros_(module.num_batches_tracked) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.ones_(module.weight) + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + elif isinstance(module, OneFormerLoss): + init.constant_(module.logit_scale, np.log(1 / self.config.contrastive_temperature)) + empty_weight = torch.ones(module.num_classes + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + + +@auto_docstring +class OneFormerModel(OneFormerPreTrainedModel): + main_input_name = ["pixel_values", "task_inputs"] + + def __init__(self, config: OneFormerConfig): + super().__init__(config) + self.pixel_level_module = OneFormerPixelLevelModule(config) + self.transformer_module = OneFormerTransformerModule(in_features=config.conv_dim, config=config) + self.task_encoder = OneFormerTaskModel(config) + self.is_training = config.is_training + + if self.is_training: + self.text_mapper = OneFormerTextMapper(config) + else: + self.text_mapper = None + + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + task_inputs: Tensor, + text_inputs: Tensor | None = None, + pixel_mask: Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> OneFormerModelOutput: + r""" + task_inputs (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Task inputs. Task inputs can be obtained using [`AutoImageProcessor`]. See [`OneFormerProcessor.__call__`] + for details. + text_inputs (`list[torch.Tensor]`, *optional*): + Tensor of shape `(num_queries, sequence_length)` to be fed to a model + + Example: + + ```python + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import OneFormerProcessor, OneFormerModel + + >>> # download texting image + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> # load processor for preprocessing the inputs + >>> processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny") + >>> model = OneFormerModel.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny") + >>> inputs = processor(image, ["semantic"], return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> mask_predictions = outputs.transformer_decoder_mask_predictions + >>> class_predictions = outputs.transformer_decoder_class_predictions + + >>> f"👉 Mask Predictions Shape: {list(mask_predictions.shape)}, Class Predictions Shape: {list(class_predictions.shape)}" + '👉 Mask Predictions Shape: [1, 150, 128, 171], Class Predictions Shape: [1, 150, 151]' + ```""" + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + batch_size, _, height, width = pixel_values.shape + + if pixel_mask is None: + pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device) + + pixel_level_module_output = self.pixel_level_module(pixel_values, output_hidden_states) + + multi_scale_features = pixel_level_module_output.decoder_features + mask_features = pixel_level_module_output.decoder_last_feature + + task_token = self.task_encoder(task_inputs.to(self.dtype)) + + if self.is_training: + text_queries = self.text_mapper(text_inputs) + else: + text_queries = None + + transformer_module_output = self.transformer_module( + multi_scale_features=multi_scale_features, + mask_features=mask_features, + task_token=task_token, + output_attentions=output_attentions, + ) + + queries = transformer_module_output.object_queries + + encoder_hidden_states = None + pixel_decoder_hidden_states = None + transformer_decoder_hidden_states = None + + if output_hidden_states: + encoder_hidden_states = pixel_level_module_output.encoder_features + pixel_decoder_hidden_states = (pixel_level_module_output.decoder_last_feature,) + for f in pixel_level_module_output.decoder_features: + pixel_decoder_hidden_states += (f,) + transformer_decoder_hidden_states = transformer_module_output.auxiliary_predictions + + output = OneFormerModelOutput( + encoder_hidden_states=encoder_hidden_states, + pixel_decoder_hidden_states=pixel_decoder_hidden_states, + transformer_decoder_hidden_states=transformer_decoder_hidden_states, + transformer_decoder_object_queries=queries, + transformer_decoder_contrastive_queries=transformer_module_output.contrastive_logits, + transformer_decoder_mask_predictions=transformer_module_output.prediction_masks, + transformer_decoder_class_predictions=transformer_module_output.prediction_class, + transformer_decoder_auxiliary_predictions=transformer_module_output.auxiliary_predictions, + text_queries=text_queries, + task_token=task_token, + attentions=transformer_module_output.attentions, + ) + + if not return_dict: + output = tuple(v for v in output.values()) + + return output + + +@auto_docstring( + custom_intro=""" + OneFormer Model for instance, semantic and panoptic image segmentation. + """ +) +class OneFormerForUniversalSegmentation(OneFormerPreTrainedModel): + main_input_name = ["pixel_values", "task_inputs"] + + def __init__(self, config: OneFormerConfig): + super().__init__(config) + self.model = OneFormerModel(config) + + self.matcher = OneFormerHungarianMatcher( + cost_class=config.class_weight, + cost_dice=config.dice_weight, + cost_mask=config.mask_weight, + num_points=config.train_num_points, + ) + + self.weight_dict: dict[str, float] = { + "loss_cross_entropy": config.class_weight, + "loss_mask": config.mask_weight, + "loss_dice": config.dice_weight, + "loss_contrastive": config.contrastive_weight, + } + + self.criterion = OneFormerLoss( + num_classes=config.num_labels, + matcher=self.matcher, + weight_dict=self.weight_dict, + eos_coef=config.no_object_weight, + num_points=config.train_num_points, + oversample_ratio=config.oversample_ratio, + importance_sample_ratio=config.importance_sample_ratio, + contrastive_temperature=config.contrastive_temperature, + ) + + self.post_init() + + def get_loss_dict( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + contrastive_queries_logits: Tensor, + mask_labels: Tensor, + class_labels: Tensor, + text_queries: Tensor, + auxiliary_predictions: dict[str, Tensor], + calculate_contrastive_loss: bool, + ) -> dict[str, Tensor]: + loss_dict: dict[str, Tensor] = self.criterion( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + contrastive_queries_logits=contrastive_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + text_queries=text_queries, + auxiliary_predictions=auxiliary_predictions, + calculate_contrastive_loss=calculate_contrastive_loss, + ) + + # weight each loss by `self.weight_dict[]` including auxiliary losses + for key, weight in self.weight_dict.items(): + for loss_key, loss in loss_dict.items(): + if key in loss_key: + loss *= weight + + return loss_dict + + def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: + return sum(loss_dict.values()) + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + task_inputs: Tensor, + text_inputs: Tensor | None = None, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + pixel_mask: Tensor | None = None, + output_auxiliary_logits: bool | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> OneFormerForUniversalSegmentationOutput: + r""" + task_inputs (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Task inputs. Task inputs can be obtained using [`AutoImageProcessor`]. See [`OneFormerProcessor.__call__`] + for details. + text_inputs (`list[torch.Tensor]`, *optional*): + Tensor of shape `(num_queries, sequence_length)` to be fed to a model + mask_labels (`list[torch.Tensor]`, *optional*): + List of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + output_auxiliary_logits (`bool`, *optional*): + Whether or not to output auxiliary logits. + + Example: + + Universal segmentation example: + + ```python + >>> from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> import torch + + >>> # load OneFormer fine-tuned on ADE20k for universal segmentation + >>> processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny") + >>> model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny") + + >>> url = ( + ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg" + ... ) + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> # Semantic Segmentation + >>> inputs = processor(image, ["semantic"], return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # you can pass them to processor for semantic postprocessing + >>> predicted_semantic_map = processor.post_process_semantic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0] + >>> f"👉 Semantic Predictions Shape: {list(predicted_semantic_map.shape)}" + '👉 Semantic Predictions Shape: [512, 683]' + + >>> # Instance Segmentation + >>> inputs = processor(image, ["instance"], return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # you can pass them to processor for instance postprocessing + >>> predicted_instance_map = processor.post_process_instance_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0]["segmentation"] + >>> f"👉 Instance Predictions Shape: {list(predicted_instance_map.shape)}" + '👉 Instance Predictions Shape: [512, 683]' + + >>> # Panoptic Segmentation + >>> inputs = processor(image, ["panoptic"], return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)` + >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` + >>> class_queries_logits = outputs.class_queries_logits + >>> masks_queries_logits = outputs.masks_queries_logits + + >>> # you can pass them to processor for panoptic postprocessing + >>> predicted_panoptic_map = processor.post_process_panoptic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0]["segmentation"] + >>> f"👉 Panoptic Predictions Shape: {list(predicted_panoptic_map.shape)}" + '👉 Panoptic Predictions Shape: [512, 683]' + ``` + """ + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.model( + pixel_values=pixel_values, + task_inputs=task_inputs, + text_inputs=text_inputs, + pixel_mask=pixel_mask, + output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, + output_attentions=output_attentions, + return_dict=True, + ) + + loss, loss_dict, auxiliary_predictions = None, None, None + + class_queries_logits = outputs.transformer_decoder_class_predictions + masks_queries_logits = outputs.transformer_decoder_mask_predictions + contrastive_queries_logits = outputs.transformer_decoder_contrastive_queries + auxiliary_predictions = outputs.transformer_decoder_auxiliary_predictions + text_queries = outputs.text_queries + + if mask_labels is not None and class_labels is not None: + loss_dict: dict[str, Tensor] = self.get_loss_dict( + masks_queries_logits=masks_queries_logits, + class_queries_logits=class_queries_logits, + contrastive_queries_logits=contrastive_queries_logits, + mask_labels=mask_labels, + class_labels=class_labels, + text_queries=text_queries, + auxiliary_predictions=auxiliary_predictions, + calculate_contrastive_loss=self.config.contrastive_temperature is not None, + ) + loss = self.get_loss(loss_dict) + + output_auxiliary_logits = ( + self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits + ) + if not output_auxiliary_logits: + auxiliary_predictions = None + + output = OneFormerForUniversalSegmentationOutput( + class_queries_logits=class_queries_logits, + masks_queries_logits=masks_queries_logits, + auxiliary_predictions=auxiliary_predictions, + loss=loss, + **outputs, + ) + + if not return_dict: + output = tuple(v for v in output.values()) + if loss is not None: + output = (loss) + output + return output + + +__all__ = ["OneFormerForUniversalSegmentation", "OneFormerModel", "OneFormerPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/oneformer/processing_oneformer.py b/third_party/transformers/src/transformers/models/oneformer/processing_oneformer.py new file mode 100644 index 0000000000000000000000000000000000000000..fef0b2c3b6803f728b9a952713afd96e0591dbef --- /dev/null +++ b/third_party/transformers/src/transformers/models/oneformer/processing_oneformer.py @@ -0,0 +1,181 @@ +# Copyright 2022 SHI Labs and The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Image/Text processor class for OneFormer +""" + +from ...processing_utils import ProcessorMixin +from ...utils import auto_docstring, is_torch_available + + +if is_torch_available(): + import torch + + +@auto_docstring +class OneFormerProcessor(ProcessorMixin): + def __init__( + self, image_processor=None, tokenizer=None, max_seq_length: int = 77, task_seq_length: int = 77, **kwargs + ): + r""" + max_seq_length (`int`, *optional*, defaults to `77`): + Maximum sequence length for encoding class names and text inputs. This parameter controls the + maximum number of tokens used when tokenizing class names and other text inputs for the model. + task_seq_length (`int`, *optional*, defaults to `77`): + Maximum sequence length specifically for encoding task descriptions. Task descriptions (e.g., + "the task is semantic") are tokenized with this length limit, which may differ from the general + text sequence length. + """ + self.max_seq_length = max_seq_length + self.task_seq_length = task_seq_length + + super().__init__(image_processor, tokenizer) + + def _preprocess_text(self, text_list=None, max_length=77): + if text_list is None: + raise ValueError("tokens cannot be None.") + + tokens = self.tokenizer(text_list, padding="max_length", max_length=max_length, truncation=True) + + attention_masks, input_ids = tokens["attention_mask"], tokens["input_ids"] + + token_inputs = [] + for attn_mask, input_id in zip(attention_masks, input_ids): + token = torch.tensor(attn_mask) * torch.tensor(input_id) + token_inputs.append(token.unsqueeze(0)) + + token_inputs = torch.cat(token_inputs, dim=0) + return token_inputs + + @auto_docstring + def __call__(self, images=None, task_inputs=None, segmentation_maps=None, **kwargs): + r""" + task_inputs (`str` or `list[str]`, *required*): + The task type(s) for segmentation. Must be one or more of `"semantic"`, `"instance"`, or `"panoptic"`. + Can be a single string for a single image, or a list of strings (one per image) for batch processing. + The task type determines which type of segmentation the model will perform on the input images. + segmentation_maps (`ImageInput`, *optional*): + The corresponding semantic segmentation maps with the pixel-wise annotations. + + (`bool`, *optional*, defaults to `True`): + Whether or not to pad images up to the largest image in a batch and create a pixel mask. + + If left to the default, will return a pixel mask that is: + + - 1 for pixels that are real (i.e. **not masked**), + - 0 for pixels that are padding (i.e. **masked**). + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + - **task_inputs** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + + if task_inputs is None: + raise ValueError("You have to specify the task_input. Found None.") + elif images is None: + raise ValueError("You have to specify the image. Found None.") + + if not all(task in ["semantic", "instance", "panoptic"] for task in task_inputs): + raise ValueError("task_inputs must be semantic, instance, or panoptic.") + + encoded_inputs = self.image_processor(images, task_inputs, segmentation_maps, **kwargs) + + if isinstance(task_inputs, str): + task_inputs = [task_inputs] + + if isinstance(task_inputs, list) and all(isinstance(task_input, str) for task_input in task_inputs): + task_token_inputs = [] + for task in task_inputs: + task_input = f"the task is {task}" + task_token_inputs.append(task_input) + encoded_inputs["task_inputs"] = self._preprocess_text(task_token_inputs, max_length=self.task_seq_length) + else: + raise TypeError("Task Inputs should be a string or a list of strings.") + + if hasattr(encoded_inputs, "text_inputs"): + texts_list = encoded_inputs.text_inputs + + text_inputs = [] + for texts in texts_list: + text_input_list = self._preprocess_text(texts, max_length=self.max_seq_length) + text_inputs.append(text_input_list.unsqueeze(0)) + + encoded_inputs["text_inputs"] = torch.cat(text_inputs, dim=0) + + return encoded_inputs + + def encode_inputs(self, images=None, task_inputs=None, segmentation_maps=None, **kwargs): + """ + This method forwards all its arguments to [`OneFormerImageProcessor.encode_inputs`] and then tokenizes the + task_inputs. Please refer to the docstring of this method for more information. + """ + + if task_inputs is None: + raise ValueError("You have to specify the task_input. Found None.") + elif images is None: + raise ValueError("You have to specify the image. Found None.") + + if not all(task in ["semantic", "instance", "panoptic"] for task in task_inputs): + raise ValueError("task_inputs must be semantic, instance, or panoptic.") + + encoded_inputs = self.image_processor.encode_inputs(images, task_inputs, segmentation_maps, **kwargs) + + if isinstance(task_inputs, str): + task_inputs = [task_inputs] + + if isinstance(task_inputs, list) and all(isinstance(task_input, str) for task_input in task_inputs): + task_token_inputs = [] + for task in task_inputs: + task_input = f"the task is {task}" + task_token_inputs.append(task_input) + encoded_inputs["task_inputs"] = self._preprocess_text(task_token_inputs, max_length=self.task_seq_length) + else: + raise TypeError("Task Inputs should be a string or a list of strings.") + + if hasattr(encoded_inputs, "text_inputs"): + texts_list = encoded_inputs.text_inputs + + text_inputs = [] + for texts in texts_list: + text_input_list = self._preprocess_text(texts, max_length=self.max_seq_length) + text_inputs.append(text_input_list.unsqueeze(0)) + + encoded_inputs["text_inputs"] = torch.cat(text_inputs, dim=0) + + return encoded_inputs + + def post_process_semantic_segmentation(self, *args, **kwargs): + """ + This method forwards all its arguments to [`OneFormerImageProcessor.post_process_semantic_segmentation`]. + Please refer to the docstring of this method for more information. + """ + return self.image_processor.post_process_semantic_segmentation(*args, **kwargs) + + def post_process_instance_segmentation(self, *args, **kwargs): + """ + This method forwards all its arguments to [`OneFormerImageProcessor.post_process_instance_segmentation`]. + Please refer to the docstring of this method for more information. + """ + return self.image_processor.post_process_instance_segmentation(*args, **kwargs) + + def post_process_panoptic_segmentation(self, *args, **kwargs): + """ + This method forwards all its arguments to [`OneFormerImageProcessor.post_process_panoptic_segmentation`]. + Please refer to the docstring of this method for more information. + """ + return self.image_processor.post_process_panoptic_segmentation(*args, **kwargs) + + +__all__ = ["OneFormerProcessor"] diff --git a/third_party/transformers/src/transformers/models/pe_video/__init__.py b/third_party/transformers/src/transformers/models/pe_video/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b3ac2f050915f88e8b1593c704e8b6fed8a25653 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_pe_video import * + from .modeling_pe_video import * + from .processing_pe_video import * + from .video_processing_pe_video import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/pe_video/configuration_pe_video.py b/third_party/transformers/src/transformers/models/pe_video/configuration_pe_video.py new file mode 100644 index 0000000000000000000000000000000000000000..16343b6578bb229f4b555eacf8c5415b51ed123f --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/configuration_pe_video.py @@ -0,0 +1,141 @@ +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig, PretrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig +from ..timm_wrapper import TimmWrapperConfig + + +@auto_docstring(checkpoint="facebook/pe-av-large") +@strict +class PeVideoEncoderConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import PeAudioEncoder, PeAudioEncoderConfig + + >>> # Initializing a PeAudioEncoder style configuration + >>> configuration = PeAudioEncoderConfig() + + >>> # Initializing a model from the pe-av-large style configuration + >>> model = PeAudioEncoder(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pe_video_encoder" + sub_configs = {"vision_config": TimmWrapperConfig} + base_config_key = "audio_video_config" + + _default_vision_config_kwargs = { + "architecture": "vit_pe_core_large_patch14_336", + "do_pooling": True, + "num_classes": 1024, + "global_pool": "map", + "initializer_range": 0.02, + } + + vision_config: dict | PreTrainedConfig | None = None + hidden_size: int = 1792 + intermediate_size: int = 4800 + num_hidden_layers: int = 6 + num_attention_heads: int = 14 + num_key_value_heads: int | None = None + head_dim: int = 128 + hidden_act: str = "silu" + max_position_embeddings: int = 10000 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.rope_parameters is None: + self.rope_parameters = {"rope_theta": 20000} + + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "timm_wrapper") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]].from_dict( + {**self._default_vision_config_kwargs, **self.vision_config} + ) + elif self.vision_config is None: + self.vision_config = CONFIG_MAPPING["timm_wrapper"].from_dict(self._default_vision_config_kwargs) + + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="facebook/pe-av-large") +@strict +class PeVideoConfig(PretrainedConfig): + r""" + video_config (`dict` or `PreTrainedConfig`, *optional*): + Configuration for the video encoder component. + + ```python + >>> from transformers import PeVideoModel, PeVideoConfig + + >>> # Initializing a PeVideoModel style configuration + >>> configuration = PeVideoConfig() + + >>> # Initializing a model from the pe-av-large style configuration + >>> model = PeVideoModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pe_video" + sub_configs = {"text_config": AutoConfig, "video_config": PeVideoEncoderConfig} + base_config_key = "audio_video_config" + + _default_text_config_kwargs = { + "model_type": "modernbert", + "hidden_size": 1024, + "intermediate_size": 2624, + "num_hidden_layers": 22, + "num_attention_heads": 16, + } + + text_config: dict | PreTrainedConfig | None = None + video_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "modernbert") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]]( + **{**self._default_text_config_kwargs, **self.text_config} + ) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["modernbert"](**self._default_text_config_kwargs) + + if isinstance(self.video_config, dict): + self.video_config = PeVideoEncoderConfig(**self.video_config) + elif self.video_config is None: + self.video_config = PeVideoEncoderConfig() + + super().__post_init__(**kwargs) + + +__all__ = ["PeVideoEncoderConfig", "PeVideoConfig"] diff --git a/third_party/transformers/src/transformers/models/pe_video/modeling_pe_video.py b/third_party/transformers/src/transformers/models/pe_video/modeling_pe_video.py new file mode 100644 index 0000000000000000000000000000000000000000..db0377f73c689aed1c343d135f166b3306c864f3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/modeling_pe_video.py @@ -0,0 +1,650 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pe_video/modular_pe_video.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_pe_video.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import AutoModel, AutoModelForImageClassification +from .configuration_pe_video import PeVideoConfig, PeVideoEncoderConfig + + +# TODO: not sure about the typing for text_model_output +@dataclass +# @auto_docstring +class PeVideoOutput(ModelOutput): + loss: torch.FloatTensor | None = None + logits_video_text: torch.FloatTensor | None = None + text_video_embeds: torch.FloatTensor | None = None + video_embeds: torch.FloatTensor | None = None + text_outputs: BaseModelOutputWithPooling = None + video_outputs: BaseModelOutputWithPooling = None + + def to_tuple(self) -> tuple[Any]: + return tuple( + self[k] if k not in ["text_outputs", "video_outputs"] else getattr(self, k).to_tuple() for k in self.keys() + ) + + +class PeVideoContrastiveHead(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + ) -> None: + super().__init__() + self.layer_norm = nn.LayerNorm(normalized_shape=in_dim, eps=1e-6) + self.proj = nn.Linear(in_dim, out_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.FloatTensor: + return self.proj(self.layer_norm(x)) + + +class PeVideoMaskedGroupNorm(nn.GroupNorm): + def forward(self, x, padding_mask=None): + if padding_mask is None: + return super().forward(x) + + batch_size, hidden_size, seq_len = x.shape + group_size = hidden_size // self.num_groups + grouped_shape = (batch_size, -1, group_size, seq_len) + + x_grouped = x.view(grouped_shape) + padding_mask_grouped = padding_mask.reshape(grouped_shape).bool() + + mean = torch.masked.mean(x_grouped, mask=padding_mask_grouped, dim=(2, 3), keepdim=True) + var = torch.masked.var(x_grouped, mask=padding_mask_grouped, dim=(2, 3), keepdim=True, unbiased=False) + + x_norm = (x_grouped - mean) / torch.sqrt(var + self.eps) + x_norm = x_norm.view(x.shape) + + if self.affine: + x_norm = x_norm * self.weight.view(1, -1, 1) + self.bias.view(1, -1, 1) + + return x_norm * padding_mask + + +class PeVideoConvBlock1d(nn.Module): + def __init__(self, config): + super().__init__() + self.groupnorm = PeVideoMaskedGroupNorm(num_groups=1, num_channels=config.hidden_size) + self.activation = nn.SiLU() + self.project = nn.Conv1d( + in_channels=config.hidden_size, + out_channels=config.hidden_size, + kernel_size=3, + padding="same", + ) + + def forward(self, x, padding_mask=None): + x = self.groupnorm(x, padding_mask=padding_mask) + x = self.activation(x) + return self.project(x) + + +class PeVideoResnetBlock1d(nn.Module): + def __init__(self, config): + super().__init__() + self.block1 = PeVideoConvBlock1d(config) + self.block2 = PeVideoConvBlock1d(config) + + def forward(self, hidden_states, padding_mask=None): + """ + Args: + hidden_states: (batch_size, seq_len, hidden_size) + padding_mask: (batch_size, seq_len) + Returns: + hidden_states: (batch_size, seq_len, hidden_size) + """ + # transpose for convolutions + # (batch_size, seq_len, hidden_size) -> (batch_size, hidden_size, seq_len) + hidden_states = hidden_states.transpose(1, 2) + + if padding_mask is not None: + padding_mask = padding_mask.unsqueeze(1).expand_as(hidden_states) + + residual = hidden_states + hidden_states = self.block1(hidden_states, padding_mask=padding_mask) + hidden_states = self.block2(hidden_states, padding_mask=padding_mask) + hidden_states = residual + hidden_states + + return hidden_states.transpose(1, 2) + + +class PeVideoEncoderPatchEmbedder(nn.Module): + def __init__(self, config): + super().__init__() + self.resnet_block = PeVideoResnetBlock1d(config) + self.class_embedding = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + + def forward(self, inputs_embeds, padding_mask=None): + # Embedding step: prepend class token and run the ResNet block. + hidden_states = torch.cat( + [self.class_embedding.expand(inputs_embeds.size(0), -1, -1), inputs_embeds], + dim=1, + ) + + if padding_mask is not None: + # TODO: any reason why we take padding_mask[0] and not just 1? + padding_mask = torch.cat([padding_mask[:, [0]], padding_mask], dim=1) + + hidden_states = self.resnet_block(hidden_states, padding_mask=padding_mask) + return hidden_states, padding_mask + + +class PeVideoEncoderEmbedder(nn.Module): + def __init__(self, config: PeVideoEncoderConfig): + super().__init__() + self.vision_model = AutoModelForImageClassification.from_config(config.vision_config) + self.proj = nn.Linear(config.vision_config.num_labels, config.hidden_size, bias=False) + self.data_proj = nn.Linear(config.hidden_size, config.hidden_size) + + def forward( + self, + pixel_values_videos: torch.Tensor, + padding_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = pixel_values_videos.shape + + pixel_values_videos = pixel_values_videos.view(-1, *input_shape[2:]) + vision_encoder_outputs = self.vision_model(pixel_values_videos) + + logits = vision_encoder_outputs.logits.view(*input_shape[:2], -1) + logits = F.normalize(logits, dim=-1) + + vision_features = self.proj(logits) + inputs_embeds = self.data_proj(vision_features) + + return inputs_embeds, padding_mask + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +def stack_freqs(cos: torch.Tensor, sin: torch.Tensor): + dim = cos.size(-1) + cos = cos.narrow(-1, 0, dim // 2) + sin = sin.narrow(-1, 0, dim // 2) + freqs_cis = torch.stack((cos, -sin, sin, cos), dim=-1).view(*cos.size(), 2, 2) + return freqs_cis + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + freqs_cis = stack_freqs(cos, sin) + freqs_cis = freqs_cis.unsqueeze(unsqueeze_dim) + q_ = q.reshape(*q.shape[:-1], -1, 1, 2) + k_ = k.reshape(*k.shape[:-1], -1, 1, 2) + return (q_ * freqs_cis).sum(5).flatten(3), (k_ * freqs_cis).sum(5).flatten(3) + + +@use_kernel_forward_from_hub("RMSNorm") +class PeVideoEncoderRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + PeVideoEncoderRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +@use_kernelized_func(apply_rotary_pos_emb) +class PeVideoEncoderAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config, layer_idx): + super().__init__() + self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = PeVideoEncoderRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # unlike olmo, only on the head dim! + self.k_norm = PeVideoEncoderRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # thus post q_norm does not need reshape + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class PeVideoEncoderMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class PeVideoEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = PeVideoEncoderAttention(config=config, layer_idx=layer_idx) + + self.mlp = PeVideoEncoderMLP(config) + self.input_layernorm = PeVideoEncoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = PeVideoEncoderRMSNorm(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 PeVideoPreTrainedModel(PreTrainedModel): + config: PeVideoConfig + base_model_prefix = "video_model" + supports_gradient_checkpointing = True + _no_split_modules = ["PeVideoEncoderLayer"] + _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": PeVideoEncoderLayer, + "attentions": PeVideoEncoderAttention, + } + main_input_name = "pixel_values_videos" + + def _init_weights(self, module): + super()._init_weights(module) + + if hasattr(self.config, "initializer_range"): + std = self.config.initializer_range + else: + # 0.02 is the standard default value across the library + std = getattr(self.config.get_text_config(), "initializer_range", 0.02) + + if isinstance(module, PeVideoEncoderPatchEmbedder): + embed_dim = module.class_embedding.shape[-1] + init.normal_(module.class_embedding, mean=0.0, std=embed_dim**-0.5 * std) + + +class PeVideoEncoderRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: PeVideoEncoderConfig, 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: PeVideoEncoderConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@auto_docstring( + custom_intro=""" + The PeVideo Encoder model. + """ +) +class PeVideoEncoder(PeVideoPreTrainedModel): + config: PeVideoEncoderConfig + main_input_name = "pixel_values_videos" + base_model_prefix = "video_model.video_encoder" + + def __init__(self, config: PeVideoEncoderConfig): + super().__init__(config) + self.embedder = PeVideoEncoderEmbedder(config) + self.patch_embedder = PeVideoEncoderPatchEmbedder(config) + self.layers = nn.ModuleList( + [PeVideoEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = PeVideoEncoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = PeVideoEncoderRotaryEmbedding(config=config) + self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.gradient_checkpointing = False + + self.post_init() + + @can_return_tuple + @merge_with_config_defaults + @capture_outputs + def forward( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPooling: + inputs_embeds, padding_mask = self.embedder(pixel_values_videos, padding_mask=padding_mask_videos) + inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) + + if attention_mask is not None: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + for encoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + hidden_states = self.output(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states[:, 1:], + pooler_output=hidden_states[:, 0], + ) + + +class PeVideoModel(PeVideoPreTrainedModel): + main_input_name = "input_ids" + + def __init__(self, config: PeVideoConfig): + super().__init__(config) + self.text_model = AutoModel.from_config(config.text_config) + self.video_encoder = PeVideoEncoder(config.video_config) + + self.text_video_head = PeVideoContrastiveHead(config.text_config.hidden_size, config.text_config.hidden_size) + self.video_head = PeVideoContrastiveHead(config.video_config.hidden_size, config.text_config.hidden_size) + + self.text_video_logit_scale = nn.Parameter(torch.zeros(1)) + self.text_video_logit_bias = nn.Parameter(torch.zeros(1)) + + self.post_init() + + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + text_outputs.pooler_output = self.text_video_head(text_outputs.last_hidden_state) + return text_outputs + + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, + padding_mask_videos=padding_mask_videos, + return_dict=True, + **kwargs, + ) + video_outputs.pooler_output = self.video_head(video_outputs.pooler_output) + return video_outputs + + @can_return_tuple + def forward( + self, + input_ids: torch.Tensor, + pixel_values_videos: torch.Tensor, + attention_mask: torch.Tensor | None = None, + padding_mask_videos: torch.Tensor | None = None, + return_loss: bool | None = None, + **kwargs, + ) -> PeVideoOutput: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, padding_mask_videos=padding_mask_videos, **kwargs + ) + kwargs["output_hidden_states"] = True + text_outputs: MaskedLMOutput = self.text_model(input_ids=input_ids, attention_mask=attention_mask, **kwargs) + + video_embeds = video_outputs.pooler_output + video_embeds = self.video_head(video_embeds) + + text_video_embeds = text_outputs.hidden_states[-1][:, 0] + text_video_embeds = self.text_video_head(text_video_embeds) + + logits_video_text = video_embeds @ text_video_embeds.T + logits_video_text = logits_video_text * self.text_video_logit_scale + self.text_video_logit_bias + + loss = None + if return_loss: + labels = torch.eye(logits_video_text.shape[0], device=logits_video_text.device) + loss = -F.logsigmoid(labels * logits_video_text).sum() / logits_video_text.shape[0] + + return PeVideoOutput( + logits_video_text=logits_video_text, + text_video_embeds=text_video_embeds, + video_embeds=video_embeds, + text_outputs=text_outputs, + video_outputs=video_outputs, + loss=loss, + ) + + +__all__ = ["PeVideoEncoder", "PeVideoModel"] diff --git a/third_party/transformers/src/transformers/models/pe_video/modular_pe_video.py b/third_party/transformers/src/transformers/models/pe_video/modular_pe_video.py new file mode 100644 index 0000000000000000000000000000000000000000..2c7764e6fe251cd55c2ee777ef5401f59a6fba1c --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/modular_pe_video.py @@ -0,0 +1,237 @@ +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...masking_utils import create_bidirectional_mask +from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import AutoModel, AutoModelForImageClassification +from ..pe_audio_video.modeling_pe_audio_video import ( + PeAudioVideoContrastiveHead, + PeAudioVideoEncoder, + PeAudioVideoEncoderPatchEmbedder, + PeAudioVideoPreTrainedModel, +) +from .configuration_pe_video import PeVideoConfig, PeVideoEncoderConfig + + +# TODO: not sure about the typing for text_model_output +@dataclass +# @auto_docstring +class PeVideoOutput(ModelOutput): + loss: torch.FloatTensor | None = None + logits_video_text: torch.FloatTensor | None = None + text_video_embeds: torch.FloatTensor | None = None + video_embeds: torch.FloatTensor | None = None + text_outputs: BaseModelOutputWithPooling = None + video_outputs: BaseModelOutputWithPooling = None + + def to_tuple(self) -> tuple[Any]: + return tuple( + self[k] if k not in ["text_outputs", "video_outputs"] else getattr(self, k).to_tuple() for k in self.keys() + ) + + +class PeVideoContrastiveHead(PeAudioVideoContrastiveHead): ... + + +class PeVideoEncoderPatchEmbedder(PeAudioVideoEncoderPatchEmbedder): ... + + +class PeVideoEncoderEmbedder(nn.Module): + def __init__(self, config: PeVideoEncoderConfig): + super().__init__() + self.vision_model = AutoModelForImageClassification.from_config(config.vision_config) + self.proj = nn.Linear(config.vision_config.num_labels, config.hidden_size, bias=False) + self.data_proj = nn.Linear(config.hidden_size, config.hidden_size) + + def forward( + self, + pixel_values_videos: torch.Tensor, + padding_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = pixel_values_videos.shape + + pixel_values_videos = pixel_values_videos.view(-1, *input_shape[2:]) + vision_encoder_outputs = self.vision_model(pixel_values_videos) + + logits = vision_encoder_outputs.logits.view(*input_shape[:2], -1) + logits = F.normalize(logits, dim=-1) + + vision_features = self.proj(logits) + inputs_embeds = self.data_proj(vision_features) + + return inputs_embeds, padding_mask + + +class PeVideoPreTrainedModel(PeAudioVideoPreTrainedModel): + base_model_prefix = "video_model" + main_input_name = "pixel_values_videos" + + +@auto_docstring( + custom_intro=""" + The PeVideo Encoder model. + """ +) +class PeVideoEncoder(PeAudioVideoEncoder): + base_model_prefix = "video_model.video_encoder" + main_input_name = "pixel_values_videos" + + def __init__(self, config: PeVideoEncoderConfig): + super().__init__(config) + self.embedder = PeVideoEncoderEmbedder(config) + + @can_return_tuple + @merge_with_config_defaults + @capture_outputs + def forward( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPooling: + inputs_embeds, padding_mask = self.embedder(pixel_values_videos, padding_mask=padding_mask_videos) + inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) + + if attention_mask is not None: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + for encoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + hidden_states = self.output(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states[:, 1:], + pooler_output=hidden_states[:, 0], + ) + + +class PeVideoModel(PeVideoPreTrainedModel): + main_input_name = "input_ids" + + def __init__(self, config: PeVideoConfig): + super().__init__(config) + self.text_model = AutoModel.from_config(config.text_config) + self.video_encoder = PeVideoEncoder(config.video_config) + + self.text_video_head = PeVideoContrastiveHead(config.text_config.hidden_size, config.text_config.hidden_size) + self.video_head = PeVideoContrastiveHead(config.video_config.hidden_size, config.text_config.hidden_size) + + self.text_video_logit_scale = nn.Parameter(torch.zeros(1)) + self.text_video_logit_bias = nn.Parameter(torch.zeros(1)) + + self.post_init() + + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + text_outputs.pooler_output = self.text_video_head(text_outputs.last_hidden_state) + return text_outputs + + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, + padding_mask_videos=padding_mask_videos, + return_dict=True, + **kwargs, + ) + video_outputs.pooler_output = self.video_head(video_outputs.pooler_output) + return video_outputs + + @can_return_tuple + def forward( + self, + input_ids: torch.Tensor, + pixel_values_videos: torch.Tensor, + attention_mask: torch.Tensor | None = None, + padding_mask_videos: torch.Tensor | None = None, + return_loss: bool | None = None, + **kwargs, + ) -> PeVideoOutput: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, padding_mask_videos=padding_mask_videos, **kwargs + ) + kwargs["output_hidden_states"] = True + text_outputs: MaskedLMOutput = self.text_model(input_ids=input_ids, attention_mask=attention_mask, **kwargs) + + video_embeds = video_outputs.pooler_output + video_embeds = self.video_head(video_embeds) + + text_video_embeds = text_outputs.hidden_states[-1][:, 0] + text_video_embeds = self.text_video_head(text_video_embeds) + + logits_video_text = video_embeds @ text_video_embeds.T + logits_video_text = logits_video_text * self.text_video_logit_scale + self.text_video_logit_bias + + loss = None + if return_loss: + labels = torch.eye(logits_video_text.shape[0], device=logits_video_text.device) + loss = -F.logsigmoid(labels * logits_video_text).sum() / logits_video_text.shape[0] + + return PeVideoOutput( + logits_video_text=logits_video_text, + text_video_embeds=text_video_embeds, + video_embeds=video_embeds, + text_outputs=text_outputs, + video_outputs=video_outputs, + loss=loss, + ) + + +__all__ = [ + "PeVideoEncoder", + "PeVideoModel", +] diff --git a/third_party/transformers/src/transformers/models/pe_video/processing_pe_video.py b/third_party/transformers/src/transformers/models/pe_video/processing_pe_video.py new file mode 100644 index 0000000000000000000000000000000000000000..515311ca99413991eeef735a7784c768d297d65c --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/processing_pe_video.py @@ -0,0 +1,10 @@ +from ...processing_utils import ProcessorMixin + + +class PeVideoProcessor(ProcessorMixin): + attributes = ["video_processor", "tokenizer"] + video_processor_class = "PeVideoVideoProcessor" + tokenizer_class = "AutoTokenizer" + + +__all__ = ["PeVideoProcessor"] diff --git a/third_party/transformers/src/transformers/models/pe_video/video_processing_pe_video.py b/third_party/transformers/src/transformers/models/pe_video/video_processing_pe_video.py new file mode 100644 index 0000000000000000000000000000000000000000..82a6d0e1121e0b1c5051fed6ffb940ff1a2e18ad --- /dev/null +++ b/third_party/transformers/src/transformers/models/pe_video/video_processing_pe_video.py @@ -0,0 +1,64 @@ +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...image_processing_utils import BatchFeature +from ...image_utils import PILImageResampling +from ...processing_utils import Unpack, VideosKwargs +from ...video_processing_utils import BaseVideoProcessor, VideoMetadata +from ...video_utils import VideoInput + + +class PeVideoVideoProcessor(BaseVideoProcessor): + resample = PILImageResampling.BILINEAR + + def sample_frames( + self, + metadata: VideoMetadata, + num_frames: int | None = None, + fps: int | float | None = None, + **kwargs, + ): + if num_frames: + total_frames = metadata.total_num_frames + num_frames = num_frames if num_frames is not None else self.num_frames + assert num_frames is not None, "`num_frames` must be specified if `fixed_len_video == True`" + frame_idxs = [int(i * (total_frames - 1) / (num_frames - 1)) for i in range(num_frames)] + return torch.tensor(frame_idxs) + else: + return super().sample_frames(metadata, num_frames, fps, **kwargs) + + def _preprocess( + self, + videos: VideoInput, + **kwargs: Unpack[VideosKwargs], + ) -> BatchFeature: + # Always set `return_tensors` to `None` since it won't pad variable length videos + # We'll handle this after we call the parent' method + return_tensors = kwargs.pop("return_tensors", None) + result = super()._preprocess(videos, **kwargs) + pixels = result.pixel_values_videos + data = {"pixel_values_videos": pixels} + if return_tensors: + lengths = torch.tensor([video.size(0) for video in pixels]) + pixels = torch.nn.utils.rnn.pad_sequence(pixels, batch_first=True, padding_value=0.0) + data["pixel_values_videos"] = pixels + if lengths.unique().size(0) > 1: + mask = torch.arange(lengths.max())[None] < lengths[:, None] + data["padding_mask_videos"] = mask + return BatchFeature(data=data, tensor_type=return_tensors) + + +__all__ = ["PeVideoVideoProcessor"] diff --git a/third_party/transformers/src/transformers/models/pegasus_x/__init__.py b/third_party/transformers/src/transformers/models/pegasus_x/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d362bacf491ea32dcbc89e7b8bba4d9ae1a9261 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pegasus_x/__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_pegasus_x import * + from .modeling_pegasus_x import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/pegasus_x/configuration_pegasus_x.py b/third_party/transformers/src/transformers/models/pegasus_x/configuration_pegasus_x.py new file mode 100644 index 0000000000000000000000000000000000000000..1191b6ee227beb973cab335fc93d02a4af49597f --- /dev/null +++ b/third_party/transformers/src/transformers/models/pegasus_x/configuration_pegasus_x.py @@ -0,0 +1,86 @@ +# Copyright 2022, Google 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. +"""PEGASUS-X model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/pegasus-x-large") +@strict +class PegasusXConfig(PreTrainedConfig): + r""" + num_global_tokens (`int`, *optional*, defaults to 128): + Number of global tokens to use for the encoder + block_size (`int`, *optional*, defaults to 512): + Block size for encoder local attention. Sequence length should be an exact multiple of block size. + block_size must be a multiple of 2 if stagger_local_block is True + stagger_local_blocks (`bool`, *optional*, defaults to `True`): + Whether to stagger every other local attention by half a block + + Example: + + ```python + >>> from transformers import PegasusXConfig, PegasusXModel + + >>> # Initializing a PEGASUS google/pegasus-x-large style configuration + >>> configuration = PegasusXConfig() + + >>> # Initializing a model (with random weights) from the google/pegasus-x-large style configuration + >>> model = PegasusXModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pegasus_x" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "num_attention_heads": "encoder_attention_heads", + "hidden_size": "d_model", + "num_hidden_layers": "encoder_layers", + } + + vocab_size: int = 96103 + max_position_embeddings: int = 16384 + encoder_layers: int = 16 + encoder_ffn_dim: int = 4096 + encoder_attention_heads: int = 16 + decoder_layers: int = 16 + decoder_ffn_dim: int = 4096 + decoder_attention_heads: int = 16 + encoder_layerdrop: float | int = 0.0 + decoder_layerdrop: float | int = 0.0 + use_cache: bool = True + is_encoder_decoder: bool = True + activation_function: str = "gelu" + d_model: int = 1024 + dropout: float | int = 0.1 + attention_dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + init_std: float = 0.02 + decoder_start_token_id: int | None = 0 + scale_embedding: bool = True + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + forced_eos_token_id: int | list[int] | None = 1 + num_global_tokens: int = 32 + block_size: int = 512 + stagger_local_blocks: bool = True + tie_word_embeddings: bool = True + + +__all__ = ["PegasusXConfig"] diff --git a/third_party/transformers/src/transformers/models/pegasus_x/modeling_pegasus_x.py b/third_party/transformers/src/transformers/models/pegasus_x/modeling_pegasus_x.py new file mode 100644 index 0000000000000000000000000000000000000000..233805c47426955abbb49a8c87a03e900ef88db6 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pegasus_x/modeling_pegasus_x.py @@ -0,0 +1,1353 @@ +# Copyright 2022, Google 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 PEGASUS-X model.""" + +import math +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn +from torch.nn import CrossEntropyLoss + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils.deprecation import deprecate_kwarg +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_pegasus_x import PegasusXConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +class DimensionInfo: + """Wrapper for dimension info.""" + + batch_size: int # batch size + seq_len: int # token length + block_size: int # block size + num_heads: int # num heads + hidden_dim: int # hidden dim + dim_per_head: int # dim per head + num_blocks: int # num blocks + global_len: int # global length + padded_seq_len: int # padded token seq length + + +# Copied from transformers.models.bart.modeling_bart.shift_tokens_right +def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): + """ + Shift input ids one token to the right. + """ + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +# Copied from transformers.models.bart.modeling_bart.BartScaledWordEmbedding with Bart->PegasusX +class PegasusXScaledWordEmbedding(nn.Embedding): + """ + This module overrides nn.Embeddings' forward by multiplying with embeddings scale. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float | None = 1.0): + super().__init__(num_embeddings, embedding_dim, padding_idx) + self.embed_scale = embed_scale + + def forward(self, input_ids: torch.Tensor): + return super().forward(input_ids) * self.embed_scale + + +class PegasusXSinusoidalPositionalEmbedding(nn.Module): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, embed_dim, max_scale: int = 10000.0): + super().__init__() + self.embed_dim = embed_dim + self.max_scale = max_scale + + @deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") + @torch.no_grad() + def forward( + self, inputs_embeds: torch.Tensor, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None + ) -> torch.Tensor: + """`input_ids_shape` is expected to be [bsz x seqlen].""" + batch_size, seq_len = inputs_embeds.shape[:2] + if position_ids is None: + position_ids = torch.arange( + past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=inputs_embeds.device + )[:, None] + + pe = torch.zeros((seq_len, self.embed_dim), device=inputs_embeds.device, dtype=inputs_embeds.dtype) + half_d_feature = self.embed_dim // 2 + div_term = torch.exp( + torch.arange(half_d_feature, device=inputs_embeds.device, dtype=torch.int64).type_as(inputs_embeds) + * -(np.log(float(self.max_scale)) / (half_d_feature - 1)) + ) + pe[:, :half_d_feature] = torch.sin(position_ids * div_term) + pe[:, half_d_feature:] = torch.cos(position_ids * div_term) + return pe[None].expand(batch_size, -1, -1) + + +# Copied from transformers.models.bert.modeling_bert.eager_attention_forward +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float | None = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->PegasusX +class PegasusXAttention(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: PegasusXConfig | None = None, + layer_idx: int | None = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + self.layer_idx = layer_idx + if layer_idx is None and self.is_decoder: + logger.warning_once( + f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and " + "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2) + + is_updated = False + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + + current_states = key_value_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = curr_past_key_values.layers[self.layer_idx].keys + value_states = curr_past_key_values.layers[self.layer_idx].values + else: + key_states = self.k_proj(current_states) + value_states = self.v_proj(current_states) + key_states = key_states.view(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class PegasusXGlobalLocalAttention(nn.Module): + """Global + Local attention. For use with Encoder only.""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + block_size: int, + dropout: float = 0.0, + is_decoder: bool = False, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.block_size = block_size + self.dropout = dropout + self.head_dim = embed_dim // num_heads + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + token_hidden_states: torch.Tensor, + global_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + dim = DimensionInfo( + batch_size=token_hidden_states.shape[0], + seq_len=token_hidden_states.shape[1], + block_size=self.block_size, + num_heads=self.num_heads, + hidden_dim=token_hidden_states.shape[2], + dim_per_head=self.head_dim, + num_blocks=token_hidden_states.shape[1] // self.block_size, + global_len=global_hidden_states.shape[1], + padded_seq_len=token_hidden_states.shape[1], + ) + + # [batch_size, num_heads, padded_seq_len, dim_per_head] + local_q = self._shape( + self.q_proj(token_hidden_states) * self.scaling, + seq_len=dim.padded_seq_len, + bsz=dim.batch_size, + ) + local_k = self._shape( + self.k_proj(token_hidden_states), + seq_len=dim.padded_seq_len, + bsz=dim.batch_size, + ) + local_v = self._shape( + self.v_proj(token_hidden_states), + seq_len=dim.padded_seq_len, + bsz=dim.batch_size, + ) + + # [batch_size, num_heads, global_len, dim_per_head] + global_q = self._shape( + self.q_proj(global_hidden_states) * self.scaling, + seq_len=dim.global_len, + bsz=dim.batch_size, + ) + global_k = self._shape( + self.k_proj(global_hidden_states), + seq_len=dim.global_len, + bsz=dim.batch_size, + ) + global_v = self._shape( + self.v_proj(global_hidden_states), + seq_len=dim.global_len, + bsz=dim.batch_size, + ) + + global_attn_output, global_attn_probs = self.compute_global_attention_representations( + global_q=global_q, + global_k=global_k, + global_v=global_v, + local_k=local_k, + local_v=local_v, + mask=attention_mask, + dim=dim, + ) + local_attn_output, local_attn_probs = self.compute_local_attention_representations( + global_k=global_k, + global_v=global_v, + local_q=local_q, + local_k=local_k, + local_v=local_v, + mask=attention_mask, + dim=dim, + ) + + # [batch_size, global_len, hidden_dim] + global_attn_output = ( + global_attn_output.transpose(1, 2).contiguous().view(dim.batch_size, dim.global_len, dim.hidden_dim) + ) + # [batch_size, global_len, hidden_dim] + global_attn_output = self.out_proj(global_attn_output) + # [batch_size, num_heads, block_size, num_heads, dim_per_head] + local_attn_output = local_attn_output.permute(0, 2, 3, 1, 4).contiguous() + # [batch_size, padded_seq_len, hidden_dim] + local_attn_output = local_attn_output.view(dim.batch_size, dim.padded_seq_len, dim.hidden_dim) + # [batch_size, padded_seq_len, hidden_dim] + local_attn_output = self.out_proj(local_attn_output) + + attn_probs = {"global": global_attn_probs, "local": local_attn_probs} + + return local_attn_output, global_attn_output, attn_probs + + def compute_global_attention_representations( + self, global_q, global_k, global_v, local_k, local_v, mask, dim: DimensionInfo + ): + """Compute attention representations for global tokens. + + Global tokens will attend to both global tokens as well as all input sequence tokens. Because the input + sequence tokens are arranged in blocks for local attention, we unblock them and compute attention. + + Args: + global_q (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]: + query vectors from global tokens + global_k (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]: + key vectors from global tokens + global_v (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]: + value vectors from global tokens + local_k (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]: + key vectors from local tokens + local_v (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]: + value vectors from local tokens + mask (`torch.FloatTensor`) of shape [batch_size, padded_seq_len]: attention mask + dim (DimensionInfo): DimensionInfo wrapper for dimensions + + Returns: + output of shape `[batch_sizes, length, features]`. where length will be padded to a multiple of block_size + """ + # [batch_size, num_heads, global_len+padded_seq_len, dim_per_head] + global_and_local_k = torch.cat([global_k, local_k], dim=2) + # [batch_size, num_heads, global_len+padded_seq_len, dim_per_head] + global_and_local_v = torch.cat([global_v, local_v], dim=2) + + # [batch_size, global_len+padded_seq_len] + extended_mask = nn.functional.pad(mask, pad=(dim.global_len, 0), value=0) + + # [batch_size, num_heads, global_len, global_len+padded_seq_len] + attn_weights = torch.einsum("BHGF,BHXF->BHGX", global_q, global_and_local_k) + attn_weights = attn_weights + extended_mask[:, None, None, :] + attn_probs = nn.functional.softmax(attn_weights, dim=-1) + attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training) + + # [batch_size, num_heads, global_len, F] + attn_output = torch.einsum("BHGX,BHXF->BHGF", attn_probs, global_and_local_v) + return attn_output, attn_probs + + def compute_local_attention_representations( + self, global_k, global_v, local_q, local_k, local_v, mask, dim: DimensionInfo + ): + """Compute attention representations for local tokens. + + Local tokens will attend to both global tokens as well as all other tokens within the same local block. Hence, + we need to tile and concatenate the global tokens to every local block + + Args: + global_k (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]: + key vectors from global tokens + global_v (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]: + value vectors from global tokens + local_q (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]: + query vectors from local tokens + local_k (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]: + key vectors from local tokens + local_v (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]: + value vectors from local tokens + mask (`torch.FloatTensor`) of shape [batch_size, padded_seq_len]: attention mask + dim (DimensionInfo): DimensionInfo wrapper for dimensions + + Returns: + output of shape `[batch_sizes, length, features]`. where length will be padded to a multiple of block_size + """ + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + blocked_local_q = local_q.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head) + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + blocked_local_k = local_k.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head) + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + blocked_local_v = local_v.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head) + + # [batch_size, num_blocks, global_len+block_size] + extended_mask = nn.functional.pad( + mask.view(dim.batch_size, dim.num_blocks, dim.block_size), + pad=(dim.global_len, 0), + value=0, + ) + + # [batch_size, num_heads, num_blocks, block_size, global_len] + blocked_local2global = torch.einsum("BHNKF,BHGF->BHNKG", blocked_local_q, global_k) + # [batch_size, num_heads, num_blocks, block_size, block_size] + blocked_local2local = torch.einsum("BHNKF,BHNXF->BHNKX", blocked_local_q, blocked_local_k) + + # [batch_size, num_heads, num_blocks, block_size, global_len+block_size] + attn_weights = torch.cat([blocked_local2global, blocked_local2local], dim=-1) + attn_weights = attn_weights + extended_mask[:, None, :, None, :] + attn_probs = nn.functional.softmax(attn_weights, dim=-1) + attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training) + + # [batch_size, num_heads, num_blocks, block_size, global_len] + local2global_attn_probs = attn_probs[:, :, :, :, : dim.global_len] + # [batch_size, num_heads, num_blocks, block_size, block_size] + local2local_attn_probs = attn_probs[:, :, :, :, dim.global_len :] + + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + local2global_attn_output = torch.einsum("BHNKG,BHGF->BHNKF", local2global_attn_probs, global_v) + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + local2local_attn_output = torch.einsum("BHNKX,BHNXF->BHNKF", local2local_attn_probs, blocked_local_v) + # [batch_size, num_heads, num_blocks, block_size, dim_per_head] + attn_output = local2global_attn_output + local2local_attn_output + return attn_output, attn_probs + + +class PegasusXEncoderLayer(GradientCheckpointingLayer): + def __init__(self, stagger_blocks_this_layer: bool, config: PegasusXConfig): + super().__init__() + self.embed_dim = config.d_model + self.self_attn = PegasusXGlobalLocalAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + block_size=config.block_size, + dropout=config.attention_dropout, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.global_self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + self.stagger_blocks_this_layer = stagger_blocks_this_layer + self.block_size = config.block_size + + def forward( + self, + hidden_states: torch.Tensor, + global_hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + global_hidden_states (`torch.FloatTensor`): global token hidden states + *(seq_len, num_global_tokens, embed_dim)* + attention_mask (`torch.FloatTensor`): attention mask of size + *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + global_residual = global_hidden_states + + hidden_states = self.self_attn_layer_norm(hidden_states) + global_hidden_states = self.global_self_attn_layer_norm(global_hidden_states) + + if self.stagger_blocks_this_layer: + # Pad the blocks to simulate staggering + hidden_states, attention_mask = self.pad_local_tokens( + hidden_states=hidden_states, attention_mask=attention_mask, block_size=self.block_size + ) + + hidden_states, global_hidden_states, _ = self.self_attn( + token_hidden_states=hidden_states, + global_hidden_states=global_hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + + if self.stagger_blocks_this_layer: + # Undo the padding + hidden_states = self.unpad_local_tokens(padded_hidden_states=hidden_states, block_size=self.block_size) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training) + global_hidden_states = global_residual + global_hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + global_residual = global_hidden_states + global_hidden_states = self.final_layer_norm(global_hidden_states) + global_hidden_states = self.activation_fn(self.fc1(global_hidden_states)) + global_hidden_states = nn.functional.dropout( + global_hidden_states, p=self.activation_dropout, training=self.training + ) + global_hidden_states = self.fc2(global_hidden_states) + global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training) + global_hidden_states = global_residual + global_hidden_states + + return hidden_states, global_hidden_states + + @classmethod + def pad_local_tokens(cls, hidden_states, attention_mask, block_size): + # hidden_states: [batch_size, seq_len, hidden_dim] + pad_size = block_size // 2 + mask_min_value = torch.finfo(hidden_states.dtype).min + padded_hidden_states = torch.nn.functional.pad( + hidden_states, + pad=(0, 0, pad_size, pad_size), + ) + padded_mask = torch.nn.functional.pad( + attention_mask, + pad=(pad_size, pad_size), + value=mask_min_value, + ) + return padded_hidden_states, padded_mask + + @classmethod + def unpad_local_tokens(cls, padded_hidden_states, block_size): + # padded_hidden_states: [batch_size, padded seq_len, hidden_dim] + pad_size = block_size // 2 + return padded_hidden_states[:, pad_size:-pad_size, :] + + +class PegasusXDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: PegasusXConfig, layer_idx: int | None = None): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = PegasusXAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + bias=False, + config=config, + layer_idx=layer_idx, + ) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.encoder_attn = PegasusXAttention( + self.embed_dim, + config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + bias=False, + config=config, + layer_idx=layer_idx, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim) + self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = True, + **kwargs, + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + attention_mask (`torch.FloatTensor`): attention mask of size + *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape *(seq_len, batch, embed_dim)* + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + past_key_values (`Cache`): cached past key and value projection states + use_cache: Whether to us KV cache for decoding + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring +class PegasusXPreTrainedModel(PreTrainedModel): + config: PegasusXConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = [r"PegasusXEncoderLayer", r"PegasusXDecoderLayer"] + _supports_flash_attn = True + # Flaky logits + _supports_sdpa = False + _supports_flex_attn = True + _can_compile_fullgraph = True + + +class PegasusXEncoder(PegasusXPreTrainedModel): + _can_record_outputs = {"attentions": OutputRecorder(PegasusXGlobalLocalAttention, index=2)} + + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`PegasusXEncoderLayer`]. + + Args: + config: PegasusXConfig + embed_tokens (nn.Embedding): output embedding + """ + + def __init__(self, config: PegasusXConfig): + super().__init__(config) + + self.dropout = config.dropout + self.layerdrop = config.encoder_layerdrop + + embed_dim = config.d_model + padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 + + self.embed_tokens = PegasusXScaledWordEmbedding( + config.vocab_size, embed_dim, padding_idx, embed_scale=embed_scale + ) + + self.embed_global = nn.Embedding(config.num_global_tokens, embed_dim) + self.embed_positions = PegasusXSinusoidalPositionalEmbedding(embed_dim) + self.layers = nn.ModuleList( + [ + PegasusXEncoderLayer( + stagger_blocks_this_layer=i % 2 == 1 and config.stagger_local_blocks, config=config + ) + for i in range(config.encoder_layers) + ] + ) + self.layer_norm = nn.LayerNorm(config.d_model) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def resize_position_embeddings(self, new_num_position_embeddings: int): + """ + Resizes position embeddings matrix of the model if `new_num_position_embeddings != + config.max_position_embeddings`. + + Arguments: + new_num_position_embeddings (`int`): + The number of new position embeddings. If position embeddings are learned, increasing the size will add + newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If + position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will + add correct vectors at the end following the position encoding algorithm, whereas reducing the size + will remove vectors from the end. + """ + logger.info(f"Setting `config.max_position_embeddings={new_num_position_embeddings}`...") + self.config.max_position_embeddings = new_num_position_embeddings + + self.embed_positions = PegasusXSinusoidalPositionalEmbedding(self.config.d_model) + self.embed_positions.to(self.device) + + def get_position_embeddings(self) -> nn.Embedding: + """ + Returns the position embeddings matrix + """ + return self.embed_positions + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids=None, + attention_mask=None, + inputs_embeds=None, + output_hidden_states=None, + **kwargs: Unpack[TransformersKwargs], + ): + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + """ + # We need to treat this special because it only adds the last global state which is unique + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + embed_pos = self.embed_positions(inputs_embeds) + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + batch_size, seq_len, _ = hidden_states.shape + + # Setup mask + if attention_mask is None: + attention_mask = torch.ones( + *inputs_embeds.shape[:-1], dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + attention_mask = attention_mask.to(dtype=hidden_states.dtype) + mask_min_value = torch.finfo(hidden_states.dtype).min + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), + mask_min_value, + ) + + # padding to block_size + if seq_len % self.config.block_size != 0: + pad_len = self.config.block_size - seq_len % self.config.block_size + hidden_states = nn.functional.pad(hidden_states, pad=(0, 0, 0, pad_len), value=0) + attention_mask = nn.functional.pad(attention_mask, pad=(0, pad_len), value=mask_min_value) + + # Global tokens + global_hidden_states = self.embed_global( + torch.arange(self.config.num_global_tokens, device=hidden_states.device)[None].expand(batch_size, -1) + ) + + encoder_states = () if output_hidden_states else None + for encoder_layer in self.layers: + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if to_drop: + hidden_states, global_hidden_states = (None, None) + else: + hidden_states, global_hidden_states = encoder_layer( + hidden_states, + global_hidden_states, + attention_mask, + **kwargs, + ) + + # Undo padding-to-block-size + hidden_states = hidden_states[:, :seq_len] + + hidden_states = self.layer_norm(hidden_states) + + if output_hidden_states: + encoder_states = encoder_states + ((hidden_states, global_hidden_states),) + + return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states) + + +class PegasusXDecoder(PegasusXPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`PegasusDecoderLayer`] + + Args: + config: PegasusXConfig + embed_tokens (nn.Embedding): output embedding + """ + + _can_record_outputs = { + "hidden_states": PegasusXDecoderLayer, + "attentions": OutputRecorder(PegasusXAttention, index=1, layer_name="self_attn"), + "cross_attentions": OutputRecorder(PegasusXAttention, index=1, layer_name="encoder_attn"), + } + + def __init__(self, config: PegasusXConfig): + super().__init__(config) + self.dropout = config.dropout + self.layerdrop = config.decoder_layerdrop + self.max_target_positions = config.max_position_embeddings + embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + padding_idx = config.pad_token_id + + self.embed_tokens = PegasusXScaledWordEmbedding( + config.vocab_size, config.d_model, padding_idx=padding_idx, embed_scale=embed_scale + ) + + self.embed_positions = PegasusXSinusoidalPositionalEmbedding(config.d_model) + self.layers = nn.ModuleList([PegasusXDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)]) + self.layer_norm = nn.LayerNorm(config.d_model) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + @capture_outputs + @merge_with_config_defaults + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ): + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention + of the decoder. + encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): + Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values + selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the + cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those + that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of + all `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of + shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing + `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more + control over how to convert `input_ids` indices into associated vectors than the model's internal + embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # initialize `past_key_values` + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + + batch_size, seq_length = inputs_embeds.size()[:-1] + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + + if attention_mask is None and not is_torchdynamo_compiling(): + # required mask seq length can be calculated via length of past cache + mask_seq_length = past_key_values_length + seq_length + attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) + + self_attn_cache = ( + past_key_values.self_attention_cache + if isinstance(past_key_values, EncoderDecoderCache) + else past_key_values + ) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=self_attn_cache, + ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + # embed positions + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length + position_ids = position_ids.unsqueeze(1) + position_ids = self.embed_positions(inputs_embeds, past_key_values_length, position_ids) + position_ids = position_ids.to(inputs_embeds.device) + hidden_states = inputs_embeds + position_ids + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + # decoder layers + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + hidden_states = decoder_layer( + hidden_states, + causal_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.layer_norm(hidden_states) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class PegasusXModel(PegasusXPreTrainedModel): + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + } + + def __init__(self, config: PegasusXConfig): + super().__init__(config) + + vocab_size = config.vocab_size + embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + padding_idx = config.pad_token_id + self.shared = PegasusXScaledWordEmbedding( + vocab_size, config.d_model, padding_idx=padding_idx, embed_scale=embed_scale + ) + + self.encoder = PegasusXEncoder(config) + self.decoder = PegasusXDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, value): + self.shared = value + self.encoder.embed_tokens = self.shared + self.decoder.embed_tokens = self.shared + + def resize_position_embeddings(self, new_num_position_embeddings: int): + """ + Resizes position embeddings matrix of the model if `new_num_position_embeddings != + config.max_position_embeddings`. + + Arguments: + new_num_position_embeddings (`int`): + The number of new position embeddings. If position embeddings are learned, increasing the size will add + newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If + position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will + add correct vectors at the end following the position encoding algorithm, whereas reducing the size + will remove vectors from the end. + """ + self.config.max_position_embeddings = new_num_position_embeddings + self.encoder.resize_position_embeddings(new_num_position_embeddings) + self.decoder.resize_position_embeddings(new_num_position_embeddings) + + def get_position_embeddings(self) -> tuple[nn.Embedding]: + """ + Returns the position embeddings matrix + """ + return (self.encoder.get_position_embeddings(), self.decoder.get_position_embeddings()) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + encoder_outputs: tuple[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Seq2SeqModelOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + PEGASUS-X uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If + `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + + Example: + + ```python + >>> from transformers import AutoTokenizer, PegasusModel + + >>> tokenizer = AutoTokenizer.from_pretrained("google/pegasus-x-large") + >>> model = PegasusModel.from_pretrained("google/pegasus-x-large") + + >>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt") + >>> decoder_inputs = tokenizer("Studies show that", return_tensors="pt") + >>> outputs = model(input_ids=inputs.input_ids, decoder_input_ids=decoder_inputs.input_ids) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 4, 1024] + ```""" + + if encoder_outputs is None: + encoder_outputs: BaseModelOutput = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + **kwargs, + ) + elif not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs[0], + encoder_attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The PEGASUS-X for conditional generation (e.g. summarization). + """ +) +class PegasusXForConditionalGeneration(PegasusXPreTrainedModel, GenerationMixin): + base_model_prefix = "model" + _tied_weights_keys = { + "lm_head.weight": "model.shared.weight", + } + + def __init__(self, config: PegasusXConfig): + super().__init__(config) + self.model = PegasusXModel(config) + self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def resize_position_embeddings(self, new_num_position_embeddings: int): + """ + Resizes position embeddings matrix of the model if `new_num_position_embeddings != + config.max_position_embeddings`. + + Arguments: + new_num_position_embeddings (`int`): + The number of new position embeddings. If position embeddings are learned, increasing the size will add + newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If + position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will + add correct vectors at the end following the position encoding algorithm, whereas reducing the size + will remove vectors from the end. + """ + self.config.max_position_embeddings = new_num_position_embeddings + self.model.encoder.resize_position_embeddings(new_num_position_embeddings) + self.model.decoder.resize_position_embeddings(new_num_position_embeddings) + + def get_position_embeddings(self) -> tuple[nn.Embedding]: + """ + Returns the position embeddings matrix + """ + return (self.model.encoder.get_position_embeddings(), self.model.decoder.get_position_embeddings()) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + encoder_outputs: tuple[torch.FloatTensor] | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Seq2SeqLMOutput: + r""" + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + PEGASUS-X uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If + `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + """ + if labels is not None: + if use_cache: + logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.") + use_cache = False + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + + outputs: Seq2SeqModelOutput = self.model( + input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + lm_logits = self.lm_head(outputs.last_hidden_state) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1)) + + return Seq2SeqLMOutput( + loss=masked_lm_loss, + logits=lm_logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): + return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id) + + +# Copied from transformers.models.bart.modeling_bart.BartDecoderWrapper with Bart->PegasusX +class PegasusXDecoderWrapper(PegasusXPreTrainedModel): + """ + This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is + used in combination with the [`EncoderDecoderModel`] framework. + """ + + def __init__(self, config): + super().__init__(config) + self.decoder = PegasusXDecoder(config) + self.post_init() + + def forward(self, *args, **kwargs): + return self.decoder(*args, **kwargs) + + +__all__ = ["PegasusXForConditionalGeneration", "PegasusXModel", "PegasusXPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/__init__.py b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b09dc0a5d3439eef8bfedf6278db94f79cb6770d --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_pp_ocrv5_server_rec import * + from .image_processing_pp_ocrv5_server_rec import * + from .modeling_pp_ocrv5_server_rec import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modeling_pp_ocrv5_server_rec.py b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modeling_pp_ocrv5_server_rec.py new file mode 100644 index 0000000000000000000000000000000000000000..9bffac55b779f92cb25f0f8c9029c77a5f48f208 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modeling_pp_ocrv5_server_rec.py @@ -0,0 +1,390 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_pp_ocrv5_server_rec.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from ...activations import ACT2FN +from ...backbone_utils import load_backbone +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithNoAttention +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_pp_ocrv5_server_rec import PPOCRV5ServerRecConfig + + +class PPOCRV5ServerRecBlock(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = PPOCRV5ServerRecAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = PPOCRV5ServerRecMLP( + config=config, + in_features=self.embed_dim, + hidden_features=int(self.embed_dim * config.mlp_ratio), + ) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +# Adapted from transformers.models.siglip.modeling_siglip.eager_attention_forward -> BLIP doesn't cast attn weights to fp32 +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs, +): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class PPOCRV5ServerRecAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.is_causal = False + self.attention_dropout = config.attention_dropout + + # small tweak here compared to CLIP, no bias here + self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=False) + + if config.qkv_bias: + q_bias = nn.Parameter(torch.zeros(self.embed_dim)) + v_bias = nn.Parameter(torch.zeros(self.embed_dim)) + else: + q_bias = None + v_bias = None + + if q_bias is not None: + qkv_bias = torch.cat((q_bias, torch.zeros_like(v_bias, requires_grad=False), v_bias)) + self.qkv.bias = nn.Parameter(qkv_bias) + + self.projection = nn.Linear(self.embed_dim, self.embed_dim) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, # Not used but kept for signature matching in downstream modules + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + bsz, tgt_len, embed_dim = hidden_states.size() + + mixed_qkv = self.qkv(hidden_states) + + mixed_qkv = mixed_qkv.reshape(bsz, tgt_len, 3, self.num_heads, embed_dim // self.num_heads).permute( + 2, 0, 3, 1, 4 + ) + query_states, key_states, value_states = mixed_qkv[0], mixed_qkv[1], mixed_qkv[2] + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scale, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.projection(attn_output) + + return attn_output, attn_weights + + +class PPOCRV5ServerRecConvLayer(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: tuple[int, int] = (3, 3), + stride: int = 1, + activation: str = "silu", + ): + super().__init__() + self.convolution = nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=(kernel_size[0] // 2, kernel_size[1] // 2), + bias=False, + ) + self.normalization = nn.BatchNorm2d(out_channels) + self.activation = ACT2FN[activation] if activation is not None else nn.Identity() + + def forward(self, input: Tensor) -> Tensor: + hidden_state = self.convolution(input) + hidden_state = self.normalization(hidden_state) + hidden_state = self.activation(hidden_state) + return hidden_state + + +@auto_docstring +class PPOCRV5ServerRecPreTrainedModel(PreTrainedModel): + config: PPOCRV5ServerRecConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["PPOCRV5ServerRecBlock"] + _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": PPOCRV5ServerRecBlock, + } + main_input_name = "pixel_values" + input_modalities = ("image",) + + +class PPOCRV5ServerRecHead(nn.Module): + def __init__(self, config): + super().__init__() + + self.encoder = PPOCRV5ServerRecEncoderWithSVTR(config) + self.head = nn.Linear(config.hidden_size, config.head_out_channels) + + def forward(self, hidden_states: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + outputs = self.encoder(hidden_states, **kwargs) + hidden_states = self.head(outputs.last_hidden_state) + hidden_states = F.softmax(hidden_states, dim=2, dtype=torch.float32).to(hidden_states.dtype) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=outputs.hidden_states) + + +class PPOCRV5ServerRecMLP(nn.Module): + def __init__(self, config, in_features, hidden_features=None, out_features=None, drop=0.0): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.activation = ACT2FN[config.hidden_act] + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, hidden_state): + hidden_state = self.fc1(hidden_state) + hidden_state = self.activation(hidden_state) + hidden_state = self.drop(hidden_state) + hidden_state = self.fc2(hidden_state) + hidden_state = self.drop(hidden_state) + return hidden_state + + +class PPOCRV5ServerRecEncoderWithSVTR(PPOCRV5ServerRecPreTrainedModel): + """ + SVTR: Scene Text Recognition with a Single Visual Model + https://www.paddleocr.ai/v2.10.0/en/algorithm/text_recognition/algorithm_rec_svtr.html + """ + + def __init__( + self, + config, + ): + super().__init__(config) + in_channels = config.backbone_config.stage_out_channels[-1] + hidden_size = config.hidden_size + + self.conv_block = nn.ModuleList( + [ + PPOCRV5ServerRecConvLayer( + in_channels=in_channels, + out_channels=in_channels // 8, + kernel_size=config.conv_kernel_size, + ), + PPOCRV5ServerRecConvLayer(in_channels=in_channels // 8, out_channels=hidden_size, kernel_size=(1, 1)), + PPOCRV5ServerRecConvLayer(in_channels=hidden_size, out_channels=in_channels, kernel_size=(1, 1)), + PPOCRV5ServerRecConvLayer( + in_channels=2 * in_channels, + out_channels=in_channels // 8, + kernel_size=config.conv_kernel_size, + ), + PPOCRV5ServerRecConvLayer(in_channels=in_channels // 8, out_channels=hidden_size, kernel_size=(1, 1)), + ] + ) + self.svtr_block = nn.ModuleList() + for _ in range(config.depth): + self.svtr_block.append(PPOCRV5ServerRecBlock(config=config)) + + self.norm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward(self, hidden_states: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + residual = hidden_states + + hidden_states = self.conv_block[0](hidden_states) + hidden_states = self.conv_block[1](hidden_states) + + batch_size, channels, height, width = hidden_states.shape + hidden_states = hidden_states.flatten(2).transpose(1, 2) + for block in self.svtr_block: + hidden_states = block(hidden_states) + + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.view(batch_size, height, width, channels).permute(0, 3, 1, 2) + hidden_states = self.conv_block[2](hidden_states) + hidden_states = self.conv_block[3](torch.cat((residual, hidden_states), dim=1)) + hidden_states = self.conv_block[4](hidden_states) + hidden_states = hidden_states.squeeze(2).transpose(1, 2) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states) + + +@auto_docstring(custom_intro="PPOCRV5ServerRec model, consisting of Backbone and Head networks.") +class PPOCRV5ServerRecModel(PPOCRV5ServerRecPreTrainedModel): + def __init__(self, config: PPOCRV5ServerRecConfig): + super().__init__(config) + self.backbone = load_backbone(config) + + self.gradient_checkpointing = False + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention: + outputs = self.backbone(pixel_values, **kwargs) + hidden_state = outputs.feature_maps[-1] + hidden_state = F.avg_pool2d(hidden_state, (3, 2)) + + return BaseModelOutputWithNoAttention( + last_hidden_state=hidden_state, + hidden_states=outputs.hidden_states, + ) + + +@dataclass +@auto_docstring +class PPOCRV5ServerRecForTextRecognitionOutput(BaseModelOutputWithNoAttention): + r""" + head_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Hidden-states of the PPOCRV5ServerRecHead at the output of each layer plus the optional initial embedding outputs. + """ + + head_hidden_states: torch.FloatTensor | None = None + + +@auto_docstring(custom_intro="PPOCRV5ServerRec model for text recognition tasks.") +class PPOCRV5ServerRecForTextRecognition(PPOCRV5ServerRecPreTrainedModel): + _keys_to_ignore_on_load_missing = ["num_batches_tracked"] + + def __init__(self, config: PPOCRV5ServerRecConfig): + super().__init__(config) + self.model = PPOCRV5ServerRecModel(config) + self.head = PPOCRV5ServerRecHead(config) + + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention: + outputs = self.model(pixel_values, **kwargs) + head_outputs = self.head(outputs.last_hidden_state, **kwargs) + + return PPOCRV5ServerRecForTextRecognitionOutput( + last_hidden_state=head_outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + head_hidden_states=head_outputs.hidden_states, + ) + + +__all__ = [ + "PPOCRV5ServerRecForTextRecognition", + "PPOCRV5ServerRecModel", + "PPOCRV5ServerRecPreTrainedModel", + "PPOCRV5ServerRecEncoderWithSVTR", +] diff --git a/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.py b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.py new file mode 100644 index 0000000000000000000000000000000000000000..54e7d5f423eb0148d187fbb7a0c5577aeab1184d --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.py @@ -0,0 +1,469 @@ +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms.v2.functional as tvF +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config, load_backbone +from ...configuration_utils import PreTrainedConfig +from ...feature_extraction_utils import BatchFeature +from ...image_processing_backends import TorchvisionBackend +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + PILImageResampling, + SizeDict, +) +from ...modeling_outputs import BaseModelOutputWithNoAttention +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, + requires_backends, +) +from ...utils.constants import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, +) +from ...utils.generic import TensorType, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import AutoConfig +from ..blip_2.modeling_blip_2 import Blip2Attention +from ..clip.modeling_clip import CLIPEncoderLayer +from ..focalnet.modeling_focalnet import FocalNetMlp +from ..paddleocr_vl.modeling_paddleocr_vl import PaddleOCRVLPreTrainedModel +from ..resnet.modeling_resnet import ResNetConvLayer + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="PaddlePaddle/PP-OCRv5_server_rec_safetensors") +@strict +class PPOCRV5ServerRecConfig(PreTrainedConfig): + r""" + head_out_channels (`int`, *optional*, defaults to 18385): + The number of output channels from the PPOCRV5ServerRecHead, responsible for final classification. + """ + + model_type = "pp_ocrv5_server_rec" + sub_configs = {"backbone_config": AutoConfig} + + hidden_act: str = "silu" + backbone_config: dict | PreTrainedConfig | None = None + hidden_size: int = 120 + mlp_ratio: float = 2.0 + depth: int = 2 + head_out_channels: int = 18385 + conv_kernel_size: list | None = None + qkv_bias: bool = True + num_attention_heads: int = 8 + attention_dropout: float | int = 0.0 + layer_norm_eps: float = 1e-6 + + def __post_init__(self, **kwargs): + if self.conv_kernel_size is None: + self.conv_kernel_size = [1, 3] + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="hgnet_v2", + default_config_kwargs={ + "arch": "L", + "return_idx": [0, 1, 2, 3], + "freeze_stem_only": True, + "freeze_at": 0, + "freeze_norm": True, + "lr_mult_list": [1.0, 1.0, 1.0, 1.0, 1.0], + "out_features": ["stage1", "stage2", "stage3", "stage4"], + "stage_downsample": [True, True, True, True], + "stem_strides": [2, 1, 1, 1, 1], + "stage_downsample_strides": [[2, 1], [1, 2], [2, 1], [2, 1]], + }, + **kwargs, + ) + super().__post_init__(**kwargs) + + +class PPOCRV5ServerRecImageProcessorKwargs(ImagesKwargs, total=False): + r""" + max_image_width (`int`, *optional*, defaults to `3200`): + Maximum image width used during resizing. + character_list (`list`, *optional*, defaults to `[]`): + Vocabulary list used for text recognition. + """ + + max_image_width: int + character_list: str + + +@auto_docstring +class PPOCRV5ServerRecImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 48, "width": 320} + pad_size = {"height": 48, "width": 320} + do_resize = True + do_rescale = True + do_convert_rgb = True + do_normalize = True + do_pad = True + max_image_width = 3200 + character_list = [] + valid_kwargs = PPOCRV5ServerRecImageProcessorKwargs + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool | None, + pad_size: SizeDict | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + + # [Key Change] Use get_target_size to calculate target_size for resizing. + shape_list = list(grouped_images.keys()) + target_size = self.get_target_size(shape_list) + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images.float(), size=target_size, resample=resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + if do_pad and target_size.width < pad_size.width: + processed_images = self.pad(processed_images, pad_size=pad_size, disable_grouping=disable_grouping) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + def get_target_size(self, shape_list: list[torch.Size]): + """ + Calculate the width and height from the widest image in the batch. + """ + max_width = -1 + max_height = -1 + for height, width in shape_list: + # We need the width and height of the widest image in the batch + if width > max_width: + max_width = width + max_height = height + + default_height, default_width = self.size["height"], self.size["width"] + ratio = max(max_width / max_height, default_width / default_height) + + target_width = int(default_height * ratio) + target_height = default_height + + if target_width > self.max_image_width: + target_width = self.max_image_width + else: + ratio = max_width / float(max_height) + if target_width >= math.ceil(default_height * ratio): + target_width = int(math.ceil(default_height * ratio)) + + return SizeDict(height=target_height, width=target_width) + + def post_process_text_recognition( + self, + predictions, + ) -> list[dict[str, str | float]]: + """ + Post-processes raw model logits to decode the recognized text and its confidence score. + + Args: + predictions: Model outputs with `logits` attribute (probability maps of shape `(batch_size, height, vocab_size)`). + + Returns: + A list of dictionaries, where each dictionary corresponds to an image in the batch. + Each dictionary contains: + - "text" (str): The decoded text string. + - "score" (float): The average confidence score of the characters in the decoded text. + """ + requires_backends(self, ["torch"]) + + logits = predictions.last_hidden_state + batch_size = logits.shape[0] + + preds_prob, preds_idx = logits.max(dim=-1) + results = [] + for idx in range(batch_size): + selection = torch.ones(len(preds_idx[idx]), dtype=torch.bool, device=preds_idx.device) + + # remove_duplicate + selection[1:] = preds_idx[idx][1:] != preds_idx[idx][:-1] + # ignore blank token + selection &= preds_idx[idx] != 0 + + character_list = [] + for text_id in preds_idx[idx][selection]: + character_list.append(self.character_list[text_id]) + + results.append( + { + "text": "".join(character_list), + "score": preds_prob[idx][selection].mean().item(), + } + ) + + return results + + +class PPOCRV5ServerRecBlock(CLIPEncoderLayer): + def __init__(self, config): + super().__init__() + self.mlp = PPOCRV5ServerRecMLP( + config=config, + in_features=self.embed_dim, + hidden_features=int(self.embed_dim * config.mlp_ratio), + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + super().forward(hidden_states, attention_mask, **kwargs) + + +class PPOCRV5ServerRecAttention(Blip2Attention): + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, # Not used but kept for signature matching in downstream modules + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + super().forward(hidden_states, **kwargs) + + +class PPOCRV5ServerRecConvLayer(ResNetConvLayer): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: tuple[int, int] = (3, 3), + stride: int = 1, + activation: str = "silu", + ): + super().__init__() + self.convolution = nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=(kernel_size[0] // 2, kernel_size[1] // 2), + bias=False, + ) + + +class PPOCRV5ServerRecPreTrainedModel(PaddleOCRVLPreTrainedModel): + _no_split_modules = ["PPOCRV5ServerRecBlock"] + main_input_name = "pixel_values" + input_modalities = ("image",) + _can_record_outputs = { + "hidden_states": PPOCRV5ServerRecBlock, + } + + def _init_weights(self, module): + raise AttributeError("No override needed here") + + +class PPOCRV5ServerRecHead(nn.Module): + def __init__(self, config): + super().__init__() + + self.encoder = PPOCRV5ServerRecEncoderWithSVTR(config) + self.head = nn.Linear(config.hidden_size, config.head_out_channels) + + def forward(self, hidden_states: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + outputs = self.encoder(hidden_states, **kwargs) + hidden_states = self.head(outputs.last_hidden_state) + hidden_states = F.softmax(hidden_states, dim=2, dtype=torch.float32).to(hidden_states.dtype) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=outputs.hidden_states) + + +class PPOCRV5ServerRecMLP(FocalNetMlp): + pass + + +class PPOCRV5ServerRecEncoderWithSVTR(PPOCRV5ServerRecPreTrainedModel): + """ + SVTR: Scene Text Recognition with a Single Visual Model + https://www.paddleocr.ai/v2.10.0/en/algorithm/text_recognition/algorithm_rec_svtr.html + """ + + def __init__( + self, + config, + ): + super().__init__(config) + in_channels = config.backbone_config.stage_out_channels[-1] + hidden_size = config.hidden_size + + self.conv_block = nn.ModuleList( + [ + PPOCRV5ServerRecConvLayer( + in_channels=in_channels, + out_channels=in_channels // 8, + kernel_size=config.conv_kernel_size, + ), + PPOCRV5ServerRecConvLayer(in_channels=in_channels // 8, out_channels=hidden_size, kernel_size=(1, 1)), + PPOCRV5ServerRecConvLayer(in_channels=hidden_size, out_channels=in_channels, kernel_size=(1, 1)), + PPOCRV5ServerRecConvLayer( + in_channels=2 * in_channels, + out_channels=in_channels // 8, + kernel_size=config.conv_kernel_size, + ), + PPOCRV5ServerRecConvLayer(in_channels=in_channels // 8, out_channels=hidden_size, kernel_size=(1, 1)), + ] + ) + self.svtr_block = nn.ModuleList() + for _ in range(config.depth): + self.svtr_block.append(PPOCRV5ServerRecBlock(config=config)) + + self.norm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps) + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward(self, hidden_states: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + residual = hidden_states + + hidden_states = self.conv_block[0](hidden_states) + hidden_states = self.conv_block[1](hidden_states) + + batch_size, channels, height, width = hidden_states.shape + hidden_states = hidden_states.flatten(2).transpose(1, 2) + for block in self.svtr_block: + hidden_states = block(hidden_states) + + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.view(batch_size, height, width, channels).permute(0, 3, 1, 2) + hidden_states = self.conv_block[2](hidden_states) + hidden_states = self.conv_block[3](torch.cat((residual, hidden_states), dim=1)) + hidden_states = self.conv_block[4](hidden_states) + hidden_states = hidden_states.squeeze(2).transpose(1, 2) + + return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states) + + +@auto_docstring(custom_intro="PPOCRV5ServerRec model, consisting of Backbone and Head networks.") +class PPOCRV5ServerRecModel(PPOCRV5ServerRecPreTrainedModel): + def __init__(self, config: PPOCRV5ServerRecConfig): + super().__init__(config) + self.backbone = load_backbone(config) + + self.gradient_checkpointing = False + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention: + outputs = self.backbone(pixel_values, **kwargs) + hidden_state = outputs.feature_maps[-1] + hidden_state = F.avg_pool2d(hidden_state, (3, 2)) + + return BaseModelOutputWithNoAttention( + last_hidden_state=hidden_state, + hidden_states=outputs.hidden_states, + ) + + +@dataclass +@auto_docstring +class PPOCRV5ServerRecForTextRecognitionOutput(BaseModelOutputWithNoAttention): + r""" + head_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Hidden-states of the PPOCRV5ServerRecHead at the output of each layer plus the optional initial embedding outputs. + """ + + head_hidden_states: torch.FloatTensor | None = None + + +@auto_docstring(custom_intro="PPOCRV5ServerRec model for text recognition tasks.") +class PPOCRV5ServerRecForTextRecognition(PPOCRV5ServerRecPreTrainedModel): + _keys_to_ignore_on_load_missing = ["num_batches_tracked"] + + def __init__(self, config: PPOCRV5ServerRecConfig): + super().__init__(config) + self.model = PPOCRV5ServerRecModel(config) + self.head = PPOCRV5ServerRecHead(config) + + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention: + outputs = self.model(pixel_values, **kwargs) + head_outputs = self.head(outputs.last_hidden_state, **kwargs) + + return PPOCRV5ServerRecForTextRecognitionOutput( + last_hidden_state=head_outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + head_hidden_states=head_outputs.hidden_states, + ) + + +__all__ = [ + "PPOCRV5ServerRecForTextRecognition", + "PPOCRV5ServerRecImageProcessor", + "PPOCRV5ServerRecConfig", + "PPOCRV5ServerRecModel", + "PPOCRV5ServerRecPreTrainedModel", + "PPOCRV5ServerRecEncoderWithSVTR", +] diff --git a/third_party/transformers/src/transformers/models/qwen3_vl_moe/__init__.py b/third_party/transformers/src/transformers/models/qwen3_vl_moe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4000cb272723d9920136a6c78465e8413a8b4d1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen3_vl_moe/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_qwen3_vl_moe import * + from .modeling_qwen3_vl_moe import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/third_party/transformers/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..dee0a4e04420bec4738b81a422c46849f3f0eafb --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -0,0 +1,188 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen3_vl_moe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeTextConfig(PreTrainedConfig): + r""" + decoder_sparse_step (`int`, *optional*, defaults to 1): + The frequency of the MoE layer. + mlp_only_layers (`List[int]`, *optional*, defaults to `[]`): + Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock + The list contains layer index, from 0 to num_layers-1 if we have num_layers layers + If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. + + ```python + >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig + + >>> # Initializing a Qwen3VLMoe style configuration + >>> configuration = Qwen3VLMoeConfig() + + >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration + >>> model = Qwen3VLMoeForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen3_vl_moe_text" + keys_to_ignore_at_inference = ["past_key_values"] + + attribute_map = { + "num_experts": "num_local_experts", + } + # Default tensor parallel plan for base model `Qwen3VLMoe` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 151936 + hidden_size: int = 2048 + + intermediate_size: int = 5632 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + hidden_act: str = "silu" + max_position_embeddings: int = 128000 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + decoder_sparse_step: int = 1 + moe_intermediate_size: int = 1408 + num_experts_per_tok: int = 4 + num_experts: int = 60 + router_aux_loss_coef: float = 0.001 + mlp_only_layers: list[int] | None = None + pad_token_id: int | None = None + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + base_config_key = "text_config" + default_theta = 500000.0 + ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} + head_dim: int | None = None + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads + self.sliding_window = None + self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeVisionConfig(PreTrainedConfig): + r""" + out_hidden_size (`int`, *optional*, defaults to 3584): + The output hidden size of the vision model. + num_position_embeddings (`int`, *optional*, defaults to 2304): + The maximum sequence length that this model might ever be used with + deepstack_visual_indexes (`list[int]`, *optional*, defaults to `[8, 16, 24]`): + Indexed of layers for deepstack embeddings. + """ + + model_type = "qwen3_vl_moe" + base_config_key = "vision_config" + + depth: int = 27 + hidden_size: int = 1152 + hidden_act: str = "gelu_pytorch_tanh" + intermediate_size: int = 4304 + num_heads: int = 16 + in_channels: int = 3 + patch_size: int | list[int] | tuple[int, int] = 16 + spatial_merge_size: int = 2 + temporal_patch_size: int | list[int] | tuple[int, int] = 2 + out_hidden_size: int = 3584 + num_position_embeddings: int = 2304 + deepstack_visual_indexes: list[int] | tuple[int, ...] = (8, 16, 24) + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig + + >>> # Initializing a Qwen3-VL-MOE style configuration + >>> configuration = Qwen3VLMoeConfig() + + >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration + >>> model = Qwen3VLMoeForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen3_vl_moe" + sub_configs = {"vision_config": Qwen3VLMoeVisionConfig, "text_config": Qwen3VLMoeTextConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**self.vision_config) + elif self.vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + if isinstance(self.text_config, dict): + self.text_config = self.sub_configs["text_config"](**self.text_config) + elif self.text_config is None: + self.text_config = self.sub_configs["text_config"]() + + super().__post_init__(**kwargs) + + +__all__ = ["Qwen3VLMoeConfig", "Qwen3VLMoeTextConfig"] diff --git a/third_party/transformers/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/third_party/transformers/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..6d4c68c1a7524ae9c693b71ff98d5a83809288e1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -0,0 +1,1972 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen3_vl_moe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput, MoeModelOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_qwen3_vl_moe import Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class Qwen3VLMoeTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Qwen3VLMoeTextRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +@use_experts_implementation +class Qwen3VLMoeTextExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) + self.act_fn = ACT2FN[config.hidden_act] + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states + + +class Qwen3VLMoeTextTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) + router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1) + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + router_top_value = router_top_value.to(router_logits.dtype) + router_scores = router_top_value + return router_logits, router_scores, router_indices + + +class Qwen3VLMoeTextSparseMoeBlock(nn.Module): + def __init__(self, config: Qwen3VLMoeTextConfig): + super().__init__() + self.experts = Qwen3VLMoeTextExperts(config) + self.gate = Qwen3VLMoeTextTopKRouter(config) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_reshaped = hidden_states.view(-1, hidden_dim) + _, routing_weights, selected_experts = self.gate(hidden_states_reshaped) + final_hidden_states = self.experts(hidden_states_reshaped, selected_experts, routing_weights) + return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +@use_kernelized_func(apply_rotary_pos_emb) +class Qwen3VLMoeTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3VLMoeTextConfig, layer_idx: int): + super().__init__() + self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3VLMoeTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # unlike olmo, only on the head dim! + self.k_norm = Qwen3VLMoeTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # thus post q_norm does not need reshape + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3VLMoeTextMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Qwen3VLMoeTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3VLMoeTextConfig, layer_idx: int): + super().__init__() + self.self_attn = Qwen3VLMoeTextAttention(config, layer_idx) + if (layer_idx not in config.mlp_only_layers) and ( + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = Qwen3VLMoeTextSparseMoeBlock(config) + else: + self.mlp = Qwen3VLMoeTextMLP(config, intermediate_size=config.intermediate_size) + self.input_layernorm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hidden_size = config.hidden_size + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class Qwen3VLMoePreTrainedModel(PreTrainedModel): + config: Qwen3VLMoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, index=0), + "hidden_states": Qwen3VLMoeTextDecoderLayer, + "attentions": Qwen3VLMoeTextAttention, + } + input_modalities = ("text", "image", "video") + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if hasattr(self.config, "initializer_range"): + std = self.config.initializer_range + else: + std = getattr(self.config.get_text_config(), "initializer_range", 0.02) + if isinstance(module, Qwen3VLMoeTextExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, Qwen3VLMoeTextTopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + elif isinstance(module, Qwen3VLMoeVisionRotaryEmbedding): + inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim)) + init.copy_(module.inv_freq, inv_freq) + + +class Qwen3VLMoeVisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.theta = theta + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +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 Qwen3VLMoeVisionAttention(nn.Module): + def __init__(self, config: Qwen3VLMoeVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 # needed for eager attention + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + if is_flash_attention_requested(self.config): + # Flash Attention: Use cu_seqlens for variable length attention + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen3VLMoeVisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + +class Qwen3VLMoeVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3VLMoeVisionAttention(config=config) + self.mlp = Qwen3VLMoeVisionMLP(config=config) + + @auto_docstring + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + r""" + cu_seqlens (`torch.Tensor`): + Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels. + rotary_pos_emb (`torch.Tensor`, *optional*): + Precomputed rotary positional embeddings applied to the vision attention query/key states. + """ + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + +class Qwen3VLMoeVisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class Qwen3VLMoeVisionPatchMerger(nn.Module): + def __init__(self, config: Qwen3VLMoeVisionConfig, use_postshuffle_norm=False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + +class Qwen3VLMoeVisionModel(Qwen3VLMoePreTrainedModel): + config: Qwen3VLMoeVisionConfig + input_modalities = ("image", "video") + _no_split_modules = ["Qwen3VLMoeVisionBlock"] + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0), + "hidden_states": Qwen3VLMoeVisionBlock, + "attentions": Qwen3VLMoeVisionAttention, + } + + def __init__(self, config, *inputs, **kwargs) -> None: + super().__init__(config, *inputs, **kwargs) + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen3VLMoeVisionPatchEmbed( + config=config, + ) + + self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen3VLMoeVisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([Qwen3VLMoeVisionBlock(config) for _ in range(config.depth)]) + self.merger = Qwen3VLMoeVisionPatchMerger( + config=config, + use_postshuffle_norm=False, + ) + + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nn.ModuleList( + [ + Qwen3VLMoeVisionPatchMerger( + config=config, + use_postshuffle_norm=True, + ) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + self.gradient_checkpointing = False + + self.post_init() + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + merge_size = self.spatial_merge_size + grid_thw_list = grid_thw.tolist() + + max_hw = max(max(h, w) for _, h, w in grid_thw_list) + freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2) + device = freq_table.device + + total_tokens = sum(t * h * w for t, h, w in grid_thw_list) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + + offset = 0 + for num_frames, height, width in grid_thw_list: + merged_h, merged_w = height // merge_size, width // merge_size + + block_rows = torch.arange(merged_h, device=device) # block row indices + block_cols = torch.arange(merged_w, device=device) # block col indices + intra_row = torch.arange(merge_size, device=device) # intra-block row offsets + intra_col = torch.arange(merge_size, device=device) # intra-block col offsets + + # Compute full-resolution positions + row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] + col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] + + row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + + coords = torch.stack((row_idx, col_idx), dim=-1) + + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + + num_tokens = coords.shape[0] + pos_ids[offset : offset + num_tokens] = coords + offset += num_tokens + + embeddings = freq_table[pos_ids] # lookup rotary embeddings + embeddings = embeddings.flatten(1) + return embeddings + + def fast_pos_embed_interpolate(self, grid_thw): + grid_thw_list = grid_thw.tolist() + grid_ts = [row[0] for row in grid_thw_list] + grid_hs = [row[1] for row in grid_thw_list] + grid_ws = [row[2] for row in grid_thw_list] + device = self.pos_embed.weight.device + + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + + for t, h, w in grid_thw_list: + h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w) + + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor + + base_h = h_idxs_floor * self.num_grid_per_side + base_h_ceil = h_idxs_ceil * self.num_grid_per_side + + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] + + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device) + weight_tensor = torch.tensor(weight_list, dtype=self.pos_embed.weight.dtype, device=device) + pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] + + patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)]) + + patch_pos_embeds_permute = [] + merge_size = self.config.spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) + pos_embed = ( + pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + patch_pos_embeds_permute.append(pos_embed) + patch_pos_embeds = torch.cat(patch_pos_embeds_permute) + return patch_pos_embeds + + @merge_with_config_defaults + @capture_outputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + """ + Args: + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + deepstack_feature_lists = [] + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + **kwargs, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)]( + hidden_states + ) + deepstack_feature_lists.append(deepstack_feature) + + merged_hidden_states = self.merger(hidden_states) + + return BaseModelOutputWithDeepstackFeatures( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + deepstack_features=deepstack_feature_lists, + ) + + +class Qwen3VLMoeTextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Qwen3VLMoeTextConfig, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + self.mrope_section = config.rope_parameters.get("mrope_section", [24, 20, 20]) + + @staticmethod + def compute_default_rope_parameters( + config: Qwen3VLMoeTextConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + # In contrast to other models, Qwen3VLMoe has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def apply_interleaved_mrope(self, freqs, mrope_section): + """Apply interleaved MRoPE to 3D rotary embeddings. + Reorganizes frequency layout from chunked [TTT...HHH...WWW] to + interleaved [THWTHWTHW...TT], preserving frequency continuity. + args: + x: (3, bs, seq_len, head_dim // 2) + mrope_section: (3,) + returns: + x_t: (bs, seq_len, head_dim // 2) + """ + freqs_t = freqs[0] # just overwrite the first dimension T + for dim, offset in enumerate((1, 2), start=1): # H, W + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + +@auto_docstring( + custom_intro=( + "Text part of Qwen3VLMoe, " + "not a pure text-only model, as DeepStack integrates visual features into the early hidden states." + ) +) +class Qwen3VLMoeTextModel(Qwen3VLMoePreTrainedModel): + config: Qwen3VLMoeTextConfig + input_modalities = ("text",) + _no_split_modules = ["Qwen3VLMoeTextDecoderLayer"] + + def __init__(self, config: Qwen3VLMoeTextConfig): + 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( + [Qwen3VLMoeTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3VLMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLMoeTextRotaryEmbedding(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, + # args for deepstack + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | MoeModelOutputWithPast: + r""" + visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*): + The mask of the visual positions. + deepstack_visual_embeds (`list[torch.Tensor]`, *optional*): + The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). + The feature is extracted from the different visual encoder layers, and fed to the decoder + hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334). + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # the hard coded `4` is for text, temporal, height and width. + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(4, 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 + + attention_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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + for layer_idx, decoder_layer in enumerate(self.layers): + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = layer_outputs + + # add visual features to the hidden states of first several layers + if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Qwen3VLTextModel + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + def _deepstack_process( + self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor + ): + visual_pos_masks = visual_pos_masks.to(hidden_states.device) + visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) + hidden_states = hidden_states.clone() + local_this = hidden_states[visual_pos_masks, :] + visual_embeds + hidden_states[visual_pos_masks, :] = local_this + return hidden_states + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Llava outputs, with hidden states and attentions. + """ +) +class Qwen3VLMoeModelOutputWithPast(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 + router_logits: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Qwen3VLMoe causal language model (or autoregressive) outputs. + """ +) +class Qwen3VLMoeCausalLMOutputWithPast(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 + router_logits: tuple[torch.FloatTensor] | None = None + aux_loss: torch.FloatTensor | None = None + + +@auto_docstring +class Qwen3VLMoeModel(Qwen3VLMoePreTrainedModel): + base_model_prefix = "model" + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLMoeConfig + _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen3VLMoeVisionModel._from_config(config.vision_config) + self.language_model = Qwen3VLMoeTextModel._from_config(config.text_config) + self.rope_deltas = None # cache rope_deltas here + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + + def get_rope_index( + self, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Difference from Qwen2VL/Qwen2.5VL's get_rope_index: + - Since Qwen3.5 use timestamps to seperate videos, like , the video_grid_thw should also be split too. + + 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)`) + """ + + # Separate video grid thw into multiple grids because timestamps are used to seperate videos. + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + video_grid_thw[:, 0] = 1 + 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 + + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + # Same implementation as for images + return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs) + + @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 | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + pixel_values = pixel_values.type(self.visual.dtype) + vision_output: BaseModelOutputWithDeepstackFeatures = self.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs + ) + image_embeds = vision_output.pooler_output + split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() + image_embeds = torch.split(image_embeds, split_sizes) + vision_output.pooler_output = image_embeds + + return vision_output + + def get_placeholder_mask( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_features: torch.FloatTensor | None = None, + video_features: torch.FloatTensor | None = None, + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + special_video_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_video_mask = special_video_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + special_video_mask = input_ids == self.config.video_token_id + + n_image_tokens = special_image_mask.sum() + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if image_features is not None: + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}", + ) + + n_video_tokens = special_video_mask.sum() + special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if video_features is not None: + torch_compilable_check( + inputs_embeds[special_video_mask].numel() == video_features.numel(), + f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}", + ) + return special_image_mask, special_video_mask + + def compute_3d_position_ids( + self, + input_ids: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + image_grid_thw: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + ) -> torch.Tensor | None: + past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() + has_multimodal = image_grid_thw is not None or video_grid_thw is not None + if has_multimodal and mm_token_type_ids is None and input_ids is not None: + raise ValueError( + "Multimodal data was passed (via `image_grid_thw` or `video_grid_thw`) but `mm_token_type_ids` is " + "missing. Please pass `mm_token_type_ids` to the model so that multimodal RoPE (M-RoPE) can be " + "computed correctly. `mm_token_type_ids` is returned by the processor alongside `input_ids`." + ) + can_compute_mrope = input_ids is not None and mm_token_type_ids is not None and has_multimodal + + if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, + ) + self.rope_deltas = rope_deltas + # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental + # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids + # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape + # mismatches from stale rope_deltas (e.g., training forward pass after generation). + elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None): + batch_size, seq_length, _ = inputs_embeds.shape + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids = position_ids.masked_fill(attention_mask == 0, 0) + position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device) + else: + position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) + position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) + delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) + position_ids = position_ids + delta.to(device=inputs_embeds.device) + else: + # Can't build correct 3D positions. Let the model infer it + position_ids = None + return position_ids + + @auto_docstring + @can_return_tuple + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen3VLMoeModelOutputWithPast: + r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + image_mask = None + video_mask = None + + if pixel_values is not None: + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + deepstack_image_embeds = image_outputs.deepstack_features + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_outputs: BaseModelOutputWithDeepstackFeatures = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ) + video_embeds = video_outputs.pooler_output + deepstack_video_embeds = video_outputs.deepstack_features + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + visual_pos_masks = None + deepstack_visual_embeds = None + if image_mask is not None and video_mask is not None: + # aggregate visual_pos_masks and deepstack_visual_embeds + image_mask = image_mask[..., 0] + video_mask = video_mask[..., 0] + visual_pos_masks = image_mask | video_mask + deepstack_visual_embeds = [] + image_mask_joint = image_mask[visual_pos_masks] + video_mask_joint = video_mask[visual_pos_masks] + for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds): + embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device) + embed_joint[image_mask_joint, :] = img_embed + embed_joint[video_mask_joint, :] = vid_embed + deepstack_visual_embeds.append(embed_joint) + elif image_mask is not None: + image_mask = image_mask[..., 0] + visual_pos_masks = image_mask + deepstack_visual_embeds = deepstack_image_embeds + elif video_mask is not None: + video_mask = video_mask[..., 0] + visual_pos_masks = video_mask + deepstack_visual_embeds = deepstack_video_embeds + + if position_ids is None: + position_ids = self.compute_3d_position_ids( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + ) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + **kwargs, + ) + + return Qwen3VLMoeModelOutputWithPast( + **outputs, + rope_deltas=self.rope_deltas, + ) + + +def load_balancing_loss_func( + gate_logits: torch.Tensor | tuple[torch.Tensor] | None, + num_experts: int | None = None, + top_k=2, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor | int: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +class Qwen3VLMoeForConditionalGeneration(Qwen3VLMoePreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLMoeConfig + + def __init__(self, config): + super().__init__(config) + self.model = Qwen3VLMoeModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) + + @can_return_tuple + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen3VLMoeCausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + + Example: + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration + + >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + }, + {"type": "text", "text": "Describe this image in short."}, + ], + } + ] + + >>> # Preparation for inference + >>> inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt" + ) + >>> inputs = inputs.to(model.device) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=128) + >>> generated_ids_trimmed = [ + out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) + ] + >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background." + ```""" + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + hidden_states = outputs[0] + + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size) + + aux_loss = None + if kwargs.get("output_router_logits", False): + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.config.text_config.num_experts, + self.config.text_config.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.config.text_config.router_aux_loss_coef * aux_loss.to( + loss.device + ) # make sure to reside in the same device + + return Qwen3VLMoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=outputs.rope_deltas, + router_logits=outputs.router_logits, + ) + + 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 -- Qwen3VLMoe use timestamps and remove second_per_grid_ts + # Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"] + + 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) + ) + + # video_nums: (batch_size,) + # since video_nums is the number of videos in the input dependent on the input_ids(vision_start), + # but Qwen3VLMoe append vision_start to each frame of each video, so we need to recover the real video_nums according to video_grid_thw + if video_grid_thw is not None: + cumulative_frame_counts = torch.cumsum(video_grid_thw[:, 0], dim=0) + cumulative_token_video_counts = torch.cumsum(video_nums, dim=0) + # Find video boundaries in cumulative_frame_counts + video_boundary_indices = torch.searchsorted(cumulative_frame_counts, cumulative_token_video_counts) + # example: video_boundary_indices = [3, 5] means video_nums = [4, 2] + video_nums = torch.diff(torch.cat([-video_boundary_indices.new_ones(1), video_boundary_indices])) + + 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 + ) + 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__ = [ + "Qwen3VLMoeVisionModel", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3VLMoeModel", + "Qwen3VLMoePreTrainedModel", + "Qwen3VLMoeTextModel", +] diff --git a/third_party/transformers/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/third_party/transformers/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..fa840e0685fe95fd5ccbb5a6b7b4d92799cd0d4d --- /dev/null +++ b/third_party/transformers/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -0,0 +1,459 @@ +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Qwen3-VL-MOE model.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import MoeModelOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.output_capturing import OutputRecorder +from ..qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig +from ..qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeDecoderLayer, + Qwen3MoeExperts, + Qwen3MoePreTrainedModel, + Qwen3MoeRMSNorm, + Qwen3MoeSparseMoeBlock, + load_balancing_loss_func, +) +from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig +from ..qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLCausalLMOutputWithPast, + Qwen3VLForConditionalGeneration, + Qwen3VLModelOutputWithPast, + Qwen3VLTextAttention, + Qwen3VLTextModel, + Qwen3VLVisionAttention, + Qwen3VLVisionBlock, + Qwen3VLVisionModel, + Qwen3VLVisionRotaryEmbedding, +) + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeTextConfig(Qwen3MoeConfig): + r""" + decoder_sparse_step (`int`, *optional*, defaults to 1): + The frequency of the MoE layer. + mlp_only_layers (`List[int]`, *optional*, defaults to `[]`): + Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock + The list contains layer index, from 0 to num_layers-1 if we have num_layers layers + If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. + + ```python + >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig + + >>> # Initializing a Qwen3VLMoe style configuration + >>> configuration = Qwen3VLMoeConfig() + + >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration + >>> model = Qwen3VLMoeForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen3_vl_moe_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 500000.0 + # Default tensor parallel plan for base model `Qwen3VLMoe` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} + + intermediate_size: int = 5632 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + max_position_embeddings: int = 128000 + moe_intermediate_size: int = 1408 + num_experts_per_tok: int = 4 + num_experts: int = 60 + head_dim: int | None = None + tie_word_embeddings: bool = True + + norm_topk_prob = AttributeError() + output_router_logits = AttributeError() + use_sliding_window = AttributeError() + sliding_window = AttributeError() + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads + super().__post_init__(**kwargs) + self.sliding_window = None + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeVisionConfig(Qwen3VLVisionConfig): + pass + + +@auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") +@strict +class Qwen3VLMoeConfig(Qwen3VLConfig): + r""" + Example: + + ```python + >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig + + >>> # Initializing a Qwen3-VL-MOE style configuration + >>> configuration = Qwen3VLMoeConfig() + + >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration + >>> model = Qwen3VLMoeForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + pass + + +class Qwen3VLMoeTextRMSNorm(Qwen3MoeRMSNorm): + pass + + +class Qwen3VLMoeTextExperts(Qwen3MoeExperts): + pass + + +class Qwen3VLMoeTextTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) + router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1) + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + router_top_value = router_top_value.to(router_logits.dtype) + router_scores = router_top_value + return router_logits, router_scores, router_indices + + +class Qwen3VLMoeTextSparseMoeBlock(Qwen3MoeSparseMoeBlock): + pass + + +class Qwen3VLMoeTextAttention(Qwen3VLTextAttention): + pass + + +class Qwen3VLMoeTextDecoderLayer(Qwen3MoeDecoderLayer): + pass + + +class Qwen3VLMoePreTrainedModel(Qwen3MoePreTrainedModel): + config: Qwen3VLMoeConfig + input_modalities = ("text", "image", "video") + _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + PreTrainedModel._init_weights(self, module) + if hasattr(self.config, "initializer_range"): + std = self.config.initializer_range + else: + std = getattr(self.config.get_text_config(), "initializer_range", 0.02) + if isinstance(module, Qwen3VLMoeTextExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, Qwen3VLMoeTextTopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + elif isinstance(module, Qwen3VLMoeVisionRotaryEmbedding): + inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim)) + init.copy_(module.inv_freq, inv_freq) + + +class Qwen3VLMoeVisionRotaryEmbedding(Qwen3VLVisionRotaryEmbedding): + pass + + +class Qwen3VLMoeVisionAttention(Qwen3VLVisionAttention): + pass + + +class Qwen3VLMoeVisionBlock(Qwen3VLVisionBlock): + pass + + +class Qwen3VLMoeVisionModel(Qwen3VLVisionModel): + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0), + "hidden_states": Qwen3VLMoeVisionBlock, + "attentions": Qwen3VLMoeVisionAttention, + } + + +class Qwen3VLMoeTextModel(Qwen3VLTextModel): + 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, + # args for deepstack + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | MoeModelOutputWithPast: + r""" + visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*): + The mask of the visual positions. + deepstack_visual_embeds (`list[torch.Tensor]`, *optional*): + The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). + The feature is extracted from the different visual encoder layers, and fed to the decoder + hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334). + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # the hard coded `4` is for text, temporal, height and width. + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(4, 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 + + attention_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 + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + for layer_idx, decoder_layer in enumerate(self.layers): + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = layer_outputs + + # add visual features to the hidden states of first several layers + if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Qwen3VLTextModel + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class Qwen3VLMoeModelOutputWithPast(Qwen3VLModelOutputWithPast): + router_logits: tuple[torch.FloatTensor] | None = None + + +class Qwen3VLMoeCausalLMOutputWithPast(Qwen3VLCausalLMOutputWithPast): + router_logits: tuple[torch.FloatTensor] | None = None + aux_loss: torch.FloatTensor | None = None + + +class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration): + @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: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Qwen3VLMoeCausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + + Example: + ```python + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration + + >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + }, + {"type": "text", "text": "Describe this image in short."}, + ], + } + ] + + >>> # Preparation for inference + >>> inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt" + ) + >>> inputs = inputs.to(model.device) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=128) + >>> generated_ids_trimmed = [ + out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) + ] + >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background." + ```""" + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + hidden_states = outputs[0] + + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size) + + aux_loss = None + if kwargs.get("output_router_logits", False): + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.config.text_config.num_experts, + self.config.text_config.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.config.text_config.router_aux_loss_coef * aux_loss.to( + loss.device + ) # make sure to reside in the same device + + return Qwen3VLMoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=outputs.rope_deltas, + router_logits=outputs.router_logits, + ) + + +__all__ = [ + "Qwen3VLMoeConfig", + "Qwen3VLMoeTextConfig", + "Qwen3VLMoeVisionModel", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3VLMoeModel", # noqa + "Qwen3VLMoePreTrainedModel", + "Qwen3VLMoeTextModel", +] diff --git a/third_party/transformers/src/transformers/models/recurrent_gemma/__init__.py b/third_party/transformers/src/transformers/models/recurrent_gemma/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ab9335fc4beff8b178336ccb546bdecb4cd45171 --- /dev/null +++ b/third_party/transformers/src/transformers/models/recurrent_gemma/__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_recurrent_gemma import * + from .modeling_recurrent_gemma import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/third_party/transformers/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..56610ebdee47922c2db3cc499d05342dfa1e4530 --- /dev/null +++ b/third_party/transformers/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py @@ -0,0 +1,103 @@ +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""RecurrentGemma model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="google/recurrentgemma-2b") +@strict +class RecurrentGemmaConfig(PreTrainedConfig): + r""" + lru_width (`int` or `None`, *optional*): + Dimension of the hidden representations of the RG-LRU. If `None` + this will be set to `hidden_size`. + Whether to scale the output of the embeddings by `sqrt(hidden_size)`. + attention_window_size (`int`, *optional*, defaults to 2048): + The size of the attention window used in the attention block. + conv1d_width (`int`, *optional*, defaults to 4): + The kernel size of conv1d layers used in the recurrent blocks. + logits_soft_cap (`float`, *optional*, defaults to 30.0): + The value at which the logits should be soft-capped to after the transformer and LM-head computation in the Causal LM architecture. + block_types (`list[str]`, *optional*, defaults to `('recurrent', 'recurrent', 'attention')`): + List of aleternating blocks that will be repeated to initialize the `temporal_block` layer. + w_init_variance_scale (`float`, *optional*, defaults to 0.01): + weight initialization variance. + + ```python + >>> from transformers import RecurrentGemmaModel, RecurrentGemmaConfig + + >>> # Initializing a RecurrentGemma recurrentgemma-2b style configuration + >>> configuration = RecurrentGemmaConfig() + + >>> # Initializing a model from the recurrentgemma-2b style configuration + >>> model = RecurrentGemmaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "recurrent_gemma" + attribute_map = {"sliding_window": "attention_window_size"} + + num_hidden_layers: int = 26 + vocab_size: int = 256000 + hidden_size: int = 2560 + intermediate_size: int = 3 * 2560 + num_attention_heads: int = 10 + lru_width: int | None = None + attention_window_size: int = 2048 + conv1d_width: int = 4 + logits_soft_cap: float = 30.0 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + bos_token_id: int | None = 2 + hidden_activation: str = "gelu_pytorch_tanh" + rope_parameters: RopeParameters | dict | None = None + block_types: list[str] | tuple[str, ...] | None = ("recurrent", "recurrent", "attention") + attention_dropout: float | int = 0.0 + num_key_value_heads: int | None = None + attention_bias: bool = False + w_init_variance_scale: float = 0.01 + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + self.lru_width = self.lru_width if self.lru_width is not None else self.hidden_size + self.block_types = list(self.block_types) + self.head_dim = self.hidden_size // self.num_attention_heads + self.num_key_value_heads = ( + self.num_key_value_heads if self.num_key_value_heads is not None else self.num_attention_heads + ) + self.final_w_init_variance_scale = 2.0 / self.num_hidden_layers + kwargs.setdefault("partial_rotary_factor", 0.5) # assign default for BC + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.num_key_value_heads > self.num_attention_heads: + raise ValueError("The number of `num_key_value_heads` must be smaller than `num_attention_heads`") + + @property + def layers_block_type(self): + return (self.block_types * 100)[: self.num_hidden_layers] + + +__all__ = ["RecurrentGemmaConfig"] diff --git a/third_party/transformers/src/transformers/models/recurrent_gemma/convert_recurrent_gemma_to_hf.py b/third_party/transformers/src/transformers/models/recurrent_gemma/convert_recurrent_gemma_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..084fa9b5da79c6adb761445981f9d3fe0540291c --- /dev/null +++ b/third_party/transformers/src/transformers/models/recurrent_gemma/convert_recurrent_gemma_to_hf.py @@ -0,0 +1,214 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os +import warnings + +import torch + +from transformers import GemmaTokenizer, RecurrentGemmaConfig, RecurrentGemmaForCausalLM + + +try: + from transformers import GemmaTokenizerFast +except ImportError as e: + warnings.warn(e) + warnings.warn( + "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion" + ) + GemmaTokenizerFast = None + +import regex as re + + +""" +Sample usage: + +``` +python src/transformers/models/gemma/convert_gemma_weights_to_hf.py \ + --input_dir /path/to/downloaded/gemma/weights --model_size 7B --output_dir /output/path +``` + +Thereafter, models can be loaded via: + +```py +from transformers import GemmaForCausalLM, GemmaTokenizerFast + +model = GemmaForCausalLM.from_pretrained("/output/path") +tokenizer = GemmaTokenizerFast.from_pretrained("/output/path") +``` + +Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions +come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM). +""" + +gemma_2b_config = RecurrentGemmaConfig( + num_attention_heads=10, + num_key_value_heads=1, + hidden_size=2560, + intermediate_size=15360, + vocab_size=256000, + num_hidden_layers=26, +) + +gemma_7b_config = RecurrentGemmaConfig() + +CONFIG_MAPPING = {"2B": gemma_2b_config, "7B": gemma_7b_config} +LAYER_NAME_MAPPING = {"embedder.weight": "model.embed_tokens.weight"} + + +def write_model(save_path, input_base_path, config, push_to_hub=False, dtype=torch.float32): + print(f"Fetching all parameters from the checkpoint at '{input_base_path}'") + model_state_dict = torch.load(input_base_path, map_location="cpu", weights_only=True) + + REPLACEMENT = { + "blocks.": "layers.", + ".ffw_down.b": ".down_proj.b", + ".ffw_down.w": ".down_proj.w", + ".ffw_up.b": ".up_proj.bias", + ".ffw_up.w": ".up_proj.weight", + "recurrent_block": "temporal_block", + "attention_block": "temporal_block", + "temporal_block.proj_final": "temporal_block.out_proj", + "norm.scale": "norm.weight", + ".proj_k": ".k_proj", + ".proj_q": ".q_proj", + ".proj_v": ".v_proj", + ".proj_final": ".o_proj", + "embedder.input_embedding": "embed_tokens.weight", + "conv_1d.w": "conv_1d.weight", + "conv_1d.b": "conv_1d.bias", + "input_gate.w": "input_gate.weight", + "input_gate.b": "input_gate.bias", + "a_param": "recurrent_param", + "a_gate.b": "recurrent_gate.bias", + "a_gate.w": "recurrent_gate.weight", + } + + state_dict = {} + for k, v in model_state_dict.items(): + k = "model." + k + pattern = re.compile("|".join(map(re.escape, REPLACEMENT.keys()))) + key = pattern.sub(lambda match: REPLACEMENT[match.group(0)], k) + if "conv_1d.weight" in key: + v = v[:, None, :].transpose(0, 2) + if "up_proj.weight" in key: + state_dict[key.replace("up_proj", "gate_proj")] = v[0].T.contiguous() + v = v[1].T.contiguous() + if "up_proj.bias" in key: + state_dict[key.replace("up_proj", "gate_proj")] = v[0, 0, 0].clone() + v = v[1, 0, 0].contiguous() + if "recurrent_gate.bias" in key: + state_dict[key.replace("gate.", "gate_")] = v.contiguous().clone() + elif "recurrent_gate.weight" in key: + state_dict[key.replace("gate.", "gate_")] = v.contiguous().clone() + elif "input_gate.b" in key: + state_dict[key.replace("gate.", "gate_")] = v.contiguous().clone() + elif "input_gate.w" in key: + state_dict[key.replace("gate.", "gate_")] = v.contiguous().clone() + elif "embed_tokens" in key: + state_dict[key] = v[: config.vocab_size, :].contiguous().clone() + state_dict["lm_head.weight"] = v[: config.vocab_size, :].contiguous().clone() + else: + state_dict[key] = v.contiguous() + + torch.set_default_dtype(dtype) + + print("Loading the checkpoint in a Gemma model.") + with torch.device("meta"): + model = RecurrentGemmaForCausalLM(config) + model.load_state_dict(state_dict, assign=True, strict=True) + + model.config.dtype = torch.float32 + del model.config._name_or_path + print("Saving in the Transformers format.") + + if push_to_hub: + print(f"pushing the model to {save_path}") + else: + model.save_pretrained(save_path) + + +def write_tokenizer(input_tokenizer_path, save_path, push_to_hub=False): + # Initialize the tokenizer based on the `spm` model + tokenizer_class = GemmaTokenizer if GemmaTokenizerFast is None else GemmaTokenizerFast + print(f"Saving a {tokenizer_class.__name__} to {save_path}.") + tokenizer = tokenizer_class(input_tokenizer_path) + if push_to_hub: + tokenizer.push_to_hub(save_path) + else: + tokenizer.save_pretrained(save_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_checkpoint", + help="Absolute path to the target Gemma weights.", + default="/home/arthur/transformers_recurrentgemma/google/recurrent-gemma-2b-it/ToBeDeleted/2b-it.pt", + ) + parser.add_argument( + "--tokenizer_checkpoint", + help="Location of Gemma tokenizer model", + ) + parser.add_argument( + "--model_size", + default="2B", + choices=["2B", "7B", "tokenizer_only"], + help="'f' models correspond to the finetuned versions, and are specific to the Gemma2 official release. For more details on Gemma2, check out the original repo: https://huggingface.co/google/gemma-7b", + ) + parser.add_argument( + "--output_dir", + default="google/recurrent-gemma-2b-it-hf", + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--convert_tokenizer", + help="Whether or not to convert the tokenizer as well.", + action="store_true", + default=False, + ) + parser.add_argument( + "--push_to_hub", + help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally.", + action="store_true", + default=False, + ) + parser.add_argument( + "--dtype", + default="float32", + help="Target dtype of the converted model", + ) + args = parser.parse_args() + + if args.convert_tokenizer: + if args.tokenizer_checkpoint is None: + raise ValueError("Path to the tokenizer is required when passing --convert_tokenizer") + + spm_path = os.path.join(args.tokenizer_checkpoint) + write_tokenizer(spm_path, args.output_dir, args.push_to_hub) + + config = CONFIG_MAPPING[args.model_size] + dtype = getattr(torch, args.dtype) + write_model( + config=config, + input_base_path=args.input_checkpoint, + save_path=args.output_dir, + push_to_hub=args.push_to_hub, + dtype=dtype, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/third_party/transformers/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..6e9c072b88606d812ddd8789706c6044a4357e9c --- /dev/null +++ b/third_party/transformers/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -0,0 +1,790 @@ +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch RecurrentGemma model.""" + +import math +from collections.abc import Callable +from typing import Optional + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_sliding_window_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithNoAttention, CausalLMOutput +from ...modeling_rope_utils import dynamic_rope_update +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import maybe_autocast +from .configuration_recurrent_gemma import RecurrentGemmaConfig + + +logger = logging.get_logger(__name__) +_MAX_SQRT_GRADIENT = 1000.0 + + +# Copied from transformers.models.gemma.modeling_gemma.GemmaRMSNorm with Gemma->RecurrentGemma +class RecurrentGemmaRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst RecurrentGemma is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->RecurrentGemma +class RecurrentGemmaRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + # Ignore copy + def __init__(self, config: RecurrentGemmaConfig, device=None): + super().__init__() + + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + raise ValueError( + f"RecurrentGemmaRotaryEmbedding does not support RoPE types other than `default` but got {self.rope_type}" + ) + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + @staticmethod + # Ignore copy + def compute_default_rope_parameters( + config: RecurrentGemmaConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class RecurrentGemmaSdpaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: RecurrentGemmaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads + self.partial_rotary_factor = config.partial_rotary_factor + + self.q_proj = nn.Linear(self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_attention_heads * self.head_dim, self.hidden_size, bias=True) + self.rotary_emb = RecurrentGemmaRotaryEmbedding(config=config) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_attention_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = self.rotary_emb(value_states, position_ids) + + # Partial rotary embedding + query_rot, query_pass = torch.chunk(query_states, int(1 / self.partial_rotary_factor), dim=-1) + key_rot, key_pass = torch.chunk(key_states, int(1 / self.partial_rotary_factor), dim=-1) + query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin) + query_states = torch.cat((query_rot, query_pass), dim=-1) + key_states = torch.cat((key_rot, key_pass), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states.contiguous(), + key_states.contiguous(), + value_states.contiguous(), + attn_mask=attention_mask, # pretty much a must for sliding window backend! + dropout_p=self.attention_dropout if self.training else 0.0, + scale=self.head_dim**-0.5, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, self.hidden_size) + attn_output = self.o_proj(attn_output) + return attn_output + + +class SqrtBoundDerivative(torch.autograd.Function): + """Computes a square root with a gradient clipped at `_MAX_SQRT_GRADIENT`.""" + + @staticmethod + def forward(ctx, x: torch.Tensor) -> torch.Tensor: + """The forward pass, which is a normal `sqrt`.""" + ctx.save_for_backward(x) + return torch.sqrt(x) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + """The backward pass, which clips the `sqrt` gradient.""" + (x,) = ctx.saved_tensors + clipped_x_times_4 = torch.clip(4.0 * x, min=1 / (_MAX_SQRT_GRADIENT**2)) + return grad_output / torch.sqrt(clipped_x_times_4) + + +class RecurrentGemmaRglru(nn.Module): + """A Real-Gated Linear Recurrent Unit (RG-LRU) layer.""" + + def __init__(self, config): + super().__init__() + self.num_attention_heads = config.num_attention_heads + self.block_width = config.lru_width // self.num_attention_heads + + self.recurrent_param = nn.Parameter(torch.empty([config.lru_width])) + self.input_gate_weight = nn.Parameter( + torch.empty([self.num_attention_heads, self.block_width, self.block_width]) + ) + self.input_gate_bias = nn.Parameter(torch.empty([self.num_attention_heads, self.block_width])) + + self.recurrent_gate_weight = nn.Parameter( + torch.empty([self.num_attention_heads, self.block_width, self.block_width]) + ) + self.recurrent_gate_bias = nn.Parameter(torch.empty([self.num_attention_heads, self.block_width])) + self.recurrent_states = None + + def forward( + self, + activations: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, seq_len, lru_width = activations.shape + reset = position_ids[:, :, None] == 0 + + reshape_act = activations.reshape(batch_size * seq_len, self.num_attention_heads, self.block_width) + reshape_act = reshape_act.permute(1, 0, 2) + + res = torch.baddbmm(self.input_gate_bias[:, None, :], reshape_act, self.input_gate_weight) + input_gate = torch.sigmoid(res.transpose(0, 1).reshape(batch_size, seq_len, lru_width)) + + res = torch.baddbmm(self.recurrent_gate_bias[:, None, :], reshape_act, self.recurrent_gate_weight) + recurrent_gate = torch.sigmoid(res.transpose(0, 1).reshape(batch_size, seq_len, lru_width)) + + # Compute the parameter `A` of the recurrence. + log_recurrent_gate = -8.0 * recurrent_gate * nn.functional.softplus(self.recurrent_param) + recurrent_gate = torch.exp(log_recurrent_gate) + a_square = torch.exp(2 * log_recurrent_gate) + + # Gate the input. + gated_inputs = activations * input_gate + + # Apply gamma normalization to the input. We need to clip the derivatives of + # `sqrt` in order to prevent NaNs during training in bfloat16. TODO a bit annoying + multiplier = SqrtBoundDerivative.apply(1 - a_square) + multiplier = reset + ~reset * multiplier + normalized_x = gated_inputs * multiplier.type(activations.dtype) + + hidden_states, recurrent_states = self._rnn_scan( + hidden_states=normalized_x, + recurrent_gate=recurrent_gate, + reset=reset, + recurrent_states=self.recurrent_states, + ) + self.recurrent_states = recurrent_states + return hidden_states + + # TODO refactor + def _rnn_scan( + self, + hidden_states: torch.Tensor, + recurrent_gate: torch.Tensor, + reset: torch.Tensor, + recurrent_states: torch.Tensor | None, + acc_dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Runs the recurrence of a linear RNN. + + Args: + hidden_states: The input sequence. + recurrent_gate: The diagonal of the recurrence matrix `A`. + reset: Indicator of document boundaries, e.g. when to reset the hidden state + of the RNN. + recurrent_states: The initial hidden state. + acc_dtype: The data type for the accumulation. + + Returns: + The output of the linear recurrence. + """ + # Multiply `a` by the reset. + recurrent_gate = recurrent_gate * ~reset + + if hidden_states.shape[1] == 1: + # Using scan in sampling mode. + if recurrent_states is None: # same here, when decoding you always have cache + return hidden_states, hidden_states[:, 0].type(acc_dtype) + + else: + contextualized_states = recurrent_gate.type(acc_dtype) * recurrent_states[:, None].to( + recurrent_gate.device + ) + contextualized_states += hidden_states.type(acc_dtype) + return contextualized_states.type(hidden_states.dtype), contextualized_states[:, -1] + + else: + # Using scan in linear mode. + if recurrent_states is None: + recurrent_states = torch.zeros(hidden_states[:, 0].shape, dtype=acc_dtype, device=hidden_states.device) + + contextualized_states = torch.zeros_like(hidden_states) + for t in range(hidden_states.shape[1]): + recurrent_states = recurrent_gate[:, t].type(acc_dtype) * recurrent_states.to(recurrent_gate.device) + recurrent_states = recurrent_states + hidden_states[:, t].type(acc_dtype) + contextualized_states[:, t] = recurrent_states.type(hidden_states.dtype) + + return contextualized_states, recurrent_states + + +class RecurrentGemmaRecurrentBlock(nn.Module): + """Griffin and Hawk's recurrent block.""" + + def __init__(self, config: RecurrentGemmaConfig, layer_idx: int): + super().__init__() + self.lru_width = config.lru_width + self.hidden_size = config.hidden_size + self.linear_y = nn.Linear(in_features=config.hidden_size, out_features=config.lru_width) + self.linear_x = nn.Linear(in_features=config.hidden_size, out_features=config.lru_width) + self.linear_out = nn.Linear(in_features=config.lru_width, out_features=config.hidden_size) + self.conv1d_width = config.conv1d_width + self.conv_1d = nn.Conv1d( + config.lru_width, + config.lru_width, + kernel_size=config.conv1d_width, + groups=config.lru_width, + padding=config.conv1d_width - 1, + ) + self.rg_lru = RecurrentGemmaRglru(config) + self.act_fn = ACT2FN[config.hidden_activation] + + self.conv1d_state = None + + def forward( + self, + input_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + use_cache: bool = True, + **kwargs, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + _, seq_len, _ = input_states.shape + batch_size = input_states.shape[0] + + y_branch = self.linear_y(input_states) + y_branch = self.act_fn(y_branch) + + x_branch = self.linear_x(input_states) + x_branch = x_branch.transpose(1, 2) + + if use_cache: + # Check if cache needs initialization (None or batch size mismatch) + if self.conv1d_state is None or self.conv1d_state.shape[0] != batch_size: + self.conv1d_state = torch.zeros( + (batch_size, self.hidden_size, self.conv1d_width - 1), + device=input_states.device, + dtype=input_states.dtype, + ) + self.rg_lru.recurrent_states = torch.zeros( + (batch_size, self.lru_width), device=input_states.device, dtype=torch.float32 + ) + + if position_ids.shape[1] != 1: # prefill + self.conv1d_state = nn.functional.pad(x_branch, (self.conv1d_width - x_branch.shape[-1] - 1, 0)) + x_branch = self.conv_1d(x_branch)[..., :seq_len] + else: # decoding + conv_state = torch.cat((self.conv1d_state, x_branch), -1) + x_branch = torch.sum(conv_state * self.conv_1d.weight[:, 0, :], dim=-1) + self.conv_1d.bias + x_branch = x_branch.unsqueeze(-1) + self.conv1d_state = conv_state[:, :, 1:] + else: + self.conv1d_state = None + self.rg_lru.recurrent_states = None + x_branch = self.conv_1d(x_branch)[..., :seq_len] + + x_branch = self.rg_lru(x_branch.transpose(1, 2), position_ids) + + hidden_states = x_branch * y_branch + hidden_states = self.linear_out(hidden_states) + return hidden_states + + def _setup_cache(self, batch, device, dtype): + # recurrent_states always computed in full precision + self.rg_lru.recurrent_states = torch.zeros((batch, self.lru_width), device=device, dtype=torch.float32) + self.conv1d_state = torch.zeros((batch, self.hidden_size, self.conv1d_width - 1), device=device, dtype=dtype) + + +TEMPORAL_BLOCK_CLASSES = {"recurrent": RecurrentGemmaRecurrentBlock, "attention": RecurrentGemmaSdpaAttention} + + +class RecurrentGemmaMlp(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size // 2 + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + self.act_fn = ACT2FN[config.hidden_activation] + + def forward(self, hidden_states): + gate = self.act_fn(self.gate_proj(hidden_states)) + return self.down_proj(gate * self.up_proj(hidden_states)) + + +class RecurrentGemmaDecoderLayer(GradientCheckpointingLayer): + """Griffin and Hawk's residual block.""" + + def __init__(self, config, layer_idx): + super().__init__() + self.temporal_pre_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.temporal_block = TEMPORAL_BLOCK_CLASSES[config.layers_block_type[layer_idx]](config, layer_idx) + self.channel_pre_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp_block = RecurrentGemmaMlp(config) + + def forward( + self, + activations: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + use_cache: bool | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + raw_activations = activations + inputs_normalized = self.temporal_pre_norm(raw_activations) # RMSNorm introduces slight slight differences + + hidden_states = self.temporal_block( + inputs_normalized, + position_ids, + attention_mask, + use_cache=use_cache, + past_key_values=past_key_values, + **kwargs, + ) + + residual = hidden_states + raw_activations + + hidden_states = self.channel_pre_norm(residual) + hidden_states = self.mlp_block(hidden_states) + + hidden_states = hidden_states + residual + return hidden_states + + +@auto_docstring +class RecurrentGemmaPreTrainedModel(PreTrainedModel): + config: RecurrentGemmaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["RecurrentGemmaDecoderLayer"] + _skip_keys_device_placement = ["cache"] + _supports_flash_attn = False + _supports_sdpa = False # we can't compare with eager for now + + @torch.no_grad() + def _init_weights(self, module): + std = math.sqrt(self.config.w_init_variance_scale / self.config.conv1d_width) + if isinstance(module, nn.Conv1d): + init.normal_(module.weight, mean=0.0, std=std) + init.zeros_(module.bias) + elif isinstance(module, RecurrentGemmaSdpaAttention): + init.normal_(module.q_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size)) + init.normal_(module.k_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size)) + init.normal_(module.v_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size)) + + std = math.sqrt(self.config.final_w_init_variance_scale / self.config.hidden_size) + init.normal_(module.o_proj.weight, mean=0.0, std=std) + elif isinstance(module, RecurrentGemmaRecurrentBlock): + init.zeros_(module.linear_x.bias) + init.normal_(module.linear_x.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size)) + + init.zeros_(module.linear_y.bias) + init.normal_(module.linear_y.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size)) + + std = math.sqrt(self.config.final_w_init_variance_scale / self.config.lru_width) + init.normal_(module.linear_out.weight, mean=0.0, std=std) + init.zeros_(module.linear_out.bias) + elif isinstance(module, RecurrentGemmaRglru): + std = math.sqrt( + self.config.w_init_variance_scale / (self.config.lru_width // self.config.num_attention_heads) + ) + init.normal_(module.input_gate_weight, mean=0.0, std=std) + init.normal_(module.recurrent_gate_weight, mean=0.0, std=std) + init.zeros_(module.input_gate_bias) + init.zeros_(module.recurrent_gate_bias) + + recurrent_param = torch.empty_like(module.recurrent_param).uniform_(0.9**2 + 1e-8, 0.999**2 + 1e-8) + recurrent_param.log_().mul_(0.5).neg_().exp_().sub_(1.0).log_() + init.copy_(module.recurrent_param, recurrent_param) + elif isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=std) + if getattr(module, "bias", None) is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight) + elif isinstance(module, RecurrentGemmaRMSNorm): + init.zeros_(module.weight) + elif isinstance(module, RecurrentGemmaModel): + init.constant_(module.normalizer, module.config.hidden_size**0.5) + elif isinstance(module, RecurrentGemmaRotaryEmbedding): + buffer_value, _ = module.compute_default_rope_parameters(module.config) + init.copy_(module.inv_freq, buffer_value) + init.copy_(module.original_inv_freq, buffer_value) + + def _setup_cache(self, config, batch, device, dtype): + layers = getattr(self, "model", self).layers + for layer in layers: + if hasattr(layer.temporal_block, "_setup_cache"): + layer.temporal_block._setup_cache(batch, device, dtype) + + +def _get_seq_length(self, layer_idx: int = 0) -> int: + return self.layers[self.first_attention_layer].get_seq_length() + + +def _get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: + return self.layers[self.first_attention_layer].get_mask_sizes(query_length) + + +@auto_docstring +class RecurrentGemmaModel(RecurrentGemmaPreTrainedModel): + def __init__(self, config: RecurrentGemmaConfig): + 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( + [RecurrentGemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.final_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.gradient_checkpointing = False + + self.register_buffer( + "normalizer", torch.tensor(self.config.hidden_size**0.5, dtype=torch.bfloat16), persistent=False + ) + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithNoAttention: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + if use_cache and past_key_values is None: + self._setup_cache(self.config, hidden_states.shape[0], hidden_states.device, hidden_states.dtype) + past_key_values = DynamicCache(config=self.config) + + # Hack because the mamba layer indices will stay empty in `past_key_values`, and we want `get_seq_length` and + # `get_mask_sizes` to use the first attention layer by default for the mask function to create correct masks + if past_key_values is not None: + past_key_values.first_attention_layer = self.config.layers_block_type.index("attention") + # bound new methods to this instance only + past_key_values.get_seq_length = _get_seq_length.__get__(past_key_values) + past_key_values.get_mask_sizes = _get_mask_sizes.__get__(past_key_values) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_sliding_window_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = hidden_states * self.normalizer.type(hidden_states.dtype) + + all_hidden_states = () if output_hidden_states else None + for i, residual_block in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + hidden_states = residual_block( + hidden_states, position_ids, causal_mask, use_cache, past_key_values, **kwargs + ) + + hidden_states = self.final_norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states] if v is not None) + + return BaseModelOutputWithNoAttention( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + ) + + +# TODO: re-enable check: Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->RECURRENTGEMMA,Llama->RecurrentGemma,llama->gemma +@auto_docstring +class RecurrentGemmaForCausalLM(RecurrentGemmaPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = RecurrentGemmaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + # Ignore copy + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, RecurrentGemmaForCausalLM + + >>> model = RecurrentGemmaForCausalLM.from_pretrained("google/recurrentgemma-2b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/recurrentgemma-2b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + 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 + output_hidden_states = True + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + # Soft-cap the logits TODO remove if always done. + # if self.config.logits_soft_cap is not None: + cap = self.config.logits_soft_cap + logits = nn.functional.tanh(logits / cap) * cap + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + ) + + +__all__ = ["RecurrentGemmaForCausalLM", "RecurrentGemmaModel", "RecurrentGemmaPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/slanext/__init__.py b/third_party/transformers/src/transformers/models/slanext/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fcd21c1731d44ad05207078647116c7ee86e8502 --- /dev/null +++ b/third_party/transformers/src/transformers/models/slanext/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_slanext import * + from .image_processing_slanext import * + from .modeling_slanext import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/slanext/configuration_slanext.py b/third_party/transformers/src/transformers/models/slanext/configuration_slanext.py new file mode 100644 index 0000000000000000000000000000000000000000..aac603aa4a01583a9350cb82bfe7fac1ce6b2f80 --- /dev/null +++ b/third_party/transformers/src/transformers/models/slanext/configuration_slanext.py @@ -0,0 +1,103 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/slanext/modular_slanext.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_slanext.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="PaddlePaddle/SLANeXt_wired_safetensors") +@strict +class SLANeXtVisionConfig(PreTrainedConfig): + r""" + output_channels (`int`, *optional*, defaults to 256): + Dimensionality of the output channels in the Patch Encoder. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. + window_size (`int`, *optional*, defaults to 14): + Window size for relative position. + global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. + """ + + base_config_key = "vision_config" + hidden_size: int = 768 + output_channels: int = 256 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int = 512 + patch_size: int | list[int] | tuple[int, int] = 16 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-06 + attention_dropout: float | int = 0.0 + initializer_range: float = 1e-10 + qkv_bias: bool = True + use_abs_pos: bool = True + use_rel_pos: bool = True + window_size: int = 14 + global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11) + mlp_dim: int = 3072 + + +@auto_docstring(checkpoint="PaddlePaddle/SLANeXt_wired_safetensors") +@strict +class SLANeXtConfig(PreTrainedConfig): + r""" + vision_config (`dict` or [`SLANeXtVisionConfig`], *optional*): + Configuration for the vision encoder. If `None`, a default [`SLANeXtVisionConfig`] is used. + post_conv_in_channels (`int`, *optional*, defaults to 256): + Number of input channels for the post-encoder convolution layer. + post_conv_out_channels (`int`, *optional*, defaults to 512): + Number of output channels for the post-encoder convolution layer. + out_channels (`int`, *optional*, defaults to 50): + Vocabulary size for the table structure token prediction head, i.e., the number of distinct structure + tokens the model can predict. + hidden_size (`int`, *optional*, defaults to 512): + Dimensionality of the hidden states in the attention GRU cell and the structure/location prediction heads. + max_text_length (`int`, *optional*, defaults to 500): + Maximum number of autoregressive decoding steps (tokens) for the structure and location decoder. + """ + + model_type = "slanext" + sub_configs = {"vision_config": SLANeXtVisionConfig} + + vision_config: dict | SLANeXtVisionConfig | None = None + post_conv_in_channels: int = 256 + post_conv_out_channels: int = 512 + out_channels: int = 50 + hidden_size: int = 512 + max_text_length: int = 500 + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = SLANeXtVisionConfig() + elif isinstance(self.vision_config, dict): + self.vision_config = SLANeXtVisionConfig(**self.vision_config) + super().__post_init__(**kwargs) + + +__all__ = ["SLANeXtConfig"] diff --git a/third_party/transformers/src/transformers/models/slanext/image_processing_slanext.py b/third_party/transformers/src/transformers/models/slanext/image_processing_slanext.py new file mode 100644 index 0000000000000000000000000000000000000000..f87aed72137f8871b1a0212bb42b6fef194454ef --- /dev/null +++ b/third_party/transformers/src/transformers/models/slanext/image_processing_slanext.py @@ -0,0 +1,257 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/slanext/modular_slanext.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_slanext.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torchvision.transforms.v2.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, SizeDict +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import auto_docstring, is_torchdynamo_compiling, logging +from ...utils.generic import TensorType +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + + +@auto_docstring +@requires(backends=("torch",)) +class SLANeXtImageProcessor(TorchvisionBackend): + resample = 2 # PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"height": 512, "width": 512} + pad_size = {"height": 512, "width": 512} + do_convert_rgb = True + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + + def _resize( + self, + image: "torch.Tensor", + size: SizeDict, + ) -> "torch.Tensor": + batch_size, channels, height, width = image.shape + image = image.view(batch_size * channels, height, width) + + device = image.device + + scale = max(size.height, size.width) / max(height, width) + target_height = round(height * scale) + target_width = round(width * scale) + + target_col = torch.arange(target_width, dtype=torch.float32, device=device) + src_col = (target_col + 0.5) * (float(width) / float(target_width)) - 0.5 + src_col_floor = src_col.floor().to(torch.int32) + src_col_frac = src_col - src_col_floor.float() + # boundary handling + src_col_frac = torch.where(src_col_floor < 0, torch.zeros_like(src_col_frac), src_col_frac) + src_col_floor = torch.where(src_col_floor < 0, torch.zeros_like(src_col_floor), src_col_floor) + src_col_frac = torch.where(src_col_floor >= width - 1, torch.ones_like(src_col_frac), src_col_frac) + src_col_floor = torch.where( + src_col_floor >= width - 1, torch.full_like(src_col_floor, width - 2), src_col_floor + ) + # fixed-point weights + weight_right = (src_col_frac * 2048 + 0.5).floor().to(torch.int32) # round-to-nearest + weight_left = 2048 - weight_right # (target_w,) + # --- row coordinate tables --- + target_row = torch.arange(target_height, dtype=torch.float32, device=device) + src_row = (target_row + 0.5) * (float(height) / float(target_height)) - 0.5 + src_row_floor = src_row.floor().to(torch.int32) + src_row_frac = src_row - src_row_floor.float() + src_row_frac = torch.where(src_row_floor < 0, torch.zeros_like(src_row_frac), src_row_frac) + src_row_floor = torch.where(src_row_floor < 0, torch.zeros_like(src_row_floor), src_row_floor) + src_row_frac = torch.where(src_row_floor >= height - 1, torch.ones_like(src_row_frac), src_row_frac) + src_row_floor = torch.where( + src_row_floor >= height - 1, torch.full_like(src_row_floor, height - 2), src_row_floor + ) + weight_bottom = (src_row_frac * 2048 + 0.5).floor().to(torch.int32) + weight_top = 2048 - weight_bottom # (target_h,) + + image_uint8 = image.clamp(0, 255).to(torch.uint8) # (C, H, W) + image_int32 = image_uint8.to(torch.int32) # (C, H, W) + col_left = src_col_floor.long() # (target_w,) + col_right = (src_col_floor + 1).long() # (target_w,) safe: src_col_floor <= width-2 + row_top = src_row_floor.long() # (target_h,) + row_bottom = (src_row_floor + 1).long() # (target_h,) + # gather 4 neighbours: (C, target_h, target_w) + pixel_top_left = image_int32[:, row_top[:, None], col_left[None, :]] + pixel_top_right = image_int32[:, row_top[:, None], col_right[None, :]] + pixel_bottom_left = image_int32[:, row_bottom[:, None], col_left[None, :]] + pixel_bottom_right = image_int32[:, row_bottom[:, None], col_right[None, :]] + # fixed-point bilinear: weights broadcast over (C, target_h, target_w) + weight_bottom_3d = weight_bottom.view(1, target_height, 1) + weight_top_3d = weight_top.view(1, target_height, 1) + weight_right_3d = weight_right.view(1, 1, target_width) + weight_left_3d = weight_left.view(1, 1, target_width) + interp = weight_top_3d * ( + weight_left_3d * pixel_top_left + weight_right_3d * pixel_top_right + ) + weight_bottom_3d * (weight_left_3d * pixel_bottom_left + weight_right_3d * pixel_bottom_right) + interp = (interp + (1 << 21)) >> 22 + result = interp.clamp(0, 255).to(torch.uint8) # (B*C, target_h, target_w) + + return result.view(batch_size, channels, target_height, target_width).to(dtype=image.dtype) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "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: + if resample is not None and not is_torchdynamo_compiling(): + logger.warning_once("Resampling is not supported in SLANeXt") + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self._resize(image=stacked_images, size=size) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + if do_pad: + processed_images = self.pad(processed_images, pad_size=pad_size, disable_grouping=disable_grouping) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + def __init__(self, **kwargs: Unpack[ImagesKwargs]): + super().__init__(**kwargs) + self.init_decoder() + + def init_decoder(self): + """ + Initialize the decoder vocabulary for table structure recognition. + + Builds a character dictionary mapping HTML table structure tokens (e.g., ``, ``, ``, colspan/ + rowspan attributes) to integer indices. The dictionary includes special `"sos"` (start-of-sequence) and + `"eos"` (end-of-sequence) tokens. Merged `` tokens are used in place of standalone `` tokens + when applicable. + """ + dict_character = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + ] + dict_character += [f' colspan="{i + 2}"' for i in range(19)] + dict_character += [f' rowspan="{i + 2}"' for i in range(19)] + + if "" not in dict_character: + dict_character.append("") + if "" in dict_character: + dict_character.remove("") + + dict_character = ["sos"] + dict_character + ["eos"] + self.dict = {char: i for i, char in enumerate(dict_character)} + self.character = dict_character + self.td_token = ["", ""] + self.bos_id = self.dict["sos"] + self.eos_id = self.dict["eos"] + + def post_process_table_recognition(self, outputs): + """ + Post-process the raw model outputs to decode the predicted table structure into an HTML token sequence. + + Converts the model's predicted probability distributions over the structure vocabulary into a sequence of + HTML tokens representing the table structure. The decoded tokens are wrapped with ``, ``, and + `` tags to form a complete HTML table structure. + + Args: + outputs ([`SLANeXtForTableRecognitionOutput`]): + Raw outputs from the SLANeXt model. The `last_hidden_state` field contains the predicted probability + distributions over the structure vocabulary at each decoding step, with shape + `(batch_size, max_text_length, num_classes)`. + + Returns: + `dict`: A dictionary containing: + - **structure** (`list[str]`): The predicted HTML table structure as a list of tokens, wrapped with + ``, ``, and `
` tags. + - **structure_score** (`float`): The mean confidence score across all predicted tokens. + """ + self.pred = outputs.last_hidden_state + structure_probs = self.pred[0:1] + ignored_tokens = [int(self.bos_id), int(self.eos_id)] + end_idx = int(self.eos_id) + + structure_idx = structure_probs.argmax(dim=2) + structure_probs = structure_probs.max(dim=2).values + + structure_str_list = [] + batch_size = structure_idx.shape[0] + for batch_index in range(batch_size): + structure_list = [] + score_list = [] + for position in range(structure_idx.shape[1]): + char_idx = int(structure_idx[batch_index, position]) + if position > 0 and char_idx == end_idx: + break + if char_idx in ignored_tokens: + continue + text = self.character[char_idx] + structure_list.append(text) + score_list.append(structure_probs[batch_index, position]) + structure_str_list.append(structure_list) + structure_score = torch.stack(score_list).mean().item() + + structure = ["", "", "
"] + structure_str_list[0] + ["
", "", ""] + return {"structure": structure, "structure_score": structure_score} + + +__all__ = ["SLANeXtImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/slanext/modeling_slanext.py b/third_party/transformers/src/transformers/models/slanext/modeling_slanext.py new file mode 100644 index 0000000000000000000000000000000000000000..a71b17e670e984453f6351b53baf39b40119d16e --- /dev/null +++ b/third_party/transformers/src/transformers/models/slanext/modeling_slanext.py @@ -0,0 +1,643 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/slanext/modular_slanext.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_slanext.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import collections +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ... import initialization as init +from ...activations import ACT2CLS, ACT2FN +from ...backbone_utils import filter_output_hidden_states +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_slanext import SLANeXtConfig, SLANeXtVisionConfig + + +class SLANeXtVisionAttention(nn.Module): + """Multi-head Attention block with relative position embeddings.""" + + def __init__(self, config, window_size): + super().__init__() + input_size = ( + (config.image_size // config.patch_size, config.image_size // config.patch_size) + if window_size == 0 + else (window_size, window_size) + ) + + self.num_attention_heads = config.num_attention_heads + head_dim = config.hidden_size // config.num_attention_heads + self.scale = head_dim**-0.5 + self.dropout = config.attention_dropout + + self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias) + self.proj = nn.Linear(config.hidden_size, config.hidden_size) + + self.use_rel_pos = config.use_rel_pos + if self.use_rel_pos: + if input_size is None: + raise ValueError("Input size must be provided if using relative positional encoding.") + + # initialize relative positional embeddings + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) + + def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: + """ + Get relative positional embeddings according to the relative positions of + query and key sizes. + + Args: + q_size (int): + size of the query. + k_size (int): + size of key k. + rel_pos (`torch.Tensor`): + relative position embeddings (L, channel). + + Returns: + Extracted positional embeddings according to relative positions. + """ + max_rel_dist = int(2 * max(q_size, k_size) - 1) + # Interpolate rel pos. + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + ) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + + # Scale the coords with short length if shapes for q and k are different. + q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + def get_decomposed_rel_pos( + self, + query: torch.Tensor, + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + q_size: tuple[int, int], + k_size: tuple[int, int], + ) -> torch.Tensor: + """ + Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. + https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py + + Args: + query (`torch.Tensor`): + query q in the attention layer with shape (batch_size, query_height * query_width, channel). + rel_pos_h (`torch.Tensor`): + relative position embeddings (Lh, channel) for height axis. + rel_pos_w (`torch.Tensor`): + relative position embeddings (Lw, channel) for width axis. + q_size (tuple): + spatial sequence size of query q with (query_height, query_width). + k_size (tuple): + spatial sequence size of key k with (key_height, key_width). + + Returns: + decomposed_rel_pos (`torch.Tensor`): + decomposed relative position embeddings. + """ + query_height, query_width = q_size + key_height, key_width = k_size + relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h) + relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w) + + batch_size, _, dim = query.shape + reshaped_query = query.reshape(batch_size, query_height, query_width, dim) + rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height) + rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width) + + decomposed_rel_pos = rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] + + return decomposed_rel_pos + + def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, height, width, _ = hidden_states.shape + # qkv with shape (3, batch_size, nHead, height * width, channel) + qkv = ( + self.qkv(hidden_states) + .reshape(batch_size, height * width, 3, self.num_attention_heads, -1) + .permute(2, 0, 3, 1, 4) + ) + # q, k, v with shape (batch_size * nHead, height * width, channel) + query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0) + + attn_weights = (query * self.scale) @ key.transpose(-2, -1) + + if self.use_rel_pos: + decomposed_rel_pos = self.get_decomposed_rel_pos( + query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width) + ) + decomposed_rel_pos = decomposed_rel_pos.reshape_as(attn_weights) + attn_weights = attn_weights + decomposed_rel_pos + + attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype) + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1) + attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1) + + attn_output = self.proj(attn_output) + return attn_output, attn_weights + + +class SLANeXtAttentionGRUCell(nn.Module): + def __init__(self, input_size, hidden_size, num_embeddings): + super().__init__() + + self.input_to_hidden = nn.Linear(input_size, hidden_size, bias=False) + self.hidden_to_hidden = nn.Linear(hidden_size, hidden_size) + self.score = nn.Linear(hidden_size, 1, bias=False) + + self.rnn = nn.GRUCell(input_size + num_embeddings, hidden_size) + + def forward( + self, + prev_hidden: torch.FloatTensor, + batch_hidden: torch.FloatTensor, + char_onehots: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ): + batch_hidden_proj = self.input_to_hidden(batch_hidden) + prev_hidden_proj = self.hidden_to_hidden(prev_hidden).unsqueeze(1) + + attention_scores = batch_hidden_proj + prev_hidden_proj + attention_scores = torch.tanh(attention_scores) + attention_scores = self.score(attention_scores) + + attn_weights = F.softmax(attention_scores, dim=1, dtype=torch.float32).to(attention_scores.dtype) + attn_weights = attn_weights.transpose(1, 2) + context = torch.matmul(attn_weights, batch_hidden).squeeze(1) + concat_context = torch.cat([context, char_onehots], 1) + hidden_states = self.rnn(concat_context, prev_hidden) + + return hidden_states, attn_weights + + +class SLANeXtMLP(nn.Module): + def __init__(self, hidden_size, out_channels, activation=None): + super().__init__() + self.fc1 = nn.Linear(hidden_size, hidden_size) + self.fc2 = nn.Linear(hidden_size, out_channels) + self.act_fn = nn.Identity() if activation is None else ACT2CLS[activation]() + + def forward(self, hidden_states): + hidden_states = self.fc1(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = self.act_fn(hidden_states) + return hidden_states + + +class SLANeXtPreTrainedModel(PreTrainedModel): + config: SLANeXtConfig + base_model_prefix = "backbone" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _keep_in_fp32_modules_strict = ["structure_attention_cell", "structure_generator"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + + # Initialize positional embeddings to zero (SLANeXtVisionEncoder holds pos_embed) + if isinstance(module, SLANeXtVisionEncoder): + if module.pos_embed is not None: + init.constant_(module.pos_embed, 0.0) + + # Initialize relative positional embeddings to zero (SLANeXtVisionAttention holds rel_pos_h/w) + if isinstance(module, SLANeXtVisionAttention): + if module.use_rel_pos: + init.constant_(module.rel_pos_h, 0.0) + init.constant_(module.rel_pos_w, 0.0) + + # Initialize GRUCell (replicates PyTorch default reset_parameters) + if isinstance(module, nn.GRUCell): + std = 1.0 / math.sqrt(module.hidden_size) if module.hidden_size > 0 else 0 + init.uniform_(module.weight_ih, -std, std) + init.uniform_(module.weight_hh, -std, std) + if module.bias_ih is not None: + init.uniform_(module.bias_ih, -std, std) + if module.bias_hh is not None: + init.uniform_(module.bias_hh, -std, std) + + # Initialize SLAHead layers + if isinstance(module, SLANeXtSLAHead): + std = 1.0 / math.sqrt(self.config.hidden_size * 1.0) + # Initialize structure_generator and loc_generator layers + for generator in (module.structure_generator,): + for layer in generator.children(): + if isinstance(layer, nn.Linear): + init.uniform_(layer.weight, -std, std) + if layer.bias is not None: + init.uniform_(layer.bias, -std, std) + + +class SLANeXtMLPBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim) + self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size) + self.act = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.lin1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.lin2(hidden_states) + return hidden_states + + +class SLANeXtVisionLayer(GradientCheckpointingLayer): + def __init__(self, config, window_size): + super().__init__() + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attn = SLANeXtVisionAttention(config, window_size) + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = SLANeXtMLPBlock(config) + self.window_size = window_size + + def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> tuple[torch.Tensor, tuple[int, int]]: + """ + Args: + Partition into non-overlapping windows with padding if needed. + hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window + size. + + Returns: + windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel]. + (pad_height, pad_width): padded height and width before partition + """ + batch_size, height, width, channel = hidden_states.shape + + pad_h = (window_size - height % window_size) % window_size + pad_w = (window_size - width % window_size) % window_size + hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h)) + pad_height, pad_width = height + pad_h, width + pad_w + + hidden_states = hidden_states.reshape( + batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel + ) + windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel) + return windows, (pad_height, pad_width) + + def window_unpartition( + self, windows: torch.Tensor, window_size: int, padding_shape: tuple[int, int], original_shape: tuple[int, int] + ) -> torch.Tensor: + """ + Args: + Window unpartition into original sequences and removing padding. + hidden_states (tensor): + input tokens with [batch_size * num_windows, window_size, window_size, channel]. + window_size (int): + window size. + padding_shape (Tuple): + padded height and width (pad_height, pad_width). + original_shape (Tuple): original height and width (height, width) before padding. + + Returns: + hidden_states: unpartitioned sequences with [batch_size, height, width, channel]. + """ + pad_height, pad_width = padding_shape + height, width = original_shape + batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size) + hidden_states = windows.reshape( + batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1 + ) + hidden_states = ( + hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1) + ) + + hidden_states = hidden_states[:, :height, :width, :].contiguous() + return hidden_states + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.FloatTensor]: + residual = hidden_states + hidden_states = self.layer_norm1(hidden_states) + # Window partition + if self.window_size > 0: + height, width = hidden_states.shape[1], hidden_states.shape[2] + hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size) + + hidden_states, attn_weights = self.attn( + hidden_states=hidden_states, + ) + # Reverse window partition + if self.window_size > 0: + hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width)) + + hidden_states = residual + hidden_states + layernorm_output = self.layer_norm2(hidden_states) + hidden_states = hidden_states + self.mlp(layernorm_output) + return hidden_states + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for slanext vision model's outputs that also contains image embeddings obtained by applying the projection + layer to the pooler_output. + """ +) +class SLANeXtVisionEncoderOutput(ModelOutput): + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + """ + + image_embeds: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +class SLANeXtPatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values): + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + if height != self.image_size[0] or width != self.image_size[1]: + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})." + ) + embeddings = self.projection(pixel_values).permute(0, 2, 3, 1) + return embeddings + + +class SLANeXtLayerNorm(nn.LayerNorm): + r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, + width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs): + super().__init__(normalized_shape, eps=eps, **kwargs) + if data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError(f"Unsupported data format: {data_format}") + self.data_format = data_format + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """ + Args: + features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels) + """ + if self.data_format == "channels_first": + features = features.permute(0, 2, 3, 1) + features = super().forward(features) + features = features.permute(0, 3, 1, 2) + else: + features = super().forward(features) + return features + + +class SLANeXtVisionNeck(nn.Module): + def __init__(self, config: SLANeXtVisionConfig): + super().__init__() + self.config = config + + self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False) + self.layer_norm1 = SLANeXtLayerNorm(config.output_channels, data_format="channels_first") + self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False) + self.layer_norm2 = SLANeXtLayerNorm(config.output_channels, data_format="channels_first") + + def forward(self, hidden_states): + hidden_states = hidden_states.permute(0, 3, 1, 2) + hidden_states = self.conv1(hidden_states) + hidden_states = self.layer_norm1(hidden_states) + + hidden_states = self.conv2(hidden_states) + hidden_states = self.layer_norm2(hidden_states) + return hidden_states + + +class SLANeXtVisionEncoder(SLANeXtPreTrainedModel): + _can_record_outputs = {"hidden_states": SLANeXtVisionLayer, "attentions": SLANeXtVisionAttention} + input_modalities = ("image",) + + def __init__(self, config: SLANeXtVisionConfig): + super().__init__(config) + self.config = config + self.image_size = config.image_size + self.patch_embed = SLANeXtPatchEmbeddings(config) + + self.pos_embed = None + if config.use_abs_pos: + # Initialize absolute positional embedding with pretrain image size. + self.pos_embed = nn.Parameter( + torch.zeros( + 1, + config.image_size // config.patch_size, + config.image_size // config.patch_size, + config.hidden_size, + ) + ) + + self.layers = nn.ModuleList() + for i in range(config.num_hidden_layers): + layer = SLANeXtVisionLayer( + config, + window_size=config.window_size if i not in config.global_attn_indexes else 0, + ) + self.layers.append(layer) + + self.neck = SLANeXtVisionNeck(config) + + self.gradient_checkpointing = False + self.post_init() + + def get_input_embeddings(self): + return self.patch_embed + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + def forward( + self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | SLANeXtVisionEncoderOutput: + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + hidden_states = self.patch_embed(pixel_values) + if self.pos_embed is not None: + hidden_states = hidden_states + self.pos_embed + for layer_module in self.layers: + hidden_states = layer_module(hidden_states) + hidden_states = self.neck(hidden_states) + return SLANeXtVisionEncoderOutput( + last_hidden_state=hidden_states, + ) + + +class SLANeXtBackbone(SLANeXtPreTrainedModel): + def __init__( + self, + config: dict | None = None, + **kwargs, + ): + super().__init__(config) + self.vision_tower = SLANeXtVisionEncoder(config.vision_config) + self.post_conv = nn.Conv2d( + config.post_conv_in_channels, config.post_conv_out_channels, kernel_size=3, stride=2, padding=1, bias=False + ) + self.post_init() + + def forward(self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs]): + vision_output = self.vision_tower(hidden_states, **kwargs) + hidden_states = self.post_conv(vision_output.last_hidden_state) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=vision_output.hidden_states, + attentions=vision_output.attentions, + ) + + +class SLANeXtSLAHead(SLANeXtPreTrainedModel): + _can_record_outputs = { + "attentions": SLANeXtAttentionGRUCell, + } + + def __init__( + self, + config: dict | None = None, + **kwargs, + ): + super().__init__(config) + + self.structure_attention_cell = SLANeXtAttentionGRUCell( + config.post_conv_out_channels, config.hidden_size, config.out_channels + ) + self.structure_generator = SLANeXtMLP(config.hidden_size, config.out_channels) + + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @filter_output_hidden_states + def forward( + self, + hidden_states: torch.FloatTensor, + targets: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + features = torch.zeros( + (hidden_states.shape[0], self.config.hidden_size), dtype=torch.float32, device=hidden_states.device + ) + predicted_chars = torch.zeros(size=[hidden_states.shape[0]], dtype=torch.long, device=hidden_states.device) + + structure_preds_list = [] + structure_ids_list = [] + for _ in range(self.config.max_text_length + 1): + embedding_feature = F.one_hot(predicted_chars, self.config.out_channels).float() + features, _ = self.structure_attention_cell(features, hidden_states.float(), embedding_feature) + structure_step = self.structure_generator(features) + predicted_chars = structure_step.argmax(dim=1) + + structure_preds_list.append(structure_step) + structure_ids_list.append(predicted_chars) + if torch.stack(structure_ids_list, dim=1).eq(self.config.out_channels - 1).any(-1).all(): + break + structure_preds = F.softmax(torch.stack(structure_preds_list, dim=1), dim=-1, dtype=torch.float32).to( + hidden_states.dtype + ) + + return BaseModelOutput(last_hidden_state=structure_preds, hidden_states=structure_preds_list) + + +@dataclass +@auto_docstring +class SLANeXtForTableRecognitionOutput(BaseModelOutput): + r""" + head_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Hidden-states of the SLANeXtSLAHead at each prediction step, varies up to max `self.config.max_text_length` states (depending on early exits). + head_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Attentions of the SLANeXtSLAHead at each prediction step, varies up to max `self.config.max_text_length` attentions (depending on early exits). + """ + + head_hidden_states: torch.FloatTensor | None = None + head_attentions: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + SLANeXt Table Recognition model for table recognition tasks. Wraps the core SLANeXtPreTrainedModel + and returns outputs compatible with the Transformers table recognition API. + """ +) +class SLANeXtForTableRecognition(SLANeXtPreTrainedModel): + def __init__(self, config: SLANeXtConfig): + super().__init__(config) + self.backbone = SLANeXtBackbone(config=config) + self.head = SLANeXtSLAHead(config=config) + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple[torch.FloatTensor] | SLANeXtForTableRecognitionOutput: + backbone_outputs = self.backbone(pixel_values, **kwargs) + head_outputs = self.head(backbone_outputs.last_hidden_state, **kwargs) + return SLANeXtForTableRecognitionOutput( + last_hidden_state=head_outputs.last_hidden_state, + hidden_states=backbone_outputs.hidden_states, + attentions=backbone_outputs.attentions, + head_hidden_states=head_outputs.hidden_states, + head_attentions=head_outputs.attentions, + ) + + +__all__ = ["SLANeXtSLAHead", "SLANeXtBackbone", "SLANeXtForTableRecognition", "SLANeXtPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/slanext/modular_slanext.py b/third_party/transformers/src/transformers/models/slanext/modular_slanext.py new file mode 100644 index 0000000000000000000000000000000000000000..965e0da7fbe64b7bf6011c62fb8c6fbfab1cdb21 --- /dev/null +++ b/third_party/transformers/src/transformers/models/slanext/modular_slanext.py @@ -0,0 +1,539 @@ +# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms.v2.functional as tvF +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...activations import ACT2CLS +from ...backbone_utils import filter_output_hidden_states +from ...configuration_utils import PreTrainedConfig +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, SizeDict +from ...modeling_outputs import BaseModelOutput +from ...modeling_utils import PreTrainedModel +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils.generic import TensorType, merge_with_config_defaults +from ...utils.import_utils import requires +from ...utils.output_capturing import capture_outputs +from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig +from ..got_ocr2.modeling_got_ocr2 import ( + GotOcr2VisionAttention, + GotOcr2VisionEncoder, +) + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="PaddlePaddle/SLANeXt_wired_safetensors") +@strict +class SLANeXtVisionConfig(GotOcr2VisionConfig): + image_size: int = 512 + + +class SLANeXtVisionAttention(GotOcr2VisionAttention): + pass + + +@auto_docstring(checkpoint="PaddlePaddle/SLANeXt_wired_safetensors") +@strict +class SLANeXtConfig(PreTrainedConfig): + r""" + vision_config (`dict` or [`SLANeXtVisionConfig`], *optional*): + Configuration for the vision encoder. If `None`, a default [`SLANeXtVisionConfig`] is used. + post_conv_in_channels (`int`, *optional*, defaults to 256): + Number of input channels for the post-encoder convolution layer. + post_conv_out_channels (`int`, *optional*, defaults to 512): + Number of output channels for the post-encoder convolution layer. + out_channels (`int`, *optional*, defaults to 50): + Vocabulary size for the table structure token prediction head, i.e., the number of distinct structure + tokens the model can predict. + hidden_size (`int`, *optional*, defaults to 512): + Dimensionality of the hidden states in the attention GRU cell and the structure/location prediction heads. + max_text_length (`int`, *optional*, defaults to 500): + Maximum number of autoregressive decoding steps (tokens) for the structure and location decoder. + """ + + model_type = "slanext" + sub_configs = {"vision_config": SLANeXtVisionConfig} + + vision_config: dict | SLANeXtVisionConfig | None = None + post_conv_in_channels: int = 256 + post_conv_out_channels: int = 512 + out_channels: int = 50 + hidden_size: int = 512 + max_text_length: int = 500 + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = SLANeXtVisionConfig() + elif isinstance(self.vision_config, dict): + self.vision_config = SLANeXtVisionConfig(**self.vision_config) + super().__post_init__(**kwargs) + + +class SLANeXtAttentionGRUCell(nn.Module): + def __init__(self, input_size, hidden_size, num_embeddings): + super().__init__() + + self.input_to_hidden = nn.Linear(input_size, hidden_size, bias=False) + self.hidden_to_hidden = nn.Linear(hidden_size, hidden_size) + self.score = nn.Linear(hidden_size, 1, bias=False) + + self.rnn = nn.GRUCell(input_size + num_embeddings, hidden_size) + + def forward( + self, + prev_hidden: torch.FloatTensor, + batch_hidden: torch.FloatTensor, + char_onehots: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ): + batch_hidden_proj = self.input_to_hidden(batch_hidden) + prev_hidden_proj = self.hidden_to_hidden(prev_hidden).unsqueeze(1) + + attention_scores = batch_hidden_proj + prev_hidden_proj + attention_scores = torch.tanh(attention_scores) + attention_scores = self.score(attention_scores) + + attn_weights = F.softmax(attention_scores, dim=1, dtype=torch.float32).to(attention_scores.dtype) + attn_weights = attn_weights.transpose(1, 2) + context = torch.matmul(attn_weights, batch_hidden).squeeze(1) + concat_context = torch.cat([context, char_onehots], 1) + hidden_states = self.rnn(concat_context, prev_hidden) + + return hidden_states, attn_weights + + +class SLANeXtMLP(nn.Module): + def __init__(self, hidden_size, out_channels, activation=None): + super().__init__() + self.fc1 = nn.Linear(hidden_size, hidden_size) + self.fc2 = nn.Linear(hidden_size, out_channels) + self.act_fn = nn.Identity() if activation is None else ACT2CLS[activation]() + + def forward(self, hidden_states): + hidden_states = self.fc1(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = self.act_fn(hidden_states) + return hidden_states + + +class SLANeXtPreTrainedModel(PreTrainedModel): + config: SLANeXtConfig + base_model_prefix = "backbone" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _keep_in_fp32_modules_strict = ["structure_attention_cell", "structure_generator"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + + # Initialize positional embeddings to zero (SLANeXtVisionEncoder holds pos_embed) + if isinstance(module, SLANeXtVisionEncoder): + if module.pos_embed is not None: + init.constant_(module.pos_embed, 0.0) + + # Initialize relative positional embeddings to zero (SLANeXtVisionAttention holds rel_pos_h/w) + if isinstance(module, SLANeXtVisionAttention): + if module.use_rel_pos: + init.constant_(module.rel_pos_h, 0.0) + init.constant_(module.rel_pos_w, 0.0) + + # Initialize GRUCell (replicates PyTorch default reset_parameters) + if isinstance(module, nn.GRUCell): + std = 1.0 / math.sqrt(module.hidden_size) if module.hidden_size > 0 else 0 + init.uniform_(module.weight_ih, -std, std) + init.uniform_(module.weight_hh, -std, std) + if module.bias_ih is not None: + init.uniform_(module.bias_ih, -std, std) + if module.bias_hh is not None: + init.uniform_(module.bias_hh, -std, std) + + # Initialize SLAHead layers + if isinstance(module, SLANeXtSLAHead): + std = 1.0 / math.sqrt(self.config.hidden_size * 1.0) + # Initialize structure_generator and loc_generator layers + for generator in (module.structure_generator,): + for layer in generator.children(): + if isinstance(layer, nn.Linear): + init.uniform_(layer.weight, -std, std) + if layer.bias is not None: + init.uniform_(layer.bias, -std, std) + + +class SLANeXtVisionEncoder(GotOcr2VisionEncoder): + pass + + +class SLANeXtBackbone(SLANeXtPreTrainedModel): + def __init__( + self, + config: dict | None = None, + **kwargs, + ): + super().__init__(config) + self.vision_tower = SLANeXtVisionEncoder(config.vision_config) + self.post_conv = nn.Conv2d( + config.post_conv_in_channels, config.post_conv_out_channels, kernel_size=3, stride=2, padding=1, bias=False + ) + self.post_init() + + def forward(self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs]): + vision_output = self.vision_tower(hidden_states, **kwargs) + hidden_states = self.post_conv(vision_output.last_hidden_state) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=vision_output.hidden_states, + attentions=vision_output.attentions, + ) + + +class SLANeXtSLAHead(SLANeXtPreTrainedModel): + _can_record_outputs = { + "attentions": SLANeXtAttentionGRUCell, + } + + def __init__( + self, + config: dict | None = None, + **kwargs, + ): + super().__init__(config) + + self.structure_attention_cell = SLANeXtAttentionGRUCell( + config.post_conv_out_channels, config.hidden_size, config.out_channels + ) + self.structure_generator = SLANeXtMLP(config.hidden_size, config.out_channels) + + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @filter_output_hidden_states + def forward( + self, + hidden_states: torch.FloatTensor, + targets: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + features = torch.zeros( + (hidden_states.shape[0], self.config.hidden_size), dtype=torch.float32, device=hidden_states.device + ) + predicted_chars = torch.zeros(size=[hidden_states.shape[0]], dtype=torch.long, device=hidden_states.device) + + structure_preds_list = [] + structure_ids_list = [] + for _ in range(self.config.max_text_length + 1): + embedding_feature = F.one_hot(predicted_chars, self.config.out_channels).float() + features, _ = self.structure_attention_cell(features, hidden_states.float(), embedding_feature) + structure_step = self.structure_generator(features) + predicted_chars = structure_step.argmax(dim=1) + + structure_preds_list.append(structure_step) + structure_ids_list.append(predicted_chars) + if torch.stack(structure_ids_list, dim=1).eq(self.config.out_channels - 1).any(-1).all(): + break + structure_preds = F.softmax(torch.stack(structure_preds_list, dim=1), dim=-1, dtype=torch.float32).to( + hidden_states.dtype + ) + + return BaseModelOutput(last_hidden_state=structure_preds, hidden_states=structure_preds_list) + + +@dataclass +@auto_docstring +class SLANeXtForTableRecognitionOutput(BaseModelOutput): + r""" + head_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Hidden-states of the SLANeXtSLAHead at each prediction step, varies up to max `self.config.max_text_length` states (depending on early exits). + head_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Attentions of the SLANeXtSLAHead at each prediction step, varies up to max `self.config.max_text_length` attentions (depending on early exits). + """ + + head_hidden_states: torch.FloatTensor | None = None + head_attentions: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + SLANeXt Table Recognition model for table recognition tasks. Wraps the core SLANeXtPreTrainedModel + and returns outputs compatible with the Transformers table recognition API. + """ +) +class SLANeXtForTableRecognition(SLANeXtPreTrainedModel): + def __init__(self, config: SLANeXtConfig): + super().__init__(config) + self.backbone = SLANeXtBackbone(config=config) + self.head = SLANeXtSLAHead(config=config) + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple[torch.FloatTensor] | SLANeXtForTableRecognitionOutput: + backbone_outputs = self.backbone(pixel_values, **kwargs) + head_outputs = self.head(backbone_outputs.last_hidden_state, **kwargs) + return SLANeXtForTableRecognitionOutput( + last_hidden_state=head_outputs.last_hidden_state, + hidden_states=backbone_outputs.hidden_states, + attentions=backbone_outputs.attentions, + head_hidden_states=head_outputs.hidden_states, + head_attentions=head_outputs.attentions, + ) + + +@auto_docstring +@requires(backends=("torch",)) +class SLANeXtImageProcessor(TorchvisionBackend): + resample = 2 # PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"height": 512, "width": 512} + pad_size = {"height": 512, "width": 512} + do_convert_rgb = True + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + + def _resize( + self, + image: "torch.Tensor", + size: SizeDict, + ) -> "torch.Tensor": + batch_size, channels, height, width = image.shape + image = image.view(batch_size * channels, height, width) + + device = image.device + + scale = max(size.height, size.width) / max(height, width) + target_height = round(height * scale) + target_width = round(width * scale) + + target_col = torch.arange(target_width, dtype=torch.float32, device=device) + src_col = (target_col + 0.5) * (float(width) / float(target_width)) - 0.5 + src_col_floor = src_col.floor().to(torch.int32) + src_col_frac = src_col - src_col_floor.float() + # boundary handling + src_col_frac = torch.where(src_col_floor < 0, torch.zeros_like(src_col_frac), src_col_frac) + src_col_floor = torch.where(src_col_floor < 0, torch.zeros_like(src_col_floor), src_col_floor) + src_col_frac = torch.where(src_col_floor >= width - 1, torch.ones_like(src_col_frac), src_col_frac) + src_col_floor = torch.where( + src_col_floor >= width - 1, torch.full_like(src_col_floor, width - 2), src_col_floor + ) + # fixed-point weights + weight_right = (src_col_frac * 2048 + 0.5).floor().to(torch.int32) # round-to-nearest + weight_left = 2048 - weight_right # (target_w,) + # --- row coordinate tables --- + target_row = torch.arange(target_height, dtype=torch.float32, device=device) + src_row = (target_row + 0.5) * (float(height) / float(target_height)) - 0.5 + src_row_floor = src_row.floor().to(torch.int32) + src_row_frac = src_row - src_row_floor.float() + src_row_frac = torch.where(src_row_floor < 0, torch.zeros_like(src_row_frac), src_row_frac) + src_row_floor = torch.where(src_row_floor < 0, torch.zeros_like(src_row_floor), src_row_floor) + src_row_frac = torch.where(src_row_floor >= height - 1, torch.ones_like(src_row_frac), src_row_frac) + src_row_floor = torch.where( + src_row_floor >= height - 1, torch.full_like(src_row_floor, height - 2), src_row_floor + ) + weight_bottom = (src_row_frac * 2048 + 0.5).floor().to(torch.int32) + weight_top = 2048 - weight_bottom # (target_h,) + + image_uint8 = image.clamp(0, 255).to(torch.uint8) # (C, H, W) + image_int32 = image_uint8.to(torch.int32) # (C, H, W) + col_left = src_col_floor.long() # (target_w,) + col_right = (src_col_floor + 1).long() # (target_w,) safe: src_col_floor <= width-2 + row_top = src_row_floor.long() # (target_h,) + row_bottom = (src_row_floor + 1).long() # (target_h,) + # gather 4 neighbours: (C, target_h, target_w) + pixel_top_left = image_int32[:, row_top[:, None], col_left[None, :]] + pixel_top_right = image_int32[:, row_top[:, None], col_right[None, :]] + pixel_bottom_left = image_int32[:, row_bottom[:, None], col_left[None, :]] + pixel_bottom_right = image_int32[:, row_bottom[:, None], col_right[None, :]] + # fixed-point bilinear: weights broadcast over (C, target_h, target_w) + weight_bottom_3d = weight_bottom.view(1, target_height, 1) + weight_top_3d = weight_top.view(1, target_height, 1) + weight_right_3d = weight_right.view(1, 1, target_width) + weight_left_3d = weight_left.view(1, 1, target_width) + interp = weight_top_3d * ( + weight_left_3d * pixel_top_left + weight_right_3d * pixel_top_right + ) + weight_bottom_3d * (weight_left_3d * pixel_bottom_left + weight_right_3d * pixel_bottom_right) + interp = (interp + (1 << 21)) >> 22 + result = interp.clamp(0, 255).to(torch.uint8) # (B*C, target_h, target_w) + + return result.view(batch_size, channels, target_height, target_width).to(dtype=image.dtype) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "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: + if resample is not None and not is_torchdynamo_compiling(): + logger.warning_once("Resampling is not supported in SLANeXt") + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self._resize(image=stacked_images, size=size) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + if do_pad: + processed_images = self.pad(processed_images, pad_size=pad_size, disable_grouping=disable_grouping) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + def __init__(self, **kwargs: Unpack[ImagesKwargs]): + super().__init__(**kwargs) + self.init_decoder() + + def init_decoder(self): + """ + Initialize the decoder vocabulary for table structure recognition. + + Builds a character dictionary mapping HTML table structure tokens (e.g., ``, ``, ``, colspan/ + rowspan attributes) to integer indices. The dictionary includes special `"sos"` (start-of-sequence) and + `"eos"` (end-of-sequence) tokens. Merged `` tokens are used in place of standalone `` tokens + when applicable. + """ + dict_character = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + ] + dict_character += [f' colspan="{i + 2}"' for i in range(19)] + dict_character += [f' rowspan="{i + 2}"' for i in range(19)] + + if "" not in dict_character: + dict_character.append("") + if "" in dict_character: + dict_character.remove("") + + dict_character = ["sos"] + dict_character + ["eos"] + self.dict = {char: i for i, char in enumerate(dict_character)} + self.character = dict_character + self.td_token = ["", ""] + self.bos_id = self.dict["sos"] + self.eos_id = self.dict["eos"] + + def post_process_table_recognition(self, outputs): + """ + Post-process the raw model outputs to decode the predicted table structure into an HTML token sequence. + + Converts the model's predicted probability distributions over the structure vocabulary into a sequence of + HTML tokens representing the table structure. The decoded tokens are wrapped with ``, ``, and + `` tags to form a complete HTML table structure. + + Args: + outputs ([`SLANeXtForTableRecognitionOutput`]): + Raw outputs from the SLANeXt model. The `last_hidden_state` field contains the predicted probability + distributions over the structure vocabulary at each decoding step, with shape + `(batch_size, max_text_length, num_classes)`. + + Returns: + `dict`: A dictionary containing: + - **structure** (`list[str]`): The predicted HTML table structure as a list of tokens, wrapped with + ``, ``, and `
` tags. + - **structure_score** (`float`): The mean confidence score across all predicted tokens. + """ + self.pred = outputs.last_hidden_state + structure_probs = self.pred[0:1] + ignored_tokens = [int(self.bos_id), int(self.eos_id)] + end_idx = int(self.eos_id) + + structure_idx = structure_probs.argmax(dim=2) + structure_probs = structure_probs.max(dim=2).values + + structure_str_list = [] + batch_size = structure_idx.shape[0] + for batch_index in range(batch_size): + structure_list = [] + score_list = [] + for position in range(structure_idx.shape[1]): + char_idx = int(structure_idx[batch_index, position]) + if position > 0 and char_idx == end_idx: + break + if char_idx in ignored_tokens: + continue + text = self.character[char_idx] + structure_list.append(text) + score_list.append(structure_probs[batch_index, position]) + structure_str_list.append(structure_list) + structure_score = torch.stack(score_list).mean().item() + + structure = ["", "", "
"] + structure_str_list[0] + ["
", "", ""] + return {"structure": structure, "structure_score": structure_score} + + +__all__ = [ + "SLANeXtImageProcessor", + "SLANeXtConfig", + "SLANeXtSLAHead", + "SLANeXtBackbone", + "SLANeXtForTableRecognition", + "SLANeXtPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/udop/__init__.py b/third_party/transformers/src/transformers/models/udop/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..852ad5e3c81ad5979157b7783b21579b8dc22efb --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_udop import * + from .modeling_udop import * + from .processing_udop import * + from .tokenization_udop import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/udop/configuration_udop.py b/third_party/transformers/src/transformers/models/udop/configuration_udop.py new file mode 100644 index 0000000000000000000000000000000000000000..58a2315604554eae79cb4bf7a799d13874146913 --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/configuration_udop.py @@ -0,0 +1,96 @@ +# Copyright 2024 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""UDOP model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/udop-large") +@strict +class UdopConfig(PreTrainedConfig): + r""" + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. + relative_attention_max_distance (`int`, *optional*, defaults to 128): + The maximum distance of the longer sequences for the bucket separation. + relative_bias_args (`list[dict]`, *optional*, defaults to `[{'type': '1d'}, {'type': 'horizontal'}, {'type': 'vertical'}]`): + A list of dictionaries containing the arguments for the relative bias layers. + feed_forward_proj (`string`, *optional*, defaults to `"relu"`): + Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. Udopv1.1 uses the + `"gated-gelu"` feed forward projection. Original Udop uses `"relu"`. + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum absolute position embeddings for relative position encoding. + """ + + model_type = "udop" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"} + + vocab_size: int = 33201 + d_model: int = 1024 + d_kv: int = 64 + d_ff: int = 4096 + num_layers: int = 24 + num_decoder_layers: int | None = None + num_heads: int = 16 + relative_attention_num_buckets: int = 32 + relative_attention_max_distance: int = 128 + relative_bias_args: list[dict] | None = None + dropout_rate: float | int = 0.1 + layer_norm_epsilon: float = 1e-6 + initializer_factor: float = 1.0 + feed_forward_proj: str = "relu" + is_encoder_decoder: bool = True + use_cache: bool = True + pad_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + max_2d_position_embeddings: int = 1024 + image_size: int | list[int] | tuple[int, int] = 224 + patch_size: int | list[int] | tuple[int, int] = 16 + num_channels: int = 3 + is_decoder: bool = False + add_cross_attention: bool = False + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + if self.relative_bias_args is None: + self.relative_bias_args = [{"type": "1d"}, {"type": "horizontal"}, {"type": "vertical"}] + + self.num_decoder_layers = ( + self.num_decoder_layers if self.num_decoder_layers is not None else self.num_layers + ) # default = symmetry + + act_info = self.feed_forward_proj.split("-") + self.dense_act_fn = act_info[-1] + self.is_gated_act = act_info[0] == "gated" + + kwargs.pop("tie_word_embeddings", None) + self.tie_word_embeddings = True + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + act_info = self.feed_forward_proj.split("-") + if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2: + raise ValueError( + f"`feed_forward_proj`: {self.feed_forward_proj} is not a valid activation function of the dense layer." + "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. " + "'gated-gelu' or 'relu'" + ) + + +__all__ = ["UdopConfig"] diff --git a/third_party/transformers/src/transformers/models/udop/convert_udop_to_hf.py b/third_party/transformers/src/transformers/models/udop/convert_udop_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..b138a6ae2c5534278ced1626a3850abd302f9a0c --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/convert_udop_to_hf.py @@ -0,0 +1,225 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert UDOP checkpoints from the original repository. URL: https://github.com/microsoft/i-Code/tree/main/i-Code-Doc""" + +import argparse + +import torch +from huggingface_hub import hf_hub_download +from PIL import Image +from torchvision import transforms as T + +from transformers import ( + LayoutLMv3ImageProcessor, + UdopConfig, + UdopForConditionalGeneration, + UdopProcessor, + UdopTokenizer, +) +from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD + + +def original_transform(image, image_size=224): + transform = T.Compose( + [ + T.Resize([image_size, image_size]), + T.ToTensor(), + T.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD), + ] + ) + + image = transform(image) + return image + + +def get_image(): + filepath = hf_hub_download( + repo_id="hf-internal-testing/fixtures_docvqa", filename="document_2.png", repo_type="dataset" + ) + image = Image.open(filepath).convert("RGB") + + return image + + +def prepare_dummy_inputs(tokenizer, image_processor): + prompt = "Question answering. What is the name of the company?" + prompt = "Question answering. In which year is the report made?" + prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) + + image = get_image() + # words, boxes = apply_tesseract(image, lang=None) + # fmt: off + words = ['7', 'ITC', 'Limited', 'REPORT', 'AND', 'ACCOUNTS', '2013', 'ITC’s', 'Brands:', 'An', 'Asset', 'for', 'the', 'Nation', 'The', 'consumer', 'needs', 'and', 'aspirations', 'they', 'fulfil,', 'the', 'benefit', 'they', 'generate', 'for', 'millions', 'across', 'ITC’s', 'value', 'chains,', 'the', 'future-ready', 'capabilities', 'that', 'support', 'them,', 'and', 'the', 'value', 'that', 'they', 'create', 'for', 'the', 'country,', 'have', 'made', 'ITC’s', 'brands', 'national', 'assets,', 'adding', 'to', 'India’s', 'competitiveness.', 'It', 'is', 'ITC’s', 'aspiration', 'to', 'be', 'the', 'No', '1', 'FMCG', 'player', 'in', 'the', 'country,', 'driven', 'by', 'its', 'new', 'FMCG', 'businesses.', 'A', 'recent', 'Nielsen', 'report', 'has', 'highlighted', 'that', "ITC's", 'new', 'FMCG', 'businesses', 'are', 'the', 'fastest', 'growing', 'among', 'the', 'top', 'consumer', 'goods', 'companies', 'operating', 'in', 'India.', 'ITC', 'takes', 'justifiable', 'pride', 'that,', 'along', 'with', 'generating', 'economic', 'value,', 'these', 'celebrated', 'Indian', 'brands', 'also', 'drive', 'the', 'creation', 'of', 'larger', 'societal', 'capital', 'through', 'the', 'virtuous', 'cycle', 'of', 'sustainable', 'and', 'inclusive', 'growth.', 'DI', 'WILLS', '*', ';', 'LOVE', 'DELIGHTFULLY', 'SOFT', 'SKIN?', 'aia', 'Ans', 'Source:', 'https://www.industrydocuments.ucsf.edu/docs/snbx0223'] + boxes = [[0, 45, 67, 80], [72, 56, 109, 67], [116, 56, 189, 67], [198, 59, 253, 66], [257, 59, 285, 66], [289, 59, 365, 66], [372, 59, 407, 66], [74, 136, 161, 158], [175, 137, 306, 158], [318, 137, 363, 158], [374, 137, 472, 158], [483, 136, 529, 158], [540, 137, 593, 158], [608, 137, 717, 158], [73, 194, 100, 203], [106, 196, 177, 203], [183, 194, 227, 203], [233, 194, 259, 203], [265, 194, 344, 205], [74, 211, 104, 222], [109, 210, 141, 221], [147, 211, 169, 220], [175, 210, 223, 220], [229, 211, 259, 222], [265, 211, 329, 222], [334, 210, 352, 220], [74, 227, 127, 236], [133, 229, 180, 236], [187, 227, 221, 236], [226, 227, 264, 236], [270, 227, 320, 237], [327, 227, 349, 236], [74, 243, 161, 254], [166, 243, 249, 254], [254, 243, 281, 252], [286, 244, 342, 254], [74, 260, 112, 270], [119, 260, 145, 269], [151, 260, 174, 269], [179, 260, 217, 269], [222, 260, 249, 269], [254, 260, 285, 271], [290, 260, 335, 269], [340, 259, 359, 269], [74, 276, 95, 284], [101, 276, 156, 287], [164, 276, 198, 284], [203, 276, 244, 284], [251, 275, 285, 284], [291, 276, 340, 284], [74, 292, 129, 301], [135, 292, 185, 302], [192, 292, 242, 303], [248, 292, 261, 301], [267, 292, 312, 301], [74, 308, 195, 319], [75, 335, 82, 344], [88, 335, 98, 344], [105, 335, 138, 344], [144, 335, 214, 346], [220, 336, 233, 344], [239, 335, 256, 344], [262, 335, 283, 344], [290, 335, 309, 344], [316, 335, 320, 344], [74, 351, 119, 360], [126, 352, 170, 362], [176, 352, 186, 360], [192, 352, 214, 360], [220, 352, 276, 362], [282, 352, 326, 360], [333, 352, 349, 362], [74, 368, 89, 377], [95, 370, 124, 377], [129, 367, 175, 377], [181, 368, 266, 377], [272, 368, 283, 376], [289, 368, 333, 377], [74, 384, 126, 393], [134, 385, 175, 395], [181, 384, 206, 393], [212, 384, 292, 395], [298, 384, 325, 393], [330, 384, 366, 393], [74, 403, 103, 409], [109, 400, 154, 409], [161, 401, 241, 409], [247, 403, 269, 409], [275, 401, 296, 409], [302, 400, 349, 409], [74, 417, 131, 428], [137, 419, 186, 428], [192, 417, 214, 426], [219, 417, 242, 428], [248, 419, 319, 426], [74, 433, 119, 444], [125, 433, 204, 444], [210, 433, 278, 444], [285, 433, 295, 441], [302, 433, 340, 442], [75, 449, 98, 458], [104, 449, 142, 458], [146, 449, 215, 460], [221, 449, 258, 460], [263, 449, 293, 459], [300, 449, 339, 460], [74, 466, 101, 474], [108, 466, 185, 476], [191, 466, 261, 474], [267, 466, 309, 476], [315, 466, 354, 474], [74, 482, 151, 491], [158, 482, 201, 491], [208, 482, 258, 491], [263, 482, 292, 491], [298, 482, 333, 491], [338, 482, 360, 491], [74, 498, 131, 507], [137, 498, 150, 507], [156, 498, 197, 509], [202, 498, 257, 507], [263, 498, 310, 509], [74, 515, 128, 525], [134, 515, 156, 523], [161, 515, 218, 523], [223, 515, 261, 525], [267, 514, 280, 523], [74, 531, 156, 540], [162, 531, 188, 540], [195, 531, 257, 540], [263, 531, 315, 542], [871, 199, 878, 202], [883, 199, 908, 202], [894, 251, 904, 257], [841, 268, 841, 270], [784, 373, 811, 378], [816, 373, 896, 378], [784, 381, 811, 387], [815, 381, 847, 387], [645, 908, 670, 915], [692, 908, 712, 915], [220, 984, 285, 993], [293, 983, 779, 996]] + # fmt: on + text_list = [] + bbox_list = [] + for text, box in zip(words, boxes): + if text == "": + continue + sub_tokens = tokenizer.tokenize(text) + for sub_token in sub_tokens: + text_list.append(sub_token) + bbox_list.append(box) + + input_ids = tokenizer.convert_tokens_to_ids(text_list) + + input_ids = prompt_ids + input_ids + bbox = [[0, 0, 0, 0]] * len(prompt_ids) + bbox_list + + pixel_values = image_processor(image, return_tensors="pt").pixel_values + original_pixel_values = original_transform(image, image_size=image_processor.size["height"]).unsqueeze(0) + # verify pixel values + assert torch.allclose(original_pixel_values, pixel_values) + print("Pixel values are ok!") + + return torch.tensor(input_ids).unsqueeze(0), torch.tensor(bbox).unsqueeze(0).float(), pixel_values + + +def convert_udop_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False): + # model_name to checkpoint_path + name_to_checkpoint_path = { + "udop-large": "/Users/nielsrogge/Documents/UDOP/udop-unimodel-large-224/pytorch_model.bin", + "udop-large-512": "/Users/nielsrogge/Documents/UDOP/udop-unimodel-large-512/pytorch_model.bin", + "udop-large-512-300k": "/Users/nielsrogge/Documents/UDOP/udop-unimodel-large-512-300k-steps/pytorch_model.bin", + } + + # load original state dict + checkpoint_path = name_to_checkpoint_path[model_name] + state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + print("Checkpoint path:", checkpoint_path) + + # create HF model + image_size = 512 if "512" in model_name else 224 + config = UdopConfig(decoder_start_token_id=0, image_size=image_size) + model = UdopForConditionalGeneration(config) + model.eval() + + # rename keys + state_dict = {k.replace("cell2dembedding", "cell_2d_embedding"): v for k, v in state_dict.items()} + + # load weights + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + assert missing_keys == ["encoder.embed_patches.proj.weight", "encoder.embed_patches.proj.bias"] + assert unexpected_keys == ["pos_embed"] + + # Add extra_ids to the special token list + # NOTE special tokens have a unique order + # see https://github.com/huggingface/transformers/issues/29591 for details + # fmt: off + additional_special_tokens = ['', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', ''] + # fmt: on + + tokenizer = UdopTokenizer.from_pretrained( + "/Users/nielsrogge/Documents/UDOP/udop-unimodel-large-512", + legacy=True, + additional_special_tokens=additional_special_tokens, + ) + size = {"height": image_size, "width": image_size} + image_processor = LayoutLMv3ImageProcessor( + image_mean=IMAGENET_DEFAULT_MEAN, image_std=IMAGENET_DEFAULT_STD, size=size + ) + processor = UdopProcessor(image_processor=image_processor, tokenizer=tokenizer) + + # prepare dummy inputs + input_ids, bbox, image = prepare_dummy_inputs(tokenizer, image_processor) + prompt = "Question answering. In which year is the report made?" + encoding = processor(images=get_image(), text=prompt, return_tensors="pt") + + input_ids = encoding.input_ids + try: + EXPECTED_INPUT_IDS = torch.tensor([[11860, 18243, 5, 86, 84, 215, 19, 8, 934, 263, 58, 1, 489, 27, 3838, 7363, 4083, 14536, 3430, 5686, 5911, 17161, 134, 2038, 27, 3838, 22, 7, 4688, 7, 10, 389, 18202, 21, 8, 11046, 37, 3733, 523, 11, 38, 2388, 1628, 3, 13133, 23334, 6, 8, 1656, 79, 3806, 21, 4040, 640, 27, 3838, 22, 7, 701, 16534, 6, 8, 3, 76, 2693, 18, 23015, 5644, 24, 380, 3, 6015, 6, 11, 8, 701, 24, 79, 482, 21, 3, 88, 684, 6, 43, 263, 27, 3838, 22, 7, 3635, 1157, 4089, 6, 2651, 12, 1547, 22, 7, 3265, 655, 5, 19, 27, 3838, 22, 7, 38, 2388, 257, 12, 36, 8, 465, 209, 13409, 12150, 1959, 16, 8, 684, 6, 6737, 57, 165, 126, 13409, 12150, 1623, 5, 71, 1100, 30298, 934, 65, 12566, 24, 27, 3838, 31, 7, 126, 13409, 12150, 1623, 33, 8, 10391, 1710, 859, 8, 420, 3733, 4968, 688, 2699, 16, 1547, 5, 27, 3838, 1217, 131, 99, 23, 179, 6064, 24, 6, 590, 28, 3, 11600, 1456, 701, 6, 175, 9443, 2557, 3635, 92, 1262, 8, 3409, 13, 2186, 3, 27908, 1784, 190, 8, 3, 5771, 17, 13281, 4005, 13, 5086, 11, 13066, 1170, 5, 10826, 16309, 134, 3, 2, 276, 26, 3, 55, 391, 13570, 5, 10315, 309, 3577, 19114, 371, 4254, 5121, 5055, 6245, 3, 10047, 3162, 58, 3, 9, 61, 1713, 2703, 476, 667, 25158, 301, 6058, 6038, 476, 3765, 9149, 10, 4893, 1303, 1986, 5, 13580, 7, 8224, 28244, 7, 5, 76, 75, 7, 89, 5, 15, 1259, 87, 7171, 7, 87, 7, 29, 115, 226, 4305, 2773, 1]]) # fmt: skip + torch.testing.assert_close(EXPECTED_INPUT_IDS, input_ids) + bbox = encoding.bbox.float() + pixel_values = encoding.pixel_values + except Exception: + print("Input_ids don't match, preparing dummy inputs") + input_ids, bbox, pixel_values = prepare_dummy_inputs(tokenizer, image_processor) + + # Verify single forward pass + print("Testing single forward pass..") + with torch.no_grad(): + decoder_input_ids = torch.tensor([[101]]) + outputs = model(input_ids=input_ids, bbox=bbox, pixel_values=pixel_values, decoder_input_ids=decoder_input_ids) + print("Shape of logits:", outputs.logits.shape) + print("First values of logits:", outputs.logits[0, :3, :3]) + + # tensor([[-18.5262, 1.5087, -15.7051]]) on linux + # tensor([[-19.4976, 0.8515, -17.1873]]) on mac + try: + assert torch.allclose(outputs.logits[0, :3, :3], torch.tensor([[-18.5262, 1.5087, -15.7051]]), atol=1e-4) + print("Looks ok!") + except Exception: + print("logits don't match let's try to generate") + + # Verify autoregressive decoding + print("Testing generation...") + model_kwargs = {"bbox": bbox, "pixel_values": pixel_values} + outputs = model.generate(input_ids=input_ids, **model_kwargs, max_new_tokens=20) + + print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) + + # autoregressive decoding with original input data + print("Testing generation with original inputs...") + filepath = hf_hub_download(repo_id="nielsr/test-image", filename="input_ids_udop.pt", repo_type="dataset") + input_ids = torch.load(filepath, weights_only=True) + filepath = hf_hub_download(repo_id="nielsr/test-image", filename="bbox_udop.pt", repo_type="dataset") + bbox = torch.load(filepath, weights_only=True) + pixel_values_filename = "pixel_values_udop_512.pt" if "512" in model_name else "pixel_values_udop_224.pt" + filepath = hf_hub_download(repo_id="nielsr/test-image", filename=pixel_values_filename, repo_type="dataset") + pixel_values = torch.load(filepath, weights_only=True) + + print("Decoded input ids:", tokenizer.decode(input_ids[0], skip_special_tokens=True)) + print("Bbox shape:", bbox.shape) + + model_kwargs = {"bbox": bbox, "pixel_values": pixel_values} + outputs = model.generate(input_ids=input_ids, **model_kwargs, max_new_tokens=20) + generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0] + print("Generated:", generated_text) + + if pytorch_dump_folder_path is not None: + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + model.push_to_hub(f"microsoft/{model_name}") + processor.push_to_hub(f"microsoft/{model_name}") + # BIG note here: to save the fast tokenizer files in the repo on the hub, you need to do the following: + # see https://discuss.huggingface.co/t/convert-slow-xlmrobertatokenizer-to-fast-one/20876 + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="udop-large", + type=str, + choices=["udop-large", "udop-large-512", "udop-large-512-300k"], + help=("Name of the UDOP model you'd like to convert."), + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory." + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + help="Whether or not to push the converted model to the Hugging Face hub.", + ) + + args = parser.parse_args() + convert_udop_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/udop/modeling_udop.py b/third_party/transformers/src/transformers/models/udop/modeling_udop.py new file mode 100644 index 0000000000000000000000000000000000000000..3cf52f9b2f6a3555705a2d06348ff1a14a4a71a9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/modeling_udop.py @@ -0,0 +1,1726 @@ +# Copyright 2024 Microsoft Research and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch UDOP model.""" + +import collections +import logging +import math +import random +from abc import ABC, abstractmethod +from collections.abc import Sequence +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor, nn +from torch.nn import CrossEntropyLoss + +from transformers import UdopConfig +from transformers.modeling_outputs import ( + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_utils import PreTrainedModel +from ...utils import ( + ModelOutput, + auto_docstring, + is_torchdynamo_compiling, +) + + +logger = logging.getLogger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for the model's outputs that may also contain a past key/values (to speed up sequential decoding). Includes + an additional attention mask. + """ +) +class BaseModelOutputWithAttentionMask(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. If `past_key_values` is used only + the last hidden-state of the sequences of shape `(batch_size, 1, hidden_size)` is output. + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Attention mask used in the model's forward pass 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**. + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the + self-attention blocks and optionally if `config.is_encoder_decoder=True` in the cross-attention blocks) + that can be used (see `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of + the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in + the self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, + used to compute the weighted average in the cross-attention heads. + """ + + last_hidden_state: torch.FloatTensor | None = None + attention_mask: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + cross_attentions: tuple[torch.FloatTensor] | None = None + + +def get_visual_bbox(image_size=224, patch_size=16): + image_feature_pool_shape = [image_size // patch_size, image_size // patch_size] + visual_bbox_x = torch.arange(0, 1.0 * (image_feature_pool_shape[1] + 1), 1.0) + visual_bbox_x /= image_feature_pool_shape[1] + + visual_bbox_y = torch.arange(0, 1.0 * (image_feature_pool_shape[0] + 1), 1.0) + visual_bbox_y /= image_feature_pool_shape[0] + + visual_bbox_input = torch.stack( + [ + visual_bbox_x[:-1].repeat(image_feature_pool_shape[0], 1), + visual_bbox_y[:-1].repeat(image_feature_pool_shape[1], 1).transpose(0, 1), + visual_bbox_x[1:].repeat(image_feature_pool_shape[0], 1), + visual_bbox_y[1:].repeat(image_feature_pool_shape[1], 1).transpose(0, 1), + ], + dim=-1, + ) + + visual_bbox_input = visual_bbox_input.view(-1, 4) + + return visual_bbox_input + + +def pad_sequence(seq, target_len, pad_value=0): + if isinstance(seq, torch.Tensor): + n = seq.shape[0] + else: + n = len(seq) + seq = torch.tensor(seq) + m = target_len - n + if m > 0: + ret = torch.stack([pad_value] * m).to(seq) + seq = torch.cat([seq, ret], dim=0) + return seq[:target_len] + + +def combine_image_text_embeddings( + image_embeddings, + inputs_embeds, + bbox, + visual_bbox, + attention_mask=None, + num_patches=14, + max_len=0, + image_size=224, + patch_size=16, +): + """ + Combine the image and text embeddings for the input to the encoder/decoder of UDOP. + + First, the image embeddings are created by checking for each visual patch if it is inside the bounding box of a + token. If it is, the visual patch is combined with the token embedding. Then, the visual bounding boxes are combined + with the text bounding boxes. Finally, the visual bounding boxes are combined with the text attention mask. + """ + + sequence_length = num_patches + ocr_points_x = torch.clip( + torch.floor((bbox[:, :, 0] + bbox[:, :, 2]) / 2.0 * sequence_length).long(), 0, sequence_length - 1 + ) + ocr_points_y = ( + torch.clip(torch.floor((bbox[:, :, 1] + bbox[:, :, 3]) / 2.0 * sequence_length).long(), 0, sequence_length - 1) + * sequence_length + ) + ocr_points = ocr_points_x + ocr_points_y + # make sure bounding boxes are of type float to calculate means + bbox = bbox.to(torch.float64) + target_seg = (bbox.mean(-1) == 0.0) | (bbox.mean(-1) == 1.0) + repeated_vision_embeds = torch.gather( + image_embeddings, 1, ocr_points.unsqueeze(-1).repeat(1, 1, image_embeddings.size(-1)) + ) + repeated_vision_embeds[target_seg] = 0.0 + inputs_embeds += repeated_vision_embeds + + patch_inds = torch.full_like(image_embeddings[:, :, 0], True).bool() + ind = torch.cat( + [ + torch.arange(len(ocr_points))[:, None].repeat(1, ocr_points.size(-1))[:, :, None].to(ocr_points), + ocr_points[:, :, None], + ], + dim=-1, + ) + ind = ind.flatten(0, 1) + rows, cols = zip(*ind) + patch_inds[rows, cols] = False + + input_vision_patches = [image_embeddings[i][patch_inds[i]] for i in range(len(patch_inds))] + + if visual_bbox is None: + visual_bbox = get_visual_bbox(image_size=image_size, patch_size=patch_size) + visual_bbox = visual_bbox.unsqueeze(0).repeat(image_embeddings.size(0), 1, 1) + visual_bbox = visual_bbox.to(image_embeddings.device) + + visual_bbox = [visual_bbox[i][patch_inds[i]] for i in range(len(patch_inds))] + + if attention_mask is not None: + visual_attention_mask = [ + torch.ones(item.size(0), dtype=attention_mask.dtype, device=attention_mask.device) for item in visual_bbox + ] + + if max_len == 0: + max_len = image_embeddings.size(1) + else: + max_len = max_len - inputs_embeds.size(1) + inputs_vision_patches = torch.stack( + [pad_sequence(item, max_len, torch.zeros_like(image_embeddings[0, 0])) for item in input_vision_patches] + ) + visual_bbox = torch.stack([pad_sequence(item, max_len, torch.zeros_like(bbox[0, 0])) for item in visual_bbox]) + if attention_mask is not None: + visual_attention_mask = torch.stack( + [pad_sequence(item, max_len, torch.zeros_like(attention_mask[0, 0])) for item in visual_attention_mask] + ) + + inputs_embeds = torch.cat([inputs_embeds, inputs_vision_patches], 1) + bbox = torch.cat([bbox, visual_bbox], 1) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask, visual_attention_mask], 1) + return inputs_embeds, bbox, attention_mask + + +class UdopPatchEmbeddings(nn.Module): + """2D Image to Patch Embeddings""" + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.proj = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values): + batch_size, num_channels, height, width = pixel_values.shape + if height != self.image_size[0] or width != self.image_size[1]: + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})." + ) + embeddings = self.proj(pixel_values) + embeddings = embeddings.flatten(2).transpose(1, 2) + return embeddings + + +@auto_docstring +class UdopPreTrainedModel(PreTrainedModel): + config: UdopConfig + base_model_prefix = "transformer" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + + _can_compile_fullgraph = False + _keep_in_fp32_modules = ["wo"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor # Used for testing weights initialization + if isinstance(module, UdopLayerNorm): + init.constant_(module.weight, factor * 1.0) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=factor) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + elif isinstance(module, nn.Conv2d): + init.trunc_normal_(module.weight, mean=0.0, std=factor) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, RelativePositionBiasBase): + factor = self.config.initializer_factor + d_model = self.config.d_model + init.normal_(module.relative_attention_bias.weight, mean=0.0, std=factor * ((d_model) ** -0.5)) + elif isinstance(module, UdopModel): + init.normal_(module.shared.weight, mean=0.0, std=factor * 1.0) + elif isinstance(module, UdopForConditionalGeneration): + if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: + init.normal_(module.lm_head.weight, mean=0.0, std=factor * 1.0) + elif isinstance(module, UdopDenseActDense): + init.normal_(module.wi.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi, "bias") and module.wi.bias is not None: + init.zeros_(module.wi.bias) + init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + init.zeros_(module.wo.bias) + elif isinstance(module, UdopDenseGatedActDense): + init.normal_(module.wi_0.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: + init.zeros_(module.wi_0.bias) + init.normal_(module.wi_1.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) + if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: + init.zeros_(module.wi_1.bias) + init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + init.zeros_(module.wo.bias) + elif isinstance(module, UdopAttention): + d_model = self.config.d_model + key_value_proj_dim = self.config.d_kv + n_heads = self.config.num_heads + init.normal_(module.q.weight, mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) + init.normal_(module.k.weight, mean=0.0, std=factor * (d_model**-0.5)) + init.normal_(module.v.weight, mean=0.0, std=factor * (d_model**-0.5)) + init.normal_(module.o.weight, mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) + if module.has_relative_attention_bias: + init.normal_(module.relative_attention_bias.weight, mean=0.0, std=factor * ((d_model) ** -0.5)) + + # Copied from transformers.models.prophetnet.modeling_prophetnet.ProphetNetPreTrainedModel._shift_right with ProphetNet->Udop + def _shift_right(self, input_ids): + decoder_start_token_id = self.config.decoder_start_token_id + pad_token_id = self.config.pad_token_id + + assert decoder_start_token_id is not None, ( + "self.model.config.decoder_start_token_id has to be defined. In Udop it is usually set to the" + " pad_token_id. See Udop docs for more information" + ) + + # shift inputs to the right + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = decoder_start_token_id + + assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined." + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + assert torch.all(shifted_input_ids >= 0).item(), "Verify that `shifted_input_ids` has only positive values" + + return shifted_input_ids + + +# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->Udop +class UdopLayerNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Construct a layernorm module in the Udop style. No bias and no subtraction of mean. + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # Udop uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean + # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated + # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for + # half-precision inputs is done in fp32 + + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->Udop +class UdopDenseActDense(nn.Module): + def __init__(self, config: UdopConfig): + super().__init__() + self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_states = self.wi(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.dropout(hidden_states) + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + hidden_states = self.wo(hidden_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->Udop +class UdopDenseGatedActDense(nn.Module): + def __init__(self, config: UdopConfig): + super().__init__() + self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->Udop +class UdopLayerFF(nn.Module): + def __init__(self, config: UdopConfig): + super().__init__() + if config.is_gated_act: + self.DenseReluDense = UdopDenseGatedActDense(config) + else: + self.DenseReluDense = UdopDenseActDense(config) + + self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, hidden_states): + forwarded_states = self.layer_norm(hidden_states) + forwarded_states = self.DenseReluDense(forwarded_states) + hidden_states = hidden_states + self.dropout(forwarded_states) + return hidden_states + + +# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->Udop +class UdopAttention(nn.Module): + def __init__( + self, + config: UdopConfig, + has_relative_attention_bias=False, + layer_idx: int | None = None, + ): + super().__init__() + self.is_decoder = config.is_decoder + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + self.d_model = config.d_model + self.key_value_proj_dim = config.d_kv + self.n_heads = config.num_heads + self.dropout = config.dropout_rate + self.inner_dim = self.n_heads * self.key_value_proj_dim + self.layer_idx = layer_idx + if layer_idx is None and self.is_decoder: + logger.warning_once( + f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and " + "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) + + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) + + self.gradient_checkpointing = False + + @staticmethod + def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + + Translate relative position to a bucket number for relative attention. The relative position is defined as + memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to + position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for + small absolute relative_position and larger buckets for larger absolute relative_positions. All relative + positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. + This should allow for more graceful generalization to longer sequences than the model has been trained on + + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + + Returns: + a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) + """ + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) + # now relative_position is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) + ) + + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) + return relative_buckets + + def compute_bias(self, query_length, key_length, device=None, past_seen_tokens=0): + """Compute binned relative position bias""" + if device is None: + device = self.relative_attention_bias.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + past_seen_tokens + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] + relative_position = memory_position - context_position # shape (query_length, key_length) + relative_position_bucket = self._relative_position_bucket( + relative_position, # shape (query_length, key_length) + bidirectional=(not self.is_decoder), + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) + values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) + return values + + def forward( + self, + hidden_states, + mask=None, + key_value_states=None, + position_bias=None, + past_key_values=None, + output_attentions=False, + **kwargs, + ): + """ + Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). + """ + # Input is (batch_size, seq_length, dim) + # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder) + batch_size, seq_length = hidden_states.shape[:2] + past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0 + # We clone here for StaticCache, as we get the value before updating it, but use it after and it's the same ref + past_seen_tokens = past_seen_tokens.clone() if isinstance(past_seen_tokens, torch.Tensor) else past_seen_tokens + + # if key_value_states are provided this layer is used as a cross-attention layer for the decoder + is_cross_attention = key_value_states is not None + + q_input_shape = (batch_size, seq_length, -1, self.key_value_proj_dim) + query_states = self.q(hidden_states).view(*q_input_shape).transpose(1, 2) + + # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` + is_updated = False + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache + else: + curr_past_key_values = past_key_values + + current_states = key_value_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = curr_past_key_values.layers[self.layer_idx].keys + value_states = curr_past_key_values.layers[self.layer_idx].values + else: + kv_input_shape = (batch_size, current_states.shape[1], -1, self.key_value_proj_dim) + key_states = self.k(current_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v(current_states).view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + + # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 + scores = torch.matmul(query_states, key_states.transpose(3, 2)) + + if position_bias is None: + key_length = key_states.shape[-2] + if not self.has_relative_attention_bias: + position_bias = torch.zeros( + (1, query_states.shape[1], seq_length, key_length), device=scores.device, dtype=scores.dtype + ) + if self.gradient_checkpointing and self.training: + position_bias.requires_grad = True + else: + position_bias = self.compute_bias( + seq_length, key_length, device=scores.device, past_seen_tokens=past_seen_tokens + ) + + if mask is not None: + causal_mask = mask[:, :, :, : key_states.shape[-2]] + position_bias = position_bias + causal_mask + + position_bias_masked = position_bias + scores += position_bias_masked + + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(batch_size, seq_length, -1) + attn_output = self.o(attn_output) + + outputs = (attn_output, position_bias) + + if output_attentions: + outputs = outputs + (attn_weights,) + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->Udop +class UdopLayerSelfAttention(nn.Module): + def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): + super().__init__() + self.SelfAttention = UdopAttention( + config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx + ) + self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + past_key_values=None, + use_cache=False, + output_attentions=False, + **kwargs, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.SelfAttention( + normed_hidden_states, + mask=attention_mask, + position_bias=position_bias, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = hidden_states + self.dropout(attention_output[0]) + outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->Udop +class UdopLayerCrossAttention(nn.Module): + def __init__(self, config, layer_idx: int | None = None): + super().__init__() + self.EncDecAttention = UdopAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + key_value_states, + attention_mask=None, + position_bias=None, + past_key_values=None, + output_attentions=False, + **kwargs, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.EncDecAttention( + normed_hidden_states, + mask=attention_mask, + key_value_states=key_value_states, + position_bias=position_bias, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + layer_output = hidden_states + self.dropout(attention_output[0]) + outputs = (layer_output,) + attention_output[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5Block with T5->Udop +class UdopBlock(GradientCheckpointingLayer): + def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): + super().__init__() + self.is_decoder = config.is_decoder + self.layer = nn.ModuleList() + self.layer.append( + UdopLayerSelfAttention( + config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx + ) + ) + if self.is_decoder: + self.layer.append(UdopLayerCrossAttention(config, layer_idx=layer_idx)) + + self.layer.append(UdopLayerFF(config)) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + encoder_decoder_position_bias=None, + past_key_values=None, + use_cache=False, + output_attentions=False, + return_dict=True, + **kwargs, + ): + self_attention_outputs = self.layer[0]( + hidden_states, + attention_mask=attention_mask, + position_bias=position_bias, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = self_attention_outputs[0] + attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + do_cross_attention = self.is_decoder and encoder_hidden_states is not None + if do_cross_attention: + cross_attention_outputs = self.layer[1]( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_bias=encoder_decoder_position_bias, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + hidden_states = cross_attention_outputs[0] + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + # Keep cross-attention outputs and relative position weights + attention_outputs = attention_outputs + cross_attention_outputs[1:] + + # Apply Feed Forward layer + hidden_states = self.layer[-1](hidden_states) + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + return ( + outputs + attention_outputs + ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + + +class UdopCellEmbeddings(nn.Module): + def __init__(self, max_2d_position_embeddings=501, hidden_size=1024): + super().__init__() + self.max_2d_position_embeddings = max_2d_position_embeddings + + self.x_position_embeddings = nn.Embedding(max_2d_position_embeddings, hidden_size) + self.y_position_embeddings = nn.Embedding(max_2d_position_embeddings, hidden_size) + + def forward(self, bbox): + bbox = torch.clip(bbox, 0.0, 1.0) + bbox = (bbox * (self.max_2d_position_embeddings - 1)).long() + left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0]) + upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1]) + right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2]) + lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3]) + + embeddings = ( + left_position_embeddings + + upper_position_embeddings + + right_position_embeddings + + lower_position_embeddings + ) + + return embeddings + + +# get function for bucket computation +# protected member access seems to be lesser evil than copy paste whole function +get_relative_position_bucket = UdopAttention._relative_position_bucket +AUGMENTATION_RANGE = (0.80, 1.25) + + +class RelativePositionBiasBase(nn.Module, ABC): + """ + Base class of relative biases. + + Args: + num_heads (`int`): + Number of attention heads in the model, it will create embeddings of size `num_heads`, which will be added to the scores of each token pair. + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + Pair token metric (distance in the sequence, distance in pixels etc.) will be bucketed, parameter is defining number of such + buckets. + bidirectional (`bool`, *optional*, defaults to `True`): + Whether the distance should be bidirectional for a pair of tokens. If `False`, then distance(tok1, tok2) == distance(tok2, tok1). + scaling_factor (`int`, *optional*, defaults to 1): + Defining factor which will be used to scale relative distance. + max_distance (`int`, *optional*, defaults to 128): + All distances above this value will end up in the one/same bucket. + augmentation (`bool`, *optional*, defaults to `False`): + Whether to multiply relative distances by a random scalar. + expand (`bool`, *optional*, defaults to `False`): + Whether to expand an existing pretrained model with subsequent additions of prefix_bucket. + """ + + def __init__( + self, + num_heads=None, + relative_attention_num_buckets=32, + bidirectional=True, + scaling_factor=1, + max_distance=128, + level="tokens", + augmentation=False, + prefix_bucket=False, + expand=False, + ): + super().__init__() + self.prefix_bucket = prefix_bucket + self.augmentation = augmentation + self.level = level + self.max_distance = max_distance + self.scaling_factor = scaling_factor + self.bidirectional = bidirectional + self.num_heads = num_heads + self.expand = expand + self.relative_attention_num_buckets = relative_attention_num_buckets + extra_head = 2 if prefix_bucket and not self.expand else 0 + self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets + extra_head, self.num_heads) + + @abstractmethod + def prepare_input( + self, + attention_mask: Tensor | None = None, + bbox: dict[str, Any] | None = None, + ) -> Tensor: + pass + + def get_bucket(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> Tensor: + relative_position = self.prepare_input(attention_mask, bbox) + rp_bucket: Tensor = get_relative_position_bucket( + relative_position, + bidirectional=self.bidirectional, + num_buckets=self.relative_attention_num_buckets, + max_distance=self.max_distance, + ) + return rp_bucket + + def get_relative_position(self, positions): + context_position = positions[:, :, None] + memory_position = positions[:, None, :] + relative_position = memory_position - context_position + if self.augmentation and self.training: + relative_position *= random.uniform(*AUGMENTATION_RANGE) + relative_position *= self.scaling_factor + + return relative_position.to(torch.long) + + def forward(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> Tensor: + # re-using pretrained model with subsequent addition of prefix_bucket + if self.expand and self.prefix_bucket: + new_bias = nn.Embedding(self.relative_attention_num_buckets + 2, self.num_heads) + new_bias.weight.data[: self.relative_attention_num_buckets] = self.relative_attention_bias.weight.data + new_bias.weight.data[self.relative_attention_num_buckets :] = 0.1 + self.relative_attention_bias = new_bias + self.expand = False + + rp_bucket = self.get_bucket(attention_mask, bbox) + + if self.prefix_bucket: + if rp_bucket.size(0) == 1 and attention_mask.size(0) > 1: + rp_bucket = rp_bucket.repeat(attention_mask.size(0), 1, 1) + # based on assumption that prefix bboxes are negative + is_prefix = bbox[:, :, 1] < 0 + num_prefix = is_prefix.sum(-1) + for idx, num_prefix_row in enumerate(num_prefix.cpu().numpy()): + rp_bucket[idx, :num_prefix_row, num_prefix_row:] = self.relative_attention_num_buckets + rp_bucket[idx, num_prefix_row:, :num_prefix_row] = self.relative_attention_num_buckets + 1 + + values: Tensor = self.relative_attention_bias(rp_bucket) + if values.dim() != 4: + raise ValueError("Wrong dimension of values tensor") + values = values.permute([0, 3, 1, 2]) + + return values + + +class RelativePositionBias1D(RelativePositionBiasBase): + def __init__(self, scaling_factor=1, max_distance=128, **kwargs): + """ + Reimplementation of T5 relative position bias. Distance between given tokens is their distance in the sequence. + Parameters are the same as in base class + """ + super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs) + + def prepare_input(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> Tensor: + if self.scaling_factor != 1: + raise ValueError("No need to scale 1d features") + relative_position = self.get_relative_position( + torch.arange(attention_mask.size(1), dtype=torch.long, device=attention_mask.device)[None, :] + ) + + return relative_position + + +class RelativePositionBiasHorizontal(RelativePositionBiasBase): + def __init__(self, scaling_factor=100, max_distance=100, **kwargs): + """ + Represents in the bucket embeddings horizontal distance between two tokens. Parameters are the same as in base + class + """ + super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs) + + def prepare_input(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> Tensor: + if not self.scaling_factor > 1.0: + raise ValueError("Need to scale the values of bboxes, as there are in small (0,1) range") + if bbox is None: + raise ValueError("Bbox is required for horizontal relative position bias") + # get x positions of left point of bbox + horizontal_position: Tensor = bbox[:, :, [0, 2]].mean(dim=-1) + + return self.get_relative_position(horizontal_position) + + +class RelativePositionBiasVertical(RelativePositionBiasBase): + def __init__(self, scaling_factor=100, max_distance=100, **kwargs): + """ + Represents in the bucket embeddings vertical distance between two tokens. Parameters are the same as in base + class + """ + super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs) + + def prepare_input(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> Tensor: + if not self.scaling_factor > 1.0: + raise ValueError("Need to scale the values of bboxes, as there are in small (0,1) range") + if bbox is None: + raise ValueError("Bbox is required for vertical relative position bias") + # get y positions of middle of bbox + vertical_position: Tensor = bbox[:, :, [1, 3]].mean(dim=-1) + + return self.get_relative_position(vertical_position) + + +class RelativePositionBiasAggregated(nn.Module): + def __init__(self, modules: Sequence[RelativePositionBiasBase]): + """ + Class which sums up various computed biases. + + Args: + modules (Sequence[RelativePositionBiasBase]): + List of relative bias modules. + """ + super().__init__() + self.biases = nn.ModuleList(modules) + + def forward(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] | None = None) -> float | Tensor: + output = 0.0 + for bias in self.biases: # type: ignore + output = bias(attention_mask, bbox) + output + + return output + + +BIAS_CLASSES = { + "1d": RelativePositionBias1D, + "horizontal": RelativePositionBiasHorizontal, + "vertical": RelativePositionBiasVertical, +} + + +def create_relative_bias(config: UdopConfig) -> Sequence[RelativePositionBiasBase]: + """ + Creates empty list or one/multiple relative biases. + + :param config: Model's configuration :return: Sequence with created bias modules. + """ + bias_list = [] + if hasattr(config, "relative_bias_args"): + for bias_kwargs_org in config.relative_bias_args: + bias_kwargs = deepcopy(bias_kwargs_org) + bias_type = bias_kwargs.pop("type") + model_num_heads = config.num_heads if hasattr(config, "num_heads") else config.num_attention_heads + if "num_heads" in bias_kwargs: + if bias_kwargs["num_heads"] != model_num_heads: + raise ValueError("Number of heads must match num of heads in the model") + else: + bias_kwargs["num_heads"] = model_num_heads + bias_list.append(BIAS_CLASSES[bias_type](**bias_kwargs)) # type: ignore + + return bias_list + + +class UdopStack(UdopPreTrainedModel): + """ + This class is based on `T5Stack`, but modified to take into account the image modality as well as 2D position + embeddings. + """ + + def __init__(self, config): + super().__init__(config) + # text and image embeddings + self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model) + self.embed_patches = UdopPatchEmbeddings(config) + self.is_decoder = config.is_decoder + self.num_layers = config.num_layers + + self.block = nn.ModuleList( + [UdopBlock(config, has_relative_attention_bias=bool(i == 0), layer_idx=i) for i in range(self.num_layers)] + ) + self.final_layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) + + self.dropout = nn.Dropout(config.dropout_rate) + + if not self.is_decoder: + self.cell_2d_embedding = UdopCellEmbeddings(config.max_2d_position_embeddings, config.hidden_size) + + # get weights from encoder position bias + self.relative_bias = self._get_relative_bias(config) + self.post_init() + + @staticmethod + def _get_relative_bias(config: UdopConfig) -> RelativePositionBiasAggregated: + relative_bias_list = create_relative_bias(config) + return RelativePositionBiasAggregated(relative_bias_list) + + def get_output_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, new_embeddings): + self.embed_tokens = new_embeddings + + def forward( + self, + input_ids=None, + attention_mask=None, + bbox=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + inputs_embeds=None, + pixel_values=None, + visual_bbox=None, + image_embeddings=None, + position_bias=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ) -> tuple | BaseModelOutputWithAttentionMask: + use_cache = use_cache if use_cache is not None else self.config.use_cache + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # input embeddings processing + + if input_ids is not None and inputs_embeds is not None: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError( + f"You cannot specify both {err_msg_prefix}inputs and {err_msg_prefix}inputs_embeds at the same time" + ) + elif input_ids is not None and torch.numel(input_ids) > 0: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is None and input_ids is not None and torch.numel(input_ids) == 0: + input_ids = torch.full((4, 1024), self.config.pad_token_id, device=input_ids.device, dtype=input_ids.dtype) + attention_mask = torch.zeros((4, 1024), device=input_ids.device, dtype=input_ids.dtype) + bbox = torch.zeros((4, 1024, 4), device=input_ids.device, dtype=input_ids.dtype) + input_shape = input_ids.size() + position_bias = torch.zeros_like(self.get_extended_attention_mask(attention_mask, input_shape)) + # encoder_attention_mask = attention_mask + logger.warning("Empty batch") + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError(f"You have to specify either {err_msg_prefix}inputs or {err_msg_prefix}inputs_embeds") + + if inputs_embeds is None: + if self.embed_tokens is None: + raise ValueError("You have to initialize the model with valid token embeddings") + inputs_embeds = self.embed_tokens(input_ids) + + if pixel_values is not None: + image_embeddings = self.embed_patches(pixel_values) + + if image_embeddings is not None: + # combine visual and OCR text embeddings + num_patches = self.config.image_size // self.config.patch_size + inputs_embeds, bbox, attention_mask = combine_image_text_embeddings( + image_embeddings, + inputs_embeds, + bbox, + visual_bbox, + attention_mask, + num_patches, + 0, + self.config.image_size, + self.config.patch_size, + ) + input_shape = inputs_embeds.size()[:-1] + + if not self.is_decoder and bbox is not None: + inputs_embeds += self.cell_2d_embedding(bbox) + + batch_size, seq_length = input_shape + + if use_cache is True: + assert self.is_decoder, f"`use_cache` can only be set to `True` if {self} is used as a decoder" + + if self.is_decoder: + if use_cache and past_key_values is None: + if self.config.is_encoder_decoder: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + else: + past_key_values = DynamicCache(config=self.config) + elif not self.is_decoder: + # do not pass cache object down the line for encoder stack + # it messes indexing later in decoder-stack because cache object is modified in-place + past_key_values = None + + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if attention_mask is None and not is_torchdynamo_compiling(): + # required mask seq length can be calculated via length of past cache + mask_seq_length = past_key_values_length + seq_length + attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) + + if self.config.is_decoder: + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + ) + else: + causal_mask = attention_mask[:, None, None, :] + causal_mask = causal_mask.to(dtype=inputs_embeds.dtype) + causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min + + if self.is_decoder and encoder_attention_mask is not None: + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + all_cross_attentions = () if (output_attentions and self.is_decoder) else None + + if self.is_decoder: # modified lines + position_bias = None + else: + position_bias = self.relative_bias(attention_mask=attention_mask, bbox=bbox) + position_bias = position_bias + causal_mask + encoder_decoder_position_bias = None + + hidden_states = inputs_embeds + + hidden_states = self.dropout(hidden_states) + + for i, layer_module in enumerate(self.block): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_outputs = layer_module( + hidden_states, + causal_mask, + position_bias, + encoder_hidden_states, + encoder_extended_attention_mask, + encoder_decoder_position_bias, # as a positional argument for gradient checkpointing + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + # We share the position biases between the layers - the first layer store them + # layer_outputs = hidden-states, key-value-states (self-attention weights), + # (self-attention position bias), (cross-attention weights), (cross-attention position bias) + + position_bias = layer_outputs[1] + if self.is_decoder and encoder_hidden_states is not None: + encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[2],) # We keep only self-attention weights for now + if self.is_decoder: + all_cross_attentions = all_cross_attentions + (layer_outputs[4],) + + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + # Add last layer + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + attention_mask, + past_key_values, + all_hidden_states, + all_attentions, + all_cross_attentions, + ] + if v is not None + ) + + return BaseModelOutputWithAttentionMask( + last_hidden_state=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attentions, + cross_attentions=all_cross_attentions, + ) + + +@auto_docstring +class UdopModel(UdopPreTrainedModel): + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + "encoder.embed_patches.proj.weight": "patch_embed.proj.weight", + "encoder.embed_patches.proj.bias": "patch_embed.proj.bias", + } + + def __init__(self, config): + super().__init__(config) + + # text and image embeddings + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.patch_embed = UdopPatchEmbeddings(config) + + encoder_config = deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + self.encoder = UdopStack(encoder_config) + + decoder_config = deepcopy(config) + decoder_config.is_decoder = True + decoder_config.num_layers = config.num_decoder_layers + self.decoder = UdopStack(decoder_config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + bbox: dict[str, Any] | None = None, + pixel_values: Tensor | None = None, + visual_bbox: dict[str, Any] | None = None, + decoder_input_ids: Tensor | None = None, + decoder_attention_mask: Tensor | None = None, + inputs_embeds: Tensor | None = None, + encoder_outputs: Tensor | None = None, + past_key_values: Cache | None = None, + decoder_inputs_embeds: Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | Seq2SeqModelOutput: + r""" + bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + visual_bbox (`torch.LongTensor` of shape `(batch_size, patch_sequence_length, 4)`, *optional*): + Bounding boxes of each patch in the image. If not provided, bounding boxes are created in the model. + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using + [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting + token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last + `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare + `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + + Example: + + ```python + >>> from transformers import AutoProcessor, AutoModel + >>> from datasets import load_dataset + >>> import torch + + >>> # load model and processor + >>> # in this case, we already have performed OCR ourselves + >>> # so we initialize the processor with `apply_ocr=False` + >>> processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=False) + >>> model = AutoModel.from_pretrained("microsoft/udop-large") + + >>> # load an example image, along with the words and coordinates + >>> # which were extracted using an OCR engine + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + >>> inputs = processor(image, words, boxes=boxes, return_tensors="pt") + + >>> decoder_input_ids = torch.tensor([[model.config.decoder_start_token_id]]) + + >>> # forward pass + >>> outputs = model(**inputs, decoder_input_ids=decoder_input_ids) + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 1, 1024] + ```""" + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + bbox=bbox, + pixel_values=pixel_values, + visual_bbox=visual_bbox, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + encoder_attention_mask = encoder_outputs.attention_mask if return_dict else encoder_outputs[1] + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if not return_dict: + # we filter out the attention mask + decoder_outputs = tuple(value for idx, value in enumerate(decoder_outputs) if idx != 1) + encoder_outputs = tuple(value for idx, value in enumerate(encoder_outputs) if idx != 1) + return decoder_outputs + encoder_outputs + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The UDOP encoder-decoder Transformer with a language modeling head on top, enabling to generate text given document + images and an optional prompt. + + This class is based on [`T5ForConditionalGeneration`], extended to deal with images and layout (2D) data. + """ +) +class UdopForConditionalGeneration(UdopPreTrainedModel, GenerationMixin): + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + "encoder.embed_patches.proj.weight": "patch_embed.proj.weight", + "encoder.embed_patches.proj.bias": "patch_embed.proj.bias", + "encoder.relative_bias.biases.0.relative_attention_bias.weight": "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", + "decoder.relative_bias.biases.0.relative_attention_bias.weight": "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", + "lm_head.weight": "shared.weight", + } + + def __init__(self, config): + super().__init__(config) + + # text and image embeddings + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.patch_embed = UdopPatchEmbeddings(config) + + encoder_config = deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + self.encoder = UdopStack(encoder_config) + + decoder_config = deepcopy(config) + decoder_config.is_decoder = True + decoder_config.num_layers = config.num_decoder_layers + self.decoder = UdopStack(decoder_config) + + # The weights of the language modeling head are shared with those of the encoder and decoder + self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + bbox: dict[str, Any] | None = None, + pixel_values: Tensor | None = None, + visual_bbox: dict[str, Any] | None = None, + decoder_input_ids: Tensor | None = None, + decoder_attention_mask: Tensor | None = None, + inputs_embeds: Tensor | None = None, + encoder_outputs: Tensor | None = None, + past_key_values: Cache | None = None, + decoder_inputs_embeds: Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: Tensor | None = None, + **kwargs, + ) -> tuple | Seq2SeqLMOutput: + r""" + bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + visual_bbox (`torch.LongTensor` of shape `(batch_size, patch_sequence_length, 4)`, *optional*): + Bounding boxes of each patch in the image. If not provided, bounding boxes are created in the model. + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using + [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting + token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last + `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare + `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size - + 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., + config.vocab_size]`. + + Examples: + + ```python + >>> from transformers import AutoProcessor, UdopForConditionalGeneration + >>> from datasets import load_dataset + + >>> # load model and processor + >>> # in this case, we already have performed OCR ourselves + >>> # so we initialize the processor with `apply_ocr=False` + >>> processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=False) + >>> model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large") + + >>> # load an example image, along with the words and coordinates + >>> # which were extracted using an OCR engine + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + + >>> # one can use the various task prefixes (prompts) used during pre-training + >>> # e.g. the task prefix for DocVQA is "Question answering. " + >>> question = "Question answering. What is the date on the form?" + >>> encoding = processor(image, question, text_pair=words, boxes=boxes, return_tensors="pt") + + >>> # autoregressive generation + >>> predicted_ids = model.generate(**encoding) + >>> print(processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]) + 9/30/92 + ```""" + + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if decoder_input_ids is None and labels is not None: + decoder_input_ids = self._shift_right(labels) + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + bbox=bbox, + visual_bbox=visual_bbox, + pixel_values=pixel_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + encoder_attention_mask = encoder_outputs.attention_mask if return_dict else encoder_outputs[1] + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = decoder_outputs[0] + + if self.config.tie_word_embeddings: + sequence_output = sequence_output * (self.config.d_model**-0.5) + + lm_logits = self.lm_head(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss(ignore_index=-100) + loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) + + if not return_dict: + output = (lm_logits,) + decoder_outputs[2:] + (encoder_outputs[0],) + encoder_outputs[2:] + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@auto_docstring +class UdopEncoderModel(UdopPreTrainedModel): + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "encoder.embed_patches.proj.weight": "patch_embed.proj.weight", + "encoder.embed_patches.proj.bias": "patch_embed.proj.bias", + "encoder.relative_bias.biases.0.relative_attention_bias.weight": "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", + } + + def __init__(self, config: UdopConfig): + super().__init__(config) + + # text and image embeddings + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.patch_embed = UdopPatchEmbeddings(config) + + encoder_config = deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = UdopStack(encoder_config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + + @auto_docstring + def forward( + self, + input_ids: Tensor | None = None, + bbox: dict[str, Any] | None = None, + attention_mask: Tensor | None = None, + pixel_values: Tensor | None = None, + visual_bbox: dict[str, Any] | None = None, + inputs_embeds: Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor] | BaseModelOutputWithAttentionMask: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). + bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. + + Note that `sequence_length = token_sequence_length + patch_sequence_length + 1` where `1` is for [CLS] + token. See `pixel_values` for `patch_sequence_length`. + visual_bbox (`torch.LongTensor` of shape `(batch_size, patch_sequence_length, 4)`, *optional*): + Bounding boxes of each patch in the image. If not provided, bounding boxes are created in the model. + + Example: + + ```python + >>> from transformers import AutoProcessor, UdopEncoderModel + >>> from huggingface_hub import hf_hub_download + >>> from datasets import load_dataset + + >>> # load model and processor + >>> # in this case, we already have performed OCR ourselves + >>> # so we initialize the processor with `apply_ocr=False` + >>> processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=False) + >>> model = UdopEncoderModel.from_pretrained("microsoft/udop-large") + + >>> # load an example image, along with the words and coordinates + >>> # which were extracted using an OCR engine + >>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train") + >>> example = dataset[0] + >>> image = example["image"] + >>> words = example["tokens"] + >>> boxes = example["bboxes"] + >>> encoding = processor(image, words, boxes=boxes, return_tensors="pt") + + >>> outputs = model(**encoding) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + encoder_outputs = self.encoder( + input_ids=input_ids, + bbox=bbox, + visual_bbox=visual_bbox, + pixel_values=pixel_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + return encoder_outputs + + +__all__ = ["UdopForConditionalGeneration", "UdopPreTrainedModel", "UdopModel", "UdopEncoderModel"] diff --git a/third_party/transformers/src/transformers/models/udop/processing_udop.py b/third_party/transformers/src/transformers/models/udop/processing_udop.py new file mode 100644 index 0000000000000000000000000000000000000000..707b5693a2d56ba7b46887a8ceed6e298eb45f8f --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/processing_udop.py @@ -0,0 +1,168 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for UDOP. +""" + +from transformers import logging + +from ...image_processing_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +logger = logging.get_logger(__name__) + + +class UdopTextKwargs(TextKwargs, total=False): + word_labels: list[int] | list[list[int]] | None + boxes: list[list[int]] | list[list[list[int]]] | None + + +class UdopProcessorKwargs(ProcessingKwargs, total=False): + text_kwargs: UdopTextKwargs + _defaults = { + "text_kwargs": { + "add_special_tokens": True, + "padding": False, + "truncation": False, + "stride": 0, + "return_overflowing_tokens": False, + "return_special_tokens_mask": False, + "return_offsets_mapping": False, + "return_length": False, + "verbose": True, + }, + } + + +@auto_docstring +class UdopProcessor(ProcessorMixin): + r""" + Constructs a UDOP processor which combines a LayoutLMv3 image processor and a UDOP tokenizer into a single processor. + + [`UdopProcessor`] offers all the functionalities you need to prepare data for the model. + + It first uses [`LayoutLMv3ImageProcessor`] to resize, rescale and normalize document images, and optionally applies OCR + to get words and normalized bounding boxes. These are then provided to [`UdopTokenizer`], + which turns the words and bounding boxes into token-level `input_ids`, `attention_mask`, `token_type_ids`, `bbox`. + Optionally, one can provide integer `word_labels`, which are turned into token-level `labels` for token + classification tasks (such as FUNSD, CORD). + + Additionally, it also supports passing `text_target` and `text_pair_target` to the tokenizer, which can be used to + prepare labels for language modeling tasks. + """ + + def __init__(self, image_processor, tokenizer): + super().__init__(image_processor, tokenizer) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[UdopProcessorKwargs], + ) -> BatchFeature: + # verify input + output_kwargs = self._merge_kwargs( + UdopProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + boxes = output_kwargs["text_kwargs"].pop("boxes", None) + word_labels = output_kwargs["text_kwargs"].pop("word_labels", None) + text_pair = output_kwargs["text_kwargs"].pop("text_pair", None) + return_overflowing_tokens = output_kwargs["text_kwargs"].get("return_overflowing_tokens", False) + return_offsets_mapping = output_kwargs["text_kwargs"].get("return_offsets_mapping", False) + text_target = output_kwargs["text_kwargs"].get("text_target", None) + + if self.image_processor.apply_ocr and (boxes is not None): + raise ValueError( + "You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True." + ) + + if self.image_processor.apply_ocr and (word_labels is not None): + raise ValueError( + "You cannot provide word labels if you initialized the image processor with apply_ocr set to True." + ) + + if return_overflowing_tokens and not return_offsets_mapping: + raise ValueError("You cannot return overflowing tokens without returning the offsets mapping.") + + if text_target is not None: + # use the processor to prepare the targets of UDOP + return self.tokenizer( + **output_kwargs["text_kwargs"], + ) + + else: + # use the processor to prepare the inputs of UDOP + # first, apply the image processor + features = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + features_words = features.pop("words", None) + features_boxes = features.pop("boxes", None) + + output_kwargs["text_kwargs"].pop("text_target", None) + output_kwargs["text_kwargs"].pop("text_pair_target", None) + output_kwargs["text_kwargs"]["text_pair"] = text_pair + output_kwargs["text_kwargs"]["boxes"] = boxes if boxes is not None else features_boxes + output_kwargs["text_kwargs"]["word_labels"] = word_labels + + # second, apply the tokenizer + if text is not None and self.image_processor.apply_ocr and text_pair is None: + if isinstance(text, str): + text = [text] # add batch dimension (as the image processor always adds a batch dimension) + output_kwargs["text_kwargs"]["text_pair"] = features_words + + encoded_inputs = self.tokenizer( + text=text if text is not None else features_words, + **output_kwargs["text_kwargs"], + ) + + # add pixel values + if return_overflowing_tokens is True: + features["pixel_values"] = self.get_overflowing_images( + features["pixel_values"], encoded_inputs["overflow_to_sample_mapping"] + ) + features.update(encoded_inputs) + + return features + + # Copied from transformers.models.layoutlmv3.processing_layoutlmv3.LayoutLMv3Processor.get_overflowing_images + def get_overflowing_images(self, images, overflow_to_sample_mapping): + # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image + images_with_overflow = [] + for sample_idx in overflow_to_sample_mapping: + images_with_overflow.append(images[sample_idx]) + + if len(images_with_overflow) != len(overflow_to_sample_mapping): + raise ValueError( + "Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got" + f" {len(images_with_overflow)} and {len(overflow_to_sample_mapping)}" + ) + + return images_with_overflow + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + + return list(tokenizer_input_names + image_processor_input_names + ["bbox"]) + + +__all__ = ["UdopProcessor"] diff --git a/third_party/transformers/src/transformers/models/udop/tokenization_udop.py b/third_party/transformers/src/transformers/models/udop/tokenization_udop.py new file mode 100644 index 0000000000000000000000000000000000000000..cf49ebfd7e552b646a789bcc74754269747cbd02 --- /dev/null +++ b/third_party/transformers/src/transformers/models/udop/tokenization_udop.py @@ -0,0 +1,1031 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License +"""Tokenization classes for UDOP model.""" + +from tokenizers import Tokenizer, decoders, pre_tokenizers, processors +from tokenizers.models import Unigram + +from ...tokenization_utils_base import ( + BatchEncoding, + EncodedInput, + PreTokenizedInput, + TextInput, + TextInputPair, + TruncationStrategy, +) +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import PaddingStrategy, TensorType, add_end_docstrings, logging + + +VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"} + + +logger = logging.get_logger(__name__) + +UDOP_ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to encode the sequences with the special tokens relative to their model. + padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + stride (`int`, *optional*, defaults to 0): + If set to a number along with `max_length`, the overflowing tokens returned when + `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence + returned to provide some overlap between truncated and overflowing sequences. The value of this + argument defines the number of overlapping tokens. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). + return_tensors (`str` or [`~file_utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) + + - **bbox** -- List of bounding boxes to be fed to a model. + + - **token_type_ids** -- List of token type ids to be fed to a model (when `return_token_type_ids=True` or + if *"token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) + + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) + + - **labels** -- List of labels to be fed to a model. (when `word_labels` is specified). + - **overflowing_tokens** -- List of overflowing tokens sequences (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **num_truncated_tokens** -- Number of tokens truncated (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **special_tokens_mask** -- List of 0s and 1s, with 1 specifying added special tokens and 0 specifying + regular sequence tokens (when `add_special_tokens=True` and `return_special_tokens_mask=True`). + - **length** -- The length of the inputs (when `return_length=True`). +""" + + +class UdopTokenizer(TokenizersBackend): + """ + Construct a "fast" UDOP tokenizer (backed by HuggingFace's *tokenizers* library). Adapted from + [`LayoutXLMTokenizer`] and [`T5Tokenizer`]. Based on + [BPE](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=BPE#models). + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + + + + When building a sequence using special tokens, this is not the token that is used for the end of sequence. + The token used is the `sep_token`. + + + + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + sep_token_box (`list[int]`, *optional*, defaults to `[1000, 1000, 1000, 1000]`): + The bounding box to use for the special [SEP] token. + pad_token_box (`list[int]`, *optional*, defaults to `[0, 0, 0, 0]`): + The bounding box to use for the special [PAD] token. + pad_token_label (`int`, *optional*, defaults to -100): + The label to use for padding tokens. Defaults to -100, which is the `ignore_index` of PyTorch's + CrossEntropyLoss. + only_label_first_subword (`bool`, *optional*, defaults to `True`): + Whether or not to only label the first subword, in case word labels are provided. + extra_special_tokens (`list[str]`, *optional*, defaults to `["NOTUSED", "NOTUSED"]`): + Extra special tokens used by the tokenizer. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = Unigram + + def __init__( + self, + vocab: str | list[tuple[str, float]] | None = None, + eos_token="", + sep_token="", + unk_token="", + pad_token="", + sep_token_box=[1000, 1000, 1000, 1000], + pad_token_box=[0, 0, 0, 0], + pad_token_label=-100, + only_label_first_subword=True, + extra_special_tokens=None, + **kwargs, + ): + if "additional_special_tokens" in kwargs and "extra_special_tokens" not in kwargs: + kwargs["extra_special_tokens"] = kwargs.pop("additional_special_tokens") + if extra_special_tokens is not None: + kwargs["extra_special_tokens"] = extra_special_tokens + + if vocab is None: + vocab = [(str(pad_token), 0.0), (str(eos_token), 0.0), (str(unk_token), 0.0), ("▁", -2.0)] + + unk_id = 2 + for idx, (token, _) in enumerate(vocab): + if token == str(unk_token): + unk_id = idx + break + + self._tokenizer = Tokenizer( + Unigram( + vocab, + unk_id=unk_id, + byte_fallback=False, + ) + ) + + self._tokenizer.normalizer = None + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.WhitespaceSplit(), + pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True), + ] + ) + + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True) + + super().__init__( + eos_token=eos_token, + sep_token=sep_token, + unk_token=unk_token, + pad_token=pad_token, + sep_token_box=sep_token_box, + pad_token_box=pad_token_box, + pad_token_label=pad_token_label, + only_label_first_subword=only_label_first_subword, + **kwargs, + ) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.eos_token_id), + ], + ) + + self.sep_token_box = sep_token_box + self.pad_token_box = pad_token_box + self.pad_token_label = pad_token_label + self.only_label_first_subword = only_label_first_subword + self.init_kwargs["vocab"] = vocab + + self._tokenizer.encode_special_tokens = self.split_special_tokens + + @add_end_docstrings(UDOP_ENCODE_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + text_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + text_pair_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + **kwargs, + ) -> BatchEncoding: + if text is None and text_target is None: + raise ValueError("You need to specify either `text` or `text_target`.") + if text is not None: + # The context manager will send the inputs as normal texts and not text_target, but we shouldn't change the + # input mode in this case. + if not self._in_target_context_manager and hasattr(self, "_switch_to_input_mode"): + self._switch_to_input_mode() + encodings = self.call_boxes(text=text, text_pair=text_pair, boxes=boxes, word_labels=word_labels, **kwargs) + if text_target is not None: + if hasattr(self, "_switch_to_target_mode"): + self._switch_to_target_mode() + target_encodings = self._encode_plus( + text=text_target, + text_pair=text_pair_target, + **kwargs, + ) + # Leave back tokenizer in input mode + if hasattr(self, "_switch_to_input_mode"): + self._switch_to_input_mode() + + if text_target is None: + return encodings + elif text is None: + return target_encodings + else: + encodings["labels"] = target_encodings["input_ids"] + return encodings + + @add_end_docstrings(UDOP_ENCODE_KWARGS_DOCSTRING) + def call_boxes( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + text_pair: PreTokenizedInput | list[PreTokenizedInput] | None = None, + boxes: list[list[int]] | list[list[list[int]]] | None = None, + word_labels: list[int] | list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences with word-level normalized bounding boxes and optional labels. + + Args: + text (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string, a list of strings + (words of a single example or questions of a batch of examples) or a list of list of strings (batch of + words). + text_pair (`list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each sequence should be a list of strings + (pretokenized string). + boxes (`list[list[int]]`, `list[list[list[int]]]`): + Word-level bounding boxes. Each bounding box should be normalized to be on a 0-1000 scale. + word_labels (`list[int]`, `list[list[int]]`, *optional*): + Word-level integer labels (for token classification tasks such as FUNSD, CORD). + """ + + # Input type checking for clearer error + def _is_valid_text_input(t): + if isinstance(t, str): + # Strings are fine + return True + elif isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... empty + return True + elif isinstance(t[0], str): + # ... list of strings + return True + elif isinstance(t[0], (list, tuple)): + # ... list with an empty list or with a list of strings + return len(t[0]) == 0 or isinstance(t[0][0], str) + else: + return False + else: + return False + + if text_pair is not None: + # in case text + text_pair are provided, text = questions, text_pair = words + if not _is_valid_text_input(text): + raise ValueError("text input must of type `str` (single example) or `list[str]` (batch of examples). ") + if not isinstance(text_pair, (list, tuple)): + raise ValueError( + "words must of type `list[str]` (single pretokenized example), " + "or `list[list[str]]` (batch of pretokenized examples)." + ) + else: + # in case only text is provided => must be words + if not isinstance(text, (list, tuple)): + raise ValueError( + "Words must of type `list[str]` (single pretokenized example), " + "or `list[list[str]]` (batch of pretokenized examples)." + ) + + if text_pair is not None: + is_batched = isinstance(text, (list, tuple)) + else: + is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)) + + words = text if text_pair is None else text_pair + if boxes is None: + raise ValueError("You must provide corresponding bounding boxes") + if is_batched: + if len(words) != len(boxes): + raise ValueError("You must provide words and boxes for an equal amount of examples") + for words_example, boxes_example in zip(words, boxes): + if len(words_example) != len(boxes_example): + raise ValueError("You must provide as many words as there are bounding boxes") + else: + if len(words) != len(boxes): + raise ValueError("You must provide as many words as there are bounding boxes") + + if is_batched: + if text_pair is not None and len(text) != len(text_pair): + raise ValueError( + f"batch length of `text`: {len(text)} does not match batch length of `text_pair`:" + f" {len(text_pair)}." + ) + batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text + is_pair = bool(text_pair is not None) + return self.batch_encode_plus_boxes( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + else: + return self.encode_plus_boxes( + text=text, + text_pair=text_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def tokenize(self, text: str, pair: str | None = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + batched_input = [(text, pair)] if pair else [text] + + self._tokenizer.encode_special_tokens = kwargs.pop("split_special_tokens", self.split_special_tokens) + + encodings = self._tokenizer.encode_batch( + batched_input, add_special_tokens=add_special_tokens, is_pretokenized=False, **kwargs + ) + + return encodings[0].tokens + + def batch_encode_plus_boxes( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a list of sequences or a list of pairs of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + batch_text_or_text_pairs (`list[str]`, `list[tuple[str, str]]`, `list[list[str]]`, `list[tuple[list[str], list[str]]]`, and for not-fast tokenizers, also `list[list[int]]`, `list[tuple[list[int], list[int]]]`): + Batch of sequences or pair of sequences to be encoded. This can be a list of + string/string-sequences/int-sequences or a list of pair of string/string-sequences/int-sequence (see + details in `encode_plus`). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._batch_encode_plus_boxes( + batch_text_or_text_pairs=batch_text_or_text_pairs, + is_pair=is_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def _batch_encode_plus_boxes( + self, + batch_text_or_text_pairs: list[TextInput] | list[TextInputPair] | list[PreTokenizedInput], + is_pair: bool | None = None, + boxes: list[list[list[int]]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + if not isinstance(batch_text_or_text_pairs, list): + raise TypeError(f"batch_text_or_text_pairs has to be a list (got {type(batch_text_or_text_pairs)})") + + # Set the truncation and padding strategy and restore the initial configuration + self.set_truncation_and_padding( + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + ) + + if is_pair: + batch_text_or_text_pairs = [(text.split(), text_pair) for text, text_pair in batch_text_or_text_pairs] + + encodings = self._tokenizer.encode_batch( + batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + is_pretokenized=True, # we set this to True as LayoutLMv2 always expects pretokenized inputs + ) + + # Convert encoding to dict + # `Tokens` has type: tuple[ + # list[dict[str, list[list[int]]]] or list[dict[str, 2D-Tensor]], + # list[EncodingFast] + # ] + # with nested dimensions corresponding to batch, overflows, sequence length + tokens_and_encodings = [ + self._convert_encoding( + encoding=encoding, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=True + if word_labels is not None + else return_offsets_mapping, # we use offsets to create the labels + return_length=return_length, + verbose=verbose, + ) + for encoding in encodings + ] + + # Convert the output to have dict[list] from list[dict] and remove the additional overflows dimension + # From (variable) shape (batch, overflows, sequence length) to ~ (batch * overflows, sequence length) + # (we say ~ because the number of overflow varies with the example in the batch) + # + # To match each overflowing sample with the original sample in the batch + # we add an overflow_to_sample_mapping array (see below) + sanitized_tokens = {} + for key in tokens_and_encodings[0][0]: + stack = [e for item, _ in tokens_and_encodings for e in item[key]] + sanitized_tokens[key] = stack + sanitized_encodings = [e for _, item in tokens_and_encodings for e in item] + + # If returning overflowing tokens, we need to return a mapping + # from the batch idx to the original sample + if return_overflowing_tokens: + overflow_to_sample_mapping = [] + for i, (toks, _) in enumerate(tokens_and_encodings): + overflow_to_sample_mapping += [i] * len(toks["input_ids"]) + sanitized_tokens["overflow_to_sample_mapping"] = overflow_to_sample_mapping + + for input_ids in sanitized_tokens["input_ids"]: + self._eventual_warn_about_too_long_sequence(input_ids, max_length, verbose) + + # create the token boxes + token_boxes = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + token_boxes_example = [] + for id, sequence_id, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_encodings[batch_index].sequence_ids, + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if is_pair and sequence_id == 0: + token_boxes_example.append(self.pad_token_box) + else: + token_boxes_example.append(boxes[original_index][word_id]) + else: + if id == self.sep_token_id: + token_boxes_example.append(self.sep_token_box) + elif id == self.pad_token_id: + token_boxes_example.append(self.pad_token_box) + else: + raise ValueError("Id not recognized") + token_boxes.append(token_boxes_example) + + sanitized_tokens["bbox"] = token_boxes + + # optionally, create the labels + if word_labels is not None: + labels = [] + for batch_index in range(len(sanitized_tokens["input_ids"])): + if return_overflowing_tokens: + original_index = sanitized_tokens["overflow_to_sample_mapping"][batch_index] + else: + original_index = batch_index + labels_example = [] + previous_token_empty = False + for id, offset, word_id in zip( + sanitized_tokens["input_ids"][batch_index], + sanitized_tokens["offset_mapping"][batch_index], + sanitized_encodings[batch_index].word_ids, + ): + if word_id is not None: + if self.only_label_first_subword: + if offset[0] == 0 and not previous_token_empty: + # Use the real label id for the first token of the word, and padding ids for the remaining tokens + labels_example.append(word_labels[original_index][word_id]) + else: + labels_example.append(self.pad_token_label) + else: + labels_example.append(word_labels[original_index][word_id]) + if self.decode(id) == "": + previous_token_empty = True + else: + previous_token_empty = False + else: + labels_example.append(self.pad_token_label) + labels.append(labels_example) + + sanitized_tokens["labels"] = labels + # finally, remove offsets if the user didn't want them + if not return_offsets_mapping: + del sanitized_tokens["offset_mapping"] + + return BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors) + + def _encode_plus_boxes( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[int] | None = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: bool | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + # make it a batched input + # 2 options: + # 1) only text, in case text must be a list of str + # 2) text + text_pair, in which case text = str and text_pair a list of str + batched_input = [(text, text_pair)] if text_pair else [text] + batched_boxes = [boxes] + batched_word_labels = [word_labels] if word_labels is not None else None + batched_output = self._batch_encode_plus_boxes( + batched_input, + is_pair=bool(text_pair is not None), + boxes=batched_boxes, + word_labels=batched_word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + # Return tensor is None, then we can remove the leading batch axis + # Overflowing tokens are returned as a batch of output so we keep them in this case + if return_tensors is None and not return_overflowing_tokens: + batched_output = BatchEncoding( + { + key: value[0] if len(value) > 0 and isinstance(value[0], list) else value + for key, value in batched_output.items() + }, + batched_output.encodings, + ) + + self._eventual_warn_about_too_long_sequence(batched_output["input_ids"], max_length, verbose) + + return batched_output + + def encode_boxes( + self, + text: TextInput | PreTokenizedInput | EncodedInput, + text_pair: TextInput | PreTokenizedInput | EncodedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> list[int]: + """ + Args: + Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. Same as doing + `self.convert_tokens_to_ids(self.tokenize(text))`. + text (`str`, `list[str]` or `list[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + text_pair (`str`, `list[str]` or `list[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + """ + encoded_inputs = self.encode_plus_boxes( + text, + text_pair=text_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + return_tensors=return_tensors, + **kwargs, + ) + + return encoded_inputs["input_ids"] + + def encode_plus_boxes( + self, + text: TextInput | PreTokenizedInput, + text_pair: PreTokenizedInput | None = None, + boxes: list[list[int]] | None = None, + word_labels: list[list[int]] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a sequence or a pair of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + text (`str`, `list[str]` or (for non-fast tokenizers) `list[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + text_pair (`str`, `list[str]` or `list[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._encode_plus_boxes( + text=text, + text_pair=text_pair, + boxes=boxes, + word_labels=word_labels, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + **kwargs, + ) + + def _pad( + self, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + ) -> dict: + """ + Pad encoded inputs (on left/right and up to predefined length or max length in the batch) + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`list[int]`) or batch of tokenized inputs (`list[list[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + padding_strategy: PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in self.padding_side: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + required_input = encoded_inputs[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(required_input) + + if needs_to_be_padded: + difference = max_length - len(required_input) + padding_side = padding_side if padding_side is not None else self.padding_side + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = encoded_inputs["attention_mask"] + [0] * difference + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = ( + encoded_inputs["token_type_ids"] + [self.pad_token_type_id] * difference + ) + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = encoded_inputs["bbox"] + [self.pad_token_box] * difference + if "labels" in encoded_inputs: + encoded_inputs["labels"] = encoded_inputs["labels"] + [self.pad_token_label] * difference + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference + encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[ + "token_type_ids" + ] + if "bbox" in encoded_inputs: + encoded_inputs["bbox"] = [self.pad_token_box] * difference + encoded_inputs["bbox"] + if "labels" in encoded_inputs: + encoded_inputs["labels"] = [self.pad_token_label] * difference + encoded_inputs["labels"] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input + else: + raise ValueError("Invalid padding strategy:" + str(padding_side)) + + return encoded_inputs + + def build_inputs_with_special_tokens( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. An XLM-RoBERTa sequence has the following format: + + - single sequence: ` X ` + - pair of sequences: ` A B ` + + Args: + token_ids_0 (`list[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`list[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `list[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + """ + + if token_ids_1 is None: + return token_ids_0 + [self.sep_token_id] + sep = [self.sep_token_id] + return token_ids_0 + sep + token_ids_1 + sep + + def create_token_type_ids_from_sequences( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + """ + Create a mask from the two sequences passed to be used in a sequence-pair classification task. XLM-RoBERTa does + not make use of token type ids, therefore a list of zeros is returned. + + Args: + token_ids_0 (`list[int]`): + List of IDs. + token_ids_1 (`list[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `list[int]`: List of zeros. + + """ + + sep = [self.sep_token_id] + + if token_ids_1 is None: + return len(token_ids_0 + sep) * [0] + return len(token_ids_0 + sep + token_ids_1 + sep) * [0] + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + """ + Save the tokenizer vocabulary files. For TokenizersBackend, the tokenizer.json file is saved + by the base class. This method returns an empty tuple since we only use tokenizer.json. + """ + # The base class handles saving tokenizer.json in _save_pretrained + # We don't need to save vocab_file since we only use tokenizer.json + return () + + +__all__ = ["UdopTokenizer"] diff --git a/third_party/transformers/src/transformers/models/xlnet/__init__.py b/third_party/transformers/src/transformers/models/xlnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5734b98957c6e6424f41f839592fbf9c21bfc40f --- /dev/null +++ b/third_party/transformers/src/transformers/models/xlnet/__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_xlnet import * + from .modeling_xlnet import * + from .tokenization_xlnet import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/models/xlnet/configuration_xlnet.py b/third_party/transformers/src/transformers/models/xlnet/configuration_xlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..24fc0cbb6f96c152d1a6cbf310068090b0b7a2e1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/xlnet/configuration_xlnet.py @@ -0,0 +1,161 @@ +# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""XLNet configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="xlnet/xlnet-large-cased") +@strict +class XLNetConfig(PreTrainedConfig): + r""" + ff_activation (`str` or `Callable`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the If string, `"gelu"`, `"relu"`, `"silu"` and + `"gelu_new"` are supported. + attn_type (`str`, *optional*, defaults to `"bi"`): + The attention type used by the model. Set `"bi"` for XLNet, `"uni"` for Transformer-XL. + mem_len (`int` or `None`, *optional*): + The number of tokens to cache. The key/value pairs that have already been pre-computed in a previous + forward pass won't be re-computed. See the + [quickstart](https://huggingface.co/transformers/quickstart.html#using-the-past) for more information. + reuse_len (`int`, *optional*): + The number of tokens in the current batch to be cached and reused in the future. + use_mems_eval (`bool`, *optional*, defaults to `True`): + Whether or not the model should make use of the recurrent memory mechanism in evaluation mode. + use_mems_train (`bool`, *optional*, defaults to `False`): + Whether or not the model should make use of the recurrent memory mechanism in train mode. + + For pretraining, it is recommended to set `use_mems_train` to `True`. For fine-tuning, it is recommended to + set `use_mems_train` to `False` as discussed + [here](https://github.com/zihangdai/xlnet/issues/41#issuecomment-505102587). If `use_mems_train` is set to + `True`, one has to make sure that the train batches are correctly pre-processed, *e.g.* `batch_1 = [[This + line is], [This is the]]` and `batch_2 = [[ the first line], [ second line]]` and that all batches are of + equal size. + + bi_data (`bool`, *optional*, defaults to `False`): + Whether or not to use bidirectional input pipeline. Usually set to `True` during pretraining and `False` + during finetuning. + clamp_len (`int`, *optional*, defaults to -1): + Clamp all relative distances larger than clamp_len. Setting this attribute to -1 means no clamping. + same_length (`bool`, *optional*, defaults to `False`): + Whether or not to use the same attention length for each token. + summary_type (`str`, *optional*, defaults to "last"): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Has to be one of the following options: + - `"last"`: Take the last token hidden state (like XLNet). + - `"first"`: Take the first token hidden state (like BERT). + - `"mean"`: Take the mean of all tokens hidden states. + - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2). + - `"attn"`: Not implemented now, use multi-head attention. + summary_use_proj (`bool`, *optional*, defaults to `True`): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Whether or not to add a projection after the vector extraction. + summary_activation (`str`, *optional*): + Argument used when doing sequence summary. Used in the sequence classification and multiple choice models. + Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation. + summary_last_dropout (`float`, *optional*, defaults to 0.1): + Used in the sequence classification and multiple choice models. + The dropout ratio to be used after the projection and activation. + start_n_top (`int`, *optional*, defaults to 5): + Used in the SQuAD evaluation script. + end_n_top (`int`, *optional*, defaults to 5): + Used in the SQuAD evaluation script. + + Examples: + + ```python + >>> from transformers import XLNetConfig, XLNetModel + + >>> # Initializing a XLNet configuration + >>> configuration = XLNetConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = XLNetModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "xlnet" + keys_to_ignore_at_inference = ["mems"] + attribute_map = { + "n_token": "vocab_size", # Backward compatibility + "hidden_size": "d_model", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + vocab_size: int = 32000 + d_model: int = 1024 + n_layer: int = 24 + n_head: int = 16 + d_inner: int = 4096 + d_head: int | None = None + ff_activation: str = "gelu" + attn_type: str = "bi" + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + dropout: float | int = 0.1 + mem_len: int | None = 512 + reuse_len: int | None = None + use_mems_eval: bool = True + use_mems_train: bool = False + bi_data: bool = False + clamp_len: int = -1 + same_length: bool = False + summary_type: str = "last" + summary_use_proj: bool = True + summary_activation: str = "tanh" + summary_last_dropout: float | int = 0.1 + start_n_top: int = 5 + end_n_top: int = 5 + pad_token_id: int | None = 5 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + self.d_head = self.d_head or self.d_model // self.n_head + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.d_model % self.n_head != 0: + raise ValueError(f"'d_model % n_head' ({self.d_model % self.n_head}) should be equal to 0") + if self.d_head != self.d_model // self.n_head: + raise ValueError( + f"`d_head` ({self.d_head}) should be equal to `d_model // n_head` ({self.d_model // self.n_head})" + ) + + @property + def max_position_embeddings(self): + logger.info(f"The model {self.model_type} is one of the few models that has no sequence length limit.") + return -1 + + @max_position_embeddings.setter + def max_position_embeddings(self, value): + # Message copied from Transformer-XL documentation + raise NotImplementedError( + f"The model {self.model_type} is one of the few models that has no sequence length limit." + ) + + +__all__ = ["XLNetConfig"] diff --git a/third_party/transformers/src/transformers/models/xlnet/convert_xlnet_original_tf_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/xlnet/convert_xlnet_original_tf_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..676debb2c7d02bc890d431317b4665a919ab542d --- /dev/null +++ b/third_party/transformers/src/transformers/models/xlnet/convert_xlnet_original_tf_checkpoint_to_pytorch.py @@ -0,0 +1,262 @@ +# Copyright 2018 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert BERT checkpoint.""" + +import argparse +import os + +import torch + +from transformers import ( + XLNetConfig, + XLNetForQuestionAnswering, + XLNetForSequenceClassification, + XLNetLMHeadModel, +) +from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging + + +GLUE_TASKS_NUM_LABELS = { + "cola": 2, + "mnli": 3, + "mrpc": 2, + "sst-2": 2, + "sts-b": 1, + "qqp": 2, + "qnli": 2, + "rte": 2, + "wnli": 2, +} + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def build_tf_xlnet_to_pytorch_map(model, config, tf_weights=None): + """ + A map of modules from TF to PyTorch. I use a map to keep the PyTorch model as identical to the original PyTorch + model as possible. + """ + + tf_to_pt_map = {} + + if hasattr(model, "transformer"): + if hasattr(model, "lm_loss"): + # We will load also the output bias + tf_to_pt_map["model/lm_loss/bias"] = model.lm_loss.bias + if hasattr(model, "sequence_summary") and "model/sequnece_summary/summary/kernel" in tf_weights: + # We will load also the sequence summary + tf_to_pt_map["model/sequnece_summary/summary/kernel"] = model.sequence_summary.summary.weight + tf_to_pt_map["model/sequnece_summary/summary/bias"] = model.sequence_summary.summary.bias + if ( + hasattr(model, "logits_proj") + and config.finetuning_task is not None + and f"model/regression_{config.finetuning_task}/logit/kernel" in tf_weights + ): + tf_to_pt_map[f"model/regression_{config.finetuning_task}/logit/kernel"] = model.logits_proj.weight + tf_to_pt_map[f"model/regression_{config.finetuning_task}/logit/bias"] = model.logits_proj.bias + + # Now load the rest of the transformer + model = model.transformer + + # Embeddings and output + tf_to_pt_map.update( + { + "model/transformer/word_embedding/lookup_table": model.word_embedding.weight, + "model/transformer/mask_emb/mask_emb": model.mask_emb, + } + ) + + # Transformer blocks + for i, b in enumerate(model.layer): + layer_str = f"model/transformer/layer_{i}/" + tf_to_pt_map.update( + { + layer_str + "rel_attn/LayerNorm/gamma": b.rel_attn.layer_norm.weight, + layer_str + "rel_attn/LayerNorm/beta": b.rel_attn.layer_norm.bias, + layer_str + "rel_attn/o/kernel": b.rel_attn.o, + layer_str + "rel_attn/q/kernel": b.rel_attn.q, + layer_str + "rel_attn/k/kernel": b.rel_attn.k, + layer_str + "rel_attn/r/kernel": b.rel_attn.r, + layer_str + "rel_attn/v/kernel": b.rel_attn.v, + layer_str + "ff/LayerNorm/gamma": b.ff.layer_norm.weight, + layer_str + "ff/LayerNorm/beta": b.ff.layer_norm.bias, + layer_str + "ff/layer_1/kernel": b.ff.layer_1.weight, + layer_str + "ff/layer_1/bias": b.ff.layer_1.bias, + layer_str + "ff/layer_2/kernel": b.ff.layer_2.weight, + layer_str + "ff/layer_2/bias": b.ff.layer_2.bias, + } + ) + + # Relative positioning biases + if config.untie_r: + r_r_list = [] + r_w_list = [] + r_s_list = [] + seg_embed_list = [] + for b in model.layer: + r_r_list.append(b.rel_attn.r_r_bias) + r_w_list.append(b.rel_attn.r_w_bias) + r_s_list.append(b.rel_attn.r_s_bias) + seg_embed_list.append(b.rel_attn.seg_embed) + else: + r_r_list = [model.r_r_bias] + r_w_list = [model.r_w_bias] + r_s_list = [model.r_s_bias] + seg_embed_list = [model.seg_embed] + tf_to_pt_map.update( + { + "model/transformer/r_r_bias": r_r_list, + "model/transformer/r_w_bias": r_w_list, + "model/transformer/r_s_bias": r_s_list, + "model/transformer/seg_embed": seg_embed_list, + } + ) + return tf_to_pt_map + + +def load_tf_weights_in_xlnet(model, config, tf_path): + """Load tf checkpoints in a pytorch model""" + try: + import numpy as np + import tensorflow as tf + except ImportError: + logger.error( + "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions." + ) + raise + # Load weights from TF model + init_vars = tf.train.list_variables(tf_path) + tf_weights = {} + for name, shape in init_vars: + logger.info(f"Loading TF weight {name} with shape {shape}") + array = tf.train.load_variable(tf_path, name) + tf_weights[name] = array + + # Build TF to PyTorch weights loading map + tf_to_pt_map = build_tf_xlnet_to_pytorch_map(model, config, tf_weights) + + for name, pointer in tf_to_pt_map.items(): + logger.info(f"Importing {name}") + if name not in tf_weights: + logger.info(f"{name} not in tf pre-trained weights, skipping") + continue + array = tf_weights[name] + # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v + # which are not required for using pretrained model + if "kernel" in name and ("ff" in name or "summary" in name or "logit" in name): + logger.info("Transposing") + array = np.transpose(array) + if isinstance(pointer, list): + # Here we will split the TF weights + assert len(pointer) == array.shape[0], ( + f"Pointer length {len(pointer)} and array length {array.shape[0]} mismatched" + ) + for i, p_i in enumerate(pointer): + arr_i = array[i, ...] + try: + assert p_i.shape == arr_i.shape, ( + f"Pointer shape {p_i.shape} and array shape {arr_i.shape} mismatched" + ) + except AssertionError as e: + e.args += (p_i.shape, arr_i.shape) + raise + logger.info(f"Initialize PyTorch weight {name} for layer {i}") + p_i.data = torch.from_numpy(arr_i) + else: + try: + assert pointer.shape == array.shape, ( + f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" + ) + except AssertionError as e: + e.args += (pointer.shape, array.shape) + raise + logger.info(f"Initialize PyTorch weight {name}") + pointer.data = torch.from_numpy(array) + tf_weights.pop(name, None) + tf_weights.pop(name + "/Adam", None) + tf_weights.pop(name + "/Adam_1", None) + + logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}") + return model + + +def convert_xlnet_checkpoint_to_pytorch( + tf_checkpoint_path, bert_config_file, pytorch_dump_folder_path, finetuning_task=None +): + # Initialise PyTorch model + config = XLNetConfig.from_json_file(bert_config_file) + + finetuning_task = finetuning_task.lower() if finetuning_task is not None else "" + if finetuning_task in GLUE_TASKS_NUM_LABELS: + print(f"Building PyTorch XLNetForSequenceClassification model from configuration: {config}") + config.finetuning_task = finetuning_task + config.num_labels = GLUE_TASKS_NUM_LABELS[finetuning_task] + model = XLNetForSequenceClassification(config) + elif "squad" in finetuning_task: + config.finetuning_task = finetuning_task + model = XLNetForQuestionAnswering(config) + else: + model = XLNetLMHeadModel(config) + + # Load weights from tf checkpoint + load_tf_weights_in_xlnet(model, config, tf_checkpoint_path) + + # Save pytorch-model + pytorch_weights_dump_path = os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME) + pytorch_config_dump_path = os.path.join(pytorch_dump_folder_path, CONFIG_NAME) + print(f"Save PyTorch model to {os.path.abspath(pytorch_weights_dump_path)}") + torch.save(model.state_dict(), pytorch_weights_dump_path) + print(f"Save configuration file to {os.path.abspath(pytorch_config_dump_path)}") + with open(pytorch_config_dump_path, "w", encoding="utf-8") as f: + f.write(config.to_json_string()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path." + ) + parser.add_argument( + "--xlnet_config_file", + default=None, + type=str, + required=True, + help=( + "The config json file corresponding to the pre-trained XLNet model. \n" + "This specifies the model architecture." + ), + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + required=True, + help="Path to the folder to store the PyTorch model or dataset/vocab.", + ) + parser.add_argument( + "--finetuning_task", + default=None, + type=str, + help="Name of a task on which the XLNet TensorFlow model was fine-tuned", + ) + args = parser.parse_args() + print(args) + + convert_xlnet_checkpoint_to_pytorch( + args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task + ) diff --git a/third_party/transformers/src/transformers/models/xlnet/modeling_xlnet.py b/third_party/transformers/src/transformers/models/xlnet/modeling_xlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..7029893312873919575c684f2fd86e4fd58d7d82 --- /dev/null +++ b/third_party/transformers/src/transformers/models/xlnet/modeling_xlnet.py @@ -0,0 +1,2171 @@ +# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +PyTorch XLNet model. +""" + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN, get_activation +from ...generation import GenerationMixin +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_xlnet import XLNetConfig + + +logger = logging.get_logger(__name__) + + +class XLNetRelativeAttention(nn.Module): + def __init__(self, config): + super().__init__() + + if config.d_model % config.n_head != 0: + raise ValueError( + f"The hidden size ({config.d_model}) is not a multiple of the number of attention " + f"heads ({config.n_head}" + ) + + self.n_head = config.n_head + self.d_head = config.d_head + self.d_model = config.d_model + self.scale = 1 / (config.d_head**0.5) + + self.q = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head)) + self.k = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head)) + self.v = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head)) + self.o = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head)) + self.r = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head)) + + self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head)) + self.r_s_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head)) + self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head)) + self.seg_embed = nn.Parameter(torch.FloatTensor(2, self.n_head, self.d_head)) + + self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.dropout) + + @staticmethod + def rel_shift(x, klen=-1): + """perform relative shift to form the relative attention score.""" + x_size = x.shape + + x = x.reshape(x_size[1], x_size[0], x_size[2], x_size[3]) + x = x[1:, ...] + x = x.reshape(x_size[0], x_size[1] - 1, x_size[2], x_size[3]) + # x = x[:, 0:klen, :, :] + x = torch.index_select(x, 1, torch.arange(klen, device=x.device, dtype=torch.long)) + + return x + + @staticmethod + def rel_shift_bnij(x, klen=-1): + x_size = x.shape + + x = x.reshape(x_size[0], x_size[1], x_size[3], x_size[2]) + x = x[:, :, 1:, :] + x = x.reshape(x_size[0], x_size[1], x_size[2], x_size[3] - 1) + # Note: the tensor-slice form was faster in my testing than torch.index_select + # However, tracing doesn't like the nature of the slice, and if klen changes + # during the run then it'll fail, whereas index_select will be fine. + x = torch.index_select(x, 3, torch.arange(klen, device=x.device, dtype=torch.long)) + # x = x[:, :, :, :klen] + + return x + + def rel_attn_core( + self, + q_head, + k_head_h, + v_head_h, + k_head_r, + seg_mat=None, + attn_mask=None, + output_attentions=False, + ): + """Core relative positional attention operations.""" + + # content based attention score + ac = torch.einsum("ibnd,jbnd->bnij", q_head + self.r_w_bias, k_head_h) + + # position based attention score + bd = torch.einsum("ibnd,jbnd->bnij", q_head + self.r_r_bias, k_head_r) + bd = self.rel_shift_bnij(bd, klen=ac.shape[3]) + + # segment based attention score + if seg_mat is None: + ef = 0 + else: + ef = torch.einsum("ibnd,snd->ibns", q_head + self.r_s_bias, self.seg_embed) + ef = torch.einsum("ijbs,ibns->bnij", seg_mat, ef) + + # merge attention scores and perform masking + attn_score = (ac + bd + ef) * self.scale + if attn_mask is not None: + # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask + if attn_mask.dtype == torch.float16: + attn_score = attn_score - 65500 * torch.einsum("ijbn->bnij", attn_mask) + else: + attn_score = attn_score - 1e30 * torch.einsum("ijbn->bnij", attn_mask) + + # attention probability + attn_prob = nn.functional.softmax(attn_score, dim=3) + attn_prob = self.dropout(attn_prob) + + # attention output + attn_vec = torch.einsum("bnij,jbnd->ibnd", attn_prob, v_head_h) + + if output_attentions: + return attn_vec, torch.einsum("bnij->ijbn", attn_prob) + + return attn_vec + + def post_attention(self, h, attn_vec, residual=True): + """Post-attention processing.""" + # post-attention projection (back to `d_model`) + attn_out = torch.einsum("ibnd,hnd->ibh", attn_vec, self.o) + + attn_out = self.dropout(attn_out) + if residual: + attn_out = attn_out + h + output = self.layer_norm(attn_out) + + return output + + def forward( + self, + h, + g, + attn_mask_h, + attn_mask_g, + r, + seg_mat, + mems=None, + target_mapping=None, + output_attentions=False, + ): + if g is not None: + # Two-stream attention with relative positional encoding. + # content based attention score + if mems is not None and mems.dim() > 1: + cat = torch.cat([mems, h], dim=0) + else: + cat = h + + # content-based key head + k_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.k) + + # content-based value head + v_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.v) + + # position-based key head + k_head_r = torch.einsum("ibh,hnd->ibnd", r, self.r) + + # h-stream + # content-stream query head + q_head_h = torch.einsum("ibh,hnd->ibnd", h, self.q) + + # core attention ops + attn_vec_h = self.rel_attn_core( + q_head_h, + k_head_h, + v_head_h, + k_head_r, + seg_mat=seg_mat, + attn_mask=attn_mask_h, + output_attentions=output_attentions, + ) + + if output_attentions: + attn_vec_h, attn_prob_h = attn_vec_h + + # post processing + output_h = self.post_attention(h, attn_vec_h) + + # g-stream + # query-stream query head + q_head_g = torch.einsum("ibh,hnd->ibnd", g, self.q) + + # core attention ops + if target_mapping is not None: + q_head_g = torch.einsum("mbnd,mlb->lbnd", q_head_g, target_mapping) + attn_vec_g = self.rel_attn_core( + q_head_g, + k_head_h, + v_head_h, + k_head_r, + seg_mat=seg_mat, + attn_mask=attn_mask_g, + output_attentions=output_attentions, + ) + + if output_attentions: + attn_vec_g, attn_prob_g = attn_vec_g + + attn_vec_g = torch.einsum("lbnd,mlb->mbnd", attn_vec_g, target_mapping) + else: + attn_vec_g = self.rel_attn_core( + q_head_g, + k_head_h, + v_head_h, + k_head_r, + seg_mat=seg_mat, + attn_mask=attn_mask_g, + output_attentions=output_attentions, + ) + + if output_attentions: + attn_vec_g, attn_prob_g = attn_vec_g + + # post processing + output_g = self.post_attention(g, attn_vec_g) + + if output_attentions: + attn_prob = attn_prob_h, attn_prob_g + + else: + # Multi-head attention with relative positional encoding + if mems is not None and mems.dim() > 1: + cat = torch.cat([mems, h], dim=0) + else: + cat = h + + # content heads + q_head_h = torch.einsum("ibh,hnd->ibnd", h, self.q) + k_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.k) + v_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.v) + + # positional heads + # type casting for fp16 support + k_head_r = torch.einsum("ibh,hnd->ibnd", r.type(self.r.dtype), self.r) + + # core attention ops + attn_vec = self.rel_attn_core( + q_head_h, + k_head_h, + v_head_h, + k_head_r, + seg_mat=seg_mat, + attn_mask=attn_mask_h, + output_attentions=output_attentions, + ) + + if output_attentions: + attn_vec, attn_prob = attn_vec + + # post processing + output_h = self.post_attention(h, attn_vec) + output_g = None + + outputs = (output_h, output_g) + if output_attentions: + outputs = outputs + (attn_prob,) + return outputs + + +class XLNetFeedForward(nn.Module): + def __init__(self, config): + super().__init__() + self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps) + self.layer_1 = nn.Linear(config.d_model, config.d_inner) + self.layer_2 = nn.Linear(config.d_inner, config.d_model) + self.dropout = nn.Dropout(config.dropout) + if isinstance(config.ff_activation, str): + self.activation_function = ACT2FN[config.ff_activation] + else: + self.activation_function = config.ff_activation + + def forward(self, inp): + output = inp + output = self.layer_1(output) + output = self.activation_function(output) + output = self.dropout(output) + output = self.layer_2(output) + output = self.dropout(output) + output = self.layer_norm(output + inp) + return output + + +class XLNetLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.rel_attn = XLNetRelativeAttention(config) + self.ff = XLNetFeedForward(config) + self.dropout = nn.Dropout(config.dropout) + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + + def forward( + self, + output_h, + output_g, + attn_mask_h, + attn_mask_g, + r, + seg_mat, + mems=None, + target_mapping=None, + output_attentions=False, + ): + outputs = self.rel_attn( + output_h, + output_g, + attn_mask_h, + attn_mask_g, + r, + seg_mat, + mems=mems, + target_mapping=target_mapping, + output_attentions=output_attentions, + ) + output_h, output_g = outputs[:2] + + if output_g is not None: + output_g = apply_chunking_to_forward( + self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_g + ) + output_h = apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_h) + + outputs = (output_h, output_g) + outputs[2:] # Add again attentions if there are there + return outputs + + def ff_chunk(self, output_x): + output_x = self.ff(output_x) + return output_x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerStartLogits with XLM->XLNet +class XLNetPoolerStartLogits(nn.Module): + """ + Compute SQuAD start logits from sequence hidden states. + + Args: + config ([`XLNetConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: XLNetConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, 1) + + def forward(self, hidden_states: torch.FloatTensor, p_mask: torch.FloatTensor | None = None) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + Returns: + `torch.FloatTensor`: The start logits for SQuAD. + """ + x = self.dense(hidden_states).squeeze(-1) + + if p_mask is not None: + if p_mask.dtype == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerEndLogits with XLM->XLNet +class XLNetPoolerEndLogits(nn.Module): + """ + Compute SQuAD end logits from sequence hidden states. + + Args: + config ([`XLNetConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config: XLNetConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dense_1 = nn.Linear(config.hidden_size, 1) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + p_mask: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The end logits for SQuAD. + """ + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + slen, hsz = hidden_states.shape[-2:] + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions) # shape (bsz, 1, hsz) + start_states = start_states.expand(-1, slen, -1) # shape (bsz, slen, hsz) + + x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1)) + x = self.activation(x) + x = self.LayerNorm(x) + x = self.dense_1(x).squeeze(-1) + + if p_mask is not None: + if p_mask.dtype == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMPoolerAnswerClass with XLM->XLNet +class XLNetPoolerAnswerClass(nn.Module): + """ + Compute SQuAD 2.0 answer class from classification and start tokens hidden states. + + Args: + config ([`XLNetConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: XLNetConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + cls_index: torch.LongTensor | None = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The SQuAD 2.0 answer class. + """ + # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample. + hsz = hidden_states.shape[-1] + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions).squeeze(-2) # shape (bsz, hsz) + + if cls_index is not None: + cls_index = cls_index[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, hsz) + else: + cls_token_state = hidden_states[:, -1, :] # shape (bsz, hsz) + + x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1)) + x = self.activation(x) + x = self.dense_1(x).squeeze(-1) + + return x + + +# Copied from transformers.models.xlm.modeling_xlm.XLMSequenceSummary with XLM->XLNet +class XLNetSequenceSummary(nn.Module): + r""" + Compute a single vector summary of a sequence hidden states. + + Args: + config ([`XLNetConfig`]): + The config used by the model. Relevant arguments in the config class of the model are (refer to the actual + config class of your model for the default values it uses): + + - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are: + + - `"last"` -- Take the last token hidden state (like XLNet) + - `"first"` -- Take the first token hidden state (like Bert) + - `"mean"` -- Take the mean of all tokens hidden states + - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2) + - `"attn"` -- Not implemented now, use multi-head attention + + - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction. + - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes + (otherwise to `config.hidden_size`). + - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output, + another string or `None` will add no activation. + - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation. + - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation. + """ + + def __init__(self, config: XLNetConfig): + super().__init__() + + self.summary_type = getattr(config, "summary_type", "last") + if self.summary_type == "attn": + # We should use a standard multi-head attention module with absolute positional embedding for that. + # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276 + # We can probably just use the multi-head attention module of PyTorch >=1.1.0 + raise NotImplementedError + + self.summary = nn.Identity() + if hasattr(config, "summary_use_proj") and config.summary_use_proj: + if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0: + num_classes = config.num_labels + else: + num_classes = config.hidden_size + self.summary = nn.Linear(config.hidden_size, num_classes) + + activation_string = getattr(config, "summary_activation", None) + self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity() + + self.first_dropout = nn.Identity() + if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0: + self.first_dropout = nn.Dropout(config.summary_first_dropout) + + self.last_dropout = nn.Identity() + if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0: + self.last_dropout = nn.Dropout(config.summary_last_dropout) + + def forward( + self, hidden_states: torch.FloatTensor, cls_index: torch.LongTensor | None = None + ) -> torch.FloatTensor: + """ + Compute a single vector summary of a sequence hidden states. + + Args: + hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`): + The hidden states of the last layer. + cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*): + Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token. + + Returns: + `torch.FloatTensor`: The summary of the sequence hidden states. + """ + if self.summary_type == "last": + output = hidden_states[:, -1] + elif self.summary_type == "first": + output = hidden_states[:, 0] + elif self.summary_type == "mean": + output = hidden_states.mean(dim=1) + elif self.summary_type == "cls_index": + if cls_index is None: + cls_index = torch.full_like( + hidden_states[..., :1, :], + hidden_states.shape[-2] - 1, + dtype=torch.long, + ) + else: + cls_index = cls_index.unsqueeze(-1).unsqueeze(-1) + cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),)) + # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states + output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size) + elif self.summary_type == "attn": + raise NotImplementedError + + output = self.first_dropout(output) + output = self.summary(output) + output = self.activation(output) + output = self.last_dropout(output) + + return output + + +@auto_docstring +class XLNetPreTrainedModel(PreTrainedModel): + config: XLNetConfig + base_model_prefix = "transformer" + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, XLNetRelativeAttention): + for param in [ + module.q, + module.k, + module.v, + module.o, + module.r, + module.r_r_bias, + module.r_s_bias, + module.r_w_bias, + module.seg_embed, + ]: + init.normal_(param, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, XLNetModel): + init.normal_(module.mask_emb, mean=0.0, std=self.config.initializer_range) + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetModel`]. + """ +) +class XLNetModelOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_predict, hidden_size)`): + Sequence of hidden-states at the last layer of the model. + + `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict` + corresponds to `sequence_length`. + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + last_hidden_state: torch.FloatTensor + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetLMHeadModel`]. + """ +) +class XLNetLMHeadModelOutput(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, num_predict, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + + `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict` + corresponds to `sequence_length`. + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetForSequenceClassification`]. + """ +) +class XLNetForSequenceClassificationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `label` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetForTokenClassificationOutput`]. + """ +) +class XLNetForTokenClassificationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): + Classification scores (before SoftMax). + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetForMultipleChoice`]. + """ +) +class XLNetForMultipleChoiceOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`): + *num_choices* is the second dimension of the input tensors. (see *input_ids* above). + + Classification scores (before SoftMax). + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetForQuestionAnsweringSimple`]. + """ +) +class XLNetForQuestionAnsweringSimpleOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length,)`): + Span-start scores (before SoftMax). + end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length,)`): + Span-end scores (before SoftMax). + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + start_logits: torch.FloatTensor | None = None + end_logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`XLNetForQuestionAnswering`]. + """ +) +class XLNetForQuestionAnsweringOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided): + Classification loss as the sum of start token, end token (and is_impossible if provided) classification + losses. + start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top config.start_n_top start token possibilities (beam-search). + start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top config.start_n_top start token possibilities (beam-search). + end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities + (beam-search). + end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search). + cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the `is_impossible` label of the answers. + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The + token ids which have their past given to this model should not be passed as `input_ids` as they have + already been computed. + """ + + loss: torch.FloatTensor | None = None + start_top_log_probs: torch.FloatTensor | None = None + start_top_index: torch.LongTensor | None = None + end_top_log_probs: torch.FloatTensor | None = None + end_top_index: torch.LongTensor | None = None + cls_logits: torch.FloatTensor | None = None + mems: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + + +@auto_docstring +class XLNetModel(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.mem_len = config.mem_len + self.reuse_len = config.reuse_len + self.d_model = config.d_model + self.same_length = config.same_length + self.attn_type = config.attn_type + self.bi_data = config.bi_data + self.clamp_len = config.clamp_len + self.n_layer = config.n_layer + + self.word_embedding = nn.Embedding(config.vocab_size, config.d_model) + self.mask_emb = nn.Parameter(torch.FloatTensor(1, 1, config.d_model)) + self.layer = nn.ModuleList([XLNetLayer(config) for _ in range(config.n_layer)]) + self.dropout = nn.Dropout(config.dropout) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.word_embedding + + def set_input_embeddings(self, new_embeddings): + self.word_embedding = new_embeddings + + def create_mask(self, qlen, mlen): + """ + Creates causal attention mask. Float mask where 1.0 indicates masked, 0.0 indicates not-masked. + + Args: + qlen: Sequence length + mlen: Mask length + + :: + + same_length=False: same_length=True: < qlen > < qlen > + ^ [0 0 0 0 0 1 1 1 1] [0 0 0 0 0 1 1 1 1] + [0 0 0 0 0 0 1 1 1] [1 0 0 0 0 0 1 1 1] + qlen [0 0 0 0 0 0 0 1 1] [1 1 0 0 0 0 0 1 1] + [0 0 0 0 0 0 0 0 1] [1 1 1 0 0 0 0 0 1] + v [0 0 0 0 0 0 0 0 0] [1 1 1 1 0 0 0 0 0] + + """ + mask = torch.ones((qlen, qlen + mlen), device=self.device) + if self.same_length: + mask_lo = mask[:, :qlen].tril(-1) + mask.triu_(mlen + 1) + mask[:, :qlen] += mask_lo + else: + mask.triu_(mlen + 1) + + return mask + + def cache_mem(self, curr_out, prev_mem): + # cache hidden states into memory. + if self.reuse_len is not None and self.reuse_len > 0: + curr_out = curr_out[: self.reuse_len] + + if self.mem_len is None or self.mem_len == 0: + # If `use_mems` is active but no `mem_len` is defined, the model behaves like GPT-2 at inference time + # and returns all of the past and current hidden states. + cutoff = 0 + else: + # If `use_mems` is active and `mem_len` is defined, the model returns the last `mem_len` hidden + # states. This is the preferred setting for training and long-form generation. + cutoff = -self.mem_len + if prev_mem is None: + # if `use_mems` is active and `mem_len` is defined, the model + new_mem = curr_out[cutoff:] + else: + new_mem = torch.cat([prev_mem, curr_out], dim=0)[cutoff:] + + return new_mem.detach() + + @staticmethod + def positional_embedding(pos_seq, inv_freq, bsz=None): + sinusoid_inp = torch.einsum("i,d->id", pos_seq, inv_freq) + pos_emb = torch.cat([torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)], dim=-1) + pos_emb = pos_emb[:, None, :] + + if bsz is not None: + pos_emb = pos_emb.expand(-1, bsz, -1) + + return pos_emb + + def relative_positional_encoding(self, qlen, klen, bsz=None, device=None): + # create relative positional encoding. + freq_seq = torch.arange(0, self.d_model, 2.0, dtype=torch.int64, device=device).float() + inv_freq = 1 / torch.pow(10000, (freq_seq / self.d_model)) + + if self.attn_type == "bi": + # beg, end = klen - 1, -qlen + beg, end = klen, -qlen + elif self.attn_type == "uni": + # beg, end = klen - 1, -1 + beg, end = klen, -1 + else: + raise ValueError(f"Unknown `attn_type` {self.attn_type}.") + + if self.bi_data: + fwd_pos_seq = torch.arange(beg, end, -1.0, dtype=torch.int64, device=device).float() + bwd_pos_seq = torch.arange(-beg, -end, 1.0, dtype=torch.int64, device=device).float() + + if self.clamp_len > 0: + fwd_pos_seq = fwd_pos_seq.clamp(-self.clamp_len, self.clamp_len) + bwd_pos_seq = bwd_pos_seq.clamp(-self.clamp_len, self.clamp_len) + + if bsz is not None: + fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2) + bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2) + else: + fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq) + bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq) + + pos_emb = torch.cat([fwd_pos_emb, bwd_pos_emb], dim=1) + else: + fwd_pos_seq = torch.arange(beg, end, -1.0, dtype=torch.int64, device=device).float() + if self.clamp_len > 0: + fwd_pos_seq = fwd_pos_seq.clamp(-self.clamp_len, self.clamp_len) + pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz) + + return pos_emb + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete after depreciation warning is removed + ) -> tuple | XLNetModelOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + """ + 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.training: + use_mems = use_mems if use_mems is not None else self.config.use_mems_train + else: + use_mems = use_mems if use_mems is not None else self.config.use_mems_eval + + # the original code for XLNet uses shapes [len, bsz] with the batch dimension at the end + # but we want a unified interface in the library with the batch size on the first dimension + # so we move here the first dimension (batch) to the end + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_ids = input_ids.transpose(0, 1).contiguous() + qlen, bsz = input_ids.shape[0], input_ids.shape[1] + elif inputs_embeds is not None: + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + qlen, bsz = inputs_embeds.shape[0], inputs_embeds.shape[1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + token_type_ids = token_type_ids.transpose(0, 1).contiguous() if token_type_ids is not None else None + input_mask = input_mask.transpose(0, 1).contiguous() if input_mask is not None else None + attention_mask = attention_mask.transpose(0, 1).contiguous() if attention_mask is not None else None + perm_mask = perm_mask.permute(1, 2, 0).contiguous() if perm_mask is not None else None + target_mapping = target_mapping.permute(1, 2, 0).contiguous() if target_mapping is not None else None + + mlen = mems[0].shape[0] if mems is not None and mems[0] is not None else 0 + klen = mlen + qlen + + dtype_float = self.dtype + device = self.device + + # Attention mask + # causal attention mask + if self.attn_type == "uni": + attn_mask = self.create_mask(qlen, mlen) + attn_mask = attn_mask[:, :, None, None] + elif self.attn_type == "bi": + attn_mask = None + else: + raise ValueError(f"Unsupported attention type: {self.attn_type}") + + # data mask: input mask & perm mask + assert input_mask is None or attention_mask is None, "You can only use one of input_mask (uses 1 for padding) " + "or attention_mask (uses 0 for padding, added for compatibility with BERT). Please choose one." + if input_mask is None and attention_mask is not None: + input_mask = 1.0 - attention_mask + if input_mask is not None and perm_mask is not None: + data_mask = input_mask[None] + perm_mask + elif input_mask is not None and perm_mask is None: + data_mask = input_mask[None] + elif input_mask is None and perm_mask is not None: + data_mask = perm_mask + else: + data_mask = None + + if data_mask is not None: + # all mems can be attended to + if mlen > 0: + mems_mask = torch.zeros([data_mask.shape[0], mlen, bsz]).to(data_mask) + data_mask = torch.cat([mems_mask, data_mask], dim=1) + if attn_mask is None: + attn_mask = data_mask[:, :, :, None] + else: + attn_mask += data_mask[:, :, :, None] + + if attn_mask is not None: + attn_mask = (attn_mask > 0).to(dtype_float) + + if attn_mask is not None: + non_tgt_mask = -torch.eye(qlen).to(attn_mask) + if mlen > 0: + non_tgt_mask = torch.cat([torch.zeros([qlen, mlen]).to(attn_mask), non_tgt_mask], dim=-1) + non_tgt_mask = ((attn_mask + non_tgt_mask[:, :, None, None]) > 0).to(attn_mask) + else: + non_tgt_mask = None + + # Word embeddings and prepare h & g hidden states + if inputs_embeds is not None: + word_emb_k = inputs_embeds + else: + word_emb_k = self.word_embedding(input_ids) + output_h = self.dropout(word_emb_k) + if target_mapping is not None: + word_emb_q = self.mask_emb.expand(target_mapping.shape[0], bsz, -1) + # else: # We removed the inp_q input which was same as target mapping + # inp_q_ext = inp_q[:, :, None] + # word_emb_q = inp_q_ext * self.mask_emb + (1 - inp_q_ext) * word_emb_k + output_g = self.dropout(word_emb_q) + else: + output_g = None + + # Segment embedding + if token_type_ids is not None: + # Convert `token_type_ids` to one-hot `seg_mat` + if mlen > 0: + mem_pad = torch.zeros([mlen, bsz], dtype=torch.long, device=device) + cat_ids = torch.cat([mem_pad, token_type_ids], dim=0) + else: + cat_ids = token_type_ids + + # `1` indicates not in the same segment [qlen x klen x bsz] + seg_mat = (token_type_ids[:, None] != cat_ids[None, :]).long() + seg_mat = nn.functional.one_hot(seg_mat, num_classes=2).to(dtype_float) + else: + seg_mat = None + + # Positional encoding + pos_emb = self.relative_positional_encoding(qlen, klen, bsz=bsz, device=output_h.device) + pos_emb = self.dropout(pos_emb) + + new_mems = () + if mems is None: + mems = [None] * len(self.layer) + + attentions = [] if output_attentions else None + hidden_states = [] if output_hidden_states else None + for i, layer_module in enumerate(self.layer): + if use_mems: + # cache new mems + new_mems = new_mems + (self.cache_mem(output_h, mems[i]),) + if output_hidden_states: + hidden_states.append((output_h, output_g) if output_g is not None else output_h) + + outputs = layer_module( + output_h, + output_g, + attn_mask_h=non_tgt_mask, + attn_mask_g=attn_mask, + r=pos_emb, + seg_mat=seg_mat, + mems=mems[i], + target_mapping=target_mapping, + output_attentions=output_attentions, + ) + output_h, output_g = outputs[:2] + if output_attentions: + attentions.append(outputs[2]) + + # Add last hidden state + if output_hidden_states: + hidden_states.append((output_h, output_g) if output_g is not None else output_h) + + output = self.dropout(output_g if output_g is not None else output_h) + + # Prepare outputs, we transpose back here to shape [bsz, len, hidden_dim] (cf. beginning of forward() method) + output = output.permute(1, 0, 2).contiguous() + + if not use_mems: + new_mems = None + + if output_hidden_states: + if output_g is not None: + hidden_states = tuple(h.permute(1, 0, 2).contiguous() for hs in hidden_states for h in hs) + else: + hidden_states = tuple(hs.permute(1, 0, 2).contiguous() for hs in hidden_states) + + if output_attentions: + if target_mapping is not None: + # when target_mapping is provided, there are 2-tuple of attentions + attentions = tuple( + tuple(att_stream.permute(2, 3, 0, 1).contiguous() for att_stream in t) for t in attentions + ) + else: + attentions = tuple(t.permute(2, 3, 0, 1).contiguous() for t in attentions) + + if not return_dict: + return tuple(v for v in [output, new_mems, hidden_states, attentions] if v is not None) + + return XLNetModelOutput( + last_hidden_state=output, mems=new_mems, hidden_states=hidden_states, attentions=attentions + ) + + +@auto_docstring( + custom_intro=""" + XLNet Model with a language modeling head on top (linear layer with weights tied to the input embeddings). + """ +) +class XLNetLMHeadModel(XLNetPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_loss.weight": "transformer.word_embedding.weight"} + + def __init__(self, config): + super().__init__(config) + self.attn_type = config.attn_type + self.same_length = config.same_length + + self.transformer = XLNetModel(config) + self.lm_loss = nn.Linear(config.d_model, config.vocab_size, bias=True) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_loss + + def set_output_embeddings(self, new_embeddings): + self.lm_loss = new_embeddings + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, use_mems=None, is_first_iteration=False, **kwargs + ): + # Overwritten -- this model has unique input preparation + + # Add dummy token at the end (no attention on this one) + + effective_batch_size = input_ids.shape[0] + dummy_token = torch.zeros((effective_batch_size, 1), dtype=torch.long, device=input_ids.device) + + # At every pass, the attention values for the new token and the two last generated tokens + # are computed, the rest is reloaded from the `past` cache. A purely auto-regressive model would have + # offset = 1; offset = 2 seems to have slightly better computation. + offset = 2 + + if past_key_values: + input_ids = torch.cat([input_ids[:, -offset:], dummy_token], dim=1) + else: + input_ids = torch.cat([input_ids, dummy_token], dim=1) + + # Build permutation mask so that previous tokens don't see last token + sequence_length = input_ids.shape[1] + perm_mask = torch.zeros( + (effective_batch_size, sequence_length, sequence_length), dtype=torch.float, device=input_ids.device + ) + perm_mask[:, :, -1] = 1.0 + + # We'll only predict the last token + target_mapping = torch.zeros( + (effective_batch_size, 1, sequence_length), dtype=torch.float, device=input_ids.device + ) + target_mapping[:, 0, -1] = 1.0 + + model_inputs = { + "input_ids": input_ids, + "perm_mask": perm_mask, + "target_mapping": target_mapping, + "use_mems": use_mems, + } + + # if past is defined in model kwargs then use it for faster decoding + if past_key_values: + model_inputs["mems"] = tuple(layer_past[:-offset, :, :] for layer_past in past_key_values) + + # Attention mask is computed on the fly on XLNetModel.forward() + kwargs.pop("attention_mask", None) + # TODO: Ignoring use_cache should not happen, fixme. + kwargs.pop("use_cache", None) + # Forward ALL kwargs that are uninitialized (e.g. `use_cache`). + for key, value in kwargs.items(): + if key not in model_inputs: + model_inputs[key] = value + + return model_inputs + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetLMHeadModelOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + labels (`torch.LongTensor` of shape `(batch_size, num_predict)`, *optional*): + Labels for masked language modeling. `num_predict` corresponds to `target_mapping.shape[1]`. If + `target_mapping` is `None`, then `num_predict` corresponds to `sequence_length`. + + The labels should correspond to the masked input words that should be predicted and depends on + `target_mapping`. Note in order to perform standard auto-regressive language modeling a ** token has + to be added to the `input_ids` (see the `prepare_inputs_for_generation` function and examples below) + + Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored, the loss + is only computed for labels in `[0, ..., config.vocab_size]` + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, XLNetLMHeadModel + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("xlnet/xlnet-large-cased") + >>> model = XLNetLMHeadModel.from_pretrained("xlnet/xlnet-large-cased") + + >>> # We show how to setup inputs to predict a next token using a bi-directional context. + >>> input_ids = torch.tensor( + ... tokenizer.encode("Hello, my dog is very ", add_special_tokens=False) + ... ).unsqueeze( + ... 0 + ... ) # We will predict the masked token + >>> perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float) + >>> perm_mask[:, :, -1] = 1.0 # Previous tokens don't see last token + >>> target_mapping = torch.zeros( + ... (1, 1, input_ids.shape[1]), dtype=torch.float + ... ) # Shape [1, 1, seq_length] => let's predict one token + >>> target_mapping[ + ... 0, 0, -1 + ... ] = 1.0 # Our first (and only) prediction will be the last token of the sequence (the masked token) + + >>> outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping) + >>> next_token_logits = outputs[ + ... 0 + ... ] # Output has shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size] + + >>> # The same way can the XLNetLMHeadModel be used to be trained by standard auto-regressive language modeling. + >>> input_ids = torch.tensor( + ... tokenizer.encode("Hello, my dog is very ", add_special_tokens=False) + ... ).unsqueeze( + ... 0 + ... ) # We will predict the masked token + >>> labels = torch.tensor(tokenizer.encode("cute", add_special_tokens=False)).unsqueeze(0) + >>> assert labels.shape[0] == 1, "only one word will be predicted" + >>> perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float) + >>> perm_mask[ + ... :, :, -1 + ... ] = 1.0 # Previous tokens don't see last token as is done in standard auto-regressive lm training + >>> target_mapping = torch.zeros( + ... (1, 1, input_ids.shape[1]), dtype=torch.float + ... ) # Shape [1, 1, seq_length] => let's predict one token + >>> target_mapping[ + ... 0, 0, -1 + ... ] = 1.0 # Our first (and only) prediction will be the last token of the sequence (the masked token) + + >>> outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping, labels=labels) + >>> loss = outputs.loss + >>> next_token_logits = ( + ... outputs.logits + ... ) # Logits have shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size] + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + token_type_ids=token_type_ids, + input_mask=input_mask, + inputs_embeds=inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = transformer_outputs[0] + # Only compute necessary logits + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_loss(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + # Flatten the tokens + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1)) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return XLNetLMHeadModelOutput( + loss=loss, + logits=logits, + mems=transformer_outputs.mems, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache(mems: list[torch.Tensor], beam_idx: torch.Tensor) -> list[torch.Tensor]: + """ + This function is used to re-order the `mems` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `mems` with the correct beam_idx at every + generation step. + """ + return [layer_past.index_select(1, beam_idx.to(layer_past.device)) for layer_past in mems] + + +@auto_docstring( + custom_intro=""" + XLNet Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. + for GLUE tasks. + """ +) +class XLNetForSequenceClassification(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + + self.transformer = XLNetModel(config) + self.sequence_summary = XLNetSequenceSummary(config) + self.logits_proj = nn.Linear(config.d_model, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetForSequenceClassificationOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + 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). + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + token_type_ids=token_type_ids, + input_mask=input_mask, + inputs_embeds=inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + output = transformer_outputs[0] + + output = self.sequence_summary(output) + logits = self.logits_proj(output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return XLNetForSequenceClassificationOutput( + loss=loss, + logits=logits, + mems=transformer_outputs.mems, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class XLNetForTokenClassification(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = XLNetModel(config) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetForTokenClassificationOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]` + where *num_choices* is the size of the second dimension of the input tensors. (see *input_ids* above) + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks.emory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + token_type_ids=token_type_ids, + input_mask=input_mask, + inputs_embeds=inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return XLNetForTokenClassificationOutput( + loss=loss, + logits=logits, + mems=outputs.mems, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class XLNetForMultipleChoice(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.transformer = XLNetModel(config) + self.sequence_summary = XLNetSequenceSummary(config) + self.logits_proj = nn.Linear(config.d_model, 1) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetForMultipleChoiceOutput: + r""" + input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + input_mask (`torch.FloatTensor` of shape `batch_size, num_choices, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + flat_input_mask = input_mask.view(-1, input_mask.size(-1)) if input_mask is not None else None + flat_inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + transformer_outputs = self.transformer( + flat_input_ids, + token_type_ids=flat_token_type_ids, + input_mask=flat_input_mask, + attention_mask=flat_attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + inputs_embeds=flat_inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + output = transformer_outputs[0] + + output = self.sequence_summary(output) + logits = self.logits_proj(output) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels.view(-1)) + + if not return_dict: + output = (reshaped_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return XLNetForMultipleChoiceOutput( + loss=loss, + logits=reshaped_logits, + mems=transformer_outputs.mems, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + XLNet Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear + layers on top of the hidden-states output to compute `span start logits` and `span end logits`). + """ +) +class XLNetForQuestionAnsweringSimple(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = XLNetModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.Tensor | None = None, + end_positions: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetForQuestionAnsweringSimpleOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + token_type_ids=token_type_ids, + input_mask=input_mask, + inputs_embeds=inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[1:] + return ((total_loss,) + output) if total_loss is not None else output + + return XLNetForQuestionAnsweringSimpleOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + mems=outputs.mems, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class XLNetForQuestionAnswering(XLNetPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.start_n_top = config.start_n_top + self.end_n_top = config.end_n_top + + self.transformer = XLNetModel(config) + self.start_logits = XLNetPoolerStartLogits(config) + self.end_logits = XLNetPoolerEndLogits(config) + self.answer_class = XLNetPoolerAnswerClass(config) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + mems: torch.Tensor | None = None, + perm_mask: torch.Tensor | None = None, + target_mapping: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + input_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + start_positions: torch.Tensor | None = None, + end_positions: torch.Tensor | None = None, + is_impossible: torch.Tensor | None = None, + cls_index: torch.Tensor | None = None, + p_mask: torch.Tensor | None = None, + use_mems: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, # delete when `use_cache` is removed in XLNetModel + ) -> tuple | XLNetForQuestionAnsweringOutput: + r""" + mems (`list[torch.FloatTensor]` of length `config.n_layers`): + Contains pre-computed hidden-states (see `mems` output below) . Can be used to speed up sequential + decoding. The token ids which have their past given to this model should not be passed as `input_ids` as + they have already been computed. + + `use_mems` has to be set to `True` to make use of `mems`. + perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*): + Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`: + + - if `perm_mask[k, i, j] = 0`, i attend to j in batch k; + - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k. + + If not set, each token attends to all the others (full bidirectional attention). Only used during + pretraining (to define factorization order) or for sequential decoding (generation). + target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*): + Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is + on the j-th token. Only used during pretraining for partial prediction or for sequential decoding + (generation). + input_mask (`torch.FloatTensor` of shape `batch_size, sequence_length`, *optional*): + Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for + real tokens and 1 for padding which is kept for compatibility with the original code base. + + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **masked**, + - 0 for tokens that are **not masked**. + + You can only uses one of `input_mask` and `attention_mask`. + is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels whether a question has an answer or no answer (SQuAD 2.0) + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the classification token to use as input for computing plausibility of the + answer. + p_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Optional mask of tokens which can't be in answers (e.g. [CLS], [PAD], ...). 1.0 means token should be + masked. 0.0 mean token is not masked. + use_mems (`bool`, *optional*): + Whether to use memory states to speed up sequential decoding. If set to `True`, the model will use the hidden + states from previous forward passes to compute attention, which can significantly improve performance for + sequential decoding tasks. + + Example: + + ```python + >>> from transformers import AutoTokenizer, XLNetForQuestionAnswering + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("xlnet/xlnet-base-cased") + >>> model = XLNetForQuestionAnswering.from_pretrained("xlnet/xlnet-base-cased") + + >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze( + ... 0 + ... ) # Batch size 1 + >>> start_positions = torch.tensor([1]) + >>> end_positions = torch.tensor([3]) + >>> outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions) + + >>> loss = outputs.loss + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + transformer_outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + mems=mems, + perm_mask=perm_mask, + target_mapping=target_mapping, + token_type_ids=token_type_ids, + input_mask=input_mask, + inputs_embeds=inputs_embeds, + use_mems=use_mems, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + hidden_states = transformer_outputs[0] + start_logits = self.start_logits(hidden_states, p_mask=p_mask) + + outputs = transformer_outputs[1:] # Keep mems, hidden states, attentions if there are in it + + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, let's remove the dimension added by batch splitting + for x in (start_positions, end_positions, cls_index, is_impossible): + if x is not None and x.dim() > 1: + x.squeeze_(-1) + + # during training, compute the end logits based on the ground truth of the start position + end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask) + + loss_fct = CrossEntropyLoss() + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if cls_index is not None and is_impossible is not None: + # Predict answerability from the representation of CLS and START + cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index) + loss_fct_cls = nn.BCEWithLogitsLoss() + cls_loss = loss_fct_cls(cls_logits, is_impossible) + + # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss + total_loss += cls_loss * 0.5 + + if not return_dict: + return (total_loss,) + transformer_outputs[1:] + else: + return XLNetForQuestionAnsweringOutput( + loss=total_loss, + mems=transformer_outputs.mems, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + # during inference, compute the end logits based on beam search + bsz, slen, hsz = hidden_states.size() + start_log_probs = nn.functional.softmax(start_logits, dim=-1) # shape (bsz, slen) + + start_top_log_probs, start_top_index = torch.topk( + start_log_probs, self.start_n_top, dim=-1 + ) # shape (bsz, start_n_top) + start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz) + start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz) + start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz) + + hidden_states_expanded = hidden_states.unsqueeze(2).expand_as( + start_states + ) # shape (bsz, slen, start_n_top, hsz) + p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None + end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask) + end_log_probs = nn.functional.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top) + + end_top_log_probs, end_top_index = torch.topk( + end_log_probs, self.end_n_top, dim=1 + ) # shape (bsz, end_n_top, start_n_top) + end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top) + end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top) + + start_states = torch.einsum( + "blh,bl->bh", hidden_states, start_log_probs + ) # get the representation of START as weighted sum of hidden states + cls_logits = self.answer_class( + hidden_states, start_states=start_states, cls_index=cls_index + ) # Shape (batch size,): one single `cls_logits` for each sample + + if not return_dict: + outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + return outputs + transformer_outputs[1:] + else: + return XLNetForQuestionAnsweringOutput( + start_top_log_probs=start_top_log_probs, + start_top_index=start_top_index, + end_top_log_probs=end_top_log_probs, + end_top_index=end_top_index, + cls_logits=cls_logits, + mems=transformer_outputs.mems, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +__all__ = [ + "XLNetForMultipleChoice", + "XLNetForQuestionAnswering", + "XLNetForQuestionAnsweringSimple", + "XLNetForSequenceClassification", + "XLNetForTokenClassification", + "XLNetLMHeadModel", + "XLNetModel", + "XLNetPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/xlnet/tokenization_xlnet.py b/third_party/transformers/src/transformers/models/xlnet/tokenization_xlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..0b682564bffd0be1b3cab092595ca0030ed07614 --- /dev/null +++ b/third_party/transformers/src/transformers/models/xlnet/tokenization_xlnet.py @@ -0,0 +1,187 @@ +# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for XLNet model.""" + +from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import Unigram + +from ...tokenization_utils_base import _get_prepend_scheme +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"} + +SPIECE_UNDERLINE = "▁" + +# Segments (not really needed) +SEG_ID_A = 0 +SEG_ID_B = 1 +SEG_ID_CLS = 2 +SEG_ID_SEP = 3 +SEG_ID_PAD = 4 + + +class XLNetTokenizer(TokenizersBackend): + """ + Construct a XLNet tokenizer (backed by HuggingFace's *tokenizers* library). Based on + [Unigram](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=unigram#models). + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab (`list of tuples`, *optional*): + List of (token, score) tuples for Unigram model. If not provided, an empty list is used. + unk_id (`int`, *optional*, defaults to 0): + The ID of the unknown token in the vocabulary. + do_lower_case (`bool`, *optional*, defaults to `False`): + Whether to lowercase the input when tokenizing. + remove_space (`bool`, *optional*, defaults to `True`): + Whether to strip the text when tokenizing (removing excess spaces before and after the string). + keep_accents (`bool`, *optional*, defaults to `False`): + Whether to keep accents when tokenizing. + bos_token (`str`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + + + + When building a sequence using special tokens, this is not the token that is used for the beginning of + sequence. The token used is the `cls_token`. + + + + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + + + + When building a sequence using special tokens, this is not the token that is used for the end of sequence. + The token used is the `sep_token`. + + + + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + additional_special_tokens (`list[str]`, *optional*, defaults to `["", ""]`): + Additional special tokens used by the tokenizer. + """ + + vocab_files_names = VOCAB_FILES_NAMES + padding_side = "left" + model = Unigram + + def __init__( + self, + vocab: str | list[tuple[str, float]] | None = None, + unk_id: int = 0, + do_lower_case=False, + remove_space=True, + keep_accents=False, + bos_token="", + eos_token="", + unk_token="", + sep_token="", + pad_token="", + cls_token="", + mask_token="", + additional_special_tokens=None, + **kwargs, + ): + if additional_special_tokens is None: + additional_special_tokens = ["", ""] + + if vocab is not None: + self._vocab = vocab + else: + self._vocab = [(str(unk_token), 0.0)] + + self._tokenizer = Tokenizer( + Unigram( + self._vocab, + unk_id=unk_id, + byte_fallback=False, + ) + ) + + list_normalizers = [ + normalizers.Replace("``", '"'), + normalizers.Replace("''", '"'), + ] + # if not keep_accents: + list_normalizers.append(normalizers.NFKD()) + list_normalizers.append(normalizers.StripAccents()) + if do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + + list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " ")) + self._tokenizer.normalizer = normalizers.Sequence(list_normalizers) + + add_prefix_space = True + prepend_scheme = _get_prepend_scheme(add_prefix_space, self) + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.WhitespaceSplit(), + pre_tokenizers.Metaspace(replacement="▁", prepend_scheme=prepend_scheme), + ] + ) + + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme=prepend_scheme) + self._pad_token_type_id = 3 + self.do_lower_case = do_lower_case + self.remove_space = remove_space + self.keep_accents = keep_accents + mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token + super().__init__( + unk_id=unk_id, + do_lower_case=do_lower_case, + remove_space=remove_space, + keep_accents=keep_accents, + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=f"$A:0 {str(self.sep_token)}:0 {str(self.cls_token)}:2", + pair=f"$A:0 {str(self.sep_token)}:0 $B:1 {str(self.sep_token)}:1 {str(self.cls_token)}:2", + special_tokens=[ + (str(self.sep_token), self.sep_token_id), + (str(self.cls_token), self.cls_token_id), + ], + ) + + +__all__ = ["XLNetTokenizer"]