diff --git a/third_party/transformers/src/transformers/__init__.py b/third_party/transformers/src/transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..14a8548d03997fa2b9ceb63c8ac3b6e63f5d4ca7 --- /dev/null +++ b/third_party/transformers/src/transformers/__init__.py @@ -0,0 +1,852 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# When adding a new object to this init, remember to add it twice: once inside the `_import_structure` dictionary and +# once inside the `if TYPE_CHECKING` branch. The `TYPE_CHECKING` should have import statements as usual, but they are +# only there for type checking. The `_import_structure` is a dictionary submodule to list of object names, and is used +# to defer the actual importing for when the objects are requested. This way `import transformers` provides the names +# in the namespace without actually importing anything (and especially none of the backends). + +__version__ = "5.6.0.dev0" + +import importlib +import sys +import types +from pathlib import Path +from typing import TYPE_CHECKING + +# Check the dependencies satisfy the minimal versions required. +from . import dependency_versions_check +from .utils import ( + OptionalDependencyNotAvailable, + _LazyModule, + is_essentia_available, + is_g2p_en_available, + is_librosa_available, + is_mistral_common_available, + is_mlx_available, + is_numba_available, + is_pretty_midi_available, +) + +# Note: the following symbols are deliberately exported with `as` +# so that mypy, pylint or other static linters can recognize them, +# given that they are not exported using `__all__` in this file. +from .utils import is_bitsandbytes_available as is_bitsandbytes_available +from .utils import is_scipy_available as is_scipy_available +from .utils import is_sentencepiece_available as is_sentencepiece_available +from .utils import is_speech_available as is_speech_available +from .utils import is_timm_available as is_timm_available +from .utils import is_tokenizers_available as is_tokenizers_available +from .utils import is_torch_available as is_torch_available +from .utils import is_torchaudio_available as is_torchaudio_available +from .utils import is_torchvision_available as is_torchvision_available +from .utils import is_vision_available as is_vision_available +from .utils import logging as logging +from .utils.import_utils import define_import_structure + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Base objects, independent of any specific backend +_import_structure = { + "audio_utils": [], + "cli": [], + "configuration_utils": ["PreTrainedConfig", "PretrainedConfig"], + "convert_slow_tokenizers_checkpoints_to_fast": [], + "data": [ + "DataProcessor", + "InputExample", + "InputFeatures", + "SingleSentenceClassificationProcessor", + "SquadExample", + "SquadFeatures", + "SquadV1Processor", + "SquadV2Processor", + "glue_compute_metrics", + "glue_convert_examples_to_features", + "glue_output_modes", + "glue_processors", + "glue_tasks_num_labels", + "squad_convert_examples_to_features", + "xnli_compute_metrics", + "xnli_output_modes", + "xnli_processors", + "xnli_tasks_num_labels", + ], + "data.data_collator": [ + "DataCollator", + "DataCollatorForLanguageModeling", + "DataCollatorForMultipleChoice", + "DataCollatorForPermutationLanguageModeling", + "DataCollatorForSeq2Seq", + "DataCollatorForSOP", + "DataCollatorForTokenClassification", + "DataCollatorForWholeWordMask", + "DataCollatorWithFlattening", + "DataCollatorWithPadding", + "DefaultDataCollator", + "default_data_collator", + ], + "data.metrics": [], + "data.processors": [], + "debug_utils": [], + "dependency_versions_check": [], + "dependency_versions_table": [], + "dynamic_module_utils": [], + "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], + "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], + "file_utils": [], + "generation": [ + "AsyncTextIteratorStreamer", + "CompileConfig", + "ContinuousBatchingConfig", + "GenerationConfig", + "TextIteratorStreamer", + "TextStreamer", + "WatermarkingConfig", + ], + "hf_argparser": ["HfArgumentParser"], + "hyperparameter_search": [], + "image_processing_utils_fast": [], + "image_transforms": [], + "integrations": [ + "is_clearml_available", + "is_comet_available", + "is_dvclive_available", + "is_neptune_available", + "is_optuna_available", + "is_ray_available", + "is_ray_tune_available", + "is_swanlab_available", + "is_tensorboard_available", + "is_trackio_available", + "is_wandb_available", + ], + "loss": [], + "pipelines": [ + "AnyToAnyPipeline", + "AudioClassificationPipeline", + "AutomaticSpeechRecognitionPipeline", + "CsvPipelineDataFormat", + "DepthEstimationPipeline", + "DocumentQuestionAnsweringPipeline", + "FeatureExtractionPipeline", + "FillMaskPipeline", + "ImageClassificationPipeline", + "ImageFeatureExtractionPipeline", + "ImageSegmentationPipeline", + "ImageTextToTextPipeline", + "JsonPipelineDataFormat", + "KeypointMatchingPipeline", + "MaskGenerationPipeline", + "NerPipeline", + "ObjectDetectionPipeline", + "PipedPipelineDataFormat", + "Pipeline", + "PipelineDataFormat", + "TableQuestionAnsweringPipeline", + "TextClassificationPipeline", + "TextGenerationPipeline", + "TextToAudioPipeline", + "TokenClassificationPipeline", + "VideoClassificationPipeline", + "ZeroShotAudioClassificationPipeline", + "ZeroShotClassificationPipeline", + "ZeroShotImageClassificationPipeline", + "ZeroShotObjectDetectionPipeline", + "pipeline", + ], + "processing_utils": [ + "AudioKwargs", + "ImagesKwargs", + "ProcessingKwargs", + "ProcessorMixin", + "TextKwargs", + "VideosKwargs", + ], + "quantizers": [], + "testing_utils": [], + "tokenization_python": ["PreTrainedTokenizer", "PythonBackend"], + "tokenization_utils": [], + "tokenization_utils_base": [ + "AddedToken", + "BatchEncoding", + "CharSpan", + "PreTrainedTokenizerBase", + "TokenSpan", + ], + "tokenization_utils_fast": [], + "tokenization_utils_sentencepiece": ["SentencePieceBackend"], + "trainer_callback": [ + "DefaultFlowCallback", + "EarlyStoppingCallback", + "PrinterCallback", + "ProgressCallback", + "TrainerCallback", + "TrainerControl", + "TrainerState", + ], + "trainer_utils": [ + "EvalPrediction", + "IntervalStrategy", + "SchedulerType", + "enable_full_determinism", + "set_seed", + ], + "training_args": ["TrainingArguments"], + "training_args_seq2seq": ["Seq2SeqTrainingArguments"], + "utils": [ + "CONFIG_NAME", + "MODEL_CARD_NAME", + "SPIECE_UNDERLINE", + "WEIGHTS_NAME", + "TensorType", + "add_end_docstrings", + "add_start_docstrings", + "is_apex_available", + "is_av_available", + "is_bitsandbytes_available", + "is_datasets_available", + "is_faiss_available", + "is_matplotlib_available", + "is_mlx_available", + "is_phonemizer_available", + "is_psutil_available", + "is_py3nvml_available", + "is_pyctcdecode_available", + "is_sacremoses_available", + "is_scipy_available", + "is_sentencepiece_available", + "is_sklearn_available", + "is_speech_available", + "is_timm_available", + "is_tokenizers_available", + "is_torch_available", + "is_torch_hpu_available", + "is_torch_mlu_available", + "is_torch_musa_available", + "is_torch_neuroncore_available", + "is_torch_npu_available", + "is_torchvision_available", + "is_torch_xla_available", + "is_torch_xpu_available", + "is_vision_available", + "logging", + ], + "utils.import_utils": ["requires_backends"], + "utils.kernel_config": ["KernelConfig"], + "utils.quantization_config": [ + "AqlmConfig", + "AutoRoundConfig", + "AwqConfig", + "BitNetQuantConfig", + "BitsAndBytesConfig", + "CompressedTensorsConfig", + "EetqConfig", + "FbgemmFp8Config", + "FineGrainedFP8Config", + "FourOverSixConfig", + "FPQuantConfig", + "GPTQConfig", + "HiggsConfig", + "HqqConfig", + "MetalConfig", + "Mxfp4Config", + "QuantoConfig", + "QuarkConfig", + "SinqConfig", + "SpQRConfig", + "TorchAoConfig", + "VptqConfig", + ], + "video_utils": [], +} + +# tokenizers-backed objects +try: + if not is_tokenizers_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_tokenizers_objects + + _import_structure["utils.dummy_tokenizers_objects"] = [ + name for name in dir(dummy_tokenizers_objects) if not name.startswith("_") + ] +else: + # Fast tokenizers structure + _import_structure["tokenization_utils_tokenizers"] = [ + "PreTrainedTokenizerFast", + "TokenizersBackend", + ] + + +try: + if not (is_sentencepiece_available() and is_tokenizers_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_sentencepiece_and_tokenizers_objects + + _import_structure["utils.dummy_sentencepiece_and_tokenizers_objects"] = [ + name for name in dir(dummy_sentencepiece_and_tokenizers_objects) if not name.startswith("_") + ] +else: + _import_structure["convert_slow_tokenizer"] = [ + "SLOW_TO_FAST_CONVERTERS", + "convert_slow_tokenizer", + ] + +try: + if not (is_mistral_common_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_mistral_common_objects + + _import_structure["utils.dummy_mistral_common_objects"] = [ + name for name in dir(dummy_mistral_common_objects) if not name.startswith("_") + ] +else: + _import_structure["tokenization_mistral_common"] = ["MistralCommonBackend"] + +# Vision-specific objects +try: + if not is_vision_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_vision_objects + + _import_structure["utils.dummy_vision_objects"] = [ + name for name in dir(dummy_vision_objects) if not name.startswith("_") + ] +else: + _import_structure["image_processing_backends"] = ["PilBackend"] + _import_structure["image_processing_base"] = ["ImageProcessingMixin"] + _import_structure["image_processing_utils"] = ["BaseImageProcessor"] + _import_structure["image_utils"] = ["ImageFeatureExtractionMixin"] + +try: + if not is_torchvision_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_torchvision_objects + + _import_structure["utils.dummy_torchvision_objects"] = [ + name for name in dir(dummy_torchvision_objects) if not name.startswith("_") + ] +else: + _import_structure.setdefault("image_processing_backends", []) + _import_structure["image_processing_backends"] += ["TorchvisionBackend"] + _import_structure["video_processing_utils"] = ["BaseVideoProcessor"] + +# PyTorch-backed objects +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_pt_objects + + _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")] +else: + _import_structure["model_debugging_utils"] = [ + "model_addition_debugger_context", + ] + _import_structure["activations"] = [] + _import_structure["cache_utils"] = [ + "CacheLayerMixin", + "DynamicLayer", + "StaticLayer", + "StaticSlidingWindowLayer", + "QuantoQuantizedLayer", + "HQQQuantizedLayer", + "Cache", + "DynamicCache", + "EncoderDecoderCache", + "QuantizedCache", + "StaticCache", + ] + _import_structure["data.datasets"] = [ + "GlueDataset", + "GlueDataTrainingArguments", + "SquadDataset", + "SquadDataTrainingArguments", + ] + _import_structure["generation"].extend( + [ + "AlternatingCodebooksLogitsProcessor", + "BayesianDetectorConfig", + "BayesianDetectorModel", + "ClassifierFreeGuidanceLogitsProcessor", + "ContinuousBatchingManager", + "ContinuousMixin", + "EncoderNoRepeatNGramLogitsProcessor", + "EncoderRepetitionPenaltyLogitsProcessor", + "EosTokenCriteria", + "EpsilonLogitsWarper", + "MinPLogitsWarper", + "EtaLogitsWarper", + "ExponentialDecayLengthPenalty", + "ForcedBOSTokenLogitsProcessor", + "ForcedEOSTokenLogitsProcessor", + "GenerationMixin", + "InfNanRemoveLogitsProcessor", + "LogitNormalization", + "LogitsProcessor", + "LogitsProcessorList", + "MaxLengthCriteria", + "MaxTimeCriteria", + "MinLengthLogitsProcessor", + "MinNewTokensLengthLogitsProcessor", + "NoBadWordsLogitsProcessor", + "NoRepeatNGramLogitsProcessor", + "PrefixConstrainedLogitsProcessor", + "RepetitionPenaltyLogitsProcessor", + "SequenceBiasLogitsProcessor", + "StoppingCriteria", + "StoppingCriteriaList", + "StopStringCriteria", + "SuppressTokensAtBeginLogitsProcessor", + "SuppressTokensLogitsProcessor", + "SynthIDTextWatermarkDetector", + "SynthIDTextWatermarkingConfig", + "SynthIDTextWatermarkLogitsProcessor", + "TemperatureLogitsWarper", + "TopHLogitsWarper", + "TopKLogitsWarper", + "TopPLogitsWarper", + "TypicalLogitsWarper", + "UnbatchedClassifierFreeGuidanceLogitsProcessor", + "WatermarkDetector", + "WatermarkLogitsProcessor", + "WhisperTimeStampLogitsProcessor", + ] + ) + + # PyTorch domain libraries integration + _import_structure["integrations.executorch"] = [ + "TorchExportableModuleWithStaticCache", + "convert_and_export_with_cache", + ] + + _import_structure["core_model_loading"] = [ + "Chunk", + "Concatenate", + "ConversionOps", + "MergeModulelist", + "PermuteForRope", + "SplitModulelist", + "WeightConverter", + ] + _import_structure["modeling_flash_attention_utils"] = [] + _import_structure["modeling_layers"] = ["GradientCheckpointingLayer"] + _import_structure["modeling_outputs"] = [] + _import_structure["backbone_utils"] = ["BackboneConfigMixin", "BackboneMixin"] + _import_structure["modeling_rope_utils"] = ["ROPE_INIT_FUNCTIONS", "dynamic_rope_update", "RopeParameters"] + _import_structure["modeling_utils"] = ["PreTrainedModel", "AttentionInterface"] + _import_structure["masking_utils"] = ["AttentionMaskInterface"] + _import_structure["optimization"] = [ + "Adafactor", + "get_constant_schedule", + "get_constant_schedule_with_warmup", + "get_cosine_schedule_with_warmup", + "get_cosine_with_hard_restarts_schedule_with_warmup", + "get_cosine_with_min_lr_schedule_with_warmup", + "get_cosine_with_min_lr_schedule_with_warmup_lr_rate", + "get_greedy_schedule", + "get_inverse_sqrt_schedule", + "get_linear_schedule_with_warmup", + "get_polynomial_decay_schedule_with_warmup", + "get_reduce_on_plateau_schedule", + "get_scheduler", + "get_wsd_schedule", + "GreedyLR", + ] + _import_structure["pytorch_utils"] = ["Conv1D", "apply_chunking_to_forward"] + _import_structure["time_series_utils"] = [] + _import_structure["trainer"] = ["Trainer"] + _import_structure["trainer_pt_utils"] = ["torch_distributed_zero_first"] + _import_structure["trainer_seq2seq"] = ["Seq2SeqTrainer"] + + +# Direct imports for type-checking +if TYPE_CHECKING: + # All modeling imports + # Models + from .backbone_utils import BackboneConfigMixin, BackboneMixin + from .cache_utils import Cache as Cache + from .cache_utils import DynamicCache as DynamicCache + from .cache_utils import DynamicLayer as DynamicLayer + from .cache_utils import EncoderDecoderCache as EncoderDecoderCache + from .cache_utils import HQQQuantizedLayer as HQQQuantizedLayer + from .cache_utils import QuantizedCache as QuantizedCache + from .cache_utils import QuantoQuantizedLayer as QuantoQuantizedLayer + from .cache_utils import StaticCache as StaticCache + from .cache_utils import StaticLayer as StaticLayer + from .cache_utils import StaticSlidingWindowLayer as StaticSlidingWindowLayer + from .configuration_utils import PreTrainedConfig as PreTrainedConfig + from .configuration_utils import PretrainedConfig as PretrainedConfig + from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS as SLOW_TO_FAST_CONVERTERS + from .convert_slow_tokenizer import convert_slow_tokenizer as convert_slow_tokenizer + from .core_model_loading import Chunk as Chunk + from .core_model_loading import Concatenate as Concatenate + from .core_model_loading import ConversionOps as ConversionOps + from .core_model_loading import MergeModulelist as MergeModulelist + from .core_model_loading import PermuteForRope as PermuteForRope + from .core_model_loading import SplitModulelist as SplitModulelist + from .core_model_loading import WeightConverter as WeightConverter + + # Data + from .data import DataProcessor as DataProcessor + from .data import InputExample as InputExample + from .data import InputFeatures as InputFeatures + from .data import SingleSentenceClassificationProcessor as SingleSentenceClassificationProcessor + from .data import SquadExample as SquadExample + from .data import SquadFeatures as SquadFeatures + from .data import SquadV1Processor as SquadV1Processor + from .data import SquadV2Processor as SquadV2Processor + from .data import glue_compute_metrics as glue_compute_metrics + from .data import glue_convert_examples_to_features as glue_convert_examples_to_features + from .data import glue_output_modes as glue_output_modes + from .data import glue_processors as glue_processors + from .data import glue_tasks_num_labels as glue_tasks_num_labels + from .data import squad_convert_examples_to_features as squad_convert_examples_to_features + from .data import xnli_compute_metrics as xnli_compute_metrics + from .data import xnli_output_modes as xnli_output_modes + from .data import xnli_processors as xnli_processors + from .data import xnli_tasks_num_labels as xnli_tasks_num_labels + from .data.data_collator import DataCollator as DataCollator + from .data.data_collator import DataCollatorForLanguageModeling as DataCollatorForLanguageModeling + from .data.data_collator import DataCollatorForMultipleChoice as DataCollatorForMultipleChoice + from .data.data_collator import ( + DataCollatorForPermutationLanguageModeling as DataCollatorForPermutationLanguageModeling, + ) + from .data.data_collator import DataCollatorForSeq2Seq as DataCollatorForSeq2Seq + from .data.data_collator import DataCollatorForSOP as DataCollatorForSOP + from .data.data_collator import DataCollatorForTokenClassification as DataCollatorForTokenClassification + from .data.data_collator import DataCollatorForWholeWordMask as DataCollatorForWholeWordMask + from .data.data_collator import DataCollatorWithFlattening as DataCollatorWithFlattening + from .data.data_collator import DataCollatorWithPadding as DataCollatorWithPadding + from .data.data_collator import DefaultDataCollator as DefaultDataCollator + from .data.data_collator import default_data_collator as default_data_collator + from .data.datasets import GlueDataset as GlueDataset + from .data.datasets import GlueDataTrainingArguments as GlueDataTrainingArguments + from .data.datasets import SquadDataset as SquadDataset + from .data.datasets import SquadDataTrainingArguments as SquadDataTrainingArguments + from .feature_extraction_sequence_utils import SequenceFeatureExtractor as SequenceFeatureExtractor + + # Feature Extractor + from .feature_extraction_utils import BatchFeature as BatchFeature + from .feature_extraction_utils import FeatureExtractionMixin as FeatureExtractionMixin + + # Generation + from .generation import AlternatingCodebooksLogitsProcessor as AlternatingCodebooksLogitsProcessor + from .generation import AsyncTextIteratorStreamer as AsyncTextIteratorStreamer + from .generation import BayesianDetectorConfig as BayesianDetectorConfig + from .generation import BayesianDetectorModel as BayesianDetectorModel + from .generation import ClassifierFreeGuidanceLogitsProcessor as ClassifierFreeGuidanceLogitsProcessor + from .generation import CompileConfig as CompileConfig + from .generation import ContinuousBatchingConfig as ContinuousBatchingConfig + from .generation import ContinuousBatchingManager as ContinuousBatchingManager + from .generation import ContinuousMixin as ContinuousMixin + from .generation import EncoderNoRepeatNGramLogitsProcessor as EncoderNoRepeatNGramLogitsProcessor + from .generation import EncoderRepetitionPenaltyLogitsProcessor as EncoderRepetitionPenaltyLogitsProcessor + from .generation import EosTokenCriteria as EosTokenCriteria + from .generation import EpsilonLogitsWarper as EpsilonLogitsWarper + from .generation import EtaLogitsWarper as EtaLogitsWarper + from .generation import ExponentialDecayLengthPenalty as ExponentialDecayLengthPenalty + from .generation import ForcedBOSTokenLogitsProcessor as ForcedBOSTokenLogitsProcessor + from .generation import ForcedEOSTokenLogitsProcessor as ForcedEOSTokenLogitsProcessor + from .generation import GenerationConfig as GenerationConfig + from .generation import GenerationMixin as GenerationMixin + from .generation import InfNanRemoveLogitsProcessor as InfNanRemoveLogitsProcessor + from .generation import LogitNormalization as LogitNormalization + from .generation import LogitsProcessor as LogitsProcessor + from .generation import LogitsProcessorList as LogitsProcessorList + from .generation import MaxLengthCriteria as MaxLengthCriteria + from .generation import MaxTimeCriteria as MaxTimeCriteria + from .generation import MinLengthLogitsProcessor as MinLengthLogitsProcessor + from .generation import MinNewTokensLengthLogitsProcessor as MinNewTokensLengthLogitsProcessor + from .generation import MinPLogitsWarper as MinPLogitsWarper + from .generation import NoBadWordsLogitsProcessor as NoBadWordsLogitsProcessor + from .generation import NoRepeatNGramLogitsProcessor as NoRepeatNGramLogitsProcessor + from .generation import PrefixConstrainedLogitsProcessor as PrefixConstrainedLogitsProcessor + from .generation import RepetitionPenaltyLogitsProcessor as RepetitionPenaltyLogitsProcessor + from .generation import SequenceBiasLogitsProcessor as SequenceBiasLogitsProcessor + from .generation import StoppingCriteria as StoppingCriteria + from .generation import StoppingCriteriaList as StoppingCriteriaList + from .generation import StopStringCriteria as StopStringCriteria + from .generation import SuppressTokensAtBeginLogitsProcessor as SuppressTokensAtBeginLogitsProcessor + from .generation import SuppressTokensLogitsProcessor as SuppressTokensLogitsProcessor + from .generation import SynthIDTextWatermarkDetector as SynthIDTextWatermarkDetector + from .generation import SynthIDTextWatermarkingConfig as SynthIDTextWatermarkingConfig + from .generation import SynthIDTextWatermarkLogitsProcessor as SynthIDTextWatermarkLogitsProcessor + from .generation import TemperatureLogitsWarper as TemperatureLogitsWarper + from .generation import TextIteratorStreamer as TextIteratorStreamer + from .generation import TextStreamer as TextStreamer + from .generation import TopHLogitsWarper as TopHLogitsWarper + from .generation import TopKLogitsWarper as TopKLogitsWarper + from .generation import TopPLogitsWarper as TopPLogitsWarper + from .generation import TypicalLogitsWarper as TypicalLogitsWarper + from .generation import ( + UnbatchedClassifierFreeGuidanceLogitsProcessor as UnbatchedClassifierFreeGuidanceLogitsProcessor, + ) + from .generation import WatermarkDetector as WatermarkDetector + from .generation import WatermarkingConfig as WatermarkingConfig + from .generation import WatermarkLogitsProcessor as WatermarkLogitsProcessor + from .generation import WhisperTimeStampLogitsProcessor as WhisperTimeStampLogitsProcessor + from .hf_argparser import HfArgumentParser as HfArgumentParser + from .image_processing_backends import PilBackend as PilBackend + from .image_processing_backends import TorchvisionBackend as TorchvisionBackend + from .image_processing_base import ImageProcessingMixin as ImageProcessingMixin + from .image_processing_utils import BaseImageProcessor as BaseImageProcessor + from .image_utils import ImageFeatureExtractionMixin as ImageFeatureExtractionMixin + + # Integrations + from .integrations import is_clearml_available as is_clearml_available + from .integrations import is_comet_available as is_comet_available + from .integrations import is_dvclive_available as is_dvclive_available + from .integrations import is_neptune_available as is_neptune_available + from .integrations import is_optuna_available as is_optuna_available + from .integrations import is_ray_available as is_ray_available + from .integrations import is_ray_tune_available as is_ray_tune_available + from .integrations import is_swanlab_available as is_swanlab_available + from .integrations import is_tensorboard_available as is_tensorboard_available + from .integrations import is_trackio_available as is_trackio_available + from .integrations import is_wandb_available as is_wandb_available + from .integrations.executorch import TorchExportableModuleWithStaticCache as TorchExportableModuleWithStaticCache + from .integrations.executorch import convert_and_export_with_cache as convert_and_export_with_cache + from .masking_utils import AttentionMaskInterface as AttentionMaskInterface + from .model_debugging_utils import model_addition_debugger_context as model_addition_debugger_context + from .modeling_layers import GradientCheckpointingLayer as GradientCheckpointingLayer + from .modeling_rope_utils import ROPE_INIT_FUNCTIONS as ROPE_INIT_FUNCTIONS + from .modeling_rope_utils import RopeParameters as RopeParameters + from .modeling_rope_utils import dynamic_rope_update as dynamic_rope_update + from .modeling_utils import AttentionInterface as AttentionInterface + from .modeling_utils import PreTrainedModel as PreTrainedModel + from .models import * + from .models.timm_wrapper import TimmWrapperImageProcessor as TimmWrapperImageProcessor + + # Optimization + from .optimization import Adafactor as Adafactor + from .optimization import GreedyLR as GreedyLR + from .optimization import get_constant_schedule as get_constant_schedule + from .optimization import get_constant_schedule_with_warmup as get_constant_schedule_with_warmup + from .optimization import get_cosine_schedule_with_warmup as get_cosine_schedule_with_warmup + from .optimization import ( + get_cosine_with_hard_restarts_schedule_with_warmup as get_cosine_with_hard_restarts_schedule_with_warmup, + ) + from .optimization import ( + get_cosine_with_min_lr_schedule_with_warmup as get_cosine_with_min_lr_schedule_with_warmup, + ) + from .optimization import ( + get_cosine_with_min_lr_schedule_with_warmup_lr_rate as get_cosine_with_min_lr_schedule_with_warmup_lr_rate, + ) + from .optimization import get_greedy_schedule as get_greedy_schedule + from .optimization import get_inverse_sqrt_schedule as get_inverse_sqrt_schedule + from .optimization import get_linear_schedule_with_warmup as get_linear_schedule_with_warmup + from .optimization import get_polynomial_decay_schedule_with_warmup as get_polynomial_decay_schedule_with_warmup + from .optimization import get_scheduler as get_scheduler + from .optimization import get_wsd_schedule as get_wsd_schedule + + # Pipelines + from .pipelines import AnyToAnyPipeline as AnyToAnyPipeline + from .pipelines import AudioClassificationPipeline as AudioClassificationPipeline + from .pipelines import AutomaticSpeechRecognitionPipeline as AutomaticSpeechRecognitionPipeline + from .pipelines import CsvPipelineDataFormat as CsvPipelineDataFormat + from .pipelines import DepthEstimationPipeline as DepthEstimationPipeline + from .pipelines import DocumentQuestionAnsweringPipeline as DocumentQuestionAnsweringPipeline + from .pipelines import FeatureExtractionPipeline as FeatureExtractionPipeline + from .pipelines import FillMaskPipeline as FillMaskPipeline + from .pipelines import ImageClassificationPipeline as ImageClassificationPipeline + from .pipelines import ImageFeatureExtractionPipeline as ImageFeatureExtractionPipeline + from .pipelines import ImageSegmentationPipeline as ImageSegmentationPipeline + from .pipelines import ImageTextToTextPipeline as ImageTextToTextPipeline + from .pipelines import JsonPipelineDataFormat as JsonPipelineDataFormat + from .pipelines import KeypointMatchingPipeline as KeypointMatchingPipeline + from .pipelines import MaskGenerationPipeline as MaskGenerationPipeline + from .pipelines import NerPipeline as NerPipeline + from .pipelines import ObjectDetectionPipeline as ObjectDetectionPipeline + from .pipelines import PipedPipelineDataFormat as PipedPipelineDataFormat + from .pipelines import Pipeline as Pipeline + from .pipelines import PipelineDataFormat as PipelineDataFormat + from .pipelines import TableQuestionAnsweringPipeline as TableQuestionAnsweringPipeline + from .pipelines import TextClassificationPipeline as TextClassificationPipeline + from .pipelines import TextGenerationPipeline as TextGenerationPipeline + from .pipelines import TextToAudioPipeline as TextToAudioPipeline + from .pipelines import TokenClassificationPipeline as TokenClassificationPipeline + from .pipelines import VideoClassificationPipeline as VideoClassificationPipeline + from .pipelines import ZeroShotAudioClassificationPipeline as ZeroShotAudioClassificationPipeline + from .pipelines import ZeroShotClassificationPipeline as ZeroShotClassificationPipeline + from .pipelines import ZeroShotImageClassificationPipeline as ZeroShotImageClassificationPipeline + from .pipelines import ZeroShotObjectDetectionPipeline as ZeroShotObjectDetectionPipeline + from .pipelines import pipeline as pipeline + from .processing_utils import AudioKwargs as AudioKwargs + from .processing_utils import ImagesKwargs as ImagesKwargs + from .processing_utils import ProcessingKwargs as ProcessingKwargs + from .processing_utils import ProcessorMixin as ProcessorMixin + from .processing_utils import TextKwargs as TextKwargs + from .processing_utils import VideosKwargs as VideosKwargs + from .pytorch_utils import Conv1D as Conv1D + from .pytorch_utils import apply_chunking_to_forward as apply_chunking_to_forward + + # Tokenization + from .tokenization_python import PreTrainedTokenizer as PreTrainedTokenizer + from .tokenization_python import PythonBackend as PythonBackend + from .tokenization_utils_base import AddedToken as AddedToken + from .tokenization_utils_base import BatchEncoding as BatchEncoding + from .tokenization_utils_base import CharSpan as CharSpan + from .tokenization_utils_base import PreTrainedTokenizerBase as PreTrainedTokenizerBase + from .tokenization_utils_base import TokenSpan as TokenSpan + + # Tokenization + from .tokenization_utils_sentencepiece import SentencePieceBackend as SentencePieceBackend + from .tokenization_utils_tokenizers import PreTrainedTokenizerFast as PreTrainedTokenizerFast + from .tokenization_utils_tokenizers import ( + TokenizersBackend as TokenizersBackend, + ) + + # Trainer + from .trainer import Trainer as Trainer + from .trainer_callback import DefaultFlowCallback as DefaultFlowCallback + from .trainer_callback import EarlyStoppingCallback as EarlyStoppingCallback + from .trainer_callback import PrinterCallback as PrinterCallback + from .trainer_callback import ProgressCallback as ProgressCallback + from .trainer_callback import TrainerCallback as TrainerCallback + from .trainer_callback import TrainerControl as TrainerControl + from .trainer_callback import TrainerState as TrainerState + from .trainer_pt_utils import torch_distributed_zero_first as torch_distributed_zero_first + from .trainer_seq2seq import Seq2SeqTrainer as Seq2SeqTrainer + from .trainer_utils import EvalPrediction as EvalPrediction + from .trainer_utils import IntervalStrategy as IntervalStrategy + from .trainer_utils import SchedulerType as SchedulerType + from .trainer_utils import enable_full_determinism as enable_full_determinism + from .trainer_utils import set_seed as set_seed + from .training_args import TrainingArguments as TrainingArguments + from .training_args_seq2seq import Seq2SeqTrainingArguments as Seq2SeqTrainingArguments + + # Files and general utilities + from .utils import CONFIG_NAME as CONFIG_NAME + from .utils import MODEL_CARD_NAME as MODEL_CARD_NAME + from .utils import SPIECE_UNDERLINE as SPIECE_UNDERLINE + from .utils import WEIGHTS_NAME as WEIGHTS_NAME + from .utils import TensorType as TensorType + from .utils import add_end_docstrings as add_end_docstrings + from .utils import add_start_docstrings as add_start_docstrings + from .utils import is_apex_available as is_apex_available + from .utils import is_av_available as is_av_available + from .utils import is_datasets_available as is_datasets_available + from .utils import is_faiss_available as is_faiss_available + from .utils import is_matplotlib_available as is_matplotlib_available + from .utils import is_phonemizer_available as is_phonemizer_available + from .utils import is_psutil_available as is_psutil_available + from .utils import is_py3nvml_available as is_py3nvml_available + from .utils import is_pyctcdecode_available as is_pyctcdecode_available + from .utils import is_sacremoses_available as is_sacremoses_available + from .utils import is_sklearn_available as is_sklearn_available + from .utils import is_torch_hpu_available as is_torch_hpu_available + from .utils import is_torch_mlu_available as is_torch_mlu_available + from .utils import is_torch_musa_available as is_torch_musa_available + from .utils import is_torch_neuroncore_available as is_torch_neuroncore_available + from .utils import is_torch_npu_available as is_torch_npu_available + from .utils import is_torch_xla_available as is_torch_xla_available + from .utils import is_torch_xpu_available as is_torch_xpu_available + from .utils.import_utils import requires_backends + from .utils.kernel_config import KernelConfig as KernelConfig + + # Quantization config + from .utils.quantization_config import AqlmConfig as AqlmConfig + from .utils.quantization_config import AutoRoundConfig as AutoRoundConfig + from .utils.quantization_config import AwqConfig as AwqConfig + from .utils.quantization_config import BitNetQuantConfig as BitNetQuantConfig + from .utils.quantization_config import BitsAndBytesConfig as BitsAndBytesConfig + from .utils.quantization_config import CompressedTensorsConfig as CompressedTensorsConfig + from .utils.quantization_config import EetqConfig as EetqConfig + from .utils.quantization_config import FbgemmFp8Config as FbgemmFp8Config + from .utils.quantization_config import FineGrainedFP8Config as FineGrainedFP8Config + from .utils.quantization_config import FourOverSixConfig as FourOverSixConfig + from .utils.quantization_config import FPQuantConfig as FPQuantConfig + from .utils.quantization_config import GPTQConfig as GPTQConfig + from .utils.quantization_config import HiggsConfig as HiggsConfig + from .utils.quantization_config import HqqConfig as HqqConfig + from .utils.quantization_config import MetalConfig as MetalConfig + from .utils.quantization_config import QuantoConfig as QuantoConfig + from .utils.quantization_config import QuarkConfig as QuarkConfig + from .utils.quantization_config import SinqConfig as SinqConfig + from .utils.quantization_config import SpQRConfig as SpQRConfig + from .utils.quantization_config import TorchAoConfig as TorchAoConfig + from .utils.quantization_config import VptqConfig as VptqConfig + from .video_processing_utils import BaseVideoProcessor as BaseVideoProcessor +else: + _import_structure = {k: set(v) for k, v in _import_structure.items()} + + import_structure = define_import_structure(Path(__file__).parent / "models", prefix="models") + import_structure[frozenset({})].update(_import_structure) + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + import_structure, + module_spec=__spec__, + extra_objects={"__version__": __version__}, + ) + + def _create_module_alias(alias: str, target: str) -> None: + """ + Lazily redirect legacy module paths to their replacements without importing heavy deps. + """ + module = types.ModuleType(alias) + module.__doc__ = f"Alias module for backward compatibility with `{target}`." + # Set __file__ explicitly so that inspect.py's hasattr(module, '__file__') check + # never falls through to __getattr__ and triggers a premature (possibly circular) import. + module.__file__ = None + + def _get_target(): + return importlib.import_module(target, __name__) + + module.__getattr__ = lambda name: getattr(_get_target(), name) + module.__dir__ = lambda: dir(_get_target()) + + sys.modules[alias] = module + setattr(sys.modules[__name__], alias.rsplit(".", 1)[-1], module) + + _create_module_alias(f"{__name__}.tokenization_utils_fast", ".tokenization_utils_tokenizers") + _create_module_alias(f"{__name__}.tokenization_utils", ".tokenization_utils_sentencepiece") + _create_module_alias(f"{__name__}.image_processing_utils_fast", ".image_processing_backends") + + for _proc_file in sorted((Path(__file__).parent / "models").rglob("image_processing_*.py")): + _model = _proc_file.parent.name + _module = _proc_file.stem + _target = f".models.{_model}.{_module}" + _create_module_alias(f"{__name__}.models.{_model}.{_module}_fast", _target) + + # Also map XImageProcessorFast -> XImageProcessor for backward compat with old class names. + def getattr_factory(target): + def _getattr(name): + new_name = name.removesuffix("Fast") + logger.warning( + "Accessing `%s` from `%s`. Returning `%s` instead. Behavior may be " + "different and this alias will be removed in future versions.", + name, + target, + new_name, + ) + return getattr(importlib.import_module(target, __name__), new_name) + + return _getattr + + sys.modules[f"{__name__}.models.{_model}.{_module}_fast"].__getattr__ = getattr_factory(_target) + +if not is_torch_available(): + logger.warning_advice( + "PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used." + ) diff --git a/third_party/transformers/src/transformers/debug_utils.py b/third_party/transformers/src/transformers/debug_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..38ff0399641bf1b006125153db5296e1554928bb --- /dev/null +++ b/third_party/transformers/src/transformers/debug_utils.py @@ -0,0 +1,348 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import collections + +from .utils import ExplicitEnum, is_torch_available, logging + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +class DebugUnderflowOverflow: + """ + This debug class helps detect and understand where the model starts getting very large or very small, and more + importantly `nan` or `inf` weight and activation elements. + + There are 2 working modes: + + 1. Underflow/overflow detection (default) + 2. Specific batch absolute min/max tracing without detection + + Mode 1: Underflow/overflow detection + + To activate the underflow/overflow detection, initialize the object with the model : + + ```python + debug_overflow = DebugUnderflowOverflow(model) + ``` + + then run the training as normal and if `nan` or `inf` gets detected in at least one of the weight, input or output + elements this module will throw an exception and will print `max_frames_to_save` frames that lead to this event, + each frame reporting + + 1. the fully qualified module name plus the class name whose `forward` was run + 2. the absolute min and max value of all elements for each module weights, and the inputs and output + + For example, here is the header and the last few frames in detection report for `google/mt5-small` run in fp16 + mixed precision : + + ``` + Detected inf/nan during batch_number=0 + Last 21 forward frames: + abs min abs max metadata + [...] + encoder.block.2.layer.1.DenseReluDense.wi_0 Linear + 2.17e-07 4.50e+00 weight + 1.79e-06 4.65e+00 input[0] + 2.68e-06 3.70e+01 output + encoder.block.2.layer.1.DenseReluDense.wi_1 Linear + 8.08e-07 2.66e+01 weight + 1.79e-06 4.65e+00 input[0] + 1.27e-04 2.37e+02 output + encoder.block.2.layer.1.DenseReluDense.wo Linear + 1.01e-06 6.44e+00 weight + 0.00e+00 9.74e+03 input[0] + 3.18e-04 6.27e+04 output + encoder.block.2.layer.1.DenseReluDense T5DenseGatedGeluDense + 1.79e-06 4.65e+00 input[0] + 3.18e-04 6.27e+04 output + encoder.block.2.layer.1.dropout Dropout + 3.18e-04 6.27e+04 input[0] + 0.00e+00 inf output + ``` + + You can see here, that `T5DenseGatedGeluDense.forward` resulted in output activations, whose absolute max value was + around 62.7K, which is very close to fp16's top limit of 64K. In the next frame we have `Dropout` which + renormalizes the weights, after it zeroed some of the elements, which pushes the absolute max value to more than + 64K, and we get an overflow. + + As you can see it's the previous frames that we need to look into when the numbers start going into very large for + fp16 numbers. + + The tracking is done in a forward hook, which gets invoked immediately after `forward` has completed. + + By default the last 21 frames are printed. You can change the default to adjust for your needs. For example : + + ```python + debug_overflow = DebugUnderflowOverflow(model, max_frames_to_save=100) + ``` + + To validate that you have set up this debugging feature correctly, and you intend to use it in a training that + may take hours to complete, first run it with normal tracing enabled for one of a few batches as explained in + the next section. + + + Mode 2. Specific batch absolute min/max tracing without detection + + The second work mode is per-batch tracing with the underflow/overflow detection feature turned off. + + Let's say you want to watch the absolute min and max values for all the ingredients of each `forward` call of a + given batch, and only do that for batches 1 and 3. Then you instantiate this class as : + + ```python + debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3]) + ``` + + And now full batches 1 and 3 will be traced using the same format as explained above. Batches are 0-indexed. + + This is helpful if you know that the program starts misbehaving after a certain batch number, so you can + fast-forward right to that area. + + + Early stopping: + + You can also specify the batch number after which to stop the training, with : + + ```python + debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3], abort_after_batch_num=3) + ``` + + This feature is mainly useful in the tracing mode, but you can use it for any mode. + + + **Performance**: + + As this module measures absolute `min`/``max` of each weight of the model on every forward it'll slow the training + down. Therefore remember to turn it off once the debugging needs have been met. + + Args: + model (`nn.Module`): + The model to debug. + max_frames_to_save (`int`, *optional*, defaults to 21): + How many frames back to record + trace_batch_nums(`list[int]`, *optional*, defaults to `[]`): + Which batch numbers to trace (turns detection off) + abort_after_batch_num (`int``, *optional*): + Whether to abort after a certain batch number has finished + """ + + def __init__(self, model, max_frames_to_save=21, trace_batch_nums=None, abort_after_batch_num=None): + if trace_batch_nums is None: + trace_batch_nums = [] + self.model = model + self.trace_batch_nums = trace_batch_nums + self.abort_after_batch_num = abort_after_batch_num + + # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence + self.frames = collections.deque([], max_frames_to_save) + self.frame = [] + self.batch_number = 0 + self.total_calls = 0 + self.detected_overflow = False + self.prefix = " " + + self.analyse_model() + + self.register_forward_hook() + + def save_frame(self, frame=None): + if frame is not None: + self.expand_frame(frame) + self.frames.append("\n".join(self.frame)) + self.frame = [] # start a new frame + + def expand_frame(self, line): + self.frame.append(line) + + def trace_frames(self): + print("\n".join(self.frames)) + self.frames = [] + + def reset_saved_frames(self): + self.frames = [] + + def dump_saved_frames(self): + print(f"\nDetected inf/nan during batch_number={self.batch_number}") + print(f"Last {len(self.frames)} forward frames:") + print(f"{'abs min':8} {'abs max':8} metadata") + print("\n".join(self.frames)) + print("\n\n") + self.frames = [] + + def analyse_model(self): + # extract the fully qualified module names, to be able to report at run time. e.g.: + # encoder.block.2.layer.0.SelfAttention.o + # + # for shared weights only the first shared module name will be registered + self.module_names = {m: name for name, m in self.model.named_modules()} + # self.longest_module_name = max(len(v) for v in self.module_names.values()) + + def analyse_variable(self, var, ctx): + if torch.is_tensor(var): + self.expand_frame(get_abs_min_max(var, ctx)) + if detect_overflow(var, ctx): + self.detected_overflow = True + elif var is None: + self.expand_frame(f"{'None':>17} {ctx}") + else: + self.expand_frame(f"{'not a tensor':>17} {ctx}") + + def batch_start_frame(self): + self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***") + self.expand_frame(f"{'abs min':8} {'abs max':8} metadata") + + def batch_end_frame(self): + self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number - 1} ***\n\n") + + def create_frame(self, module, input, output): + self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}") + + # params + for name, p in module.named_parameters(recurse=False): + self.analyse_variable(p, name) + + # inputs + if isinstance(input, tuple): + for i, x in enumerate(input): + self.analyse_variable(x, f"input[{i}]") + else: + self.analyse_variable(input, "input") + + # outputs + if isinstance(output, tuple): + for i, x in enumerate(output): + # possibly a tuple of tuples + if isinstance(x, tuple): + for j, y in enumerate(x): + self.analyse_variable(y, f"output[{i}][{j}]") + else: + self.analyse_variable(x, f"output[{i}]") + else: + self.analyse_variable(output, "output") + + self.save_frame() + + def register_forward_hook(self): + self.model.apply(self._register_forward_hook) + + def _register_forward_hook(self, module): + module.register_forward_hook(self.forward_hook) + + def forward_hook(self, module, input, output): + # - input is a tuple of packed inputs (could be non-Tensors) + # - output could be a Tensor or a tuple of Tensors and non-Tensors + + last_frame_of_batch = False + + trace_mode = self.batch_number in self.trace_batch_nums + if trace_mode: + self.reset_saved_frames() + + if self.total_calls == 0: + self.batch_start_frame() + self.total_calls += 1 + + # count batch numbers - the very first forward hook of the batch will be called when the + # batch completes - i.e. it gets called very last - we know this batch has finished + if module == self.model: + self.batch_number += 1 + last_frame_of_batch = True + + self.create_frame(module, input, output) + + # if last_frame_of_batch: + # self.batch_end_frame() + + if trace_mode: + self.trace_frames() + + if last_frame_of_batch: + self.batch_start_frame() + + if self.detected_overflow and not trace_mode: + self.dump_saved_frames() + + # now we can abort, as it's pointless to continue running + raise ValueError( + "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. " + "Please scroll up above this traceback to see the activation values prior to this event." + ) + + # abort after certain batch if requested to do so + if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num: + raise ValueError( + f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to" + f" `abort_after_batch_num={self.abort_after_batch_num}` arg" + ) + + +def get_abs_min_max(var, ctx): + abs_var = var.abs() + return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}" + + +def detect_overflow(var, ctx): + """ + Report whether the tensor contains any `nan` or `inf` entries. + + This is useful for detecting overflows/underflows and best to call right after the function that did some math that + modified the tensor in question. + + This function contains a few other helper features that you can enable and tweak directly if you want to track + various other things. + + Args: + var: the tensor variable to check + ctx: the message to print as a context + + Return: + `True` if `inf` or `nan` was detected, `False` otherwise + """ + detected = False + if torch.isnan(var).any().item(): + detected = True + print(f"{ctx} has nans") + if torch.isinf(var).any().item(): + detected = True + print(f"{ctx} has infs") + + # if needed to monitor large elements can enable the following + if 0: # and detected: + n100 = var[torch.ge(var.abs(), 100)] + if n100.numel() > 0: + print(f"{ctx}: n100={n100.numel()}") + n1000 = var[torch.ge(var.abs(), 1000)] + if n1000.numel() > 0: + print(f"{ctx}: n1000={n1000.numel()}") + n10000 = var[torch.ge(var.abs(), 10000)] + if n10000.numel() > 0: + print(f"{ctx}: n10000={n10000.numel()}") + + if 0: + print(f"min={var.min():9.2e} max={var.max():9.2e}") + + if 0: + print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})") + + return detected + + +class DebugOption(ExplicitEnum): + UNDERFLOW_OVERFLOW = "underflow_overflow" + TPU_METRICS_DEBUG = "tpu_metrics_debug" diff --git a/third_party/transformers/src/transformers/distributed/__init__.py b/third_party/transformers/src/transformers/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba6db8358d2b126293c061b9ca99368a247c6104 --- /dev/null +++ b/third_party/transformers/src/transformers/distributed/__init__.py @@ -0,0 +1,33 @@ +# 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 + + +_import_structure = { + "configuration_utils": ["DistributedConfig"], +} + + +if TYPE_CHECKING: + from .configuration_utils import ( + DistributedConfig, + ) + +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/third_party/transformers/src/transformers/distributed/configuration_utils.py b/third_party/transformers/src/transformers/distributed/configuration_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7726d9f3290dba97fa288f5ebe54eb107f6af192 --- /dev/null +++ b/third_party/transformers/src/transformers/distributed/configuration_utils.py @@ -0,0 +1,110 @@ +# 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 copy +import json +import os +from dataclasses import dataclass +from typing import Any + + +@dataclass +class DistributedConfig: + """ + Base class for distributed configs + """ + + enable_expert_parallel: bool = False + # TODO: add tp_plan, pp_plan, device_mesh etc.. + + @classmethod + def from_dict(cls, config_dict, **kwargs): + """ + Constructs a DistributedConfig instance from a dictionary of parameters. + Args: + config_dict (Dict[str, Any]): Dictionary containing configuration parameters. + **kwargs: Additional keyword arguments to override dictionary values. + Returns: + DistributedConfig: Instance of DistributedConfig constructed from the dictionary. + """ + config = cls(**config_dict) + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + setattr(config, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + return config + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def to_json_file(self, json_file_path: str | os.PathLike): + """ + Save this instance to a JSON file. + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default + `QuantizationConfig()` is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + config_dict = self.to_dict() + json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + writer.write(json_string) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return copy.deepcopy(self.__dict__) + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ + def __iter__(self): + """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" + yield from copy.deepcopy(self.__dict__).items() + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def to_json_string(self): + """ + Serializes this instance to a JSON formatted string. + Returns: + str: JSON formatted string representing the configuration instance. + """ + return json.dumps(self.__dict__, indent=2) + "\n" + + def update(self, **kwargs): + """ + Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, + returning all the unused kwargs. + Args: + kwargs (`Dict[str, Any]`): + Dictionary of attributes to tentatively update this class. + Returns: + `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + """ + to_remove = [] + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + to_remove.append(key) + + # Remove all the attributes that were updated, without modifying the input dict + unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} + return unused_kwargs diff --git a/third_party/transformers/src/transformers/hyperparameter_search.py b/third_party/transformers/src/transformers/hyperparameter_search.py new file mode 100644 index 0000000000000000000000000000000000000000..267bb36a8e6eff68a7207a09eb32018d8ca532cd --- /dev/null +++ b/third_party/transformers/src/transformers/hyperparameter_search.py @@ -0,0 +1,123 @@ +# Copyright 2023-present 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 .integrations import ( + is_optuna_available, + is_ray_tune_available, + is_wandb_available, + run_hp_search_optuna, + run_hp_search_ray, + run_hp_search_wandb, +) +from .trainer_utils import ( + HPSearchBackend, + default_hp_space_optuna, + default_hp_space_ray, + default_hp_space_wandb, +) +from .utils import logging + + +logger = logging.get_logger(__name__) + + +class HyperParamSearchBackendBase: + name: str + pip_package: str | None = None + + @staticmethod + def is_available(): + raise NotImplementedError + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + raise NotImplementedError + + def default_hp_space(self, trial): + raise NotImplementedError + + def ensure_available(self): + if not self.is_available(): + raise RuntimeError( + f"You picked the {self.name} backend, but it is not installed. Run {self.pip_install()}." + ) + + @classmethod + def pip_install(cls): + return f"`pip install {cls.pip_package or cls.name}`" + + +class OptunaBackend(HyperParamSearchBackendBase): + name = "optuna" + + @staticmethod + def is_available(): + return is_optuna_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_optuna(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_optuna(trial) + + +class RayTuneBackend(HyperParamSearchBackendBase): + name = "ray" + pip_package = "'ray[tune]'" + + @staticmethod + def is_available(): + return is_ray_tune_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_ray(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_ray(trial) + + +class WandbBackend(HyperParamSearchBackendBase): + name = "wandb" + + @staticmethod + def is_available(): + return is_wandb_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_wandb(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_wandb(trial) + + +ALL_HYPERPARAMETER_SEARCH_BACKENDS = { + HPSearchBackend(backend.name): backend for backend in [OptunaBackend, RayTuneBackend, WandbBackend] +} + + +def default_hp_search_backend() -> str: + available_backends = [backend for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() if backend.is_available()] + if len(available_backends) > 0: + name = available_backends[0].name + if len(available_backends) > 1: + logger.info( + f"{len(available_backends)} hyperparameter search backends available. Using {name} as the default." + ) + return name + raise RuntimeError( + "No hyperparameter search backend available.\n" + + "\n".join( + f" - To install {backend.name} run {backend.pip_install()}" + for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() + ) + ) diff --git a/third_party/transformers/src/transformers/masking_utils.py b/third_party/transformers/src/transformers/masking_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..45e43fdaf3aa81e121267762c557a2764af07798 --- /dev/null +++ b/third_party/transformers/src/transformers/masking_utils.py @@ -0,0 +1,1608 @@ +# Copyright 2025 HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import itertools +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from .cache_utils import Cache +from .configuration_utils import PreTrainedConfig +from .utils import is_torch_xpu_available, logging +from .utils.deprecation import deprecate_kwarg +from .utils.generic import GeneralInterface, is_flash_attention_requested +from .utils.import_utils import is_torch_flex_attn_available, is_torch_greater_or_equal, is_tracing + + +if is_torch_flex_attn_available(): + from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size + from torch.nn.attention.flex_attention import BlockMask, create_block_mask +else: + # Register a fake type to avoid crashing for annotations and `isinstance` checks + BlockMask = torch.Tensor + +_is_torch_greater_or_equal_than_2_5 = is_torch_greater_or_equal("2.5", accept_dev=True) +_is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True) +_is_torch_xpu_available = is_torch_xpu_available() + +if _is_torch_greater_or_equal_than_2_6: + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + +logger = logging.get_logger(__name__) + + +def and_masks(*mask_functions: Callable) -> Callable: + """Returns a mask function that is the intersection of provided mask functions""" + if not all(callable(arg) for arg in mask_functions): + raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}") + + def and_mask(batch_idx, head_idx, q_idx, kv_idx): + result = q_idx.new_ones((), dtype=torch.bool) + for mask in mask_functions: + result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device) + return result + + return and_mask + + +def or_masks(*mask_functions: Callable) -> Callable: + """Returns a mask function that is the union of provided mask functions""" + if not all(callable(arg) for arg in mask_functions): + raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}") + + def or_mask(batch_idx, head_idx, q_idx, kv_idx): + result = q_idx.new_zeros((), dtype=torch.bool) + for mask in mask_functions: + result = result | mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device) + return result + + return or_mask + + +def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + """ + This creates a basic lower-diagonal causal mask. + """ + return kv_idx <= q_idx + + +def bidirectional_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + """ + This creates a full bidirectional mask. + + NOTE: It is important to keep an index-based version for non-vmap expansion. + """ + return q_idx >= 0 + + +def sliding_window_overlay(sliding_window: int) -> Callable: + """ + This is an overlay depicting a sliding window pattern. Add it on top of a causal mask for a proper sliding + window mask. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + return kv_idx > q_idx - sliding_window + + return inner_mask + + +def chunked_overlay(chunk_size: int, left_padding: torch.Tensor) -> Callable: + """ + This is an overlay depicting a chunked attention pattern. Add it on top of a causal mask for a proper chunked + attention mask. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + return (kv_idx - left_padding[batch_idx]) // chunk_size == (q_idx - left_padding[batch_idx]) // chunk_size + + return inner_mask + + +def sliding_window_causal_mask_function(sliding_window: int) -> Callable: + """ + This return the mask_function function to create a sliding window mask. + """ + return and_masks(sliding_window_overlay(sliding_window), causal_mask_function) + + +def sliding_window_bidirectional_overlay(sliding_window: int) -> Callable: + """ + This is an overlay depicting a bidirectional sliding window pattern. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + """A token can attend to any other token if their absolute distance is within + the (inclusive) sliding window size (distance <= sliding_window).""" + return abs(q_idx - kv_idx) <= sliding_window + + return inner_mask + + +def sliding_window_bidirectional_mask_function(sliding_window: int) -> Callable: + """ + This return the mask_function function to create a bidirectional sliding window mask. + """ + return and_masks(sliding_window_bidirectional_overlay(sliding_window), bidirectional_mask_function) + + +def chunked_causal_mask_function(chunk_size: int, left_padding: torch.Tensor) -> Callable: + """ + This return the mask_function function to create a chunked attention mask. + """ + return and_masks(chunked_overlay(chunk_size, left_padding), causal_mask_function) + + +def padding_mask_function(padding_mask: torch.Tensor) -> Callable: + """ + This return the mask_function function corresponding to a 2D padding mask. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + # Note that here the mask should ALWAYS be at least of the max `kv_index` size in the dimension 1. This is because + # we cannot pad it here in the mask_function as we don't know the final size, and we cannot try/except, as it is not + # vectorizable on accelerator devices + return padding_mask[batch_idx, kv_idx] + + return inner_mask + + +def packed_sequence_mask_function(packed_sequence_mask: torch.Tensor) -> Callable: + """ + This return the mask_function function corresponding to a 2D packed sequence mask. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + return packed_sequence_mask[batch_idx, q_idx] == packed_sequence_mask[batch_idx, kv_idx] + + return inner_mask + + +def add_offsets_to_mask_function(mask_function: Callable, q_offset: int, kv_offset: int) -> Callable: + """ + This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths, + not start and end indices. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + return mask_function(batch_idx, head_idx, q_idx + q_offset, kv_idx + kv_offset) + + return inner_mask + + +def prepare_padding_mask(attention_mask: torch.Tensor | None, kv_length: int, kv_offset: int) -> torch.Tensor | None: + """ + From the 2D attention mask, prepare the correct padding mask to use by potentially padding it. + """ + local_padding_mask = attention_mask + if attention_mask is not None: + # Pad it if necessary + if (padding_length := kv_length + kv_offset - attention_mask.shape[-1]) > 0: + local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length)) + return local_padding_mask + + +def _can_skip_causal_mask_xpu( + padding_mask: torch.Tensor | None, + query_length: int, + kv_length: int, + local_attention_size: int | None, +) -> bool: + """ + XPU-specific logic for determining if we can skip causal mask creation. + + For XPU devices, we have special handling: + - Single query tokens (query_length == 1) use the same logic as CUDA + - Multi-query tokens can skip if padding_mask is provided and correctly structured + The mask must have all True values in the query window and all False after + """ + + if is_tracing(padding_mask): + return False + + # Check local attention constraint (same as CUDA) + if local_attention_size is not None and kv_length >= local_attention_size: + return False + + if padding_mask is None: + # Without padding mask, can skip if single query token or full causal attention + return query_length == 1 or kv_length == query_length + + # XPU allows skipping under additional conditions when padding_mask is provided + if query_length == 1: + # Single query token: skip only if no padding tokens present + return padding_mask.all() + + # XPU-specific: check if query window is all True and rest is all False + # This allows XPU to optimize the 1st token in static cache + return padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() + + +def _ignore_causal_mask_sdpa( + padding_mask: torch.Tensor | None, + query_length: int, + kv_length: int, + kv_offset: int, + local_attention_size: int | None = None, +) -> bool: + """ + Detects whether the causal mask can be ignored in case PyTorch's SDPA is used, rather relying on SDPA's `is_causal` argument. + + In case no token is masked in the 2D `padding_mask` argument, if `query_length == 1` or + `key_value_length == query_length`, we rather rely on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is + passed). + """ + if padding_mask is not None and padding_mask.shape[-1] > kv_length: + mask_indices = torch.arange(kv_length, device=padding_mask.device) + mask_indices += kv_offset + padding_mask = padding_mask[:, mask_indices] + + if _is_torch_xpu_available: + # XPU devices have special handling for mask skipping: + # - Single query tokens use the same logic as CUDA + # - Multi-query tokens can skip if padding_mask is provided and correctly structured + # (all True in query window, all False after) + return _can_skip_causal_mask_xpu(padding_mask, query_length, kv_length, local_attention_size) + # When using `torch.export` or `torch.onnx.dynamo_export`, we must pass an example input, and `is_causal` behavior is + # hard-coded to the forward. If a user exports a model with query_length > 1, the exported model will hard-code `is_causal=True` + # which is in general wrong (see https://github.com/pytorch/pytorch/issues/108108). Thus, we only set + # `ignore_causal_mask = True` if we are not tracing + if ( + not is_tracing(padding_mask) + # only cases when lower and upper diags are the same, see https://github.com/pytorch/pytorch/issues/108108 + and (query_length == 1 or kv_length == query_length) + # in this case we need to add special patterns to the mask so cannot be skipped otherwise + and (local_attention_size is None or kv_length < local_attention_size) + # In this case, we need to add padding to the mask, so cannot be skipped otherwise + and (padding_mask is None or padding_mask.all()) + ): + return True + + return False + + +def _can_skip_bidirectional_mask_xpu( + padding_mask: torch.Tensor | None, + kv_length: int, + local_attention_size: int | None, +) -> bool: + """ + XPU-specific logic for determining if we can skip bidirectional mask creation. + + For XPU devices, we have special handling: + - Skip if no padding and no local attention constraint + """ + + if is_tracing(padding_mask): + return False + + # Check local attention constraint (same as CUDA) + if local_attention_size is not None and kv_length >= local_attention_size: + return False + + if padding_mask is None: + # Without padding mask, can always skip for full bidirectional attention + return True + + # Skip only if no padding tokens present + return padding_mask.all() + + +def _ignore_bidirectional_mask_sdpa( + padding_mask: torch.Tensor | None, + kv_length: int, + local_attention_size: int | None = None, +) -> bool: + """ + Detects whether the bidirectional mask can be ignored in case PyTorch's SDPA is used. + + In case no token is masked in the 2D `padding_mask` argument and no local attention constraint applies + (i.e. `local_attention_size` is None or `kv_length < local_attention_size`), we skip mask creation, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is + passed). + """ + if _is_torch_xpu_available: + # XPU devices have special handling for mask skipping: + # - Skip if no padding and no local attention constraint + return _can_skip_bidirectional_mask_xpu(padding_mask, kv_length, local_attention_size) + + # When using `torch.export` or `torch.onnx.dynamo_export`, we need to avoid to check the contents of the mask; + # otherwise, we will encounter dynamic control flows + if ( + not is_tracing(padding_mask) + and (padding_mask is None or padding_mask.all()) + # in this case we need to add special patterns to the mask so cannot be skipped otherwise + and (local_attention_size is None or kv_length < local_attention_size) + ): + return True + + return False + + +def _vmap_expansion_sdpa(mask_function: Callable) -> Callable: + """ + Used to vmap our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs. + Using vmap here allows us to keep the performance of vectorized ops, while having a single set of primitive + functions between attention interfaces (i.e. between flex and sdpa/eager, FA2 being a bit different). + """ + # We vmap the function over all 4 dimensions, broadcasting [b_idx, h_idx, q_idx, kv_idx] + dimensions = [(None, None, None, 0), (None, None, 0, None), (None, 0, None, None), (0, None, None, None)] + for dims in dimensions: + mask_function = torch.vmap(mask_function, in_dims=dims, out_dims=0) + return mask_function + + +def _non_vmap_expansion_sdpa( + batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor, kv_indices: torch.Tensor +): + """ + Used to broadcast our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs. + Allows the usage of any index-based mask function without relying on vmap. + + NOTE: This is limited to index based functions only and is not guaranteed to work otherwise. + + Reference: + - https://github.com/huggingface/optimum-onnx/blob/c123e8f4fab61b54a8e0e31ce74462bcacca576e/optimum/exporters/onnx/model_patcher.py#L362-L365 + """ + batch_indices = batch_indices[:, None, None, None] + head_indices = head_indices[None, :, None, None] + q_indices = q_indices[None, None, :, None] + kv_indices = kv_indices[None, None, None, :] + return batch_indices, head_indices, q_indices, kv_indices + + +def sdpa_mask( + batch_size: int, + q_length: int, + kv_length: int, + q_offset: int = 0, + kv_offset: int = 0, + mask_function: Callable = causal_mask_function, + attention_mask: torch.Tensor | None = None, + local_size: int | None = None, + allow_is_causal_skip: bool = True, + allow_is_bidirectional_skip: bool = False, + allow_torch_fix: bool = True, + use_vmap: bool = False, + device: torch.device | str = "cpu", + **kwargs, +) -> torch.Tensor | None: + """ + Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that + the element should take part in the attention computation, and False that it should not. + This function can only be used with torch>=2.5, as the context manager is otherwise not available. + + Args: + batch_size (`int`): + The batch size of the input sequence. + q_length (`int`): + The size that the query states will have during the attention computation. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + kv_offset (`int`, optional): + An optional offset to indicate at which first position the key and values states will refer to. + q_offset (`int`, optional): + An optional offset to indicate at which first position the query states will refer to. + mask_function (`Callable`): + The mask factory function describing the mask pattern. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length) + local_size (`int`, optional): + The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True` + to try to skip mask creation if possible. + allow_is_causal_skip (`bool`, optional): + Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in + `torch.sdpa` instead. Default to `True`. + allow_is_bidirectional_skip (`bool`, optional): + Whether to allow to return `None` for the mask under conditions where we do not have to add any bias, + i.e. full attention without any padding. Default to `False`. + allow_torch_fix (`bool`, optional): + Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older + versions. We need an arg to skip it when using eager. By default `True`. + use_vmap (`bool`, optional): + Whether to use `vmap` during the mask construction or not. Allows powerful custom patterns that may not be + index-based (for the cost of speed performance). By default `False`. + device (`torch.device` or `str`, optional): + An optional device to create the mask on. + + + ## Creating a simple causal mask: + + To create the following causal mask: + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ■ ■ ■ ■ ⬚ + 4 ■ ■ ■ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, q_length=5, kv_length=5) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [ True, True, True, True, False], + [ True, True, True, True, True]]]]) + ``` + + ## Creating a sliding window mask: + + To create the following sliding window mask (`sliding_window=3`): + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ⬚ ■ ■ ■ ⬚ + 4 ⬚ ⬚ ■ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, q_length=5, kv_length=5, mask_function=sliding_window_causal_mask_function(3)) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [False, True, True, True, False], + [False, False, True, True, True]]]]) + ``` + + ## Creating a chunked attention mask + + To create the following chunked attention mask (`chunk_size=3`): + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ⬚ ⬚ ⬚ ■ ⬚ + 4 ⬚ ⬚ ⬚ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, q_length=5, kv_length=5, mask_function=chunked_causal_mask_function(3, torch.zeros(1, dtype=int))) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [False, False, False, True, False], + [False, False, False, True, True]]]]) + ``` + + """ + # For BC on `cache_positions` that used to be an arg at the position of `q_length` + if isinstance(q_length, torch.Tensor): + logger.warning_once( + "`cache_position` is deprecated as an arg, and will be removed in Transformers v5.6. Please use `q_length` and " + "`q_offset` instead, similarly to `kv_length` and `kv_offset`" + ) + q_length, q_offset = q_length.shape[0], q_length[0].to(device) + + # Potentially pad the 2D mask + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + + # Under specific conditions, we can avoid materializing the mask + # 1. Causal masks can rely on the `is_causal` argument + # 2. Bidirectional do not need any further processing (no bias) + if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): + return None + if allow_is_bidirectional_skip and _ignore_bidirectional_mask_sdpa(padding_mask, kv_length, local_size): + return None + + # Potentially add the padding 2D mask + if padding_mask is not None: + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + batch_arange = torch.arange(batch_size, device=device) + head_arange = torch.arange(1, device=device) + q_arange = torch.arange(q_length, device=device) + q_offset + kv_arange = torch.arange(kv_length, device=device) + kv_offset + + # Actual mask creation + # Option 1: Fast non-vmap mask creation (default) + if not use_vmap: + # Apply mask function element-wise through broadcasting + attention_mask = mask_function(*_non_vmap_expansion_sdpa(batch_arange, head_arange, q_arange, kv_arange)) + # Expand the mask to match batch size and query length if they weren't used in the mask function + attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length) + + # Option 2: Vmap mask creation (torch>=2.6 and custom patterns) + elif _is_torch_greater_or_equal_than_2_6: + # This creates the 4D mask easily. Note that we need this context manager as vmap cannot handle slicing a tensor from + # scalar tensor (it internally calls `.item()` which vmap does not allow, but this context works around it + # We don't need to add an offset to the mask_function either, as we vmap directly the correct indices for k and kv indices + with TransformGetItemToIndex(): + attention_mask = _vmap_expansion_sdpa(mask_function)(batch_arange, head_arange, q_arange, kv_arange) + + # Option 3: Error out since it indicates that the user did something custom, which they shouldn't have (torch<2.6) + else: + raise ValueError( + "The vmap functionality for mask creation is only supported from torch>=2.6. " + "Please update your torch version or use `use_vmap=False` with index-based masks." + ) + + # Due to a bug in versions of torch<2.5, we need to update the mask in case a query is not attending to any + # tokens (due to padding). See details in https://github.com/pytorch/pytorch/issues/110213 + if not _is_torch_greater_or_equal_than_2_5 and allow_torch_fix: + attention_mask = attention_mask | torch.all(~attention_mask, dim=-1, keepdim=True) + + return attention_mask + + +def eager_mask( + batch_size: int, + q_length: int, + kv_length: int, + q_offset: int = 0, + kv_offset: int = 0, + mask_function: Callable = causal_mask_function, + attention_mask: torch.Tensor | None = None, + dtype: torch.dtype = torch.float32, + allow_is_bidirectional_skip: bool = False, + use_vmap: bool = False, + device: torch.device | str = "cpu", + **kwargs, +) -> torch.Tensor: + """ + Create a 4D float mask of shape `(batch_size, 1, query_length, kv_length)` where a value of 0 indicates that + the element should take part in the attention computation, and -inf (minimum value for the given `dtype`) that + it should not. + + Args: + batch_size (`int`): + The batch size of the input sequence. + q_length (`int`): + The size that the query states will have during the attention computation. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + q_offset (`int`, optional): + An optional offset to indicate at which first position the query states will refer to. + kv_offset (`int`, optional): + An optional offset to indicate at which first position the key and values states will refer to. + mask_function (`Callable`): + The mask factory function describing the mask pattern. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length) + dtype (`torch.dtype`, optional): + The dtype to use for the mask. By default, `torch.float32`. + allow_is_bidirectional_skip (`bool`, optional): + Whether to allow to return `None` for the mask under conditions where we do not have to add any bias, + i.e. full attention without any padding. Default to `False`. + use_vmap (`bool`, optional): + Whether to use `vmap` during the mask construction or not. Allows powerful custom patterns that may not be + index-based (for the cost of speed performance). By default `False`. + device (`torch.device` or `str`, optional): + An optional device to create the mask on. + """ + # The masks for eager attention are simply boolean mask from sdpa, casted to 0 and -inf + _ = kwargs.pop("allow_is_causal_skip", None) + _ = kwargs.pop("allow_torch_fix", None) + mask = sdpa_mask( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_function, + attention_mask=attention_mask, + allow_is_causal_skip=False, + allow_is_bidirectional_skip=allow_is_bidirectional_skip, + allow_torch_fix=False, + use_vmap=use_vmap, + device=device, + **kwargs, + ) + # only bidirectional masks can be skipped, otherwise we convert bool -> float + if mask is not None: + min_dtype = torch.finfo(dtype).min + # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) + mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype) + return mask + + +def flash_attention_mask( + batch_size: int, + q_length: int, + kv_length: int, + q_offset: int = 0, + kv_offset: int = 0, + mask_function: Callable = causal_mask_function, + attention_mask: torch.Tensor | None = None, + **kwargs, +): + """ + Create the attention mask necessary to use FA2. Since FA2 is un-padded by definition, here we simply return + `None` if the mask is fully causal, or we return the 2D mask which will then be used to extract the seq_lens. + We just slice it in case of sliding window. + + Args: + batch_size (`int`): + The batch size of the input sequence. + q_length (`int`): + The size that the query states will have during the attention computation. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + q_offset (`int`, optional): + An optional offset to indicate at which first position the query states will refer to. + kv_offset (`int`, optional): + An optional offset to indicate at which first position the key and values states will refer to. + mask_function (`Callable`): + The mask factory function describing the mask pattern. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length) + """ + if attention_mask is not None: + # Here we need to slice from the right if using sliding or chunked (for full attention, this is equivalent to doing nothing) + attention_mask = attention_mask[:, -kv_length:] + # We only return an actual mask if there is at least 1 padding token, otherwise we return `None` and use `is_causal` in FA2 + # (note that the attention_mask is a boolean dtype here) + if attention_mask.all(): + attention_mask = None + + return attention_mask + + +def flex_attention_mask( + batch_size: int, + q_length: int, + kv_length: int, + q_offset: int = 0, + kv_offset: int = 0, + mask_function: Callable = causal_mask_function, + attention_mask: torch.Tensor | None = None, + device: torch.device | str = "cpu", + **kwargs, +) -> BlockMask: + """ + Create a 4D block mask which is a compressed representation of the full 4D block causal mask. BlockMask is essential + for performant computation of flex attention. See: https://pytorch.org/blog/flexattention/ + + Args: + batch_size (`int`): + The batch size of the input sequence. + q_length (`int`): + The size that the query states will have during the attention computation. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + q_offset (`int`, optional): + An optional offset to indicate at which first position the query states will refer to. + kv_offset (`int`, optional): + An optional offset to indicate at which first position the key and values states will refer to. + mask_function (`Callable`): + The mask factory function describing the mask pattern. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length) + device (`torch.device` or `str`, optional): + An optional device to create the mask on. + """ + # For BC on `cache_positions` that used to be an arg at the position of `q_length` + if isinstance(q_length, torch.Tensor): + logger.warning_once( + "`cache_position` is deprecated as an arg, and will be removed in Transformers v5.6. Please use `q_length` and " + "`q_offset` instead, similarly to `kv_length` and `kv_offset`" + ) + q_length, q_offset = q_length.shape[0], q_length[0].to(device) + + # Potentially add the padding 2D mask + if attention_mask is not None: + # Older torch (2.5.x) cannot handle sequences not in multiples of 128 (default block size) + # Hence we pad to multiples of this as a minimum to ensure this + pad_len = ((attention_mask.shape[1] // flex_default_block_size) + 1) * flex_default_block_size + pad_len = pad_len - attention_mask.shape[1] + if not _is_torch_greater_or_equal_than_2_6 and pad_len > 0: + attention_mask = torch.nn.functional.pad(attention_mask, value=0, pad=(0, pad_len)) + + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + # Add the offsets on top (because flex interface only allows length, not start and end indices) + mask_function = add_offsets_to_mask_function(mask_function, q_offset, kv_offset) + + # Finally create the block mask + block_mask = create_block_mask( + mask_mod=mask_function, + B=batch_size, + H=None, + Q_LEN=q_length, + KV_LEN=kv_length, + device=device, + _compile=_is_torch_greater_or_equal_than_2_6, + ) + return block_mask + + +class AttentionMaskInterface(GeneralInterface): + # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if + # a new instance is created (in order to locally override a given function) + _global_mapping = { + "sdpa": sdpa_mask, + "eager": eager_mask, + "flash_attention_2": flash_attention_mask, + "flash_attention_3": flash_attention_mask, + "flash_attention_4": flash_attention_mask, + "flex_attention": flex_attention_mask, + } + + +# Global AttentionMaskInterface shared by all models which do not need to overwrite any of the existing ones +ALL_MASK_ATTENTION_FUNCTIONS: AttentionMaskInterface = AttentionMaskInterface() + + +def find_packed_sequence_indices(position_ids: torch.Tensor) -> torch.Tensor | None: + """ + Find the indices of the sequence to which each new query token in the sequence belongs when using packed + tensor format (i.e. several sequences packed in the same batch dimension). + + Args: + position_ids (`torch.Tensor`) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + + Returns: + A 2D tensor where each similar integer indicates that the tokens belong to the same sequence. For example, if we + pack 3 sequences of 2, 3 and 1 tokens respectively along a single batch dim, this will return [[0, 0, 1, 1, 1, 2]]. + + If the there is only one sequence in each batch item (and we don't compile), then we return `None` indicating + no packed sequences. This is the same as [[0, 0, 0, 0, 0, 0]] for the example above. + """ + # What separate different sequences is when 2 consecutive positions_ids are separated by more than 1. So + # taking the diff (by prepending the first value - 1 to keep correct indexing) and applying cumsum to the result + # gives exactly the sequence indices + # Note that we assume that a single sequence cannot span several batch dimensions, i.e. 1 single sequence + # cannot be part of the end of the first batch dim and the start of the 2nd one for example + first_dummy_value = position_ids[:, :1] - 1 # We just need the diff on this first value to be 1 + position_diff = torch.diff(position_ids, prepend=first_dummy_value, dim=-1) + packed_sequence_mask = (position_diff != 1).cumsum(-1) + + # Sadly this is a dynamic control flow, so we cannot enable this check on anything compile related + if not is_tracing(packed_sequence_mask) and (packed_sequence_mask[:, -1] == 0).all(): + return None + + return packed_sequence_mask + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def _preprocess_mask_arguments( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | BlockMask | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None, + layer_idx: int | None, + encoder_hidden_states: torch.Tensor | None = None, +) -> tuple[bool, torch.Tensor | BlockMask | None, int, int]: + """ + Perform some common pre-processing of the mask arguments we get from the modeling code. Mostly determine the + key-value length and offsets, and if we should early exit or not. + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + layer_idx (`int`, optional): + If `past_key_values` is not None, this is the layer index of the cache from which to get the key-value + length and offset. Indeed, for hybrid caches, different layers may return different lengths. + encoder_hidden_states (`torch.Tensor`, optional): + The input embeddings of shape (batch_size, kv_length, hidden_dim). If provided, it is used instead of + `inputs_embeds` to infer the kv length. + + Returns: + early_exit (`bool`): + Whether we should early exit mask creation, and return the mask as-is. + attention_mask (`torch.Tensor` or `BlockMask` or `None`): + The attention mask to either return immediately, or to use in downstream mask creation. + packed_sequence_mask (`torch.Tensor`, optional): + In case we detected packed sequence format, this is a tensor where each similar integer indicates that + the tokens belong to the same sequence. + q_length (`int`): + The size that the query states will have during the attention computation. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + q_offset (`int`, optional): + An optional offset to indicate at which first position the query states will refer to. + kv_offset (`int`): + An offset to indicate at which first position the key and values states will refer to. + """ + # If the mask is already 4D, simply return as-is (it was already prepared, or it is custom) + if isinstance(attention_mask, (torch.Tensor, BlockMask)) and len(attention_mask.shape) == 4: + return True, attention_mask, None, None, None, None, None + + # For TGI/vLLM backends, or other custom attention without equivalent mask creation: we don't need a mask! + # Note: it's not ideal to check the `_global_mapping` attribute instead of the object itself, however otherwise + # full graph dynamo tracing (i.e. torch.export or compile with `fullgraph=True`) will fail on Python<3.11 + # with `torch._dynamo.exc.Unsupported: 'inline in skipfiles:Mapping.__contains__ | __contains__, skipped + # according trace_rules.lookup SKIP_DIRS'` -- can be removed when we require Python>=3.11 + if config._attn_implementation not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping: + return True, None, None, None, None, None, None + + # Move the mask to correct device, and potentially switch dtype for efficiency + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask.to(device=inputs_embeds.device, dtype=torch.bool) + + q_length = inputs_embeds.shape[1] + # If using a cache, it can give all information about mask sizes based on seen tokens + if past_key_values is not None: + q_offset = past_key_values.get_seq_length() + # To avoid graph breaks, StaticLayer return a tensor instead of int -> this has no impact on the ops, but we + # need the correct device + q_offset = q_offset.to(inputs_embeds.device) if isinstance(q_offset, torch.Tensor) else q_offset + kv_length, kv_offset = past_key_values.get_mask_sizes(q_length, layer_idx) + # Otherwise, we infer based on our input + else: + q_offset = 0 + # 1. Rely on input directly + if attention_mask is None: + # For encoder-decoders, use encoder_hidden_states to infer kv_length if provided + kv_length = encoder_hidden_states.shape[1] if encoder_hidden_states is not None else q_length + kv_offset = 0 + # 2. Rely on the mask instead - needed for special cases like prefix tuning in PEFT + # + # This is a very unique and special case where an encoder utilizes a cache and expects its length + # to be accounted for (usually, they should never use a cache). In general, the mask should always + # match with the input sizes nonetheless (i.e. it does not affect others). + # Conclusion: "prefix tuning is evil" + else: + kv_length, kv_offset = attention_mask.shape[-1], 0 + + # We check the position_ids for potential packed sequence format (only if the 2D attention mask is explicitly None, + # and we don't have past_key_values, i.e. generally a training setup) + packed_sequence_mask = None + if position_ids is not None and attention_mask is None and past_key_values is None: + batch_size = inputs_embeds.shape[0] + # The position ids are sometimes just unsqueezed, without being expanded + if batch_size != position_ids.shape[0]: + position_ids = position_ids.expand(batch_size, -1) + packed_sequence_mask = find_packed_sequence_indices(position_ids) + + return False, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + cache_position: torch.Tensor | None = None, # not used anymore but kept for BC + *, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create a standard causal mask based on the attention implementation used (stored in the config). If `past_key_values` + has an hybrid cache structure, this function will return the mask corresponding to one of the "full_attention" layers (to align + to what is needed in the `modeling_xxx.py` files). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + cache_position (`torch.Tensor`): + Deprecated and unused. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the causal mask function (by doing the union of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the causal mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + """ + # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention. + # It allows to use decoder-only models with bi-directional attention as well + if not getattr(config, "is_causal", True): + return create_bidirectional_mask( + config, + inputs_embeds, + attention_mask, + past_key_values=past_key_values, + or_mask_function=or_mask_function, + and_mask_function=and_mask_function, + ) + + # If we have an hybrid cache structure, here we want to create the mask for the full layers + if hasattr(past_key_values, "is_sliding") and False in past_key_values.is_sliding: + layer_idx = past_key_values.is_sliding.index(False) + else: + layer_idx = 0 + + early_exit, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset = ( + _preprocess_mask_arguments(config, inputs_embeds, attention_mask, past_key_values, position_ids, layer_idx) + ) + if early_exit: + return attention_mask + + batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device + mask_factory_function = causal_mask_function + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + # Defaulting to using non-vmap based mask creations except when detecting + # users passing custom mask functions (as we cannot guarantee that they + # are properly index-based as required by our implementation). + use_vmap = False + + # Do not allow skip if we are compiling (this is to match BC) + # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it + if _is_torch_xpu_available: + # Do not allow skip if we are compiling for decoding, but for prefill, we still allow skip to optimization the perf of 1st token generation + allow_is_causal_skip = not (getattr(past_key_values, "is_compileable", False) and q_length == 1) + else: + allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False) + + # Allow slight deviations from causal mask + # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, + # padding mask, etc) as the resulting mask may otherwise not be correct! + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_causal_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_causal_skip = False + use_vmap = True + + # If we detected packing format + if packed_sequence_mask is not None: + mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask)) + allow_is_causal_skip = False + + # We now create the mask + causal_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return causal_mask + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_bidirectional_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + encoder_hidden_states: torch.Tensor | None = None, + past_key_values: Cache | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create a standard bidirectional mask based on the attention implementation used (stored in the config). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is only used to infer metadata + such as the batch size, query length, dtype, and device. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, kv_length). + It can also be an already prepared 4D mask of shape (batch_size, 1, query_length, kv_length), + in which case it is returned as-is. + encoder_hidden_states (`torch.Tensor`, optional): + The input embeddings of shape (batch_size, kv_length, hidden_dim). If provided, it is used instead of + `inputs_embeds` to infer the batch size, kv length and dtype. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the base mask function (by doing the union of both). This is + useful to easily overlay another mask on top, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the base mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top, for example for image tokens handling. + """ + # We ignore a few irrelevant arguments at the end as we do not have a (growing) cache here + early_exit, attention_mask, _, q_length, kv_length, q_offset, kv_offset = _preprocess_mask_arguments( + config, inputs_embeds, attention_mask, past_key_values, None, 0, encoder_hidden_states + ) + if early_exit: + return attention_mask + + embeds = encoder_hidden_states if encoder_hidden_states is not None else inputs_embeds + batch_size, dtype, device = embeds.shape[0], embeds.dtype, embeds.device + mask_factory_function = bidirectional_mask_function + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + # Allow skipping the mask creation except we have additional masking operators (and/or masks) + allow_is_bidirectional_skip = True + # Defaulting to using non-vmap based mask creations except when detecting + # users passing custom mask functions (as we cannot guarantee that they + # are properly index-based as required by our implementation). + use_vmap = False + + # Allow slight deviations from the base mask + # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, + # padding mask, etc) as the resulting mask may otherwise not be correct! + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_bidirectional_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_bidirectional_skip = False + use_vmap = True + + # We now create the mask + attention_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + # Additional kwargs for sdpa + allow_is_causal_skip=False, + allow_is_bidirectional_skip=allow_is_bidirectional_skip, + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return attention_mask + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_sliding_window_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + cache_position: torch.Tensor | None = None, # not used anymore but kept for BC + *, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create a sliding window causal mask based on the attention implementation used (stored in the config). This type + of attention pattern was mostly democratized by Mistral. If `past_key_values` has an hybrid cache structure, this + function will return the mask corresponding to one of the "sliding_attention" layers (to align to what is needed in the + `modeling_xxx.py` files). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + cache_position (`torch.Tensor`): + Deprecated and unused. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the sliding causal mask function (by doing the union of both). This is + useful to easily overlay another mask on top of the sliding causal one, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the sliding causal mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top of the sliding causal one, for example for image tokens handling. + """ + # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention + # It allows to use decoder-only models with bi-directional attention as well + if not getattr(config, "is_causal", True): + return create_bidirectional_sliding_window_mask( + config, + inputs_embeds, + attention_mask, + past_key_values=past_key_values, + or_mask_function=or_mask_function, + and_mask_function=and_mask_function, + ) + + # If we have an hybrid cache structure, here we want to create the mask for the sliding layers + if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding: + layer_idx = past_key_values.is_sliding.index(True) + else: + layer_idx = 0 + + early_exit, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset = ( + _preprocess_mask_arguments(config, inputs_embeds, attention_mask, past_key_values, position_ids, layer_idx) + ) + if early_exit: + return attention_mask + + sliding_window = getattr(config, "sliding_window", None) + if sliding_window is None: + raise ValueError("Could not find a `sliding_window` argument in the config, or it is not set") + + batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device + mask_factory_function = sliding_window_causal_mask_function(sliding_window) + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + # Defaulting to using non-vmap based mask creations except when detecting + # users passing custom mask functions (as we cannot guarantee that they + # are properly index-based as required by our implementation). + use_vmap = False + # Do not allow skip if we are compiling (this is to match BC) + # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it + allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False) + + # Allow slight deviations from causal mask + # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, + # padding mask, etc) as the resulting mask may otherwise not be correct! + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_causal_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_causal_skip = False + use_vmap = True + + # If we detected packing format + if packed_sequence_mask is not None: + mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask)) + allow_is_causal_skip = False + + # We now create the mask + causal_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa + local_size=sliding_window, # Additional kwarg for sdpa + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return causal_mask + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_bidirectional_sliding_window_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create a standard bidirectional sliding window mask based on the attention implementation used (stored in the config). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is only used to infer metadata + such as the batch size, query length, dtype, and device. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, kv_length). + It can also be an already prepared 4D mask of shape (batch_size, 1, query_length, kv_length), + in which case it is returned as-is. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the base mask function (by doing the union of both). This is + useful to easily overlay another mask on top, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the base mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top, for example for image tokens handling. + """ + # We ignore a few irrelevant arguments at the end as we do not have a (growing) cache here + early_exit, attention_mask, _, q_length, kv_length, q_offset, kv_offset = _preprocess_mask_arguments( + config, inputs_embeds, attention_mask, past_key_values, None, 0 + ) + if early_exit: + return attention_mask + + sliding_window = getattr(config, "sliding_window", None) + if sliding_window is None: + raise ValueError("Could not find a `sliding_window` argument in the config, or it is not set") + + batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device + mask_factory_function = sliding_window_bidirectional_mask_function(sliding_window) + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + use_vmap = False + allow_is_bidirectional_skip = True + + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_bidirectional_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_bidirectional_skip = False + use_vmap = True + + attention_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + allow_is_causal_skip=False, + allow_is_bidirectional_skip=allow_is_bidirectional_skip, + local_size=sliding_window, # Additional kwarg for sdpa + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return attention_mask + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_chunked_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + cache_position: torch.Tensor | None = None, # not used anymore but kept for BC + *, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create a chunked attention causal mask based on the attention implementation used (stored in the config). This type + of attention pattern was mostly democratized by Llama4. If `past_key_values` has an hybrid cache structure, this + function will return the mask corresponding to one of the "chunked_attention" layers (to align to what is needed in the + `modeling_xxx.py` files). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + cache_position (`torch.Tensor`): + Deprecated and unused. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the chunked causal mask function (by doing the union of both). This is + useful to easily overlay another mask on top of the chunked causal one, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the chunked causal mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top of the chunked causal one, for example for image tokens handling. + """ + # If we have an hybrid cache structure, here we want to create the mask for the sliding layers + if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding: + layer_idx = past_key_values.is_sliding.index(True) + else: + layer_idx = 0 + + early_exit, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset = ( + _preprocess_mask_arguments(config, inputs_embeds, attention_mask, past_key_values, position_ids, layer_idx) + ) + if early_exit: + return attention_mask + + chunk_size = getattr(config, "attention_chunk_size", None) + if chunk_size is None: + raise ValueError("Could not find an `attention_chunk_size` argument in the config, or it is not set") + + # Raise if using chunked attention on context too large with FA + if is_flash_attention_requested(config) and kv_length + kv_offset > chunk_size: + raise ValueError( + "Flash attention cannot handle chunked attention, and the key-value length is larger than the chunk size so the " + "chunked pattern cannot be respected. You should use another `attn_implementation` when instantiating the model" + ) + + batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device + # For chunked attention and batched inputs, we need to take the number of left padding tokens into account + # to start the chunk from the actual start of the sequence for the padded sequence + if attention_mask is not None: + # Only count the left padding tokens, not all of them + left_padding_tokens = (attention_mask.cumsum(dim=-1) == torch.zeros_like(attention_mask)).sum(dim=-1) + else: + left_padding_tokens = torch.zeros(batch_size, device=device, dtype=int) + mask_factory_function = chunked_causal_mask_function(chunk_size, left_padding_tokens) + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + # Defaulting to using non-vmap based mask creations except when detecting + # users passing custom mask functions (as we cannot guarantee that they + # are properly index-based as required by our implementation). + use_vmap = False + # Do not allow skip if we are compiling (this is to match BC) + # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it + allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False) + + # Allow slight deviations from causal mask + # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, + # padding mask, etc) as the resulting mask may otherwise not be correct! + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_causal_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_causal_skip = False + use_vmap = True + + # If we detected packing format + if packed_sequence_mask is not None: + mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask)) + allow_is_causal_skip = False + + # We now create the mask + causal_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa + local_size=chunk_size, # Additional kwarg for sdpa + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return causal_mask + + +LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING = { + "full_attention": create_causal_mask, + "sliding_attention": create_sliding_window_causal_mask, + "chunked_attention": create_chunked_causal_mask, +} + + +@deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds") +def create_masks_for_generate( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, + **kwargs, +): + """ + This function mimics how we create the masks in the `modeling_xxx.py` files, and is used in places like `generate` + in order to easily create the masks in advance, when we compile the forwards with Static caches. + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the other mask function (by doing the union of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the other mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + """ + # The attribute reside in the text config for composite models + effective_config = config.get_text_config() + # Prepare the mask args + mask_kwargs = { + "config": effective_config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + "or_mask_function": or_mask_function, + "and_mask_function": and_mask_function, + } + + # If the attribute exist, we need several masks + if hasattr(effective_config, "layer_types"): + causal_masks = {} + for layer_pattern in set(effective_config.layer_types): + causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs) + return causal_masks + # In this case, all layers are sliding + elif getattr(effective_config, "sliding_window", None) is not None: + return create_sliding_window_causal_mask(**mask_kwargs) + # In this case, all layers are chunked + elif getattr(effective_config, "attention_chunk_size", None) is not None: + return create_chunked_causal_mask(**mask_kwargs) + # All layers use standard causal attention + return create_causal_mask(**mask_kwargs) + + +# Below are utilities to pretty-print the different masks +# Print the matrix with words as row labels +GREEN = "\033[92m" +YELLOW = "\033[93m" +RESET = "\033[0m" +BLACK_SQUARE = "■" +WHITE_SQUARE = "⬚" +GREY_SQUARE = "∙" +LOW_TRIANGLE = "⬕" +UPPER_TRIANGLE = "⬔" + + +def get_style(style): + if style == "majong": + BLACK_SQUARE = "🀞" # Full block (represents "on" or active) + BLACK_SQUARE = "🀙" # Full block (represents "on" or active) + WHITE_SQUARE = "🀆" # "▒" # Light shade (represents "off" or inactive) + LOW_TRIANGLE = "🀛" # Lower left triangle (stylized indication) + UPPER_TRIANGLE = "🀛" # Upper left triangle (stylized indication) + else: + BLACK_SQUARE = "█" # Full block (represents "on" or active) + WHITE_SQUARE = "░" # "▒" # Light shade (represents "off" or inactive) + LOW_TRIANGLE = "▙" # Lower left triangle (stylized indication)) + UPPER_TRIANGLE = "▜" # Upper left triangle (stylized indication) + + return BLACK_SQUARE, WHITE_SQUARE, LOW_TRIANGLE, UPPER_TRIANGLE + + +# LOW_TRIANGLE = UPPER_TRIANGLE = "⟍" # Upper right triangle (stylized indication) + +YELLOW_SQUARE = f"{YELLOW}{BLACK_SQUARE}{RESET}" +GREEN_SQUARE = f"{GREEN}{BLACK_SQUARE}{RESET}" + + +def tensor_to_mask_visual(original_tensor: torch.Tensor, grid_size=(20, 40), style="majong") -> str: + BLACK_SQUARE, WHITE_SQUARE, LOW_TRIANGLE, UPPER_TRIANGLE = get_style(style) + h, w = original_tensor.shape + max_h, max_w = grid_size + if not (h < max_h and w < max_w): + # Preserve aspect ratio within max grid size + aspect_ratio = 2 * w / h + if aspect_ratio > 1: + w = max_w + h = min(max_h, max(1, round(max_w / aspect_ratio))) + else: + h = max_h + w = max(1, round(max_h * aspect_ratio)) + + # Step 1: Rescale tensor by average pooling + tensor = original_tensor.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions + tensor = F.adaptive_avg_pool2d(tensor, output_size=(h, w))[0, 0] # Remove extra dims + else: + tensor = original_tensor + + # Step 3: Build the string representation + result = [] + for i in range(h): + row = "" + for j in range(w): + if tensor[i, j] == 1: + row += BLACK_SQUARE + elif tensor[i, j] == 0: + row += WHITE_SQUARE + else: + if j > 0: + if tensor[i, j - 1] == 1: + row += LOW_TRIANGLE + elif tensor[i, j - 1] == 0: + row += UPPER_TRIANGLE + else: + row += BLACK_SQUARE if tensor[i, j] == 1 else WHITE_SQUARE + else: + row += ( + BLACK_SQUARE + if tensor[i, j] == 1 + else ( + WHITE_SQUARE + if tensor[i, j] == 0 + else (UPPER_TRIANGLE if tensor[i, j + 1] == 1 else LOW_TRIANGLE) + ) + ) + result.append(row) + + return "\n".join(result) + + +class AttentionMask(torch.Tensor): + def __new__(cls, data, style=None): + # Create a new instance of AttentionMask as a Tensor + cls.style = style + return torch.Tensor._make_subclass(cls, data, require_grad=False) + + def __init__(self, data): + # You can initialize any additional metadata here if needed + pass + + def to_string(self, grid_size=(20, 40), limit=4): + """Returns a string representation of the block mask.""" + dense_mask = self + *batch_dims, num_rows, num_cols = dense_mask.shape + total_vis = [] + + for idx, batch_idx in enumerate(itertools.product(*[range(i) for i in batch_dims])): + if idx == limit: + total_vis.append("...") + total_vis.append("To print out more, set AttentionMask.to_string(limit=N)") + total_vis.append("You can also index (AttentionMask[batch, head]) to choose a specific batch or head") + break + block_vis = tensor_to_mask_visual(dense_mask[batch_idx], grid_size=grid_size, style=self.style) + total_vis.append(block_vis) + + total_vis.append(f"torch.Tensor(shape={tuple(self.shape)}, dtype={self.dtype})") + return "\n".join(total_vis) + + def __repr__(self): + return self.to_string() + + def __str__(self): + return self.to_string() + + @classmethod + def from_tensor(cls, tensor: torch.Tensor, style: str | None = None) -> "AttentionMask": + res = cls(tensor) + res.style = style + return res diff --git a/third_party/transformers/src/transformers/model_debugging_utils.py b/third_party/transformers/src/transformers/model_debugging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5b230836c6e7ec63a83b276e3ced23e30e6dd300 --- /dev/null +++ b/third_party/transformers/src/transformers/model_debugging_utils.py @@ -0,0 +1,455 @@ +# 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 functools +import json +import os +import re +from contextlib import contextmanager, redirect_stdout +from io import StringIO + +from .utils import logging +from .utils.import_utils import is_torch_available, requires + + +if is_torch_available(): + import torch + from safetensors.torch import save_file + + _torch_distributed_available = False + # Note to code inspectors: this toolbox is intended for people who add models to `transformers`. + if torch.distributed.is_available(): + import torch.distributed.tensor + + _torch_distributed_available = True +else: + _torch_distributed_available = False + + +logger = logging.get_logger(__name__) + + +def _is_rank_zero(): + """Return True if rank=0 or we aren't running distributed.""" + if not (_torch_distributed_available and torch.distributed.is_initialized()): + return True + return torch.distributed.get_rank() == 0 + + +MEMORY_ADDRESS_REGEX = re.compile(r"object at 0x[0-9A-Fa-f]+") + + +def _sanitize_repr_for_diff(x_str: str) -> str: + """ + Replace memory addresses in an object's repr with a stable placeholder + so that beautiful JSON diffs won't be ruined by ephemeral addresses. + """ + return MEMORY_ADDRESS_REGEX.sub("object at 0xXXXXXXXX", x_str) + + +def _dtensor_repr(x): + """Return a stable string representation for a DTensor-like object.""" + if _is_rank_zero(): + return f"DTensor (rank0) -> {repr(x._local_tensor)}" + return "DTensor(non-rank0)" + + +def _serialize_tensor_like_io( + value, debug_path: str | None = None, use_repr: bool = True, path_to_value: str | None = None +): + """ + Converts Tensors and DTensors to a JSON-serializable dictionary representation. + + Args: + value: Any Python object, often including torch Tensors, lists, dicts, etc. + debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files. + use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensor as the + `value` property in the asscoiated FULL_TENSORS.json file, or to store the full tensors in separate + SafeTensors file and store the relative path to that file in the `value` property in the dictionary. + path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full + tensor value if `use_repr=False`. + + Returns: + A nested Python structure (list, dict, or sanitized string) that is safe to json.dump. + """ + torch.set_printoptions(sci_mode=True) + + if use_repr: + value_out = _repr_to_list(value) + elif path_to_value: + if not path_to_value.endswith(".safetensors"): + path_to_value += ".safetensors" + + filepath = os.path.join(debug_path, path_to_value) if debug_path else path_to_value + save_file({"data": value.contiguous().detach().cpu()}, filepath) + value_out = f"./{path_to_value}" + else: + raise ValueError(f"{use_repr=} and {path_to_value=} cannot both be falsy.") + + out = { + "shape": repr(value.shape), + "dtype": repr(value.dtype), + "value": value_out, + } + if value.dtype in {torch.float16, torch.float32, torch.bfloat16}: + out.update( + { + "mean": _sanitize_repr_for_diff(repr(value.mean())), + "std": _sanitize_repr_for_diff(repr(value.std())), + "min": _sanitize_repr_for_diff(repr(value.min())), + "max": _sanitize_repr_for_diff(repr(value.max())), + } + ) + return out + + +def _serialize_io(value, debug_path: str | None = None, use_repr: bool = True, path_to_value: str | None = None): + """ + Recursively build a JSON-serializable Python structure from `value`. + Tensors and DTensors become either sanitized repr strings, or are saved to disk as SafeTensors files and their + relative paths are recorded in the returned Python structure. + Lists/tuples/dicts are recursed into. + All memory addresses are replaced with a stable placeholder. + + Args: + value: Any Python object, often including torch Tensors, lists, dicts, etc. + debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files. + use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the + `value` property in the asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors + files and store the relative path to that file in the `value` property. + path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full + tensor value if `use_repr=False`. + + Returns: + A nested Python structure (list, dict, or sanitized string) that is safe to json.dump. + """ + if isinstance(value, (list, tuple)): + return [ + _serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f"{path_to_value}_{i}") + for i, v in enumerate(value) + ] + + if isinstance(value, dict): + return { + k: _serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f"{path_to_value}_{k}") + for k, v in value.items() + } + + if hasattr(value, "_local_tensor"): + return _serialize_tensor_like_io( + value._local_tensor, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value + ) + + if isinstance(value, torch.Tensor): + return _serialize_tensor_like_io(value, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value) + + return _sanitize_repr_for_diff(repr(value)) + + +def _repr_to_list(value: torch.Tensor): + """ + Converts a tensor into a sanitized multi-line string representation. + + Args: + value (`torch.Tensor`): The tensor to represent. + + Returns: + `list[str]`: List of string lines representing the tensor. + """ + torch.set_printoptions(sci_mode=True, linewidth=120) + with StringIO() as buf, redirect_stdout(buf): + print(value) # to redirected stdout to avoid line splits + raw = buf.getvalue() + return _sanitize_repr_for_diff(raw).splitlines() + + +def prune_outputs_if_children(node): + # if there are children, remove this node's "outputs" + # so we only see outputs at the leaf level + if node.get("children"): + node.pop("outputs", None) + for child in node["children"]: + prune_outputs_if_children(child) + + +LAYER_SUFFIX_RE = re.compile(r"(.*)\.(\d+)$") # should be generic enough, ends with a number + + +def is_layer_block(node): + """ + Checks whether a node represents a layer block with submodules. + + Args: + node (`dict`): A node from the call tree. + + Returns: + `bool`: Whether the node is a layer block. + """ + match = LAYER_SUFFIX_RE.match(node.get("module_path", "")) + if not match or not node.get("children"): + return False + number = match.group(2) + return any(f".{number}." in child.get("module_path", "") for child in node["children"]) + + +def prune_intermediate_layers(node): + """ + Recursively removes intermediate layers from the tree to improve readability. + Keeps at least the first and last layers if many consecutive layers are present. + + Args: + node (`dict`): The root or subnode to prune recursively. + """ + if not node.get("children"): + return + layer_blocks = [(i, child) for i, child in enumerate(node["children"]) if is_layer_block(child)] + + if len(layer_blocks) > 2: + to_remove = [i for i, _ in layer_blocks[1:-1]] + node["children"] = [child for i, child in enumerate(node["children"]) if i not in to_remove] + + for child in node["children"]: + prune_intermediate_layers(child) + + +def log_model_debug_trace(debug_path: str | None, model): + if debug_path: + try: + os.makedirs(debug_path, exist_ok=True) + base = os.path.join(debug_path, model._debugger_module_dump_name + "_debug_tree") + except Exception as e: + raise ValueError(f"Unexpected or existing debug_path={debug_path}.") from e + else: + base = model._debugger_module_dump_name + "_debug_tree" + + logger.info(f"Writing model trace at {base}.json") + full_path = base + "_FULL_TENSORS.json" + summary_path = base + "_SUMMARY.json" + + prune_outputs_if_children(model._call_tree) + + with open(full_path, "w") as f: + json.dump(model._call_tree, f, indent=2) + + # summary-only version for readability - traversing the tree again #TODO optimize? + def strip_values(node): + def clean(val): + if isinstance(val, dict): + val.pop("value", None) + for v in val.values(): + clean(v) + elif isinstance(val, list): + for item in val: + clean(item) + + clean(node.get("inputs", {})) + clean(node.get("outputs", {})) + + for child in node.get("children", []): + strip_values(child) + + tree_copy = json.loads(json.dumps(model._call_tree)) # deep copy + strip_values(tree_copy) + + with open(summary_path, "w") as f: + json.dump(tree_copy, f, indent=2) + + +def _attach_debugger_logic( + model, + debug_path: str = ".", + do_prune_layers: bool = True, + use_repr: bool = True, +): + """ + Attaches a debugging wrapper to every module in the model. + + This records structured inputs and outputs during the forward pass into a call tree. + + Args: + model (`PreTrainedModel`, `nn.Module`): Model to wrap. + debug_path (`str`): Optional directory to dump debug JSON files. + do_prune_layers (`bool`, *optional*, defaults to `True`): Whether to prune intermediate layers. + use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the + `value` property in the associated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors + files and store the relative path to that file in the `value` property. + """ + class_name = model.__class__.__name__ + + # Prepare data structures on the model object + model._call_tree = {"module_path": class_name, "inputs": None, "outputs": None, "children": []} + model._debugger_model_call_stack = [] + model._debugger_module_dump_name = class_name # used for final JSON filename + + if debug_path: + try: + os.makedirs(debug_path, exist_ok=True) + except Exception as e: + raise ValueError(f"Unexpected or existing debug_path={debug_path}.") from e + + def wrap_forward(module, full_path): + orig_forward = module.forward + + @functools.wraps(orig_forward) + def wrapped_forward(*inps, **kws): + if _is_rank_zero(): + dict_inputs = {"args": inps, "kwargs": kws} + dict_inputs = {k: dict_inputs[k] for k in dict_inputs if len(dict_inputs[k]) > 0} + node = { + "module_path": full_path, + "inputs": _serialize_io( + dict_inputs, + debug_path=debug_path, + use_repr=use_repr, + path_to_value=f"{full_path}_inputs", + ), + "outputs": None, + "children": [], + } + model._debugger_model_call_stack.append(node) + with torch.no_grad(): + out = orig_forward(*inps, **kws) + + if _is_rank_zero(): + if sum(1 for _ in module.named_children()) > 0: + node["outputs"] = None + else: + node["outputs"] = _serialize_io( + out, + debug_path=debug_path, + use_repr=use_repr, + path_to_value=f"{full_path}_outputs", + ) + + finished = model._debugger_model_call_stack.pop() + # prune empty vertices here as well (mostly empty children nodes) + if not finished["children"]: + finished.pop("children") + + if model._debugger_model_call_stack: + model._debugger_model_call_stack[-1]["children"].append(finished) + return out + + module.forward = wrapped_forward + + # wrap all submodules + for name, submodule in model.named_modules(): + if name == "": + continue + wrap_forward(submodule, f"{class_name}.{name}") + + # wrap top-level forward + real_top_forward = model.forward + + @functools.wraps(real_top_forward) + def top_wrapped_forward(*inps, **kws): + if _is_rank_zero(): + top_node = { + "module_path": f"{class_name} (top-level)", + "inputs": _serialize_io( + {"args": inps, "kwargs": kws}, + debug_path=debug_path, + use_repr=use_repr, + path_to_value=f"{class_name}_inputs", + ), + "outputs": None, + "children": [], + } + model._debugger_model_call_stack.append(top_node) + + out = real_top_forward(*inps, **kws) + if _is_rank_zero() and model._debugger_model_call_stack: + top_node["outputs"] = _serialize_io( + out, + debug_path=debug_path, + use_repr=use_repr, + path_to_value=f"{class_name}_outputs", + ) + finished = model._debugger_model_call_stack.pop() + model._call_tree["inputs"] = finished["inputs"] + model._call_tree["outputs"] = finished["outputs"] + model._call_tree["children"] = finished["children"] + # prune empty stuff for visibility + [model._call_tree.pop(k, None) for k in list(model._call_tree.keys()) if not model._call_tree[k]] + + # prune layers that are not 0 or last + if do_prune_layers: + prune_intermediate_layers(model._call_tree) + # Write final JSON trace here + log_model_debug_trace(debug_path=debug_path, model=model) + return out + + model.forward = top_wrapped_forward + + +@requires(backends=("torch",)) +@contextmanager +def model_addition_debugger_context( + model, + debug_path: str | None = None, + do_prune_layers: bool = True, + use_repr: bool = True, +): + """ + # Model addition debugger - context manager for model adders + This context manager is a power user tool intended for model adders. + + It tracks all forward calls within a model forward and logs a slice of each input and output on a nested JSON file. + If `use_repr=True` (the default), the JSON file will record a `repr()`-ized version of the tensors as a list of + strings. If `use_repr=False`, the full tensors will be stored in separate SafeTensors files and the JSON file will + provide a relative path to that file. + + To note, this context manager enforces `torch.no_grad()`. + + ## Usage + + add the context manager to a model to debug + + ```python + import torch + + from PIL import Image + from transformers import LlavaProcessor, LlavaForConditionalGeneration, model_addition_debugger_context + + torch.random.manual_seed(673) + + # load pretrained model and processor + model_id = "llava-hf/llava-1.5-7b-hf" + processor = LlavaProcessor.from_pretrained(model_id) + model = LlavaForConditionalGeneration.from_pretrained(model_id) + + # create random image input + random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy()) + + # prompt + prompt = "Describe this image." + + # process inputs + inputs = processor(text=prompt, images=random_image, return_tensors="pt") + + # call forward method (not .generate!) + with model_addition_debugger_context(model, debug_path="Your_debug_path", do_prune_layers=False): + output = model.forward(**inputs) + ``` + + """ + orig_forwards = {m: m.forward for _, m in model.named_modules()} + orig_forwards[model] = model.forward + _attach_debugger_logic(model, debug_path, do_prune_layers, use_repr) + try: + yield model + finally: + for module_instance, forward_method in orig_forwards.items(): + module_instance.forward = forward_method diff --git a/third_party/transformers/src/transformers/modeling_flash_attention_utils.py b/third_party/transformers/src/transformers/modeling_flash_attention_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9211ccb19a9e10a997e42283296b394e141ac115 --- /dev/null +++ b/third_party/transformers/src/transformers/modeling_flash_attention_utils.py @@ -0,0 +1,807 @@ +# Copyright 2025 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import importlib +import inspect +import os +from collections.abc import Callable +from functools import partial +from typing import TypedDict + +import torch +import torch.nn.functional as F + +from .utils import ( + is_flash_attn_2_available, + is_flash_attn_3_available, + is_flash_attn_4_available, + is_torch_cuda_available, + is_torch_mlu_available, + is_torch_npu_available, + is_torch_xpu_available, + logging, +) +from .utils.import_utils import PACKAGE_DISTRIBUTION_MAPPING, is_tracing + + +logger = logging.get_logger(__name__) + + +# TODO Deprecate when all models have the attention interface +def flash_attn_supports_top_left_mask(): + if is_flash_attn_2_available() or is_flash_attn_3_available() or is_flash_attn_4_available(): + return False + + from .integrations.npu_flash_attention import is_npu_fa2_top_left_aligned_causal_mask + + return is_npu_fa2_top_left_aligned_causal_mask() + + +# TODO Deprecate when all models have the attention interface +def is_flash_attn_available(): + return ( + is_flash_attn_4_available() + or is_flash_attn_3_available() + or is_flash_attn_2_available() + or is_torch_npu_available() + or is_torch_xpu_available() + ) + + +# Mapping from flash attention implementations to their kernel fallback repositories +FLASH_ATTN_KERNEL_FALLBACK = { + "flash_attention_2": "kernels-community/flash-attn2", + "flash_attention_3": "kernels-community/vllm-flash-attn3", + "flash_attention_4": "kernels-community/flash-attn4", +} + + +# Meta information on each mainline FA compatibility: +# 1. The import structure and availability +# 2. Device support (with custom ones that use other workarounds, e.g. kernels) +# 3. Supported major cuda devices, e.g. Hopper, Blackwell. Mostly found in the newest FA versions +FLASH_ATTENTION_COMPATIBILITY_MATRIX = { + 2: { + "flash_attn_version": 2, + "general_availability_check": is_flash_attn_2_available, + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None + and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "supported_devices": ( + (is_torch_cuda_available, "cuda"), + (is_torch_mlu_available, "mlu"), + (is_torch_npu_available, "npu"), + (is_torch_xpu_available, "xpu"), + ), + "custom_supported_devices": ( + (is_torch_npu_available, "Detect using FlashAttention2 on Ascend NPU."), + ( + is_torch_xpu_available, + f"Detect using FlashAttention2 (via kernel `{FLASH_ATTN_KERNEL_FALLBACK['flash_attention_2']}`) on XPU.", + ), + ), + }, + 3: { + "flash_attn_version": 3, + "general_availability_check": is_flash_attn_3_available, + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None + and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], + "supported_devices": ((is_torch_cuda_available, "cuda"),), + "cuda_min_major_version": 8, # Ampere + }, + 4: { + "flash_attn_version": 4, + "general_availability_check": is_flash_attn_4_available, + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None + and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "supported_devices": ((is_torch_cuda_available, "cuda"),), + "cuda_min_major_version": 9, # Hopper + }, +} + + +# `globals()` is not compatible with dynamo, hence we have do define them in global scope ourselves +_loaded_implementation = None +_flash_fn = None +_flash_varlen_fn = None +_flash_with_kvcache_fn = None +_pad_fn = None +_unpad_fn = None + +# function that processes kwargs, generalized to handle any supported kwarg within the function +_process_flash_kwargs_fn = None +# exceptions where hf API doesn't match the original flash attention API +_hf_api_to_flash_mapping = { + "dropout": "dropout_p", + "sliding_window": "window_size", +} +# alternative names within the different flash attention APIs, e.g. for attention sinks +_flash_api_alternative_names = {"s_aux": "learnable_sink"} + + +def _lazy_imports( + implementation: str | None, attention_wrapper: Callable | None = None, allow_all_kernels: bool = False +): + """ + Lazy loads the respective flash attention implementations. + + Return: + flash_attn_func: The base flash attention function. + flash_attn_varlen_func: The flash attention function supporting variable sequence lengths, + e.g. for padding-free training. + pad_input: The function to pad inputs into one sequence and returning the respective kwargs. + unpad_input: The function to unpad outputs based on the kwargs (from pad_input). + """ + is_fa2 = is_flash_attn_2_available() + is_fa3 = is_flash_attn_3_available() + is_fa4 = is_flash_attn_4_available() + + pad_input, unpad_input = _pad_input, _unpad_input + + is_paged = implementation.startswith("paged|") + implementation = implementation.split("|")[1] if is_paged else implementation + + if (implementation == "flash_attention_2" and is_fa2) or ( + implementation is None and is_fa2 and not is_fa3 and not is_fa4 + ): + from flash_attn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + from flash_attn.bert_padding import pad_input, unpad_input + elif is_torch_npu_available(): + # Package `flash-attn` is unavailable on Ascend NPU, which will cause ImportError + # Flash-Attention2 related apis for Ascend NPU must be imported from `.integrations.npu_flash_attention` module + from .integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func + from .integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func + from .integrations.npu_flash_attention import npu_flash_attn_with_kvcache as flash_attn_with_kvcache + else: + if implementation == "flash_attention_3" or (implementation is None and is_fa3 and not is_fa4): + from flash_attn_interface import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + elif implementation == "flash_attention_4" or (implementation is None and is_fa4): + from flash_attn.cute import flash_attn_func, flash_attn_varlen_func + + flash_attn_with_kvcache = None # not supported yet + # Kernels fallback + else: + from .integrations.hub_kernels import load_and_register_attn_kernel + + # Map standard attention names to hub kernel repos + kernel_repo = FLASH_ATTN_KERNEL_FALLBACK.get(implementation, implementation) + # We want to explicitly register the name with `paged|` if found + kernel_implementation = f"paged|{implementation}" if is_paged else kernel_repo + kernel = load_and_register_attn_kernel( + kernel_implementation, attention_wrapper, allow_all_kernels=allow_all_kernels + ) + + flash_attn_func = getattr(kernel, "flash_attn_func", None) + flash_attn_varlen_func = getattr(kernel, "flash_attn_varlen_func", None) + flash_attn_with_kvcache = getattr(kernel, "flash_attn_with_kvcache", None) + if flash_attn_varlen_func is None: + raise ValueError( + f"Could not find the currently requested flash attention implementation at `{implementation}`." + "Make sure that you request a valid kernel from the hub, e.g. `kernels-community/flash-attn2`." + ) + if flash_attn_func is None: + logger.warning( + f"The loaded flash attention implementation at `{implementation}` only supports varlen, i.e. " + "it can only be used with continuous batching and does not support the full functionality for " + "the base transformers generation methods." + ) + if flash_attn_with_kvcache is None: + logger.warning( + f"The loaded flash attention implementation at `{implementation}` does not support block tables, so" + " the full performances of continuous batching will not be achieved, only the varlen path will be " + "used." + ) + + return flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache, pad_input, unpad_input + + +def _lazy_define_process_function(flash_function): + """ + Depending on the version and kernel some features are not supported. Due to limitations in + `torch.compile`, we opt to statically type which (optional) kwarg parameters are supported + within `_process_flash_attention_kwargs`. + + NOTE: While all supported kwargs are marked as `True`, everything else is marked as `False`. + This might be confusing for kwargs that we use in any case, e.g. `is_causal`. + """ + + flash_parameters = inspect.signature(flash_function).parameters + process_parameters = inspect.signature(_process_flash_attention_kwargs).parameters + + supports_mapping = {} + for param in process_parameters: + fa_param = _hf_api_to_flash_mapping.get(param, param) + supports_mapping[fa_param] = fa_param in flash_parameters + + if (fa_alternative_name := _flash_api_alternative_names.get(param, param)) != fa_param: + supports_mapping[fa_alternative_name] = fa_alternative_name in flash_parameters + + return partial(_process_flash_attention_kwargs, supports_mapping=supports_mapping) + + +def lazy_import_flash_attention( + implementation: str | None, attention_wrapper: Callable | None = None, allow_all_kernels: bool = False +): + """ + Lazily import flash attention and return the respective functions + flags. + + NOTE: For fullgraph, this needs to be called before compile, while no fullgraph can + work without preloading. See `load_and_register_attn_kernel` in `integrations.hub_kernels`. + """ + global _loaded_implementation + if implementation is None and _loaded_implementation is None: + raise ValueError("Could not find any flash attn implementation based on your environment.") + + global _flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn, _process_flash_kwargs_fn + if implementation is not None and _loaded_implementation != implementation: + _loaded_implementation = implementation + + _flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn = _lazy_imports( + implementation, attention_wrapper, allow_all_kernels=allow_all_kernels + ) + _process_flash_kwargs_fn = _lazy_define_process_function(_flash_varlen_fn) + + return (_flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn), _process_flash_kwargs_fn + + +def lazy_import_paged_flash_attention(implementation: str | None, allow_all_kernels: bool = False): + """ + Same as `lazy_import_flash_attention` but explicitly wrapping it with the paged implementation. + """ + from .integrations.flash_paged import paged_attention_forward + + (_, flash_attn_varlen_func, flash_attn_with_kvcache_fn, _, _), _ = lazy_import_flash_attention( + implementation, attention_wrapper=paged_attention_forward, allow_all_kernels=allow_all_kernels + ) + return flash_attn_varlen_func, flash_attn_with_kvcache_fn + + +def _index_first_axis(tensor, indices): + """ + A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis, + after flattening the first two dimensions of the tensor. This is functionally equivalent to + FA2's `index_first_axis` and replaces the need to import it. + """ + # The input tensor is expected to be of shape (batch, seq_len, ...). We flatten the first + # two dimensions to get (total_tokens, ...) before indexing. + reshaped_tensor = tensor.reshape(-1, *tensor.shape[2:]) + return reshaped_tensor[indices] + + +def _unpad_input(hidden_states, attention_mask, unused_mask=None): + """ + unpad_input function for flash attention variants that do not have them within their pkg themselves, e.g. fa3. + + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. + unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. + + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. + indices: (total_nnz), the indices of masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. + """ + all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask + seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) + used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + + return ( + _index_first_axis(hidden_states, indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + used_seqlens_in_batch, + ) + + +def _pad_input(hidden_states, indices, batch, seqlen): + """ + pad_input function for flash attention variants that do not have them within their pkg themselves, e.g. fa3. + + Arguments: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. + batch: int, batch size for the padded sequence. + seqlen: int, maximum sequence length for the padded sequence. + + Return: + hidden_states: (batch, seqlen, ...) + """ + dim = hidden_states.shape[1:] + output = torch.zeros((batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype) + output[indices] = hidden_states + return output.view(batch, seqlen, *dim) + + +def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Retrieves indexing data required to repad unpadded (ragged) tensors. + + Arguments: + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + indices (`torch.Tensor`): + The indices of non-masked tokens from the flattened input sequence. + cu_seqlens (`torch.Tensor`): + The cumulative sequence lengths, used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + max_seqlen_in_batch (`int`): + Maximum sequence length in batch. + """ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def _upad_input( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + attention_mask: torch.Tensor, + query_length: int, + unpad_input_func, +): + """ + Unpads query, key, and values tensors, using a single dimension for all tokens even though they belong to different batches. + This function is used instead of `flash_attn.bert_padding.unpad_input` in order to avoid the recomputation of the same intermediary + tensors for query, key, value tensors. + + Arguments: + query_layer (`torch.Tensor`): + Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + query_length (`int`): + Target length. + unpad_input_func: + The function to use for unpadding the input tensors. + + Return: + query_layer (`torch.Tensor`): + Query state without padding. Shape: (total_target_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + indices_q (`torch.Tensor`): + The indices of non-masked tokens from the flattened input target sequence. + (cu_seqlens_q, cu_seqlens_k) (`tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + + # With static caches, the k/v states may be larger than the mask -> we need to slice them to avoid generating garbage + # It's a bit of an anti-pattern, but otherwise we silently compute wrong attentions scores + if key_layer.shape[1] > (seq_len := attention_mask.shape[-1]): + key_layer, value_layer = key_layer[:, :seq_len, :, :], value_layer[:, :seq_len, :, :] + + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = _index_first_axis(key_layer, indices_k) + value_layer = _index_first_axis(value_layer, indices_k) + if query_length == kv_seq_len: + query_layer = _index_first_axis(query_layer, indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input_func(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +def prepare_fa_kwargs_from_position_ids(position_ids): + """ + This function returns all the necessary kwargs to call `flash_attn_varlen_func` extracted from position_ids. + + Arguments: + position_ids (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + (cu_seqlens_q, cu_seqlens_k) (`tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index into + ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, + `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + tensor_kwargs = {"dtype": torch.int32, "device": position_ids.device} + + position_ids = position_ids.reshape(-1) + indices_q = (position_ids == 0).nonzero().view(-1) + + cu_seq_lens_q = torch.cat( + ( + indices_q.to(**tensor_kwargs), + torch.tensor(position_ids.size(), **tensor_kwargs), + ) + ) + cu_seq_lens_k = cu_seq_lens_q + + # https://github.com/Dao-AILab/flash-attention/blob/2dd8078adc1d9b74e315ee99718c0dea0de8eeb6/flash_attn/flash_attn_interface.py#L1423-L1424 + # We should use cu_seq_lens instead of position_ids to get the max length since position_ids is not always increasing + # for some models (e.g. qwen2-vl). + max_length_q = cu_seq_lens_q.diff().max() + max_length_k = max_length_q + + return (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) + + +def _prepare_from_posids(query, key, value, position_ids): + """ + This function returns necessary arguments to call `flash_attn_varlen_func`. + All three query, key, value states will be flattened. + Cumulative lengths of each examples in the batch will be extracted from position_ids. + NOTE: ideally cumulative lengths should be prepared at the data collator stage + + Arguments: + query (`torch.Tensor`): + Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim). + key (`torch.Tensor`): + Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + value (`torch.Tensor`): + Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + position_ids (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + query (`torch.Tensor`): + Query state without padding. Shape: (total_target_length, num_heads, head_dim). + key (`torch.Tensor`): + Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + value (`torch.Tensor`): + Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + (cu_seqlens_q, cu_seqlens_k) (`tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + query = query.contiguous().view(-1, query.size(-2), query.size(-1)) + key = key.contiguous().view(-1, key.size(-2), key.size(-1)) + value = value.contiguous().view(-1, value.size(-2), value.size(-1)) + + (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = prepare_fa_kwargs_from_position_ids(position_ids) + + return (query, key, value, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k)) + + +def _is_packed_sequence(position_ids, batch_size): + """ + Check the position ids whether packed sequences are indicated or not + 1. Position ids exist + 2. Flattened sequences only are supported + 3. Compile-friendly `not (torch.diff(position_ids, dim=-1) >= 0).all()`, i.e. we have multiple increasing sequences + """ + if position_ids is None: + return False + + increasing_position_sequences = ( + torch.arange(position_ids.shape[1], device=position_ids.device) + position_ids.min() + ) + return batch_size == 1 and (increasing_position_sequences - position_ids).abs().sum().bool() + + +def fa_peft_integration_check( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + target_dtype: torch.dtype | None = None, +): + """ + PEFT usually casts the layer norms in float32 for training stability reasons + therefore the input hidden states gets silently casted in float32. Hence, we need + cast them back in float16 / bfloat16 just to be sure everything works as expected. + This might slowdown training & inference so it is recommended to not cast the LayerNorms! + """ + if target_dtype and q.dtype == torch.float32: + logger.warning_once(f"Casting fp32 inputs back to {target_dtype} for flash-attn compatibility.") + q, k, v = q.to(target_dtype), k.to(target_dtype), v.to(target_dtype) + return q, k, v + + +class FlashAttentionKwargs(TypedDict, total=False): + """ + Keyword arguments for Flash Attention with Compile. + + Attributes: + cu_seq_lens_q (`torch.LongTensor`, *optional*) + Gets cumulative sequence length for query state. + cu_seq_lens_k (`torch.LongTensor`, *optional*) + Gets cumulative sequence length for key state. + max_length_q (`int`, *optional*): + Maximum sequence length for query state. + max_length_k (`int`, *optional*): + Maximum sequence length for key state. + """ + + cu_seq_lens_q: torch.LongTensor | None + cu_seq_lens_k: torch.LongTensor | None + max_length_q: int | None + max_length_k: int | None + + +def _process_flash_attention_kwargs( + query_length: int, + key_length: int, + is_causal: bool, + dropout: float = 0.0, + softmax_scale: float | None = None, + sliding_window: int | None = None, + use_top_left_mask: bool = False, + softcap: float | None = None, + deterministic: bool | None = None, + s_aux: torch.Tensor | None = None, + max_seqlen_q: int | torch.IntTensor | None = None, + max_seqlen_k: int | torch.IntTensor | None = None, + supports_mapping: dict[str, bool] | None = None, + **kwargs, +): + """ + Returns a set of kwargs that are passed down to the according flash attention function based on + requested features and whether it is supported - depends on the version and kernel implementation + which is dynamically configured at `lazy_import_flash_attention`. The (un)supported features can be + inspected in `supports_mapping`, see `_lazy_define_process_function` for more details. + + Args: + query_length (`int`): + Length of the query states + key_length (`int`): + Length of the key states + is_causal (`bool`): + Whether we perform causal (decoder) attention or full attention. + dropout (`float`): + Attention dropout. + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to `1 / sqrt(head_dim)`. + sliding_window (`int`, *optional*): + The size of the sliding window, i.e. we look at a max of `sliding_window` tokens back. + use_top_left_mask (`bool`): + Deprecated behavior of older versions of flash attention requiring different masking. + softcap (`float`, *optional*): + Softcap for the attention logits, used e.g. in gemma2. + deterministic (`bool`, *optional*): + Determines if the deterministic option introduced in flash_attn>=2.4.1 is enabled. + s_aux (`torch.Tensor`, *optional*): + Attention sink auxiliary that adds a `bias` to the attention calculation via an additional head. + max_seqlen_q (`Union[int, torch.IntTensor]`, *optional*): + The maximum sequence length in the query tensor during a varlen forward. + max_seqlen_k (`Union[int, torch.IntTensor]`, *optional*): + The maximum sequence length in the key/value tensor during a varlen forward. + Return: + flash_kwargs (`dict`): + A dict of kwargs that are requested and supported. + """ + flash_kwargs = { + "causal": is_causal and not (use_top_left_mask and query_length == 1), + "softmax_scale": softmax_scale, + } + + if supports_mapping["dropout_p"]: + flash_kwargs["dropout_p"] = dropout + + if supports_mapping["window_size"] and sliding_window is not None and key_length > sliding_window: + # The flash attention API sets inclusive boundaries, i.e. (4, 0) would take 4 tokens to the left + # and the current token for a total size of 5. However, we usually define our window sizes by + # their total window size (when causal). Encoder models as of now seldom use SWA and when they + # do, they must align with this symmetric logic, i.e. for a total of `2*sliding_window + 1`. + flash_kwargs["window_size"] = (sliding_window - 1, sliding_window - 1) + + if supports_mapping["deterministic"]: + flash_kwargs["deterministic"] = ( + deterministic if deterministic is not None else os.getenv("FLASH_ATTENTION_DETERMINISTIC", "0") == "1" + ) + + if supports_mapping["softcap"] and softcap is not None: + flash_kwargs["softcap"] = softcap + + if ((legacy_sink_param := supports_mapping["s_aux"]) or supports_mapping["learnable_sink"]) and s_aux is not None: + if legacy_sink_param: + flash_kwargs["s_aux"] = s_aux # e.g. FA3 (vllm) + else: + flash_kwargs["learnable_sink"] = s_aux # FA4 + + # There is a limitation of the flash attention API, as the function `flash_attn_varlen_func` + # may require `max_length_q`, `max_length_k` to be passed as `int` and not `torch.Tensor`. + # + # You can either set + # - Env: `TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1` + # - Before compiling: `torch._dynamo.config.capture_scalar_outputs = True` + # to allow torch compile to handle scalar outputs in those cases. + same_max_seqlen = max_seqlen_q is max_seqlen_k # to avoid 2x device syncs + if supports_mapping["max_seqlen_q"] and max_seqlen_q is not None: + if not isinstance(max_seqlen_q, int) and is_tracing(max_seqlen_q): + max_seqlen_q = max_seqlen_q.item() + flash_kwargs["max_seqlen_q"] = max_seqlen_q + + if supports_mapping["max_seqlen_k"] and max_seqlen_k is not None: + if same_max_seqlen and flash_kwargs["max_seqlen_q"] is not None: + max_seqlen_k = flash_kwargs["max_seqlen_q"] + elif not isinstance(max_seqlen_k, int) and is_tracing(max_seqlen_k): + max_seqlen_k = max_seqlen_k.item() + flash_kwargs["max_seqlen_k"] = max_seqlen_k + + return flash_kwargs + + +def _flash_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor | None, + query_length: int, + is_causal: bool, + dropout: float = 0.0, + position_ids: torch.Tensor | None = None, + softmax_scale: float | None = None, + sliding_window: int | None = None, + use_top_left_mask: bool = False, + softcap: float | None = None, + deterministic: bool | None = None, + cu_seq_lens_q: torch.LongTensor | None = None, + cu_seq_lens_k: torch.LongTensor | None = None, + max_length_q: int | None = None, + max_length_k: int | None = None, + target_dtype: torch.dtype | None = None, + attn_implementation: str | None = None, + **kwargs, +): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + (Optional) kwargs are described further in `_process_flash_attention_kwargs` and `FlashAttentionKwargs`. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`, *optional*): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + attn_implementation (`str`, *optional*): + The attention implementation to use. If None, will default to the one based on the environment. + """ + (flash_fn, flash_varlen_fn, _, pad_fn, unpad_fn), process_flash_kwargs_fn = lazy_import_flash_attention( + attn_implementation + ) + + # PEFT possibly silently casts tensors to fp32, this potentially reconverts to correct dtype or is a no op + query_states, key_states, value_states = fa_peft_integration_check( + query_states, key_states, value_states, target_dtype + ) + + # Extract the flash attention kwargs that have been requested (and are supported by the implementation) + flash_kwargs = partial( + process_flash_kwargs_fn, + query_length=query_length, + key_length=key_states.size(1), + is_causal=is_causal, + dropout=dropout, + softmax_scale=softmax_scale, + sliding_window=sliding_window, + use_top_left_mask=use_top_left_mask, + softcap=softcap, + deterministic=deterministic, + **kwargs, + ) + + # We will use `flash_varlen_fn` to prevent cross-example attention and also allow padding free approach under two cases: + # Case 1. If position ids is provided and the position ids indicate packed sequences, see `_is_packed_sequence`. + # Case 2. Some models pass directly pre-computed `cu_seqlens` so we don't need to infer it from position ids. It is safe to + # use `flash_varlen_fn` knowing we already have all necessary the kwargs. + # + # NOTE: it is user's responsibility to take care of flattening `position_ids` if that's needed by the model. + # See #39121 for more information. + is_fa_with_position_ids = _is_packed_sequence(position_ids, batch_size=query_states.size(0)) + is_fa_with_varlen_kwargs = all( + kwarg is not None for kwarg in (cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k) + ) + + # Contains at least one padding token in the sequence + if attention_mask is not None: + q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( + query_states, key_states, value_states, attention_mask, query_length, unpad_fn + ) + + # TODO for now this is required to work with + # https://huggingface.co/kernels-community/metal-flash-sdpa/blob/main/torch-ext/metal_flash_sdpa/__init__.py + if "mps" in str(q.device): + cu_seq_lens_k = cu_seq_lens_k.clone() + + out_unpad = flash_varlen_fn( + q, + k, + v, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + **flash_kwargs(max_seqlen_q=max_length_q, max_seqlen_k=max_length_k), + ) + if isinstance(out_unpad, tuple): + out_unpad = out_unpad[0] + + out = pad_fn(out_unpad, indices_q, query_states.size(0), query_length) + + # Padding free, i.e. sequences flattened into one total sequence + elif is_fa_with_varlen_kwargs or is_fa_with_position_ids: + if cu_seq_lens_q is None or cu_seq_lens_k is None: + q, k, v, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _prepare_from_posids( + query_states, key_states, value_states, position_ids + ) + else: + q = query_states.reshape(-1, query_states.size(-2), query_states.size(-1)) + k = key_states.reshape(-1, key_states.size(-2), key_states.size(-1)) + v = value_states.reshape(-1, value_states.size(-2), value_states.size(-1)) + + # TODO for now this is required to work with + # https://huggingface.co/kernels-community/metal-flash-sdpa/blob/main/torch-ext/metal_flash_sdpa/__init__.py + if "mps" in str(q.device): + cu_seq_lens_k = cu_seq_lens_k.clone() + + out = flash_varlen_fn( + q, + k, + v, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + **flash_kwargs(max_seqlen_q=max_length_q, max_seqlen_k=max_length_k), + ) + if isinstance(out, tuple): + out = out[0] + + out = out.view(query_states.size(0), -1, out.size(-2), out.size(-1)) + + # No padding + else: + out = flash_fn(query_states, key_states, value_states, **flash_kwargs()) + if isinstance(out, tuple): + out = out[0] + + return out diff --git a/third_party/transformers/src/transformers/modeling_layers.py b/third_party/transformers/src/transformers/modeling_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..1012606fcaaf37c7104c97e9c9f587401805a858 --- /dev/null +++ b/third_party/transformers/src/transformers/modeling_layers.py @@ -0,0 +1,288 @@ +# 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 functools import partial + +import torch +import torch.nn as nn + +from .cache_utils import Cache +from .modeling_outputs import ( + BaseModelOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from .models.auto import AutoModel +from .processing_utils import Unpack +from .utils import TransformersKwargs, auto_docstring, can_return_tuple, logging + + +logger = logging.get_logger(__name__) + + +class GradientCheckpointingLayer(nn.Module): + """Base class for layers with gradient checkpointing. + + This class enables gradient checkpointing functionality for a layer. By default, gradient checkpointing is disabled + (`gradient_checkpointing = False`). When `model.set_gradient_checkpointing()` is called, gradient checkpointing is + enabled by setting `gradient_checkpointing = True` and assigning a checkpointing function to `_gradient_checkpointing_func`. + + Important: + + When using gradient checkpointing with `use_reentrant=True`, inputs that require gradients (e.g. hidden states) + must be passed as positional arguments (`*args`) rather than keyword arguments to properly propagate gradients. + + Example: + + ```python + >>> # Correct - hidden_states passed as positional arg + >>> out = self.layer(hidden_states, attention_mask=attention_mask) + + >>> # Incorrect - hidden_states passed as keyword arg + >>> out = self.layer(hidden_states=hidden_states, attention_mask=attention_mask) + ``` + """ + + gradient_checkpointing = False + + def __call__(self, *args, **kwargs): + if self.gradient_checkpointing and self.training: + do_warn = False + layer_name = self.__class__.__name__ + message = f"Caching is incompatible with gradient checkpointing in {layer_name}. Setting" + + if "use_cache" in kwargs and kwargs["use_cache"]: + kwargs["use_cache"] = False + message += " `use_cache=False`," + do_warn = True + + # different names for the same thing in different layers + # TODO cyril: this one without `S` can be removed after deprecation cycle + if "past_key_value" in kwargs and kwargs["past_key_value"] is not None: + kwargs["past_key_value"] = None + message += " `past_key_value=None`," + do_warn = True + + if "past_key_values" in kwargs and kwargs["past_key_values"] is not None: + kwargs["past_key_values"] = None + message += " `past_key_values=None`," + do_warn = True + + if "layer_past" in kwargs and kwargs["layer_past"] is not None: + kwargs["layer_past"] = None + message += " `layer_past=None`," + do_warn = True + + # warn if anything was changed + if do_warn: + message = message.rstrip(",") + "." + logger.warning_once(message) + + return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) + return super().__call__(*args, **kwargs) + + +@auto_docstring +class GenericForSequenceClassification: + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class + setattr(self, self.base_model_prefix, AutoModel.from_config(config)) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.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, + **kwargs: Unpack[TransformersKwargs], + ) -> SequenceClassifierOutputWithPast: + transformer_outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)( + 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 = transformer_outputs.last_hidden_state + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class GenericForQuestionAnswering: + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class + setattr(self, self.base_model_prefix, AutoModel.from_config(config)) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return getattr(self, self.base_model_prefix).embed_tokens + + def set_input_embeddings(self, value): + getattr(self, self.base_model_prefix).embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + start_positions: torch.LongTensor | None = None, + end_positions: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> QuestionAnsweringModelOutput: + outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + sequence_output = outputs.last_hidden_state + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + loss = None + if start_positions is not None and end_positions is not None: + loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs) + + return QuestionAnsweringModelOutput( + loss=loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class GenericForTokenClassification: + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class + setattr(self, self.base_model_prefix, AutoModel.from_config(config)) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + 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, + **kwargs: Unpack[TransformersKwargs], + ) -> TokenClassifierOutput: + outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)( + 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, + ) + sequence_output = outputs.last_hidden_state + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/third_party/transformers/src/transformers/models/__init__.py b/third_party/transformers/src/transformers/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..989be9eb114eca42c2ec18177bc268876a64b937 --- /dev/null +++ b/third_party/transformers/src/transformers/models/__init__.py @@ -0,0 +1,472 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ..utils import _LazyModule +from ..utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .afmoe import * + from .aimv2 import * + from .albert import * + from .align import * + from .altclip import * + from .apertus import * + from .arcee import * + from .aria import * + from .audio_spectrogram_transformer import * + from .audioflamingo3 import * + from .auto import * + from .autoformer import * + from .aya_vision import * + from .bamba import * + from .bark import * + from .bart import * + from .barthez import * + from .bartpho import * + from .beit import * + from .bert import * + from .bert_generation import * + from .bert_japanese import * + from .bertweet import * + from .big_bird import * + from .bigbird_pegasus import * + from .biogpt import * + from .bit import * + from .bitnet import * + from .blenderbot import * + from .blenderbot_small import * + from .blip import * + from .blip_2 import * + from .bloom import * + from .blt import * + from .bridgetower import * + from .bros import * + from .byt5 import * + from .camembert import * + from .canine import * + from .chameleon import * + from .chinese_clip import * + from .chmv2 import * + from .clap import * + from .clip import * + from .clipseg import * + from .clvp import * + from .code_llama import * + from .codegen import * + from .cohere import * + from .cohere2 import * + from .cohere2_vision import * + from .cohere_asr import * + from .colmodernvbert import * + from .colpali import * + from .colqwen2 import * + from .conditional_detr import * + from .convbert import * + from .convnext import * + from .convnextv2 import * + from .cpm import * + from .cpmant import * + from .csm import * + from .ctrl import * + from .cvt import * + from .cwm import * + from .d_fine import * + from .dab_detr import * + from .dac import * + from .data2vec import * + from .dbrx import * + from .deberta import * + from .deberta_v2 import * + from .decision_transformer import * + from .deepseek_v2 import * + from .deepseek_v3 import * + from .deepseek_vl import * + from .deepseek_vl_hybrid import * + from .deformable_detr import * + from .deit import * + from .deprecated import * + from .depth_anything import * + from .depth_pro import * + from .detr import * + from .dia import * + from .dialogpt import * + from .diffllama import * + from .dinat import * + from .dinov2 import * + from .dinov2_with_registers import * + from .dinov3_convnext import * + from .dinov3_vit import * + from .distilbert import * + from .dit import * + from .doge import * + from .donut import * + from .dots1 import * + from .dpr import * + from .dpt import * + from .edgetam import * + from .edgetam_video import * + from .efficientloftr import * + from .efficientnet import * + from .electra import * + from .emu3 import * + from .encodec import * + from .encoder_decoder import * + from .eomt import * + from .eomt_dinov3 import * + from .ernie import * + from .ernie4_5 import * + from .ernie4_5_moe import * + from .ernie4_5_vl_moe import * + from .esm import * + from .evolla import * + from .exaone4 import * + from .exaone_moe import * + from .falcon import * + from .falcon_h1 import * + from .falcon_mamba import * + from .fast_vlm import * + from .fastspeech2_conformer import * + from .flaubert import * + from .flava import * + from .flex_olmo import * + from .florence2 import * + from .fnet import * + from .focalnet import * + from .fsmt import * + from .funnel import * + from .fuyu import * + from .gemma import * + from .gemma2 import * + from .gemma3 import * + from .gemma3n import * + from .gemma4 import * + from .git import * + from .glm import * + from .glm4 import * + from .glm4_moe import * + from .glm4_moe_lite import * + from .glm4v import * + from .glm4v_moe import * + from .glm46v import * + from .glm_image import * + from .glm_moe_dsa import * + from .glm_ocr import * + from .glmasr import * + from .glpn import * + from .got_ocr2 import * + from .gpt2 import * + from .gpt_bigcode import * + from .gpt_neo import * + from .gpt_neox import * + from .gpt_neox_japanese import * + from .gpt_oss import * + from .gpt_sw3 import * + from .gptj import * + from .granite import * + from .granite_speech import * + from .granitemoe import * + from .granitemoehybrid import * + from .granitemoeshared import * + from .grounding_dino import * + from .groupvit import * + from .helium import * + from .herbert import * + from .hgnet_v2 import * + from .hiera import * + from .higgs_audio_v2 import * + from .higgs_audio_v2_tokenizer import * + from .hubert import * + from .hunyuan_v1_dense import * + from .hunyuan_v1_moe import * + from .ibert import * + from .idefics import * + from .idefics2 import * + from .idefics3 import * + from .ijepa import * + from .imagegpt import * + from .informer import * + from .instructblip import * + from .instructblipvideo import * + from .internvl import * + from .jais2 import * + from .jamba import * + from .janus import * + from .jetmoe import * + from .jina_embeddings_v3 import * + from .kosmos2 import * + from .kosmos2_5 import * + from .kyutai_speech_to_text import * + from .lasr import * + from .layoutlm import * + from .layoutlmv2 import * + from .layoutlmv3 import * + from .layoutxlm import * + from .led import * + from .levit import * + from .lfm2 import * + from .lfm2_moe import * + from .lfm2_vl import * + from .lightglue import * + from .lilt import * + from .llama import * + from .llama4 import * + from .llava import * + from .llava_next import * + from .llava_next_video import * + from .llava_onevision import * + from .longcat_flash import * + from .longformer import * + from .longt5 import * + from .luke import * + from .lw_detr import * + from .lxmert import * + from .m2m_100 import * + from .mamba import * + from .mamba2 import * + from .marian import * + from .markuplm import * + from .mask2former import * + from .maskformer import * + from .mbart import * + from .mbart50 import * + from .megatron_bert import * + from .megatron_gpt2 import * + from .metaclip_2 import * + from .mgp_str import * + from .mimi import * + from .minimax import * + from .minimax_m2 import * + from .ministral import * + from .ministral3 import * + from .mistral import * + from .mistral3 import * + from .mistral4 import * + from .mixtral import * + from .mlcd import * + from .mllama import * + from .mluke import * + from .mm_grounding_dino import * + from .mobilebert import * + from .mobilenet_v1 import * + from .mobilenet_v2 import * + from .mobilevit import * + from .mobilevitv2 import * + from .modernbert import * + from .modernbert_decoder import * + from .modernvbert import * + from .moonshine import * + from .moonshine_streaming import * + from .moshi import * + from .mpnet import * + from .mpt import * + from .mra import * + from .mt5 import * + from .musicflamingo import * + from .musicgen import * + from .musicgen_melody import * + from .mvp import * + from .myt5 import * + from .nanochat import * + from .nemotron import * + from .nemotron_h import * + from .nllb import * + from .nllb_moe import * + from .nomic_bert import * + from .nougat import * + from .nystromformer import * + from .olmo import * + from .olmo2 import * + from .olmo3 import * + from .olmo_hybrid import * + from .olmoe import * + from .omdet_turbo import * + from .oneformer import * + from .openai import * + from .opt import * + from .ovis2 import * + from .owlv2 import * + from .owlvit import * + from .paddleocr_vl import * + from .paligemma import * + from .parakeet import * + from .patchtsmixer import * + from .patchtst import * + from .pe_audio import * + from .pe_audio_video import * + from .pe_video import * + from .pegasus import * + from .pegasus_x import * + from .perceiver import * + from .perception_lm import * + from .persimmon import * + from .phi import * + from .phi3 import * + from .phi4_multimodal import * + from .phimoe import * + from .phobert import * + from .pi0 import * + from .pi0_fast import * + from .pix2struct import * + from .pixio import * + from .pixtral import * + from .plbart import * + from .poolformer import * + from .pop2piano import * + from .pp_chart2table import * + from .pp_doclayout_v2 import * + from .pp_doclayout_v3 import * + from .pp_lcnet import * + from .pp_lcnet_v3 import * + from .pp_ocrv5_mobile_det import * + from .pp_ocrv5_server_det import * + from .prompt_depth_anything import * + from .prophetnet import * + from .pvt import * + from .pvt_v2 import * + from .qwen2 import * + from .qwen2_5_omni import * + from .qwen2_5_vl import * + from .qwen2_audio import * + from .qwen2_moe import * + from .qwen2_vl import * + from .qwen3 import * + from .qwen3_5 import * + from .qwen3_5_moe import * + from .qwen3_moe import * + from .qwen3_next import * + from .qwen3_omni_moe import * + from .qwen3_vl import * + from .qwen3_vl_moe import * + from .rag import * + from .recurrent_gemma import * + from .reformer import * + from .regnet import * + from .rembert import * + from .resnet import * + from .roberta import * + from .roberta_prelayernorm import * + from .roc_bert import * + from .roformer import * + from .rt_detr import * + from .rt_detr_v2 import * + from .rwkv import * + from .sam import * + from .sam2 import * + from .sam2_video import * + from .sam3 import * + from .sam3_tracker import * + from .sam3_tracker_video import * + from .sam3_video import * + from .sam_hq import * + from .seamless_m4t import * + from .seamless_m4t_v2 import * + from .seed_oss import * + from .segformer import * + from .seggpt import * + from .sew import * + from .sew_d import * + from .shieldgemma2 import * + from .siglip import * + from .siglip2 import * + from .slanext import * + from .smollm3 import * + from .smolvlm import * + from .solar_open import * + from .speech_encoder_decoder import * + from .speech_to_text import * + from .speecht5 import * + from .splinter import * + from .squeezebert import * + from .stablelm import * + from .starcoder2 import * + from .superglue import * + from .superpoint import * + from .swiftformer import * + from .swin import * + from .swin2sr import * + from .swinv2 import * + from .switch_transformers import * + from .t5 import * + from .t5gemma import * + from .t5gemma2 import * + from .table_transformer import * + from .tapas import * + from .textnet import * + from .time_series_transformer import * + from .timesfm import * + from .timesfm2_5 import * + from .timesformer import * + from .timm_backbone import * + from .timm_wrapper import * + from .trocr import * + from .tvp import * + from .udop import * + from .umt5 import * + from .unispeech import * + from .unispeech_sat import * + from .univnet import * + from .upernet import * + from .uvdoc import * + from .vaultgemma import * + from .vibevoice_asr import * + from .video_llama_3 import * + from .video_llava import * + from .videomae import * + from .videomt import * + from .vilt import * + from .vipllava import * + from .vision_encoder_decoder import * + from .vision_text_dual_encoder import * + from .visual_bert import * + from .vit import * + from .vit_mae import * + from .vit_msn import * + from .vitdet import * + from .vitmatte import * + from .vitpose import * + from .vitpose_backbone import * + from .vits import * + from .vivit import * + from .vjepa2 import * + from .voxtral import * + from .voxtral_realtime import * + from .wav2vec2 import * + from .wav2vec2_bert import * + from .wav2vec2_conformer import * + from .wav2vec2_phoneme import * + from .wav2vec2_with_lm import * + from .wavlm import * + from .whisper import * + from .x_clip import * + from .xcodec import * + from .xglm import * + from .xlm import * + from .xlm_roberta import * + from .xlm_roberta_xl import * + from .xlnet import * + from .xlstm import * + from .xmod import * + from .yolos import * + from .yoso import * + from .youtu import * + from .zamba import * + from .zamba2 import * + from .zoedepth 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/beit/__init__.py b/third_party/transformers/src/transformers/models/beit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e679b695f621923dbed6bda044eadc5a2bf7462e --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/__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_beit import * + from .image_processing_beit import * + from .image_processing_pil_beit import * + from .modeling_beit 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/beit/configuration_beit.py b/third_party/transformers/src/transformers/models/beit/configuration_beit.py new file mode 100644 index 0000000000000000000000000000000000000000..e2410761efa68507772f51905e0f5cc1351982ab --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/configuration_beit.py @@ -0,0 +1,117 @@ +# 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. +"""BEiT model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import BackboneConfigMixin +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/beit-base-patch16-224-pt22k") +@strict +class BeitConfig(BackboneConfigMixin, PreTrainedConfig): + r""" + use_mask_token (`bool`, *optional*, defaults to `False`): + Whether to use a mask token for masked image modeling. + use_relative_position_bias (`bool`, *optional*, defaults to `False`): + Whether to use T5-style relative position embeddings in the self-attention layers. + use_shared_relative_position_bias (`bool`, *optional*, defaults to `False`): + Whether to use the same relative position embeddings across all self-attention layers of the Transformer. + use_mean_pooling (`bool`, *optional*, defaults to `True`): + Whether to mean pool the final hidden states of the patches instead of using the final hidden state of the + CLS token, before applying the classification head. + pool_scales (`tuple[int]`, *optional*, defaults to `[1, 2, 3, 6]`): + Pooling scales used in Pooling Pyramid Module applied on the last feature map. + use_auxiliary_head (`bool`, *optional*, defaults to `True`): + Whether to use an auxiliary head during training. + auxiliary_loss_weight (`float`, *optional*, defaults to 0.4): + Weight of the cross-entropy loss of the auxiliary head. + auxiliary_channels (`int`, *optional*, defaults to 256): + Number of channels to use in the auxiliary head. + auxiliary_num_convs (`int`, *optional*, defaults to 1): + Number of convolutional layers to use in the auxiliary head. + auxiliary_concat_input (`bool`, *optional*, defaults to `False`): + Whether to concatenate the output of the auxiliary head with the input before the classification layer. + add_fpn (`bool`, *optional*, defaults to `False`): + Whether to add a FPN as part of the backbone. Only relevant for [`BeitBackbone`]. + reshape_hidden_states (`bool`, *optional*, defaults to `True`): + Whether to reshape the feature maps to 4D tensors of shape `(batch_size, hidden_size, height, width)` in + case the model is used as backbone. If `False`, the feature maps will be 3D tensors of shape `(batch_size, + seq_len, hidden_size)`. Only relevant for [`BeitBackbone`]. + + Example: + + ```python + >>> from transformers import BeitConfig, BeitModel + + >>> # Initializing a BEiT beit-base-patch16-224-pt22k style configuration + >>> configuration = BeitConfig() + + >>> # Initializing a model (with random weights) from the beit-base-patch16-224-pt22k style configuration + >>> model = BeitModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "beit" + + vocab_size: int = 8192 + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float | int = 0.0 + attention_probs_dropout_prob: float | int = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + image_size: int | list[int] | tuple[int, int] = 224 + patch_size: int | list[int] | tuple[int, int] = 16 + num_channels: int = 3 + use_mask_token: bool = False + use_absolute_position_embeddings: bool = False + use_relative_position_bias: bool = False + use_shared_relative_position_bias: bool = False + layer_scale_init_value: float = 0.1 + drop_path_rate: float | int = 0.1 + use_mean_pooling: bool = True + pool_scales: list[int] | tuple[int, ...] = (1, 2, 3, 6) + use_auxiliary_head: bool = True + auxiliary_loss_weight: float = 0.4 + auxiliary_channels: int = 256 + auxiliary_num_convs: int = 1 + auxiliary_concat_input: bool = False + semantic_loss_ignore_index: int = 255 + _out_features: list[str] | None = None + _out_indices: list[int] | None = None + add_fpn: bool = False + reshape_hidden_states: bool = True + + def __post_init__(self, **kwargs): + if "segmentation_indices" in kwargs and kwargs.get("out_indices") is None: + kwargs["out_indices"] = kwargs.pop("segmentation_indices") + + # backbone attributes + self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, self.num_hidden_layers + 1)] + self.set_output_features_output_indices( + out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None) + ) + + super().__post_init__(**kwargs) + + +__all__ = ["BeitConfig"] diff --git a/third_party/transformers/src/transformers/models/beit/convert_beit_unilm_to_pytorch.py b/third_party/transformers/src/transformers/models/beit/convert_beit_unilm_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..b1b66f3552238de5e9352f0ac15bc70d000ae4e7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/convert_beit_unilm_to_pytorch.py @@ -0,0 +1,375 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert BEiT checkpoints from the unilm repository.""" + +import argparse +import json +from io import BytesIO +from pathlib import Path + +import httpx +import torch +from datasets import load_dataset +from huggingface_hub import hf_hub_download +from PIL import Image + +from transformers import ( + BeitConfig, + BeitForImageClassification, + BeitForMaskedImageModeling, + BeitForSemanticSegmentation, + BeitImageProcessor, +) +from transformers.image_utils import PILImageResampling +from transformers.utils import logging + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +# here we list all keys to be renamed (original name on the left, our name on the right) +def create_rename_keys(config, has_lm_head=False, is_semantic=False): + prefix = "backbone." if is_semantic else "" + + rename_keys = [] + for i in range(config.num_hidden_layers): + # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms + rename_keys.append((f"{prefix}blocks.{i}.norm1.weight", f"beit.encoder.layer.{i}.layernorm_before.weight")) + rename_keys.append((f"{prefix}blocks.{i}.norm1.bias", f"beit.encoder.layer.{i}.layernorm_before.bias")) + rename_keys.append( + (f"{prefix}blocks.{i}.attn.proj.weight", f"beit.encoder.layer.{i}.attention.output.dense.weight") + ) + rename_keys.append( + (f"{prefix}blocks.{i}.attn.proj.bias", f"beit.encoder.layer.{i}.attention.output.dense.bias") + ) + rename_keys.append((f"{prefix}blocks.{i}.norm2.weight", f"beit.encoder.layer.{i}.layernorm_after.weight")) + rename_keys.append((f"{prefix}blocks.{i}.norm2.bias", f"beit.encoder.layer.{i}.layernorm_after.bias")) + rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.weight", f"beit.encoder.layer.{i}.intermediate.dense.weight")) + rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.bias", f"beit.encoder.layer.{i}.intermediate.dense.bias")) + rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.weight", f"beit.encoder.layer.{i}.output.dense.weight")) + rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.bias", f"beit.encoder.layer.{i}.output.dense.bias")) + + # projection layer + position embeddings + rename_keys.extend( + [ + (f"{prefix}cls_token", "beit.embeddings.cls_token"), + (f"{prefix}patch_embed.proj.weight", "beit.embeddings.patch_embeddings.projection.weight"), + (f"{prefix}patch_embed.proj.bias", "beit.embeddings.patch_embeddings.projection.bias"), + ] + ) + + if has_lm_head: + # mask token + shared relative position bias + layernorm + rename_keys.extend( + [ + ("mask_token", "beit.embeddings.mask_token"), + ( + "rel_pos_bias.relative_position_bias_table", + "beit.encoder.relative_position_bias.relative_position_bias_table", + ), + ( + "rel_pos_bias.relative_position_index", + "beit.encoder.relative_position_bias.relative_position_index", + ), + ("norm.weight", "layernorm.weight"), + ("norm.bias", "layernorm.bias"), + ] + ) + elif is_semantic: + # semantic segmentation classification heads + rename_keys.extend( + [ + ("decode_head.conv_seg.weight", "decode_head.classifier.weight"), + ("decode_head.conv_seg.bias", "decode_head.classifier.bias"), + ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"), + ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"), + ] + ) + else: + # layernorm + classification head + rename_keys.extend( + [ + ("fc_norm.weight", "beit.pooler.layernorm.weight"), + ("fc_norm.bias", "beit.pooler.layernorm.bias"), + ("head.weight", "classifier.weight"), + ("head.bias", "classifier.bias"), + ] + ) + + return rename_keys + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_q_k_v(state_dict, config, has_lm_head=False, is_semantic=False): + for i in range(config.num_hidden_layers): + prefix = "backbone." if is_semantic else "" + # queries, keys and values + in_proj_weight = state_dict.pop(f"{prefix}blocks.{i}.attn.qkv.weight") + q_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.q_bias") + v_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.v_bias") + + state_dict[f"beit.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[ + : config.hidden_size, : + ] + state_dict[f"beit.encoder.layer.{i}.attention.attention.query.bias"] = q_bias + state_dict[f"beit.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[ + config.hidden_size : config.hidden_size * 2, : + ] + state_dict[f"beit.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[ + -config.hidden_size :, : + ] + state_dict[f"beit.encoder.layer.{i}.attention.attention.value.bias"] = v_bias + + # gamma_1 and gamma_2 + # we call them lambda because otherwise they are renamed when using .from_pretrained + gamma_1 = state_dict.pop(f"{prefix}blocks.{i}.gamma_1") + gamma_2 = state_dict.pop(f"{prefix}blocks.{i}.gamma_2") + + state_dict[f"beit.encoder.layer.{i}.lambda_1"] = gamma_1 + state_dict[f"beit.encoder.layer.{i}.lambda_2"] = gamma_2 + + # relative_position bias table + index + if not has_lm_head: + # each layer has its own relative position bias + table = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_bias_table") + index = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_index") + + state_dict[ + f"beit.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table" + ] = table + state_dict[ + f"beit.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index" + ] = index + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@torch.no_grad() +def convert_beit_checkpoint(checkpoint_url, pytorch_dump_folder_path): + """ + Copy/paste/tweak model's weights to our BEiT structure. + """ + + # define default BEiT configuration + config = BeitConfig() + has_lm_head = False + is_semantic = False + repo_id = "huggingface/label-files" + # set config parameters based on URL + if checkpoint_url[-9:-4] == "pt22k": + # masked image modeling + config.use_shared_relative_position_bias = True + config.use_mask_token = True + has_lm_head = True + elif checkpoint_url[-9:-4] == "ft22k": + # intermediate fine-tuning on ImageNet-22k + config.use_relative_position_bias = True + config.num_labels = 21841 + filename = "imagenet-22k-id2label.json" + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + # this dataset contains 21843 labels but the model only has 21841 + # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18 + del id2label[9205] + del id2label[15027] + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + elif checkpoint_url[-8:-4] == "to1k": + # fine-tuning on ImageNet-1k + config.use_relative_position_bias = True + config.num_labels = 1000 + filename = "imagenet-1k-id2label.json" + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + if "384" in checkpoint_url: + config.image_size = 384 + if "512" in checkpoint_url: + config.image_size = 512 + elif "ade20k" in checkpoint_url: + # fine-tuning + config.use_relative_position_bias = True + config.num_labels = 150 + filename = "ade20k-id2label.json" + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + config.image_size = 640 + is_semantic = True + else: + raise ValueError("Checkpoint not supported, URL should either end with 'pt22k', 'ft22k', 'to1k' or 'ade20k'") + + # size of the architecture + if "base" in checkpoint_url: + if "ade20k" in checkpoint_url: + config.out_indices = [3, 5, 7, 11] + elif "large" in checkpoint_url: + config.hidden_size = 1024 + config.intermediate_size = 4096 + config.num_hidden_layers = 24 + config.num_attention_heads = 16 + if "ade20k" in checkpoint_url: + config.image_size = 640 + config.out_indices = [7, 11, 15, 23] + else: + raise ValueError("Should either find 'base' or 'large' in checkpoint URL") + + # load state_dict of original model, remove and rename some keys + state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu", check_hash=True) + state_dict = state_dict["model"] if "ade20k" not in checkpoint_url else state_dict["state_dict"] + + rename_keys = create_rename_keys(config, has_lm_head=has_lm_head, is_semantic=is_semantic) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + read_in_q_k_v(state_dict, config, has_lm_head=has_lm_head, is_semantic=is_semantic) + if is_semantic: + # add prefix to decoder keys + for key, val in state_dict.copy().items(): + val = state_dict.pop(key) + if key.startswith("backbone.fpn"): + key = key.replace("backbone.fpn", "fpn") + state_dict[key] = val + + # load HuggingFace model + if checkpoint_url[-9:-4] == "pt22k": + model = BeitForMaskedImageModeling(config) + elif "ade20k" in checkpoint_url: + model = BeitForSemanticSegmentation(config) + else: + model = BeitForImageClassification(config) + model.eval() + model.load_state_dict(state_dict) + + # Check outputs on an image + if is_semantic: + image_processor = BeitImageProcessor(size=config.image_size, do_center_crop=False) + ds = load_dataset("hf-internal-testing/fixtures_ade20k", split="test") + image = Image.open(ds[0]["file"]) + else: + image_processor = BeitImageProcessor( + size=config.image_size, resample=PILImageResampling.BILINEAR, do_center_crop=False + ) + image = prepare_img() + + encoding = image_processor(images=image, return_tensors="pt") + pixel_values = encoding["pixel_values"] + + outputs = model(pixel_values) + logits = outputs.logits + + # verify logits + expected_shape = torch.Size([1, 1000]) + if checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k"): + expected_shape = torch.Size([1, 196, 8192]) + elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k"): + expected_shape = torch.Size([1, 196, 8192]) + elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22k"): + expected_shape = torch.Size([1, 21841]) + expected_logits = torch.tensor([2.2288, 2.4671, 0.7395]) + expected_class_idx = 2397 + elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22k"): + expected_shape = torch.Size([1, 21841]) + expected_logits = torch.tensor([1.6881, -0.2787, 0.5901]) + expected_class_idx = 2396 + elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft1k"): + expected_logits = torch.tensor([0.1241, 0.0798, -0.6569]) + expected_class_idx = 285 + elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22kto1k"): + expected_logits = torch.tensor([-1.2385, -1.0987, -1.0108]) + expected_class_idx = 281 + elif checkpoint_url[:-4].endswith("beit_base_patch16_384_pt22k_ft22kto1k"): + expected_logits = torch.tensor([-1.5303, -0.9484, -0.3147]) + expected_class_idx = 761 + elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft1k"): + expected_logits = torch.tensor([0.4610, -0.0928, 0.2086]) + expected_class_idx = 761 + elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22kto1k"): + expected_logits = torch.tensor([-0.4804, 0.6257, -0.1837]) + expected_class_idx = 761 + elif checkpoint_url[:-4].endswith("beit_large_patch16_384_pt22k_ft22kto1k"): + expected_logits = torch.tensor([[-0.5122, 0.5117, -0.2113]]) + expected_class_idx = 761 + elif checkpoint_url[:-4].endswith("beit_large_patch16_512_pt22k_ft22kto1k"): + expected_logits = torch.tensor([-0.3062, 0.7261, 0.4852]) + expected_class_idx = 761 + elif checkpoint_url[:-4].endswith("beit_base_patch16_640_pt22k_ft22ktoade20k"): + expected_shape = (1, 150, 160, 160) + expected_logits = torch.tensor( + [ + [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]], + [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]], + [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]], + ] + ) + elif checkpoint_url[:-4].endswith("beit_large_patch16_640_pt22k_ft22ktoade20k"): + expected_shape = (1, 150, 160, 160) + expected_logits = torch.tensor( + [ + [[-4.3305, -2.3049, -3.0161], [-2.9591, -1.5305, -2.2251], [-3.4198, -1.8004, -2.9062]], + [[-5.8922, -3.7435, -4.3978], [-4.2063, -2.7872, -3.4755], [-4.2791, -3.1874, -4.1681]], + [[0.9895, 4.3467, 4.7663], [4.2476, 5.6830, 6.1518], [4.5550, 6.2495, 6.5154]], + ] + ) + else: + raise ValueError("Can't verify logits as model is not supported") + + if logits.shape != expected_shape: + raise ValueError(f"Shape of logits not as expected. {logits.shape=}, {expected_shape=}") + if not has_lm_head: + if is_semantic: + if not torch.allclose(logits[0, :3, :3, :3], expected_logits, atol=1e-3): + raise ValueError("First elements of logits not as expected") + else: + print("Predicted class idx:", logits.argmax(-1).item()) + + if not torch.allclose(logits[0, :3], expected_logits, atol=1e-3): + raise ValueError("First elements of logits not as expected") + if logits.argmax(-1).item() != expected_class_idx: + raise ValueError("Predicted class index not as expected") + + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + print(f"Saving model to {pytorch_dump_folder_path}") + model.save_pretrained(pytorch_dump_folder_path) + print(f"Saving image processor to {pytorch_dump_folder_path}") + image_processor.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_url", + default="https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth", + type=str, + help="URL to the original PyTorch checkpoint (.pth file).", + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model." + ) + args = parser.parse_args() + convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path) diff --git a/third_party/transformers/src/transformers/models/beit/image_processing_beit.py b/third_party/transformers/src/transformers/models/beit/image_processing_beit.py new file mode 100644 index 0000000000000000000000000000000000000000..53053f644539bee5ce8d136b0b912ac9e7597069 --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/image_processing_beit.py @@ -0,0 +1,228 @@ +# 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 BEiT.""" + +from typing import Union + +import torch +import torch.nn.functional as F +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_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, is_torch_available + + +class BeitImageProcessorKwargs(ImagesKwargs, total=False): + r""" + 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_reduce_labels: bool + + +@auto_docstring +class BeitImageProcessor(TorchvisionBackend): + """PIL backend for BEiT with reduce_label support.""" + + valid_kwargs = BeitImageProcessorKwargs + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + default_to_square = True + crop_size = {"height": 224, "width": 224} + do_resize = True + do_center_crop = False + do_rescale = True + do_normalize = True + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[BeitImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[BeitImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps to preprocess. + """ + return super().preprocess(images, segmentation_maps, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + return_tensors: str | TensorType | None, + device: Union[str, "torch.device"] | None = None, + **kwargs, + ) -> BatchFeature: + """Handle extra inputs beyond images.""" + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + # Prepare segmentation maps if provided + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + # Process segmentation maps with do_normalize=False and do_rescale=False + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False}) + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + # Convert to int64 and squeeze channel dimension + processed_segmentation_maps = [ + processed_segmentation_map.squeeze(0).to(torch.int64) + for processed_segmentation_map in processed_segmentation_maps + ] + data["labels"] = processed_segmentation_maps + + return BatchFeature(data=data, tensor_type=return_tensors) + + def reduce_label(self, labels: list["torch.Tensor"]) -> list["torch.Tensor"]: + """Reduce label values by 1, replacing 0 with 255.""" + for idx in range(len(labels)): + label = labels[idx] + label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype, device=label.device), label) + label = label - 1 + label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype, device=label.device), label) + labels[idx] = label + return labels + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + disable_grouping: bool | None, + do_reduce_labels: bool = False, + **kwargs, + ) -> list["torch.Tensor"]: + """Custom preprocessing for BEiT.""" + if do_reduce_labels: + images = self.reduce_label(images) + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(stacked_images, size, resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Use fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + return processed_images + + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """ + Converts the output of [`BeitForSemanticSegmentation`] into semantic segmentation maps. + + Args: + outputs ([`BeitForSemanticSegmentation`]): + Raw outputs of the model. + target_sizes (`list[Tuple]` of length `batch_size`, *optional*): + List of tuples corresponding to the requested final size (height, width) of each prediction. If unset, + predictions will not be resized. + + Returns: + semantic_segmentation: `list[torch.Tensor]` of length `batch_size`, where each item is a semantic + segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is + specified). Each entry of each `torch.Tensor` correspond to a semantic class id. + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for post_process_semantic_segmentation") + + logits = outputs.logits + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + + semantic_segmentation = [] + + for idx in range(len(logits)): + resized_logits = F.interpolate( + logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + +__all__ = ["BeitImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/beit/image_processing_pil_beit.py b/third_party/transformers/src/transformers/models/beit/image_processing_pil_beit.py new file mode 100644 index 0000000000000000000000000000000000000000..e3ccf12e909b2768281de9dfab76ff1c6458c34e --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/image_processing_pil_beit.py @@ -0,0 +1,209 @@ +# 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 BEiT.""" + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring +from ...utils.import_utils import requires + + +# Adapted from transformers.models.beit.image_processing_beit.BeitImageProcessorKwargs +class BeitImageProcessorKwargs(ImagesKwargs, total=False): + r""" + 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_reduce_labels: bool + + +@auto_docstring +class BeitImageProcessorPil(PilBackend): + """PIL backend for BEiT with reduce_label support.""" + + valid_kwargs = BeitImageProcessorKwargs + + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 224, "width": 224} + default_to_square = True + crop_size = {"height": 224, "width": 224} + do_resize = True + do_center_crop = False + do_rescale = True + do_normalize = True + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[BeitImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + **kwargs: Unpack[BeitImageProcessorKwargs], + ) -> 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, + **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 + ) + images_kwargs = kwargs.copy() + images_kwargs["do_reduce_labels"] = False + data = {} + data["pixel_values"] = self._preprocess(images, **images_kwargs) + + # Prepare segmentation maps if provided + if segmentation_maps is not None: + processed_segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + + # Process segmentation maps with do_normalize=False and do_rescale=False + segmentation_maps_kwargs = kwargs.copy() + segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False}) + processed_segmentation_maps = self._preprocess( + images=processed_segmentation_maps, **segmentation_maps_kwargs + ) + + # Convert to int64 and squeeze channel dimension + data["labels"] = [ + processed_segmentation_map.squeeze(0).astype(np.int64) + for processed_segmentation_map in 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.""" + # Avoid using underflow conversion + image[image == 0] = 255 + image = image - 1 + image[image == 254] = 255 + return image + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | None, + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_reduce_labels: bool = False, + **kwargs, + ) -> list[np.ndarray]: + """Custom preprocessing for BEiT.""" + 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_center_crop: + image = self.center_crop(image, crop_size) + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + + return processed_images + + @requires(backends=("torch",)) + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None): + """ + Converts the output of [`BeitForSemanticSegmentation`] into semantic segmentation maps. + + Args: + outputs ([`BeitForSemanticSegmentation`]): + Raw outputs of the model. + target_sizes (`list[Tuple]` of length `batch_size`, *optional*): + List of tuples corresponding to the requested final size (height, width) of each prediction. If unset, + predictions will not be resized. + + Returns: + semantic_segmentation: `list[torch.Tensor]` of length `batch_size`, where each item is a semantic + segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is + specified). Each entry of each `torch.Tensor` correspond to a semantic class id. + """ + import torch + import torch.nn.functional as F + + logits = outputs.logits + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + if isinstance(target_sizes, torch.Tensor): + target_sizes = target_sizes.numpy() + + semantic_segmentation = [] + + for idx in range(len(logits)): + resized_logits = F.interpolate( + logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + +__all__ = ["BeitImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/beit/modeling_beit.py b/third_party/transformers/src/transformers/models/beit/modeling_beit.py new file mode 100644 index 0000000000000000000000000000000000000000..0a169bf55a1d993213932f56c20956a871a01faf --- /dev/null +++ b/third_party/transformers/src/transformers/models/beit/modeling_beit.py @@ -0,0 +1,1470 @@ +# Copyright 2021 Microsoft Research and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch BEiT model.""" + +import collections.abc +import math +from dataclasses import dataclass + +import torch +from torch import Tensor, nn +from torch.nn import CrossEntropyLoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import BackboneMixin, filter_output_hidden_states +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BackboneOutput, + BaseModelOutput, + BaseModelOutputWithPooling, + ImageClassifierOutput, + MaskedLMOutput, + SemanticSegmenterOutput, +) +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import auto_docstring, logging, torch_int +from ...utils.generic import can_return_tuple +from .configuration_beit import BeitConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`BeitModel`]. + """ +) +class BeitModelOutputWithPooling(BaseModelOutputWithPooling): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if + *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token + will be returned. + """ + + +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 BeitDropPath(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}" + + +# Based on timm implementation, which can be found here: +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py +class BeitEmbeddings(nn.Module): + """ + Construct the CLS token, position and patch embeddings. Optionally, also the mask token. + + """ + + def __init__(self, config: BeitConfig) -> None: + super().__init__() + + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + if config.use_mask_token: + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + else: + self.mask_token = None + self.patch_embeddings = BeitPatchEmbeddings(config) + self.patch_size = config.patch_size + self.image_size = ( + config.image_size + if isinstance(config.image_size, collections.abc.Iterable) + else (config.image_size, config.image_size) + ) + num_patches = self.patch_embeddings.num_patches + if config.use_absolute_position_embeddings: + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size)) + else: + self.position_embeddings = None + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embeddings + + class_pos_embed = self.position_embeddings[:, :1] + patch_pos_embed = self.position_embeddings[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward( + self, + pixel_values: torch.Tensor, + bool_masked_pos: torch.BoolTensor | None = None, + ) -> torch.Tensor: + _, _, height, width = pixel_values.shape + embeddings, (patch_height, patch_width) = self.patch_embeddings(pixel_values) + batch_size, seq_len, _ = embeddings.size() + + if bool_masked_pos is not None: + mask_tokens = self.mask_token.expand(batch_size, seq_len, -1) + # replace the masked visual tokens by mask_tokens + w = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens) + embeddings = embeddings * (1 - w) + mask_tokens * w + + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + if self.position_embeddings is not None: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + + embeddings = self.dropout(embeddings) + + return embeddings, (patch_height, patch_width) + + +class BeitPatchEmbeddings(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]) + patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + self.patch_shape = patch_shape + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + + embeddings = self.projection(pixel_values.to(self.projection.weight.dtype)) + patch_height, patch_width = embeddings.shape[2], embeddings.shape[3] + embeddings = embeddings.flatten(2).transpose(1, 2) + + return embeddings, (patch_height, patch_width) + + +class BeitSelfAttention(nn.Module): + def __init__(self, config: BeitConfig, window_size: tuple | None = None) -> None: + super().__init__() + self.config = config + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size {config.hidden_size} is not a multiple of the number of attention " + f"heads {config.num_attention_heads}." + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=False) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.has_relative_position_bias = bool(window_size) + if self.has_relative_position_bias: + self.relative_position_bias = BeitRelativePositionBias(config, window_size=window_size) + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + relative_position_bias: torch.Tensor | None = None, + interpolate_pos_encoding: bool = False, + resolution: tuple[int] | None = None, + ) -> tuple[torch.Tensor] | tuple[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) + + # Add relative position bias if present. + if self.has_relative_position_bias: + height, width = resolution + window_size = (height // self.config.patch_size, width // self.config.patch_size) + attention_scores = attention_scores + self.relative_position_bias( + window_size, interpolate_pos_encoding, dim_size=hidden_states.shape[1] + ) + + # Add shared relative position bias if provided. + if relative_position_bias is not None: + attention_scores = attention_scores + relative_position_bias + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(*new_context_layer_shape) + + outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + + return outputs + + +class BeitSdpaSelfAttention(BeitSelfAttention): + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + relative_position_bias: torch.Tensor | None = None, + interpolate_pos_encoding: bool = False, + resolution: tuple[int] | None = None, + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor]: + if output_attentions: + logger.warning_once( + f"{self.__class__.__name__} does not support `output_attentions=True`. The returned attention weights will " + "be `None`. If you want to get attention weights, please set `attn_implementation='eager'` when loading the model." + ) + 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) + ) + + attn_bias = None + if self.has_relative_position_bias: + height, width = resolution + window_size = (height // self.config.patch_size, width // self.config.patch_size) + attn_bias = self.relative_position_bias( + window_size, interpolate_pos_encoding, dim_size=hidden_states.shape[1] + ) + + # Add shared relative position bias if provided. + if relative_position_bias is not None: + if attn_bias is None: + attn_bias = relative_position_bias + else: + attn_bias += relative_position_bias + + scaling = 1 / math.sqrt(self.attention_head_size) + context_layer = torch.nn.functional.scaled_dot_product_attention( + query_layer, + key_layer, + value_layer, + attn_mask=attn_bias, + dropout_p=self.config.attention_probs_dropout_prob if self.training else 0.0, + is_causal=False, + scale=scaling, + ) + 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, None + + +class BeitSelfOutput(nn.Module): + """ + The residual connection is defined in BeitLayer instead of here (as is the case with other models), due to the + layernorm applied before each block. + """ + + def __init__(self, config: BeitConfig) -> None: + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor, gamma=None) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +BEIT_SELF_ATTENTION_CLASSES = { + "eager": BeitSelfAttention, + "sdpa": BeitSdpaSelfAttention, +} + + +class BeitAttention(nn.Module): + def __init__(self, config: BeitConfig, window_size: tuple | None = None) -> None: + super().__init__() + self.attention = BEIT_SELF_ATTENTION_CLASSES[config._attn_implementation](config, window_size=window_size) + self.output = BeitSelfOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + relative_position_bias: torch.Tensor | None = None, + interpolate_pos_encoding: bool = False, + resolution: tuple[int] | None = None, + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor]: + self_outputs = self.attention( + hidden_states, output_attentions, relative_position_bias, interpolate_pos_encoding, resolution + ) + + attention_output = self.output(self_outputs[0], hidden_states) + + outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + return outputs + + +class BeitIntermediate(nn.Module): + def __init__(self, config: BeitConfig) -> None: + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + 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 BeitOutput(nn.Module): + def __init__(self, config: BeitConfig) -> None: + 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) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +class BeitLayer(GradientCheckpointingLayer): + """This corresponds to the Block class in the timm implementation.""" + + def __init__(self, config: BeitConfig, window_size: tuple | None = None, drop_path_rate: float = 0.0) -> None: + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = BeitAttention(config, window_size=window_size) + self.intermediate = BeitIntermediate(config) + self.output = BeitOutput(config) + self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.drop_path = BeitDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + init_values = config.layer_scale_init_value + if init_values > 0: + self.lambda_1 = nn.Parameter(init_values * torch.ones(config.hidden_size), requires_grad=True) + self.lambda_2 = nn.Parameter(init_values * torch.ones(config.hidden_size), requires_grad=True) + else: + self.lambda_1, self.lambda_2 = None, None + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + relative_position_bias: torch.Tensor | None = None, + interpolate_pos_encoding: bool = False, + resolution: tuple[int, int] | None = None, + ) -> tuple[torch.Tensor] | tuple[torch.Tensor, torch.Tensor]: + self_attention_outputs = self.attention( + self.layernorm_before(hidden_states), # in BEiT, layernorm is applied before self-attention + output_attentions=output_attentions, + relative_position_bias=relative_position_bias, + interpolate_pos_encoding=interpolate_pos_encoding, + resolution=resolution, + ) + attention_output = self_attention_outputs[0] + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + # apply lambda_1 if present + if self.lambda_1 is not None: + attention_output = self.lambda_1 * attention_output + + # first residual connection + hidden_states = self.drop_path(attention_output) + hidden_states + + # in BEiT, layernorm is also applied after self-attention + layer_output = self.layernorm_after(hidden_states) + + layer_output = self.intermediate(layer_output) + layer_output = self.output(layer_output) + + if self.lambda_2 is not None: + layer_output = self.lambda_2 * layer_output + + # second residual connection + layer_output = self.drop_path(layer_output) + hidden_states + + outputs = (layer_output,) + outputs + + return outputs + + +class BeitRelativePositionBias(nn.Module): + def __init__(self, config: BeitConfig, window_size: tuple) -> None: + super().__init__() + self.window_size = window_size + self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3 + self.relative_position_bias_table = nn.Parameter( + torch.zeros(self.num_relative_distance, config.num_attention_heads) + ) # 2*Wh-1 * 2*Ww-1, nH + # cls to token & token 2 cls & cls to cls + + @compile_compatible_method_lru_cache(maxsize=10) + def generate_relative_position_index(self, window_size: tuple[int, int]) -> torch.Tensor: + """ + This method creates the relative position index, modified to support arbitrary window sizes, + as introduced in [MiDaS v3.1](https://huggingface.co/papers/2307.14460). + """ + num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3 + # cls to token & token 2 cls & cls to cls + # get pair-wise relative position index for each token inside the window + window_area = window_size[0] * window_size[1] + grid = torch.meshgrid(torch.arange(window_size[0]), torch.arange(window_size[1]), indexing="ij") + coords = torch.stack(grid) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * window_size[1] - 1 + relative_position_index = torch.zeros(size=(window_area + 1,) * 2, dtype=relative_coords.dtype) + relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + relative_position_index[0, 0:] = num_relative_distance - 3 + relative_position_index[0:, 0] = num_relative_distance - 2 + relative_position_index[0, 0] = num_relative_distance - 1 + return relative_position_index + + def forward(self, window_size, interpolate_pos_encoding: bool = False, dim_size=None) -> torch.Tensor: + """ + Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes. + """ + old_height = 2 * self.window_size[0] - 1 + old_width = 2 * self.window_size[1] - 1 + + new_height = 2 * window_size[0] - 1 + new_width = 2 * window_size[1] - 1 + + old_relative_position_bias_table = self.relative_position_bias_table + + old_num_relative_distance = self.num_relative_distance + new_num_relative_distance = new_height * new_width + 3 + + old_sub_table = old_relative_position_bias_table[: old_num_relative_distance - 3] + + old_sub_table = old_sub_table.reshape(1, old_width, old_height, -1).permute(0, 3, 1, 2) + new_sub_table = nn.functional.interpolate( + old_sub_table, size=(torch_int(new_height), torch_int(new_width)), mode="bilinear" + ) + new_sub_table = new_sub_table.permute(0, 2, 3, 1).reshape(new_num_relative_distance - 3, -1) + + new_relative_position_bias_table = torch.cat( + [new_sub_table, old_relative_position_bias_table[old_num_relative_distance - 3 :]] + ) + + relative_position_index = self.generate_relative_position_index(window_size) + relative_position_bias = new_relative_position_bias_table[relative_position_index.view(-1)] + + # patch_size*num_patches_height, patch_size*num_patches_width, num_attention_heads + relative_position_bias = relative_position_bias.view( + window_size[0] * window_size[1] + 1, window_size[0] * window_size[1] + 1, -1 + ) + # num_attention_heads, patch_size*num_patches_width, patch_size*num_patches_height + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + + if interpolate_pos_encoding: + relative_position_bias = nn.functional.interpolate( + relative_position_bias.unsqueeze(1), + size=(dim_size, dim_size), + mode="bilinear", + align_corners=False, + ).squeeze(1) + + return relative_position_bias.unsqueeze(0) + + +class BeitEncoder(nn.Module): + def __init__(self, config: BeitConfig, window_size: tuple | None = None) -> None: + super().__init__() + self.config = config + self.has_relative_position_bias = config.use_shared_relative_position_bias + if self.has_relative_position_bias: + self.relative_position_bias = BeitRelativePositionBias(config, window_size=window_size) + + # stochastic depth decay rule + dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers, device="cpu")] + self.layer = nn.ModuleList( + [ + BeitLayer( + config, + window_size=window_size if config.use_relative_position_bias else None, + drop_path_rate=dpr[i], + ) + for i in range(config.num_hidden_layers) + ] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + output_hidden_states: bool = False, + interpolate_pos_encoding: bool = False, + resolution: tuple[int, int] | None = None, + return_dict: bool = True, + ) -> tuple | BaseModelOutput: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.has_relative_position_bias: + height, width = resolution + window_size = (height // self.config.patch_size, width // self.config.patch_size) + relative_position_bias = self.relative_position_bias( + window_size, interpolate_pos_encoding=interpolate_pos_encoding, dim_size=hidden_states.shape[1] + ) + else: + relative_position_bias = None + + layer_outputs = layer_module( + hidden_states, + output_attentions=output_attentions, + relative_position_bias=relative_position_bias, + interpolate_pos_encoding=interpolate_pos_encoding, + resolution=resolution, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@auto_docstring +class BeitPreTrainedModel(PreTrainedModel): + config: BeitConfig + base_model_prefix = "beit" + input_modalities = ("image",) + main_input_name = "pixel_values" + supports_gradient_checkpointing = True + _no_split_modules = ["BeitLayer"] + _keys_to_ignore_on_load_unexpected = [r".*relative_position_index.*"] + _supports_sdpa = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, BeitEmbeddings): + init.zeros_(module.cls_token) + if module.mask_token is not None: + init.zeros_(module.mask_token) + if module.position_embeddings is not None: + init.zeros_(module.position_embeddings) + elif isinstance(module, BeitRelativePositionBias): + init.zeros_(module.relative_position_bias_table) + elif isinstance(module, BeitLayer): + if module.lambda_1 is not None: + init.constant_(module.lambda_1, self.config.layer_scale_init_value) + init.constant_(module.lambda_2, self.config.layer_scale_init_value) + + +@auto_docstring +class BeitModel(BeitPreTrainedModel): + def __init__(self, config: BeitConfig, add_pooling_layer: bool = True) -> None: + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + + self.embeddings = BeitEmbeddings(config) + self.encoder = BeitEncoder(config, window_size=self.embeddings.patch_embeddings.patch_shape) + + self.layernorm = ( + nn.Identity() if config.use_mean_pooling else nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + ) + self.pooler = BeitPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.patch_embeddings + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor, + bool_masked_pos: torch.BoolTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BeitModelOutputWithPooling: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). + """ + 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 + + embedding_output, _ = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos) + resolution = pixel_values.shape[2:] + + encoder_outputs = self.encoder( + embedding_output, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + resolution=resolution, + return_dict=return_dict, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,) + return head_outputs + encoder_outputs[1:] + + return BeitModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class BeitPooler(nn.Module): + def __init__(self, config: BeitConfig) -> None: + super().__init__() + self.layernorm = ( + nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if config.use_mean_pooling else None + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.layernorm is not None: + # Mean pool the final hidden states of the patch tokens + patch_tokens = hidden_states[:, 1:, :] + pooled_output = self.layernorm(patch_tokens.mean(1)) + else: + # Pool by simply taking the final hidden state of the [CLS] token + pooled_output = hidden_states[:, 0] + + return pooled_output + + +@auto_docstring( + custom_intro=""" + Beit Model transformer with a 'language' modeling head on top. BEiT does masked image modeling by predicting + visual tokens of a Vector-Quantize Variational Autoencoder (VQ-VAE), whereas other vision models like ViT and DeiT + predict RGB pixel values. As a result, this class is incompatible with [`AutoModelForMaskedImageModeling`], so you + will need to use [`BeitForMaskedImageModeling`] directly if you wish to do masked image modeling with BEiT. + """ +) +class BeitForMaskedImageModeling(BeitPreTrainedModel): + def __init__(self, config: BeitConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.beit = BeitModel(config, add_pooling_layer=False) + + # Classifier head + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return None + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + bool_masked_pos: torch.BoolTensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | MaskedLMOutput: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, BeitForMaskedImageModeling + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224-pt22k") + >>> model = BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k") + + >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2 + >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values + >>> # create random boolean mask of shape (batch_size, num_patches) + >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool() + + >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) + >>> loss, logits = outputs.loss, outputs.logits + >>> list(logits.shape) + [1, 196, 8192] + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.beit( + pixel_values, + bool_masked_pos=bool_masked_pos, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + sequence_output = self.layernorm(sequence_output) + prediction_scores = self.lm_head(sequence_output[:, 1:]) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores[bool_masked_pos], labels) + + if not return_dict: + output = (prediction_scores,) + outputs[1:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + Beit Model transformer with an image classification head on top (a linear layer on top of the average of the final + hidden states of the patch tokens) e.g. for ImageNet. + """ +) +class BeitForImageClassification(BeitPreTrainedModel): + def __init__(self, config: BeitConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.beit = BeitModel(config, add_pooling_layer=True) + + # Classifier head + self.classifier = nn.Linear(config.hidden_size, 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, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | ImageClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + outputs = self.beit( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + ) + + pooled_output = outputs.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 ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class BeitConvModule(nn.Module): + """ + A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution + layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU). + + Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + padding: int | tuple[int, int] | str = 0, + bias: bool = False, + dilation: int | tuple[int, int] = 1, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + padding=padding, + bias=bias, + dilation=dilation, + ) + self.bn = nn.BatchNorm2d(out_channels) + self.activation = nn.ReLU() + + def forward(self, input: torch.Tensor) -> torch.Tensor: + output = self.conv(input) + output = self.bn(output) + output = self.activation(output) + + return output + + +class BeitPyramidPoolingBlock(nn.Module): + def __init__(self, pool_scale: int, in_channels: int, channels: int) -> None: + super().__init__() + self.layers = [ + nn.AdaptiveAvgPool2d(pool_scale), + BeitConvModule(in_channels, channels, kernel_size=1), + ] + for i, layer in enumerate(self.layers): + self.add_module(str(i), layer) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class BeitPyramidPoolingModule(nn.Module): + """ + Pyramid Pooling Module (PPM) used in PSPNet. + + Args: + pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid + Module. + in_channels (int): Input channels. + channels (int): Channels after modules, before conv_seg. + align_corners (bool): align_corners argument of F.interpolate. + + Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. + """ + + def __init__(self, pool_scales: tuple[int, ...], in_channels: int, channels: int, align_corners: bool) -> None: + super().__init__() + self.pool_scales = pool_scales + self.align_corners = align_corners + self.in_channels = in_channels + self.channels = channels + self.blocks = [] + for i, pool_scale in enumerate(pool_scales): + block = BeitPyramidPoolingBlock(pool_scale=pool_scale, in_channels=in_channels, channels=channels) + self.blocks.append(block) + self.add_module(str(i), block) + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + ppm_outs = [] + for ppm in self.blocks: + ppm_out = ppm(x) + upsampled_ppm_out = nn.functional.interpolate( + ppm_out, size=x.size()[2:], mode="bilinear", align_corners=self.align_corners + ) + ppm_outs.append(upsampled_ppm_out) + return ppm_outs + + +class BeitUperHead(nn.Module): + """ + Unified Perceptual Parsing for Scene Understanding. This head is the implementation of + [UPerNet](https://huggingface.co/papers/1807.10221). + + Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. + """ + + def __init__(self, config: BeitConfig) -> None: + super().__init__() + + self.pool_scales = config.pool_scales # e.g. (1, 2, 3, 6) + self.in_channels = [config.hidden_size] * 4 # e.g. [768, 768, 768, 768] + self.channels = config.hidden_size + self.align_corners = False + self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1) + + # PSP Module + self.psp_modules = BeitPyramidPoolingModule( + self.pool_scales, + self.in_channels[-1], + self.channels, + align_corners=self.align_corners, + ) + self.bottleneck = BeitConvModule( + self.in_channels[-1] + len(self.pool_scales) * self.channels, + self.channels, + kernel_size=3, + padding=1, + ) + # FPN Module + self.lateral_convs = nn.ModuleList() + self.fpn_convs = nn.ModuleList() + for in_channels in self.in_channels[:-1]: # skip the top layer + l_conv = BeitConvModule(in_channels, self.channels, kernel_size=1) + fpn_conv = BeitConvModule(self.channels, self.channels, kernel_size=3, padding=1) + self.lateral_convs.append(l_conv) + self.fpn_convs.append(fpn_conv) + + self.fpn_bottleneck = BeitConvModule( + len(self.in_channels) * self.channels, + self.channels, + kernel_size=3, + padding=1, + ) + + def psp_forward(self, inputs): + x = inputs[-1] + psp_outs = [x] + psp_outs.extend(self.psp_modules(x)) + psp_outs = torch.cat(psp_outs, dim=1) + output = self.bottleneck(psp_outs) + + return output + + def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + # build laterals + laterals = [lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)] + + laterals.append(self.psp_forward(encoder_hidden_states)) + + # build top-down path + used_backbone_levels = len(laterals) + for i in range(used_backbone_levels - 1, 0, -1): + prev_shape = laterals[i - 1].shape[2:] + laterals[i - 1] = laterals[i - 1] + nn.functional.interpolate( + laterals[i], size=prev_shape, mode="bilinear", align_corners=self.align_corners + ) + + # build outputs + fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)] + # append psp feature + fpn_outs.append(laterals[-1]) + + for i in range(used_backbone_levels - 1, 0, -1): + fpn_outs[i] = nn.functional.interpolate( + fpn_outs[i], size=fpn_outs[0].shape[2:], mode="bilinear", align_corners=self.align_corners + ) + fpn_outs = torch.cat(fpn_outs, dim=1) + output = self.fpn_bottleneck(fpn_outs) + output = self.classifier(output) + + return output + + +class BeitFCNHead(nn.Module): + """ + Fully Convolution Networks for Semantic Segmentation. This head is implemented of + [FCNNet](https://huggingface.co/papers/1411.4038>). + + Args: + config (BeitConfig): Configuration. + in_channels + kernel_size (int): The kernel size for convs in the head. Default: 3. + dilation (int): The dilation rate for convs in the head. Default: 1. + + + Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. + """ + + def __init__( + self, config: BeitConfig, in_index: int = 2, kernel_size: int = 3, dilation: int | tuple[int, int] = 1 + ) -> None: + super().__init__() + self.in_channels = config.hidden_size + self.channels = config.auxiliary_channels + self.num_convs = config.auxiliary_num_convs + self.concat_input = config.auxiliary_concat_input + self.in_index = in_index + + conv_padding = (kernel_size // 2) * dilation + convs = [] + convs.append( + BeitConvModule( + self.in_channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation + ) + ) + for i in range(self.num_convs - 1): + convs.append( + BeitConvModule( + self.channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation + ) + ) + if self.num_convs == 0: + self.convs = nn.Identity() + else: + self.convs = nn.Sequential(*convs) + if self.concat_input: + self.conv_cat = BeitConvModule( + self.in_channels + self.channels, self.channels, kernel_size=kernel_size, padding=kernel_size // 2 + ) + + self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1) + + def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + # just take the relevant feature maps + hidden_states = encoder_hidden_states[self.in_index] + output = self.convs(hidden_states) + if self.concat_input: + output = self.conv_cat(torch.cat([hidden_states, output], dim=1)) + output = self.classifier(output) + return output + + +@auto_docstring +class BeitForSemanticSegmentation(BeitPreTrainedModel): + def __init__(self, config: BeitConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.beit = BeitModel(config, add_pooling_layer=False) + + # FPNs + if len(self.config.out_indices) != 4: + raise ValueError( + "BeitForSemanticSegmentation requires config.out_indices to be a list of 4 integers, " + "specifying which features to use from the backbone. One can use [3, 5, 7, 11] in case of " + "a base-sized architecture." + ) + self.fpn1 = nn.Sequential( + nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2), + nn.BatchNorm2d(config.hidden_size), + nn.GELU(), + nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2), + ) + self.fpn2 = nn.Sequential( + nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2), + ) + self.fpn3 = nn.Identity() + self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2) + + # Semantic segmentation head(s) + self.decode_head = BeitUperHead(config) + self.auxiliary_head = BeitFCNHead(config) if config.use_auxiliary_head else None + + # Initialize weights and apply final processing + self.post_init() + + def compute_loss(self, logits, auxiliary_logits, labels): + # upsample logits to the images' original size + upsampled_logits = nn.functional.interpolate( + logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + if auxiliary_logits is not None: + upsampled_auxiliary_logits = nn.functional.interpolate( + auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False + ) + # compute weighted loss + loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index) + main_loss = loss_fct(upsampled_logits, labels) + loss = main_loss + if auxiliary_logits is not None: + auxiliary_loss = loss_fct(upsampled_auxiliary_logits, labels) + loss += self.config.auxiliary_loss_weight * auxiliary_loss + + return loss + + @auto_docstring + def forward( + self, + pixel_values: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | SemanticSegmenterOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): + Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, BeitForSemanticSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-finetuned-ade-640-640") + >>> model = BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640") + + >>> inputs = image_processor(images=image, return_tensors="pt") + >>> outputs = model(**inputs) + >>> # logits are of shape (batch_size, num_labels, height, width) + >>> logits = outputs.logits + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if labels is not None and self.config.num_labels == 1: + raise ValueError("The number of labels should be greater than one") + + outputs = self.beit( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=True, # we need the intermediate hidden states + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + ) + + encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1] + + # only keep certain features, and reshape + # note that we do +1 as the encoder_hidden_states also includes the initial embeddings + features = [feature for idx, feature in enumerate(encoder_hidden_states) if idx + 1 in self.config.out_indices] + batch_size = pixel_values.shape[0] + patch_resolution = self.config.image_size // self.config.patch_size + features = [ + x[:, 1:, :].permute(0, 2, 1).reshape(batch_size, -1, patch_resolution, patch_resolution) for x in features + ] + + # apply FPNs + ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4] + for i in range(len(features)): + features[i] = ops[i](features[i]) + + logits = self.decode_head(features) + + auxiliary_logits = None + if self.auxiliary_head is not None: + auxiliary_logits = self.auxiliary_head(features) + + loss = None + if labels is not None: + loss = self.compute_loss(logits, auxiliary_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=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + BEiT backbone, to be used with frameworks like DETR and MaskFormer. + """ +) +class BeitBackbone(BackboneMixin, BeitPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)] + self.embeddings = BeitEmbeddings(config) + self.encoder = BeitEncoder(config, window_size=self.embeddings.patch_embeddings.patch_shape) + + if config.add_fpn: + if len(self.config.out_indices) != 4: + raise ValueError( + "BeitBackbone requires config.out_indices to be a list of 4 integers, " + "specifying which features to use from the backbone. One can use [3, 5, 7, 11] in case of " + "a base-sized architecture." + ) + hidden_size = config.hidden_size + self.fpn1 = nn.Sequential( + nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2), + nn.BatchNorm2d(hidden_size, eps=config.batch_norm_eps), + nn.GELU(), + nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2), + ) + + self.fpn2 = nn.Sequential(nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2)) + self.fpn3 = nn.Identity() + self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2) + + # initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.patch_embeddings + + @can_return_tuple + @filter_output_hidden_states + @auto_docstring + def forward( + self, + pixel_values: Tensor, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BackboneOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoBackbone + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224") + >>> model = AutoBackbone.from_pretrained( + ... "microsoft/beit-base-patch16-224", out_features=["stage1", "stage2", "stage3", "stage4"] + ... ) + + >>> inputs = processor(image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> feature_maps = outputs.feature_maps + >>> list(feature_maps[-1].shape) + [1, 768, 14, 14] + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + + batch_size = pixel_values.shape[0] + embedding_output, (patch_height, patch_width) = self.embeddings(pixel_values) + resolution = pixel_values.shape[2:] + + outputs = self.encoder( + embedding_output, + output_hidden_states=True, + output_attentions=output_attentions, + resolution=resolution, + return_dict=return_dict, + ) + + hidden_states = outputs.hidden_states if return_dict else outputs[1] + + feature_maps = () + for stage, hidden_state in zip(self.stage_names, hidden_states): + if stage in self.out_features: + if self.config.reshape_hidden_states: + hidden_state = hidden_state[:, 1:, :] + hidden_state = hidden_state.permute(0, 2, 1) + hidden_state = hidden_state.reshape(batch_size, -1, patch_height, patch_width) + + feature_maps += (hidden_state,) + + if self.config.add_fpn: + feature_maps = [ + self.fpn1(feature_maps[0]), + self.fpn2(feature_maps[1]), + self.fpn3(feature_maps[2]), + self.fpn4(feature_maps[3]), + ] + feature_maps = tuple(feature_maps) + + if not return_dict: + if output_hidden_states: + output = (feature_maps,) + outputs[1:] + else: + output = (feature_maps,) + outputs[2:] + return output + + return BackboneOutput( + feature_maps=feature_maps, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=outputs.attentions, + ) + + +__all__ = [ + "BeitForImageClassification", + "BeitForMaskedImageModeling", + "BeitForSemanticSegmentation", + "BeitModel", + "BeitPreTrainedModel", + "BeitBackbone", +] diff --git a/third_party/transformers/src/transformers/models/cohere2/__init__.py b/third_party/transformers/src/transformers/models/cohere2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1447f65935601f0fffd8a88dac25bc5916b35f83 --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere2/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2024 Cohere 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_cohere2 import * + from .modeling_cohere2 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/cohere2/configuration_cohere2.py b/third_party/transformers/src/transformers/models/cohere2/configuration_cohere2.py new file mode 100644 index 0000000000000000000000000000000000000000..48c2df360354c3b61451907eb55f207b70b1082e --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere2/configuration_cohere2.py @@ -0,0 +1,107 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/cohere2/modular_cohere2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_cohere2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 Cohere Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="CohereForAI/c4ai-command-r-v01") +@strict +class Cohere2Config(PreTrainedConfig): + r""" + logit_scale (`float`, *optional*, defaults to 0.0625): + The scaling factor for the output logits. + + ```python + >>> from transformers import Cohere2Model, Cohere2Config + + >>> # Initializing a Cohere Nextmodel configuration + >>> configuration = Cohere2Config() + + >>> # Initializing a model from the Cohere2 configuration + >>> model = Cohere2Model(configuration) # doctest: +SKIP + + >>> # Accessing the model configuration + >>> configuration = model.config # doctest: +SKIP + ``` + """ + + model_type = "cohere2" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 256000 + hidden_size: int = 8192 + intermediate_size: int = 22528 + logit_scale: float = 0.0625 + num_hidden_layers: int = 40 + num_attention_heads: int = 64 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 8192 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = 0 + bos_token_id: int | None = 5 + eos_token_id: int | list[int] | None = 255001 + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + sliding_window: int | None = 4096 + layer_types: list[str] | None = None + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + # Need to specify head_dim in the config so it can be used in the attention forward functions + self.head_dim = self.hidden_size // self.num_attention_heads + + # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub + if self.layer_types is None: + # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub + _sliding_window_pattern = kwargs.pop("sliding_window_pattern", 4) + self.layer_types = [ + "sliding_attention" if bool((i + 1) % _sliding_window_pattern) else "full_attention" + for i in range(self.num_hidden_layers) + ] + + super().__post_init__(**kwargs) + + +__all__ = ["Cohere2Config"] diff --git a/third_party/transformers/src/transformers/models/cohere_asr/__init__.py b/third_party/transformers/src/transformers/models/cohere_asr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..64734198de9eaa29057c34f314de949bfca1e045 --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/__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_cohere_asr import * + from .feature_extraction_cohere_asr import * + from .modeling_cohere_asr import * + from .processing_cohere_asr 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/cohere_asr/configuration_cohere_asr.py b/third_party/transformers/src/transformers/models/cohere_asr/configuration_cohere_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..15b759ec0619a5fe4f17f64883e02d35cff32077 --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/configuration_cohere_asr.py @@ -0,0 +1,101 @@ +# 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 ...utils import auto_docstring +from ..auto import CONFIG_MAPPING +from ..parakeet.configuration_parakeet import ParakeetEncoderConfig + + +@auto_docstring(checkpoint="CohereLabs/cohere-transcribe-03-2026") +@strict +class CohereAsrConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import CohereAsrForConditionalGeneration, CohereAsrConfig + + >>> configuration = CohereAsrConfig() + >>> model = CohereAsrForConditionalGeneration(configuration) + >>> configuration = model.config + ```""" + + model_type = "cohere_asr" + sub_configs = {"encoder_config": ParakeetEncoderConfig} + + _default_encoder_config_kwargs = { + "hidden_size": 1280, + "num_hidden_layers": 48, + "num_attention_heads": 8, + "intermediate_size": 5120, + "hidden_act": "silu", + "attention_bias": True, + "convolution_bias": True, + "conv_kernel_size": 9, + "subsampling_factor": 8, + "subsampling_conv_channels": 256, + "num_mel_bins": 128, + "subsampling_conv_kernel_size": 3, + "subsampling_conv_stride": 2, + "dropout": 0.0, + "dropout_positions": 0.0, + "layerdrop": 0.0, + "activation_dropout": 0.0, + "attention_dropout": 0.0, + "max_position_embeddings": 5000, + "scale_input": False, + "initializer_range": 0.02, + } + + encoder_config: dict | PreTrainedConfig | None = None + vocab_size: int = 16384 + hidden_size: int = 1024 + num_hidden_layers: int = 8 + num_attention_heads: int = 8 + num_key_value_heads: int | None = None + intermediate_size: int = 4096 + hidden_act: str = "relu" + max_position_embeddings: int = 1024 + pad_token_id: int | None = 2 + eos_token_id: int | None = 3 + bos_token_id: int | None = 4 + is_encoder_decoder: bool = True + initializer_range: float = 0.02 + attention_dropout: float | int = 0.0 + attention_bias: bool = True + decoder_start_token_id: int | None = None + tie_word_embeddings: bool = False + head_dim: int | None = None + + def __post_init__(self, **kwargs): + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if isinstance(self.encoder_config, dict): + self.encoder_config["model_type"] = self.encoder_config.get("model_type", "parakeet_encoder") + self.encoder_config = CONFIG_MAPPING[self.encoder_config["model_type"]]( + **{**self._default_encoder_config_kwargs, **self.encoder_config} + ) + elif self.encoder_config is None: + self.encoder_config = CONFIG_MAPPING["parakeet_encoder"](**self._default_encoder_config_kwargs) + + super().__post_init__(**kwargs) + + +__all__ = ["CohereAsrConfig"] diff --git a/third_party/transformers/src/transformers/models/cohere_asr/feature_extraction_cohere_asr.py b/third_party/transformers/src/transformers/models/cohere_asr/feature_extraction_cohere_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..1192be10606d4ed0c5162b38d3c0b1dbb8369a6b --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/feature_extraction_cohere_asr.py @@ -0,0 +1,374 @@ +# 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 numpy as np +import torch + +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import TensorType, is_librosa_available, logging +from ...utils.import_utils import requires + + +if is_librosa_available(): + import librosa + + +EPSILON = 1e-5 +LOG_ZERO_GUARD_VALUE = 2**-24 + + +logger = logging.get_logger(__name__) + + +@requires(backends=("torch", "librosa")) +class CohereAsrFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs a CohereAsr feature extractor. + + This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains + most of the main methods. Users should refer to this superclass for more information regarding those methods. + + This class extracts mel-filter bank features from raw speech using a custom numpy implementation of the `Short Time + Fourier Transform` which should match pytorch's `torch.stft` equivalent. + + Args: + feature_size (`int`, *optional*, defaults to 128): + The feature dimension of the extracted features. + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + hop_length (`int`, *optional*, defaults to 160): + Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients. + n_fft (`int`, *optional*, defaults to 512): + Size of the Fourier transform. + win_length (`int`, *optional*, defaults to 400): + The window length for the STFT computation. + preemphasis (`float`, *optional*, defaults to 0.97): + A preemphasis filter coefficient. 0.0 means no preemphasis filter. + padding_value (`float`, *optional*, defaults to 0.0): + Padding value used to pad the audio. Should correspond to silences. + dither (`float`, *optional*, defaults to 1e-05): + Amount of deterministic dither noise to add before feature extraction. Each sample is seeded by its + valid waveform length so that dither is batch-composition invariant. Set to 0.0 to disable. + max_audio_clip_s (`float`, *optional*, defaults to 35.0): + Maximum duration in seconds for a single audio chunk. Audio longer than + `max_audio_clip_s - overlap_chunk_second` is split at energy-based boundaries. + overlap_chunk_second (`float`, *optional*, defaults to 5.0): + Size in seconds of the boundary search window used when splitting long audio. This is not actual + overlap between chunks — it defines how far back from the chunk boundary to search for a quiet + split point. + min_energy_window_samples (`int`, *optional*, defaults to 1600): + Size in samples of the sliding window used to find the quietest point when splitting audio chunks. + """ + + model_input_names = ["input_features", "attention_mask"] + + def __init__( + self, + feature_size=128, + sampling_rate=16000, + hop_length=160, + n_fft=512, + win_length=400, + preemphasis=0.97, + padding_value=0.0, + dither=1e-5, + max_audio_clip_s=35.0, + overlap_chunk_second=5.0, + min_energy_window_samples=1600, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + + self.hop_length = hop_length + self.n_fft = n_fft + self.win_length = win_length + self.preemphasis = preemphasis + self.dither = dither + self.max_audio_clip_s = max_audio_clip_s + self.overlap_chunk_second = overlap_chunk_second + self.min_energy_window_samples = min_energy_window_samples + + # TODO: @eustlb, for now we use librosa to compute the mel filters + # indeed mel_filter_bank uses np.float64 (while librosa uses np.float32), giving numerical differences + mel_filters = librosa.filters.mel( + sr=sampling_rate, n_fft=n_fft, n_mels=feature_size, fmin=0.0, fmax=sampling_rate / 2, norm="slaney" + ) + self.mel_filters = torch.from_numpy(mel_filters).to(torch.float32) + + def _find_split_point_energy(self, waveform: torch.Tensor, start_idx: int, end_idx: int) -> int: + segment = waveform[start_idx:end_idx] + if segment.shape[0] <= self.min_energy_window_samples: + return (start_idx + end_idx) // 2 + + min_energy = float("inf") + quietest_idx = start_idx + upper = segment.shape[0] - self.min_energy_window_samples + for i in range(0, upper, self.min_energy_window_samples): + window = segment[i : i + self.min_energy_window_samples] + energy = torch.sqrt(torch.mean(window * window)).item() + if energy < min_energy: + min_energy = energy + quietest_idx = start_idx + i + return quietest_idx + + def _split_audio_chunks_energy(self, waveform: torch.Tensor) -> list[torch.Tensor]: + chunk_size = max(1, int(round(self.max_audio_clip_s * self.sampling_rate))) + boundary_context_size = max(1, int(round(self.overlap_chunk_second * self.sampling_rate))) + total_samples = waveform.shape[0] + + if total_samples <= chunk_size: + return [waveform] + + chunks_meta: list[tuple[int, int]] = [] + idx = 0 + while idx < total_samples: + if idx + chunk_size >= total_samples: + chunks_meta.append((idx, total_samples)) + break + + search_start = max(idx, idx + chunk_size - boundary_context_size) + search_end = min(idx + chunk_size, total_samples) + if search_end <= search_start: + split_point = idx + chunk_size + else: + split_point = self._find_split_point_energy(waveform, search_start, search_end) + + split_point = max(idx + 1, min(split_point, total_samples)) + chunks_meta.append((idx, split_point)) + idx = split_point + + return [waveform[start:end] for start, end in chunks_meta if end > start] + + def _apply_dither(self, waveform: torch.Tensor, audio_lengths: torch.Tensor) -> torch.Tensor: + if self.dither <= 0: + return waveform + generator = torch.Generator(device=waveform.device) + for i in range(waveform.shape[0]): + valid_samples = min(int(audio_lengths[i].item()), waveform.shape[1]) + if valid_samples <= 0: + continue + generator.manual_seed(valid_samples) + noise = torch.randn(valid_samples, dtype=waveform.dtype, device=waveform.device, generator=generator) + waveform[i, :valid_samples] += self.dither * noise + return waveform + + def _torch_extract_fbank_features(self, waveform, device="cpu"): + # spectrogram + window = torch.hann_window(self.win_length, periodic=False, device=device) + stft = torch.stft( + waveform, + self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=window, + return_complex=True, + pad_mode="constant", + ) + # Let's match original implementation + magnitudes = torch.view_as_real(stft) + magnitudes = torch.sqrt(magnitudes.pow(2).sum(-1)) + magnitudes = magnitudes.pow(2) + + # log mel spectrogram + mel_filters = self.mel_filters.to(device) + mel_spec = mel_filters @ magnitudes + mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE) + + # (batch_size, num_mel_filters, num_frames) -> (batch_size, num_frames, num_mel_filters) + mel_spec = mel_spec.permute(0, 2, 1) + + return mel_spec + + def __call__( + self, + raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]], + truncation: bool = False, + pad_to_multiple_of: int | None = None, + return_tensors: str | TensorType | None = None, + return_attention_mask: bool | None = None, + padding: str | None = "longest", + max_length: int | None = None, + sampling_rate: int | None = None, + do_normalize: bool | None = None, + device: str | None = "cpu", + return_token_timestamps: bool | None = None, + **kwargs, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for + the STFT computation if available, otherwise a slower NumPy based one. + + Args: + raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`): + The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not + stereo, i.e. single float per timestep. + truncation (`bool`, *optional*, default to `True`): + Activates truncation to cut input sequences longer than *max_length* to *max_length*. + pad_to_multiple_of (`int`, *optional*, defaults to None): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific feature_extractor's default. + + [What are attention masks?](../glossary#attention-mask) + + + + For CohereAsr models, `attention_mask` should always be passed for batched inference, to avoid subtle + bugs. + + + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + sampling_rate (`int`, *optional*): + The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition + pipeline. + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values / vectors. + do_normalize (`bool`, *optional*, defaults to `False`): + Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly + improve the performance of the model. + device (`str`, *optional*, defaults to `'cpu'`): + Specifies the device for computation of the log-mel spectrogram of audio signals in the + `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda") + return_token_timestamps (`bool`, *optional*, defaults to `None`): + Deprecated. Use `return_attention_mask` instead from which the number of frames can be inferred. + + Whether or not to return the number of frames of the input raw_speech. + These num_frames can be used by the model to compute word level timestamps. + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a" + f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input" + f" was sampled with {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + # Convert to torch tensor + if isinstance(raw_speech, np.ndarray): + raw_speech = torch.tensor(raw_speech) + elif isinstance(raw_speech, (list, tuple)) and isinstance(raw_speech[0], np.ndarray): + raw_speech = [torch.tensor(speech) for speech in raw_speech] + + is_batched_torch = isinstance(raw_speech, torch.Tensor) and len(raw_speech.shape) > 1 + if is_batched_torch and len(raw_speech.shape) > 2: + logger.warning( + f"Only mono-channel audio is supported for input to {self.__class__.__name__}. " + "We will take the mean of the channels to convert to mono." + ) + raw_speech = raw_speech.mean(-1) + + is_batched_sequence = isinstance(raw_speech, (list, tuple)) + if is_batched_sequence: + for speech in raw_speech: + if len(speech.shape) > 1: + logger.warning( + f"Only mono-channel audio is supported for input to {self.__class__.__name__}. " + "We will take the mean of the channels to convert to mono." + ) + speech = speech.mean(-1) + + if is_batched_torch or is_batched_sequence: + raw_speech = [speech.to(torch.float32) for speech in raw_speech] + else: + raw_speech = [raw_speech.to(torch.float32)] + + # Chunk long audio at energy-based boundaries + fast_path_threshold_s = max(0.0, self.max_audio_clip_s - self.overlap_chunk_second) + audio_chunk_index: list[tuple[int, int | None]] = [] + chunked_speech: list[torch.Tensor] = [] + for sample_idx, speech in enumerate(raw_speech): + duration_s = speech.shape[0] / self.sampling_rate + if duration_s <= fast_path_threshold_s: + chunked_speech.append(speech) + audio_chunk_index.append((sample_idx, None)) + else: + chunks = self._split_audio_chunks_energy(speech) + for chunk_idx, chunk in enumerate(chunks): + chunked_speech.append(chunk) + audio_chunk_index.append((sample_idx, chunk_idx)) + + raw_speech = [speech[:, None] for speech in chunked_speech] + + audio_lengths = [len(speech) for speech in raw_speech] + batched_speech = BatchFeature({"input_features": raw_speech, "audio_lengths": audio_lengths}) + + padded_inputs = self.pad( + batched_speech, + padding=padding, + max_length=max_length, + truncation=truncation, + pad_to_multiple_of=pad_to_multiple_of, + return_tensors="pt", + ) + input_features = padded_inputs.input_features.squeeze(-1) + + # dithering + input_features = self._apply_dither(input_features, padded_inputs.audio_lengths) + + # preemphasis + if self.preemphasis is not None: + timemask = torch.arange(input_features.shape[1], device=input_features.device).unsqueeze( + 0 + ) < padded_inputs.audio_lengths.unsqueeze(1) + input_features = torch.cat( + [input_features[:, :1], input_features[:, 1:] - self.preemphasis * input_features[:, :-1]], dim=1 + ) + input_features = input_features.masked_fill(~timemask, 0.0) + + input_features = self._torch_extract_fbank_features(input_features, device) + features_lengths = torch.floor_divide( + padded_inputs.audio_lengths + self.n_fft // 2 * 2 - self.n_fft, self.hop_length + ) + attention_mask = torch.arange(input_features.shape[1], device=device)[None, :] < features_lengths[:, None] + + # normalize mel features, ignoring padding + mask = attention_mask.unsqueeze(-1) + input_features_masked = input_features * mask + mean = input_features_masked.sum(dim=1) / features_lengths.unsqueeze(-1) + mean = mean.unsqueeze(1) + variance = ((input_features_masked - mean) ** 2 * mask).sum(dim=1) / (features_lengths - 1).unsqueeze(-1) + std = torch.sqrt(variance).unsqueeze(1) + input_features = (input_features - mean) / (std + EPSILON) + input_features *= mask + + result = BatchFeature( + data={ + "input_features": input_features, + "attention_mask": attention_mask, + }, + tensor_type=return_tensors, + ) + result["audio_chunk_index"] = audio_chunk_index + return result + + +__all__ = ["CohereAsrFeatureExtractor"] diff --git a/third_party/transformers/src/transformers/models/cohere_asr/modeling_cohere_asr.py b/third_party/transformers/src/transformers/models/cohere_asr/modeling_cohere_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..3ec5d7cde57760c20d7d14288c6470b309a119ac --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/modeling_cohere_asr.py @@ -0,0 +1,658 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/cohere_asr/modular_cohere_asr.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_cohere_asr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable + +import torch +import torch.nn as nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..auto.modeling_auto import AutoModel +from .configuration_cohere_asr import CohereAsrConfig + + +class CohereAsrDecoderMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# Modular automatically inherits RoPE, hence no inheritance for now +class CohereAsrSelfAttention(nn.Module): + def __init__(self, config: CohereAsrConfig, 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, + attention_mask: torch.Tensor, + past_key_values: Cache | None = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(hidden_shape).transpose(1, 2) + key_states = key_states.view(hidden_shape).transpose(1, 2) + value_states = value_states.view(hidden_shape).transpose(1, 2) + + if past_key_values is not None: + past_key_values = past_key_values.self_attention_cache + 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 + + +# Modular automatically inherits RoPE, hence no inheritance for now +class CohereAsrCrossAttention(nn.Module): + def __init__(self, config: CohereAsrConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = encoder_hidden_states.shape[1] + + q_input_shape = (bsz, tgt_len, -1, self.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 = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False + if past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = past_key_values.cross_attention_cache.layers[self.layer_idx].keys + value_states = past_key_values.cross_attention_cache.layers[self.layer_idx].values + else: + key_states = self.k_proj(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + # save all states to the cache + key_states, value_states = past_key_values.cross_attention_cache.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 + 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.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class CohereAsrDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + self.self_attn = CohereAsrSelfAttention(config=config, layer_idx=layer_idx) + self.encoder_attn = CohereAsrCrossAttention(config=config, layer_idx=layer_idx) + + self.mlp = CohereAsrDecoderMLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size) + self.final_layernorm = nn.LayerNorm(config.hidden_size) + + 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, + position_ids: torch.LongTensor | None = None, + encoder_position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class CohereAsrPreTrainedModel(PreTrainedModel): + config: CohereAsrConfig + base_model_prefix = "model" + main_input_name = "input_features" + input_modalities = "audio" + supports_gradient_checkpointing = True + _no_split_modules = ["CohereAsrEncoderLayer", "CohereAsrDecoderLayer"] + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _keys_to_ignore_on_load_unexpected = [r"preprocessor\.featurizer\..*"] + # TODO arthur, how do we separate when it cross / self coming from different layer? + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers + """ + output_conv1_length = int((input_lengths - 127) / 64 + 1) + output_conv2_length = int((output_conv1_length - 7) / 3 + 1) + output_conv3_length = int((output_conv2_length - 3) / 2 + 1) + + return output_conv3_length + + +@auto_docstring +class CohereAsrDecoder(CohereAsrPreTrainedModel): + main_input_name = "input_ids" + _can_record_outputs = { + "attentions": OutputRecorder(CohereAsrSelfAttention, index=1, layer_name="self_attn"), + "hidden_states": CohereAsrDecoderLayer, + "cross_attentions": OutputRecorder(CohereAsrCrossAttention, index=1, layer_name="encoder_attn"), + } + + 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([CohereAsrDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)]) + self.norm = nn.LayerNorm(config.hidden_size) + self.gradient_checkpointing = False + self.pos_emb = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.embedding_layernorm = nn.LayerNorm(config.hidden_size) + self.proj = nn.Linear(config.encoder_config.hidden_size, config.hidden_size, bias=True) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: + r""" + 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.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding indices in `encoder_hidden_states`. 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 = self.proj(encoder_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) + + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), 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) + + # Fixed sinusoidal position embedding added to token embeddings, then layernorm + pos_emb = self.pos_emb(position_ids.squeeze(0)) + inputs_embeds = self.embedding_layernorm(inputs_embeds + pos_emb) + + 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, + ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + causal_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +@auto_docstring +class CohereAsrModel(CohereAsrPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.encoder = AutoModel.from_config(config.encoder_config) + self.decoder = CohereAsrDecoder(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.decoder.embed_tokens + + def set_input_embeddings(self, value): + self.decoder.embed_tokens = value + + def freeze_encoder(self): + """ + Calling this function will disable the gradient computation for the CohereAsr encoder so that its parameters will + not be updated during training. + """ + self.encoder._freeze_parameters() + + def _mask_input_features(self): + """ + Masks extracted features along time axis and/or along feature axis according to + [SpecAugment](https://huggingface.co/papers/1904.08779). + """ + raise AttributeError("Not needed for CohereAsr") + + @can_return_tuple + @auto_docstring + def forward( + self, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None, + past_key_values: EncoderDecoderCache | None = None, + decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, + decoder_position_ids: tuple[torch.LongTensor] | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: + r""" + input_features (`torch.FloatTensor` of shape `(batch_size, audio_length)`): + Float values of the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a + `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or + the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for padding + and conversion into a tensor of type `torch.FloatTensor`. + decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`): + Indices of positions of each input sequence tokens in the position embeddings. + Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings` + + Example: + + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, CohereAsrModel + >>> from datasets import load_dataset + + >>> model = CohereAsrModel.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 288] + ``` + """ + # Main difference: uses `input_features` instead of `input_values` + if encoder_outputs is None: + encoder_outputs: BaseModelOutput = self.encoder(input_features, attention_mask=attention_mask, **kwargs) + + decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=encoder_outputs.attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + position_ids=decoder_position_ids, + 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, + ) + + +def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): + """ + Shift input ids one token to the right. + """ + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +@auto_docstring( + custom_intro=""" + The CohereAsr Model with a language modeling head. Can be used for automatic speech recognition. + """ +) +class CohereAsrForConditionalGeneration(CohereAsrPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"proj_out.weight": "model.decoder.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = CohereAsrModel(config) + self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=True) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.proj_out + + def set_output_embeddings(self, new_embeddings): + self.proj_out = new_embeddings + + def get_input_embeddings(self) -> nn.Module: + return self.model.get_input_embeddings() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None, + past_key_values: EncoderDecoderCache | None = None, + decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, + decoder_position_ids: tuple[torch.LongTensor] | None = None, + use_cache: bool | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: + r""" + input_features (`torch.FloatTensor` of shape `(batch_size, audio_length)`): + Float values of the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a + `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or + the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for padding + and conversion into a tensor of type `torch.FloatTensor`. + decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`): + Indices of positions of each input sequence tokens in the position embeddings. + Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings` + + Example: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, CohereAsrForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> model = CohereAsrForConditionalGeneration.from_pretrained("UsefulSensors/cohere_asr-tiny") + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + + >>> generated_ids = model.generate(input_features, max_new_tokens=100) + + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> transcription + 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' + ```""" + # Main difference: uses `input_features` instead of `input_values` + if labels is not None: + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + + outputs: Seq2SeqModelOutput = self.model( + input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + decoder_position_ids=decoder_position_ids, + use_cache=use_cache, + **kwargs, + ) + logits = self.proj_out(outputs.last_hidden_state) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size) + + return Seq2SeqLMOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + def prepare_inputs_for_generation(self, *args, audio_chunk_index=None, **kwargs): + # audio_chunk_index is returned by the processor but not used by the model, absorb it here + return super().prepare_inputs_for_generation(*args, **kwargs) + + +__all__ = ["CohereAsrPreTrainedModel", "CohereAsrModel", "CohereAsrForConditionalGeneration"] diff --git a/third_party/transformers/src/transformers/models/cohere_asr/modular_cohere_asr.py b/third_party/transformers/src/transformers/models/cohere_asr/modular_cohere_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..e6303b24bd0cd6317409877d1cb5a93a94a6fbec --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/modular_cohere_asr.py @@ -0,0 +1,525 @@ +# Copyright 2026 the HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable + +import torch +import torch.nn as nn + +from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin +from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple +from ...utils.output_capturing import OutputRecorder +from ..auto.modeling_auto import AutoModel +from ..clip.modeling_clip import CLIPMLP +from ..moonshine.modeling_moonshine import ( + MoonshineDecoder, + MoonshineForConditionalGeneration, + MoonshineModel, + MoonshinePreTrainedModel, + eager_attention_forward, + shift_tokens_right, +) +from .configuration_cohere_asr import CohereAsrConfig + + +class CohereAsrDecoderMLP(CLIPMLP): + pass + + +# Modular automatically inherits RoPE, hence no inheritance for now +class CohereAsrSelfAttention(nn.Module): + def __init__(self, config: CohereAsrConfig, 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, + attention_mask: torch.Tensor, + past_key_values: Cache | None = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(hidden_shape).transpose(1, 2) + key_states = key_states.view(hidden_shape).transpose(1, 2) + value_states = value_states.view(hidden_shape).transpose(1, 2) + + if past_key_values is not None: + past_key_values = past_key_values.self_attention_cache + 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 + + +# Modular automatically inherits RoPE, hence no inheritance for now +class CohereAsrCrossAttention(nn.Module): + def __init__(self, config: CohereAsrConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = encoder_hidden_states.shape[1] + + q_input_shape = (bsz, tgt_len, -1, self.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 = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False + if past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_states = past_key_values.cross_attention_cache.layers[self.layer_idx].keys + value_states = past_key_values.cross_attention_cache.layers[self.layer_idx].values + else: + key_states = self.k_proj(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) + + if past_key_values is not None: + # save all states to the cache + key_states, value_states = past_key_values.cross_attention_cache.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 + 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.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class CohereAsrDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + self.self_attn = CohereAsrSelfAttention(config=config, layer_idx=layer_idx) + self.encoder_attn = CohereAsrCrossAttention(config=config, layer_idx=layer_idx) + + self.mlp = CohereAsrDecoderMLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size) + self.final_layernorm = nn.LayerNorm(config.hidden_size) + + 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, + position_ids: torch.LongTensor | None = None, + encoder_position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class CohereAsrPreTrainedModel(MoonshinePreTrainedModel): + main_input_name = "input_features" + _keys_to_ignore_on_load_unexpected = [r"preprocessor\.featurizer\..*"] + + +class CohereAsrDecoder(MoonshineDecoder): + _can_record_outputs = { + "attentions": OutputRecorder(CohereAsrSelfAttention, index=1, layer_name="self_attn"), + "hidden_states": CohereAsrDecoderLayer, + "cross_attentions": OutputRecorder(CohereAsrCrossAttention, index=1, layer_name="encoder_attn"), + } + + def __init__(self, config): + super().__init__(config) + del self.rotary_emb + self.norm = nn.LayerNorm(config.hidden_size) + self.pos_emb = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.embedding_layernorm = nn.LayerNorm(config.hidden_size) + self.proj = nn.Linear(config.encoder_config.hidden_size, config.hidden_size, bias=True) + 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, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: + r""" + 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.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding indices in `encoder_hidden_states`. 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 = self.proj(encoder_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) + + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), 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) + + # Fixed sinusoidal position embedding added to token embeddings, then layernorm + pos_emb = self.pos_emb(position_ids.squeeze(0)) + inputs_embeds = self.embedding_layernorm(inputs_embeds + pos_emb) + + 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, + ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + causal_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +class CohereAsrModel(MoonshineModel): + def __init__(self, config): + super().__init__(config) + self.encoder = AutoModel.from_config(config.encoder_config) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None, + past_key_values: EncoderDecoderCache | None = None, + decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, + decoder_position_ids: tuple[torch.LongTensor] | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: + r""" + input_features (`torch.FloatTensor` of shape `(batch_size, audio_length)`): + Float values of the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a + `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or + the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for padding + and conversion into a tensor of type `torch.FloatTensor`. + decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`): + Indices of positions of each input sequence tokens in the position embeddings. + Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings` + + Example: + + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, CohereAsrModel + >>> from datasets import load_dataset + + >>> model = CohereAsrModel.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 288] + ``` + """ + # Main difference: uses `input_features` instead of `input_values` + if encoder_outputs is None: + encoder_outputs: BaseModelOutput = self.encoder(input_features, attention_mask=attention_mask, **kwargs) + + decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=encoder_outputs.attention_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + position_ids=decoder_position_ids, + 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, + ) + + +class CohereAsrForConditionalGeneration(MoonshineForConditionalGeneration): + def __init__(self, config): + super().__init__(config) + self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=True) + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.LongTensor | None = None, + decoder_input_ids: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None, + past_key_values: EncoderDecoderCache | None = None, + decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, + decoder_position_ids: tuple[torch.LongTensor] | None = None, + use_cache: bool | None = None, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: + r""" + input_features (`torch.FloatTensor` of shape `(batch_size, audio_length)`): + Float values of the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a + `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or + the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for padding + and conversion into a tensor of type `torch.FloatTensor`. + decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`): + Indices of positions of each input sequence tokens in the position embeddings. + Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings` + + Example: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, CohereAsrForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("UsefulSensors/cohere_asr-tiny") + >>> model = CohereAsrForConditionalGeneration.from_pretrained("UsefulSensors/cohere_asr-tiny") + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + + >>> generated_ids = model.generate(input_features, max_new_tokens=100) + + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> transcription + 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' + ```""" + # Main difference: uses `input_features` instead of `input_values` + if labels is not None: + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + + outputs: Seq2SeqModelOutput = self.model( + input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + decoder_position_ids=decoder_position_ids, + use_cache=use_cache, + **kwargs, + ) + logits = self.proj_out(outputs.last_hidden_state) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size) + + return Seq2SeqLMOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + def prepare_inputs_for_generation(self, *args, audio_chunk_index=None, **kwargs): + # audio_chunk_index is returned by the processor but not used by the model, absorb it here + return GenerationMixin.prepare_inputs_for_generation(self, *args, **kwargs) + + +__all__ = [ + "CohereAsrPreTrainedModel", + "CohereAsrModel", + "CohereAsrForConditionalGeneration", +] diff --git a/third_party/transformers/src/transformers/models/cohere_asr/processing_cohere_asr.py b/third_party/transformers/src/transformers/models/cohere_asr/processing_cohere_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..91618d8bcc4d6892cadcbf43e11470bfe4a7e6cc --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere_asr/processing_cohere_asr.py @@ -0,0 +1,188 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...audio_utils import AudioInput, make_list_of_audio +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, is_torch_available, logging +from ...utils.import_utils import requires + + +if is_torch_available(): + import torch + + +LANGUAGES = {"ar", "de", "el", "en", "es", "fr", "it", "ja", "ko", "nl", "pl", "pt", "vi", "zh"} +_NO_SPACE_LANGS = {"ja", "zh"} + + +logger = logging.get_logger(__name__) + + +class CohereAsrProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "audio_kwargs": { + "sampling_rate": 16000, + "padding": "longest", + "return_attention_mask": True, + }, + "text_kwargs": { + "padding": True, + "padding_side": "right", + "add_special_tokens": False, + }, + "common_kwargs": {"return_tensors": "pt"}, + } + + +@auto_docstring +@requires(backends=("torch",)) +class CohereAsrProcessor(ProcessorMixin): + def __init__(self, feature_extractor, tokenizer): + super().__init__(feature_extractor, tokenizer) + + def get_decoder_prompt_ids(self, language: str, punctuation: bool = True) -> list[int]: + """Build the decoder prompt token IDs for the given language and punctuation settings.""" + if language not in LANGUAGES: + raise ValueError( + f"Unsupported language: {language!r}. Supported languages: {', '.join(sorted(LANGUAGES))}." + ) + pnc_token = "<|pnc|>" if punctuation else "<|nopnc|>" + tokens = [ + "▁", + "<|startofcontext|>", + "<|startoftranscript|>", + "<|emo:undefined|>", + f"<|{language}|>", + f"<|{language}|>", + pnc_token, + "<|noitn|>", + "<|notimestamp|>", + "<|nodiarize|>", + ] + return self.tokenizer.convert_tokens_to_ids(tokens) + + @auto_docstring + def __call__( + self, + audio: AudioInput, + language: str, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + punctuation: bool = True, + sampling_rate: int | None = None, + **kwargs: Unpack[CohereAsrProcessorKwargs], + ): + r""" + language (`str`): + Language code (e.g. `"en"`, `"es"`, `"fr"`) used to build the decoder prompt. The processor + constructs the full decoder prompt and returns `decoder_input_ids` alongside the audio features. + punctuation (`bool`, defaults to `True`): + Whether to enable punctuation in the decoder prompt. + sampling_rate (`int`, *optional*): + The sampling rate of the input audio in Hz. This should match the sampling rate expected by the feature + extractor (defaults to 16000 Hz). If provided, it will be validated against the processor's expected + sampling rate, and an error will be raised if they don't match. If not provided, a warning will be + issued and the default sampling rate will be assumed. + """ + audio = make_list_of_audio(audio) + + output_kwargs = self._merge_kwargs( + CohereAsrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if sampling_rate is None: + logger.warning_once( + f"You've provided audio without specifying the sampling rate. It will be assumed to be {output_kwargs['audio_kwargs']['sampling_rate']}, which can result in silent errors." + ) + elif sampling_rate != output_kwargs["audio_kwargs"]["sampling_rate"]: + raise ValueError( + f"The sampling rate of the audio ({sampling_rate}) does not match the sampling rate of the processor ({output_kwargs['audio_kwargs']['sampling_rate']}). Please provide resampled the audio to the expected sampling rate." + ) + + inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"]) + + prompt_ids = self.get_decoder_prompt_ids(language=language, punctuation=punctuation) + batch_size = inputs["input_features"].shape[0] + inputs["decoder_input_ids"] = torch.tensor([prompt_ids] * batch_size, dtype=torch.long) + + if text is not None: + encodings = self.tokenizer(text, **output_kwargs["text_kwargs"]) + inputs["labels"] = encodings["input_ids"] + + return inputs + + def decode(self, *args, audio_chunk_index=None, language=None, **kwargs): + texts = self.tokenizer.decode(*args, **kwargs) + if audio_chunk_index is None: + return texts + if language is None: + raise ValueError("`language` must be provided when `audio_chunk_index` is given.") + separator = "" if language in _NO_SPACE_LANGS else " " + return self._reassemble_chunk_texts(texts, audio_chunk_index, separator) + + @staticmethod + def _reassemble_chunk_texts( + texts: list[str], + audio_chunk_index: list[tuple[int, int | None]], + separator: str = " ", + ) -> list[str]: + """Reassemble per-chunk transcription texts back into per-sample strings. + + When audio inputs are longer than the feature extractor's `max_audio_clip_s`, they are split into + overlapping chunks before being fed to the model. This means a single original audio sample can + produce multiple decoded text segments. This method reverses that chunking: it groups the decoded + texts by their original sample index using `chunk_map`, orders the chunks, and joins them + with `separator` to reconstruct one transcription string per input sample. + + Args: + texts: Decoded text strings, one per model output (i.e. one per chunk). + audio_chunk_index: List of `(sample_idx, chunk_idx)` tuples that map each entry in + `texts` back to its original sample and chunk position. A `chunk_idx` of `None` + indicates the sample was not chunked. + separator: String used to join chunks belonging to the same sample. Defaults to a + space; callers pass an empty string for languages that don't use spaces between + words (e.g. Chinese, Japanese). + + Returns: + A list of reassembled transcription strings, one per original input sample. + """ + max_sample_idx = max(sample_idx for sample_idx, _ in audio_chunk_index) + outputs = [""] * (max_sample_idx + 1) + chunked = {} + + for (sample_idx, chunk_idx), text in zip(audio_chunk_index, texts): + if chunk_idx is None: + outputs[sample_idx] = text + else: + if sample_idx not in chunked: + chunked[sample_idx] = [] + chunked[sample_idx].append((chunk_idx, text)) + + for sample_idx, chunk_items in chunked.items(): + chunk_items.sort(key=lambda item: item[0]) + non_empty = [t for _, t in chunk_items if t and t.strip()] + parts = [non_empty[0].rstrip()] + [t.strip() for t in non_empty[1:]] + outputs[sample_idx] = separator.join(parts) + + return outputs + + @property + def model_input_names(self): + feature_extractor_input_names = self.feature_extractor.model_input_names + return feature_extractor_input_names + ["labels"] + + +__all__ = ["CohereAsrProcessor"] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/__init__.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..165c706733af65344f9a1dd1726a933f9f3c0436 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_deepseek_vl_hybrid import * + from .image_processing_deepseek_vl_hybrid import * + from .image_processing_pil_deepseek_vl_hybrid import * + from .modeling_deepseek_vl_hybrid import * + from .processing_deepseek_vl_hybrid 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/deepseek_vl_hybrid/configuration_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/configuration_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..f984faf0687f250404b9a959f019d4e271d0c2d9 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/configuration_deepseek_vl_hybrid.py @@ -0,0 +1,92 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.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_deepseek_vl_hybrid.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging +from ..auto import CONFIG_MAPPING, AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="deepseek-community/deepseek-vl-7b-chat") +@strict +class DeepseekVLHybridConfig(PreTrainedConfig): + r""" + high_res_vision_config (`Union[AutoConfig, dict]`, *optional*, defaults to `SamVisionConfig`): + The config object or dictionary of the high resolution vision backbone. + + Example: + + ```python + >>> from transformers import DeepseekVLHybridConfig, DeepseekVLHybridModel + + >>> # Initializing a DeepseekVLHybrid deepseek-community/deepseek-vl-7b-chat style configuration + >>> configuration = DeepseekVLHybridConfig() + + >>> # Initializing a model (with random weights) from the deepseek-community/deepseek-vl-7b-chat style configuration + >>> model = DeepseekVLHybridModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_vl_hybrid" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig, "high_res_vision_config": AutoConfig} + + text_config: dict | PreTrainedConfig | None = None + vision_config: dict | PreTrainedConfig | None = None + image_token_id: int = 100015 + tie_word_embeddings: bool = True + + high_res_vision_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.high_res_vision_config is None: + self.high_res_vision_config = {} + logger.info("`high_res_vision_config` is `None`. Initializing the `SamVisionConfig` with default values.") + + if isinstance(self.high_res_vision_config, dict): + self.high_res_vision_config["model_type"] = self.high_res_vision_config.get( + "model_type", "sam_vision_model" + ) + self.high_res_vision_config = CONFIG_MAPPING[self.high_res_vision_config["model_type"]]( + **self.high_res_vision_config + ) + if self.text_config is None: + self.text_config = {} + logger.info("`text_config` is `None`. Initializing the `LlamaConfig` with default values.") + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "llama") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + + if self.vision_config is None: + self.vision_config = {} + logger.info("`vision_config` is `None`. Initializing the `SiglipVisionConfig` with default values.") + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "siglip_vision_model") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config) + + super().__post_init__(**kwargs) + + +__all__ = ["DeepseekVLHybridConfig"] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/convert_deepseek_vl_hybrid_weights_to_hf.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/convert_deepseek_vl_hybrid_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..2b83767f5efd51b698ff23fd15fae3a110aa0057 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/convert_deepseek_vl_hybrid_weights_to_hf.py @@ -0,0 +1,386 @@ +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import gc +import json +import os + +import regex as re +import torch +from huggingface_hub import snapshot_download +from huggingface_hub.errors import HFValidationError +from safetensors.torch import load_file + +from transformers import ( + AutoTokenizer, + DeepseekVLHybridConfig, + DeepseekVLHybridForConditionalGeneration, + DeepseekVLHybridImageProcessor, + DeepseekVLHybridProcessor, +) +from transformers.image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + PILImageResampling, +) + + +# fmt: off +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + # # Sam (High Resolution) + r"vision_model.vision_tower_high.vision_tower.pos_embed": r"model.high_res_vision_model.vision_encoder.pos_embed", + r"vision_model.vision_tower_high.vision_tower.patch_embed.proj.(weight|bias)": r"model.high_res_vision_model.vision_encoder.patch_embed.projection.\1", + r"vision_model.vision_tower_high.vision_tower.blocks.(\d+).norm(\d+).(weight|bias)": r"model.high_res_vision_model.vision_encoder.layers.\1.layer_norm\2.\3", + r"vision_model.vision_tower_high.vision_tower.blocks.(\d+).attn.rel_pos_(h|w)": r"model.high_res_vision_model.vision_encoder.layers.\1.attn.rel_pos_\2", + r"vision_model.vision_tower_high.vision_tower.blocks.(\d+).attn.qkv.(weight|bias)": r"model.high_res_vision_model.vision_encoder.layers.\1.attn.qkv.\2", + r"vision_model.vision_tower_high.vision_tower.blocks.(\d+).attn.proj.(weight|bias)": r"model.high_res_vision_model.vision_encoder.layers.\1.attn.proj.\2", + r"vision_model.vision_tower_high.vision_tower.blocks.(\d+).mlp.lin(\d+).(weight|bias)": r"model.high_res_vision_model.vision_encoder.layers.\1.mlp.lin\2.\3", + r"vision_model.vision_tower_high.vision_tower.neck.0.weight": r"model.high_res_vision_model.vision_encoder.neck.conv1.weight", + r"vision_model.vision_tower_high.vision_tower.neck.1.(weight|bias)": r"model.high_res_vision_model.vision_encoder.neck.layer_norm1.\1", + r"vision_model.vision_tower_high.vision_tower.neck.2.weight": r"model.high_res_vision_model.vision_encoder.neck.conv2.weight", + r"vision_model.vision_tower_high.vision_tower.neck.3.(weight|bias)": r"model.high_res_vision_model.vision_encoder.neck.layer_norm2.\1", + r"vision_model.vision_tower_high.vision_tower.neck_hd.0.weight": r"model.high_res_vision_neck.conv1.weight", + r"vision_model.vision_tower_high.vision_tower.neck_hd.1.(weight|bias)": r"model.high_res_vision_neck.layer_norm1.\1", + r"vision_model.vision_tower_high.vision_tower.neck_hd.2.weight": r"model.high_res_vision_neck.conv2.weight", + r"vision_model.vision_tower_high.vision_tower.neck_hd.3.(weight|bias)": r"model.high_res_vision_neck.layer_norm2.\1", + r"vision_model.vision_tower_high.vision_tower.downsamples.0.weight": r"model.high_res_vision_proj.conv1.weight", + r"vision_model.vision_tower_high.vision_tower.downsamples.1.weight": r"model.high_res_vision_proj.conv2.weight", + r"vision_model.vision_tower_high.vision_tower.hd_alpha_downsamples": r"model.high_res_vision_alpha", + + # Siglip (Low Resolution) + r"vision_model.vision_tower_low.vision_tower.pos_embed": r"model.vision_model.vision_model.embeddings.position_embedding.weight", + r"vision_model.vision_tower_low.vision_tower.patch_embed.proj.(weight|bias)": r"model.vision_model.vision_model.embeddings.patch_embedding.\1", + r"vision_model.vision_tower_low.vision_tower.blocks.(\d+).attn.qkv.(weight|bias)": r"model.vision_model.vision_model.encoder.layers.\1.self_attn.(q|k|v)_proj.\2", + r"vision_model.vision_tower_low.vision_tower.blocks.(\d+).attn.proj.(weight|bias)": r"model.vision_model.vision_model.encoder.layers.\1.self_attn.out_proj.\2", + r"vision_model.vision_tower_low.vision_tower.blocks.(\d+).norm(\d+).(weight|bias)": r"model.vision_model.vision_model.encoder.layers.\1.layer_norm\2.\3", + r"vision_model.vision_tower_low.vision_tower.blocks.(\d+).mlp.fc(\d+).(weight|bias)": r"model.vision_model.vision_model.encoder.layers.\1.mlp.fc\2.\3", + r"vision_model.vision_tower_low.vision_tower.norm.(weight|bias)": r"model.vision_model.vision_model.post_layernorm.\1", + r"vision_model.vision_tower_low.vision_tower.attn_pool.latent": r"model.vision_model.vision_model.head.probe", + r"vision_model.vision_tower_low.vision_tower.attn_pool.proj.(weight|bias)": r"model.vision_model.vision_model.head.attention.out_proj.\1", + r"vision_model.vision_tower_low.vision_tower.attn_pool.norm.(weight|bias)": r"model.vision_model.vision_model.head.layernorm.\1", + r"vision_model.vision_tower_low.vision_tower.attn_pool.mlp.fc(\d+).(weight|bias)": r"model.vision_model.vision_model.head.mlp.fc\1.\2", + + # Vision Projection + r"aligner.layers.1.(weight|bias)": r"model.aligner.proj.\1", + r"aligner.low_up_proj.(weight|bias)": r"model.aligner.vision_proj.\1", + r"aligner.high_up_proj.(weight|bias)": r"model.aligner.high_res_vision_proj.\1", + + # Llama (Text Model) + r"language_model.model.(\w+)": r"model.language_model.\1", + r"language_model.lm_head.(weight|bias)": r"lm_head.\1", +} +# fmt: on + +# Adopted from https://github.com/deepseek-ai/DeepSeek-VL/blob/main/deepseek_vl/utils/conversation.py#L80-L91 +CHAT_TEMPLATE = ( + # Define separators and initialize counter + "{% set seps = ['\n\n', '<\uff5cend\u2581of\u2581sentence\uff5c>'] %}" + "{% set i = 0 %}" + # Start with default system prompt + "You are a helpful language and vision assistant. " + "You are able to understand the visual content that the user provides, " + "and assist the user with a variety of tasks using natural language.\n\n" + # Iterate through messages + "{% for message in messages %}" + # Identify user or assistant role + "{% if message['role']|lower == 'user' %}" + "User: " + "{% elif message['role']|lower == 'assistant' %}" + "Assistant:{% if not (loop.last and not add_generation_prompt and message['content'][0]['type']=='text' and message['content'][0]['text']=='') %} {% endif %}" + "{% else %}" + "{{ message['role'].capitalize() }}: " + "{% endif %}" + # Iterate through message content (text/images) + "{% for content in message['content'] %}" + # If content is an image, replace with placeholder + "{% if content['type'] == 'image' %}" + "" + # If content is text, handle formatting + "{% elif content['type'] == 'text' %}" + "{% set text = content['text'] %}" + # Strip whitespace for first and last text blocks + "{% if loop.first %}{% set text = text.lstrip() %}{% endif %}" + "{% if loop.last %}{% set text = text.rstrip() %}{% endif %}" + # If previous content was text, add space + "{% if not loop.first and message['content'][loop.index0-1]['type'] == 'text' %}" + "{{ ' ' + text }}" + "{% else %}" + "{{ text }}" + "{% endif %}" + "{% endif %}" + "{% endfor %}" # End message content loop + # Add separators between messages + "{% if not loop.last or add_generation_prompt %}" + "{% if message['role']|lower == 'user' %}" + "{{ seps[0] }}" + "{% else %}" + "{{ seps[1] }}" + "{% endif %}" + "{% endif %}" + "{% endfor %}" # End messages loop + # Add final Assistant prompt if required + "{% if add_generation_prompt %}Assistant:{% endif %}" +) + + +def convert_old_keys_to_new_keys(state_dict_keys: dict): + output_dict = {} + + 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 get_qkv_state_dict(key, parameter): + """ + new key which looks like this + xxxx.(q|k|v).xxx (m, n) + + is converted to + xxxx.q.xxxx (m//3, n) + xxxx.k.xxxx (m//3, n) + xxxx.v.xxxx (m//3, n) + """ + qkv_state_dict = {} + placeholder = re.search(r"(\(.*?\))", key).group(1) # finds "(query|key|value)" + replacements_keys = placeholder[1:-1].split("|") # creates ['query', 'key', 'value'] + replacements_vals = torch.split( + parameter, split_size_or_sections=parameter.size(0) // len(replacements_keys), dim=0 + ) + for replacement_key, replacement_val in zip(replacements_keys, replacements_vals): + qkv_state_dict[key.replace(placeholder, replacement_key)] = replacement_val + return qkv_state_dict + + +def update_state_dict(old_state_dict): + all_keys = list(old_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] + current_parameter = old_state_dict.pop(key) + + if "qkv" in key and "vision_tower_high" not in key: + qkv_state_dict = get_qkv_state_dict(new_key, current_parameter) + state_dict.update(qkv_state_dict) + elif "pos_embed" in key: + if "vision_tower_high" not in key: + # timm implementation of siglip creates this param of size [1, 576, 1024] + # transformers implementation of siglip creates this param of size [576, 1024] + state_dict[new_key] = current_parameter.squeeze(0) + else: + state_dict[new_key] = current_parameter + else: + state_dict[new_key] = current_parameter + + return state_dict + + +def load_model_state_dict(input_path: str) -> dict: + """ + Load model state dict, handling both single and sharded files. + """ + index_path = os.path.join(input_path, "model.safetensors.index.json") + single_file_path = os.path.join(input_path, "model.safetensors") + + # Check if we have a sharded model + if os.path.exists(index_path): + print("Loading sharded model...") + state_dict = {} + with open(index_path, "r") as f: + index = json.load(f) + + # Get unique shard files and load each one only once + unique_shard_files = sorted(set(index["weight_map"].values())) + for shard_file in unique_shard_files: + print(f"Loading shard {shard_file}...") + shard_path = os.path.join(input_path, shard_file) + shard_dict = load_file(shard_path) + state_dict.update(shard_dict) + + return state_dict + + # Single file model + elif os.path.exists(single_file_path): + print("Loading single file model...") + return load_file(single_file_path, device="cpu") + + else: + raise ValueError(f"No model files found in {input_path}") + + +def convert_model( + hf_repo_id: str, + output_dir: str | None = None, + output_hub_path: str | None = None, +): + if output_dir: + os.makedirs(output_dir, exist_ok=True) + + try: + input_path = snapshot_download(hf_repo_id) + except HFValidationError: + # If the input path is not a HF repo ID, assume it's a local path + input_path = hf_repo_id + + # ------------------------------------------------------------ + # Create and save config + # ------------------------------------------------------------ + + config = DeepseekVLHybridConfig( + text_config={ + "hidden_size": 4096, + "intermediate_size": 11008, + "max_position_embeddings": 16384, + "num_attention_heads": 32, + "num_hidden_layers": 30, + "vocab_size": 102400, + }, + vision_config={ + "hidden_size": 1024, + "intermediate_size": 4096, + "image_size": 384, + "patch_size": 16, + "hidden_act": "gelu", + "vision_use_head": False, + "num_attention_heads": 16, + "num_hidden_layers": 24, + }, + high_res_vision_config={ + "hidden_size": 768, + "intermediate_size": 3072, + "image_size": 1024, + "patch_size": 16, + "num_attention_heads": 12, + "num_hidden_layers": 12, + }, + ) + + # save config + if output_dir: + config.save_pretrained(output_dir) + print("Model config saved successfully...") + + # ------------------------------------------------------------ + # Convert processor + # ------------------------------------------------------------ + + image_processor = DeepseekVLHybridImageProcessor( + image_mean=IMAGENET_STANDARD_MEAN, + image_std=IMAGENET_STANDARD_STD, + high_res_image_mean=OPENAI_CLIP_MEAN, + high_res_image_std=OPENAI_CLIP_STD, + resample=PILImageResampling.BILINEAR, + ) + + tokenizer = AutoTokenizer.from_pretrained( + input_path, + extra_special_tokens={ + "pad_token": "<｜end▁of▁sentence｜>", + "image_token": "", + }, + ) + + processor = DeepseekVLHybridProcessor( + image_processor=image_processor, + tokenizer=tokenizer, + chat_template=CHAT_TEMPLATE, + ) + + if output_dir: + print(f"Saving processor to {output_dir}...") + processor.save_pretrained(output_dir) + if output_hub_path: + print(f"Pushing processor to hub at {output_hub_path}...") + processor.push_to_hub(output_hub_path) + + # ------------------------------------------------------------ + # Convert weights + # ------------------------------------------------------------ + + print("Creating empty model...") + with torch.device("meta"): + model = DeepseekVLHybridForConditionalGeneration(config) + + # Load and convert state dict + print("Loading state dict...") + state_dict = load_model_state_dict(input_path) + state_dict = update_state_dict(state_dict) + + # Load converted state dict + print("Loading converted weights into model...") + info = model.load_state_dict(state_dict, strict=False, assign=True) + if len(info.missing_keys) > 0: + raise ValueError(f"Missing keys: {info.missing_keys}") + + # Tie weights before any device mapping + print("Tying weights...") + model.tie_weights() + + # Save the model + if output_dir: + print(f"Saving model to {output_dir}...") + model.save_pretrained(output_dir) + if output_hub_path: + print(f"Pushing model to hub at {output_hub_path}...") + model.push_to_hub(output_hub_path) + + del state_dict, model + gc.collect() + + # Validate the saved model if saved locally + if output_dir: + print("Reloading the local model to check if it's saved correctly...") + DeepseekVLHybridForConditionalGeneration.from_pretrained(output_dir, device_map="auto") + print("Local model reloaded successfully.") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--hf_repo_id", + default="deepseek-ai/deepseek-vl-7b-chat", + help="Location of official weights from DeepseekAI on HF", + ) + parser.add_argument( + "--output_dir", + default=None, + help="Location to write the converted model and processor", + ) + parser.add_argument( + "--output_hub_path", + default=None, + help="Repository ID to push model to hub (e.g. 'username/model-name')", + ) + args = parser.parse_args() + + convert_model( + hf_repo_id=args.hf_repo_id, + output_dir=args.output_dir, + output_hub_path=args.output_hub_path, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..867d61aeb230a8085e176bab4c1c935a41938280 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py @@ -0,0 +1,297 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.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_deepseek_vl_hybrid.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from typing import Union + +import torch +import torchvision.transforms.v2.functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +class DeepseekVLHybridImageProcessorKwargs(ImagesKwargs, total=False): + r""" + min_size (`int`, *optional*, defaults to 14): + The minimum allowed size for the resized image. Ensures that neither the height nor width + falls below this value after resizing. + high_res_size (`dict`, *optional*, defaults to `{"height": 1024, "width": 1024}`): + Size of the high resolution output image after resizing. Can be overridden by the `high_res_size` parameter in the `preprocess` + method. + high_res_resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `high_res_resample` parameter in the `preprocess` method. + high_res_image_mean (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_MEAN`): + Mean to use if normalizing the high resolution image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `high_res_image_mean` parameter in the `preprocess` method. + high_res_image_std (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_STD`): + Standard deviation to use if normalizing the high resolution image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `high_res_image_std` parameter in the `preprocess` method. + """ + + min_size: int + high_res_size: dict + high_res_resample: Union["PILImageResampling", int] + high_res_image_mean: float | list[float] | tuple[float, ...] + high_res_image_std: float | list[float] | tuple[float, ...] + + +@auto_docstring +class DeepseekVLHybridImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"height": 384, "width": 384} + min_size = 14 + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + valid_kwargs = DeepseekVLHybridImageProcessorKwargs + high_res_image_mean = OPENAI_CLIP_MEAN + high_res_image_std = OPENAI_CLIP_STD + high_res_size = {"height": 1024, "width": 1024} + high_res_resample = PILImageResampling.BICUBIC + model_input_names = ["pixel_values", "high_res_pixel_values"] + + def __init__(self, **kwargs: Unpack[DeepseekVLHybridImageProcessorKwargs]): + if kwargs.get("image_mean") is None: + background_color = (127, 127, 127) + else: + background_color = tuple(int(x * 255) for x in kwargs.get("image_mean")) + if kwargs.get("high_res_image_mean") is None: + high_res_background_color = (127, 127, 127) + else: + high_res_background_color = tuple(int(x * 255) for x in kwargs.get("high_res_image_mean")) + super().__init__(**kwargs) + self.background_color = tuple(background_color) + self.high_res_background_color = tuple(high_res_background_color) + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + min_size: int, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + antialias: bool = True, + **kwargs, + ) -> "torch.Tensor": + if size.height is None or size.width is None or size.height != size.width: + raise ValueError( + f"Output height and width must be the same. Got height={size['height']} and width={size['width']}" + ) + size = size.height + + height, width = image.shape[-2:] + max_size = max(height, width) + + delta = size / max_size + # Largest side becomes `size` and the other side is scaled according to the aspect ratio. + output_size_nonpadded = SizeDict( + height=max(round(height * delta), min_size), + width=max(round(width * delta), min_size), + ) + + return super().resize(image, size=output_size_nonpadded, resample=resample, antialias=antialias) + + def pad_to_square( + self, + images: "torch.Tensor", + background_color: int | tuple[int, int, int] = 0, + ) -> "torch.Tensor": + """ + Pads an image to a square based on the longest edge. + + Args: + images (`torch.Tensor`): + The images to pad. + background_color (`int` or `tuple[int, int, int]`, *optional*, defaults to 0): + The color to use for the padding. Can be an integer for single channel or a + tuple of integers representing for multi-channel images. If passed as integer + in multi-channel mode, it will default to `0` in subsequent channels. + + Returns: + `torch.Tensor`: The padded images. + """ + height, width = images.shape[-2:] + num_channels = images.shape[1] + batch_size = images.shape[0] + + if height == width: + return images + + max_dim = max(height, width) + + # Ensure background_color is the correct shape + if isinstance(background_color, int): + background_color = [background_color] + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + padded_images = torch.zeros( + (batch_size, num_channels, max_dim, max_dim), dtype=images.dtype, device=images.device + ) + for i, color in enumerate(background_color): + padded_images[:, i, :, :] = color + if width > height: + start = (max_dim - height) // 2 + padded_images[:, :, start : start + height, :] = images + else: + start = (max_dim - width) // 2 + padded_images[:, :, :, start : start + width] = images + + return padded_images + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + high_res_size: SizeDict, + min_size: int, + resample: "PILImageResampling | None", + high_res_resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + high_res_image_mean: float | list[float] | None, + high_res_image_std: float | list[float] | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + do_pad: bool = True, + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + high_res_resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_high_res_images = self.resize( + image=stacked_images, size=high_res_size, min_size=min_size, resample=high_res_resample + ) + high_res_resized_images_grouped[shape] = stacked_high_res_images + high_res_resized_images = reorder_images(high_res_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_high_res_images, grouped_high_res_images_index = group_images_by_shape( + high_res_resized_images, disable_grouping=disable_grouping + ) + high_res_padded_images = {} + high_res_processed_images_grouped = {} + for shape, stacked_high_res_images in grouped_high_res_images.items(): + if do_pad: + stacked_high_res_images = self.pad_to_square( + stacked_high_res_images, background_color=self.high_res_background_color + ) + high_res_padded_images[shape] = stacked_high_res_images + # Fused rescale and normalize + stacked_high_res_images = self.rescale_and_normalize( + stacked_high_res_images, + do_rescale, + rescale_factor, + do_normalize, + high_res_image_mean, + high_res_image_std, + ) + high_res_processed_images_grouped[shape] = stacked_high_res_images + high_res_processed_images = reorder_images(high_res_processed_images_grouped, grouped_high_res_images_index) + + resized_images_grouped = {} + for shape, stacked_high_res_padded_images in high_res_padded_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_high_res_padded_images, size=size, min_size=min_size, resample=resample + ) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_high_res_images_index) + + grouped_resized_images, grouped_resized_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + for shape, stacked_images in grouped_resized_images.items(): + if do_pad: + stacked_images = self.pad_to_square(stacked_images, background_color=self.background_color) + # 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_resized_images_index) + + return BatchFeature( + data={"pixel_values": processed_images, "high_res_pixel_values": high_res_processed_images}, + tensor_type=return_tensors, + ) + + def postprocess(self) -> "torch.Tensor": + raise AttributeError("Not needed for DeepseekVLHybrid") + + def _standardize_kwargs( + self, + size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + high_res_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + default_to_square: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + high_res_image_mean: float | list[float] | None = None, + high_res_image_std: float | list[float] | None = None, + **kwargs, + ) -> dict: + """ + Update kwargs that need further processing before being validated + Can be overridden by subclasses to customize the processing of kwargs. + """ + if kwargs is None: + kwargs = {} + if size is not None and not isinstance(size, SizeDict): + size = SizeDict(**get_size_dict(size=size, default_to_square=default_to_square)) + if high_res_size is not None and not isinstance(high_res_size, SizeDict): + high_res_size = SizeDict(**get_size_dict(size=high_res_size, default_to_square=default_to_square)) + if isinstance(image_mean, list): + image_mean = tuple(image_mean) + if isinstance(image_std, list): + image_std = tuple(image_std) + if isinstance(high_res_image_mean, list): + high_res_image_mean = tuple(high_res_image_mean) + if isinstance(high_res_image_std, list): + high_res_image_std = tuple(high_res_image_std) + + kwargs["size"] = size + kwargs["high_res_size"] = high_res_size + kwargs["image_mean"] = image_mean + kwargs["image_std"] = image_std + kwargs["high_res_image_mean"] = high_res_image_mean + kwargs["high_res_image_std"] = high_res_image_std + + return kwargs + + +__all__ = ["DeepseekVLHybridImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_pil_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_pil_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..55573c35c4231e1b667431368a8900c726adcddf --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/image_processing_pil_deepseek_vl_hybrid.py @@ -0,0 +1,261 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.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_deepseek_vl_hybrid.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from typing import Union + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import resize as np_resize +from ...image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +class DeepseekVLHybridImageProcessorKwargs(ImagesKwargs, total=False): + r""" + min_size (`int`, *optional*, defaults to 14): + The minimum allowed size for the resized image. Ensures that neither the height nor width + falls below this value after resizing. + high_res_size (`dict`, *optional*, defaults to `{"height": 1024, "width": 1024}`): + Size of the high resolution output image after resizing. Can be overridden by the `high_res_size` parameter in the `preprocess` + method. + high_res_resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `high_res_resample` parameter in the `preprocess` method. + high_res_image_mean (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_MEAN`): + Mean to use if normalizing the high resolution image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `high_res_image_mean` parameter in the `preprocess` method. + high_res_image_std (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_STD`): + Standard deviation to use if normalizing the high resolution image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `high_res_image_std` parameter in the `preprocess` method. + """ + + min_size: int + high_res_size: dict + high_res_resample: Union["PILImageResampling", int] + high_res_image_mean: float | list[float] | tuple[float, ...] + high_res_image_std: float | list[float] | tuple[float, ...] + + +@auto_docstring +class DeepseekVLHybridImageProcessorPil(PilBackend): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"height": 384, "width": 384} + min_size = 14 + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + valid_kwargs = DeepseekVLHybridImageProcessorKwargs + high_res_image_mean = OPENAI_CLIP_MEAN + high_res_image_std = OPENAI_CLIP_STD + high_res_size = {"height": 1024, "width": 1024} + high_res_resample = PILImageResampling.BICUBIC + model_input_names = ["pixel_values", "high_res_pixel_values"] + + def __init__(self, **kwargs: Unpack[DeepseekVLHybridImageProcessorKwargs]): + if kwargs.get("image_mean") is None: + background_color = (127, 127, 127) + else: + background_color = tuple(int(x * 255) for x in kwargs.get("image_mean")) + if kwargs.get("high_res_image_mean") is None: + high_res_background_color = (127, 127, 127) + else: + high_res_background_color = tuple(int(x * 255) for x in kwargs.get("high_res_image_mean")) + super().__init__(**kwargs) + self.background_color = tuple(background_color) + self.high_res_background_color = tuple(high_res_background_color) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[DeepseekVLHybridImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def resize( + self, + image: np.ndarray, + size: SizeDict, + min_size: int, + resample: PILImageResampling | None = None, + **kwargs, + ) -> np.ndarray: + """Resize so largest side becomes size, with min_size floor.""" + if size.height is None or size.width is None or size.height != size.width: + raise ValueError( + f"Output height and width must be the same. Got height={size.height} and width={size.width}" + ) + target_size = size.height + + height, width = image.shape[-2:] + max_size = max(height, width) + + delta = target_size / max_size + new_height = max(round(height * delta), min_size) + new_width = max(round(width * delta), min_size) + + return np_resize( + image, + size=(new_height, new_width), + resample=resample or self.resample, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + + def pad_to_square( + self, + image: np.ndarray, + background_color: int | tuple[int, int, int] = 0, + ) -> np.ndarray: + """Pad an image to a square based on the longest edge.""" + height, width = image.shape[-2:] + num_channels = image.shape[0] + + if height == width: + return image + + max_dim = max(height, width) + + if isinstance(background_color, int): + background_color = [background_color] + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + padded_image = np.zeros((num_channels, max_dim, max_dim), dtype=image.dtype) + for i, color in enumerate(background_color): + padded_image[i, :, :] = color + + if width > height: + start = (max_dim - height) // 2 + padded_image[:, start : start + height, :] = image + else: + start = (max_dim - width) // 2 + padded_image[:, :, start : start + width] = image + + return padded_image + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + high_res_size: SizeDict, + min_size: int, + resample: "PILImageResampling | None", + high_res_resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + high_res_image_mean: float | list[float] | None, + high_res_image_std: float | list[float] | None, + return_tensors: str | TensorType | None, + do_pad: bool = True, + **kwargs, + ) -> BatchFeature: + high_res_processed_images = [] + processed_images = [] + for image in images: + # high_res_image: resize (high) -> rescale -> normalize (high) + # low_res_image: resize (high) -> rescale -> resize (low) -> normalize (low) + high_res_image = image + if do_resize: + high_res_image = self.resize( + image=high_res_image, size=high_res_size, min_size=min_size, resample=high_res_resample + ) + if do_pad: + high_res_image = self.pad_to_square( + high_res_image, background_color=self.high_res_background_color + ) + image = self.resize(image=high_res_image, size=size, min_size=min_size, resample=resample) + if do_pad: + image = self.pad_to_square(image, background_color=self.background_color) + if do_rescale: + high_res_image = self.rescale(high_res_image, rescale_factor) + image = self.rescale(image, rescale_factor) + if do_normalize: + high_res_image = self.normalize(high_res_image, high_res_image_mean, high_res_image_std) + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + high_res_processed_images.append(high_res_image) + + return BatchFeature( + data={"pixel_values": processed_images, "high_res_pixel_values": high_res_processed_images}, + tensor_type=return_tensors, + ) + + def postprocess(self): + """Applies post-processing to the decoded image tokens by reversing transformations applied during preprocessing.""" + raise AttributeError("Not needed for DeepseekVLHybrid") + + def _standardize_kwargs( + self, + size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + high_res_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + default_to_square: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + high_res_image_mean: float | list[float] | None = None, + high_res_image_std: float | list[float] | None = None, + **kwargs, + ) -> dict: + """ + Update kwargs that need further processing before being validated + Can be overridden by subclasses to customize the processing of kwargs. + """ + if kwargs is None: + kwargs = {} + if size is not None and not isinstance(size, SizeDict): + size = SizeDict(**get_size_dict(size=size, default_to_square=default_to_square)) + if high_res_size is not None and not isinstance(high_res_size, SizeDict): + high_res_size = SizeDict(**get_size_dict(size=high_res_size, default_to_square=default_to_square)) + if isinstance(image_mean, list): + image_mean = tuple(image_mean) + if isinstance(image_std, list): + image_std = tuple(image_std) + if isinstance(high_res_image_mean, list): + high_res_image_mean = tuple(high_res_image_mean) + if isinstance(high_res_image_std, list): + high_res_image_std = tuple(high_res_image_std) + + kwargs["size"] = size + kwargs["high_res_size"] = high_res_size + kwargs["image_mean"] = image_mean + kwargs["image_std"] = image_std + kwargs["high_res_image_mean"] = high_res_image_mean + kwargs["high_res_image_std"] = high_res_image_std + + return kwargs + + +__all__ = ["DeepseekVLHybridImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..eb85a8d02a76de8fba8aaffeba1fe6942c4cfa67 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py @@ -0,0 +1,539 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.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_deepseek_vl_hybrid.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ... import initialization as init +from ...cache_utils import Cache +from ...generation import GenerationMixin +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ..auto import AutoModel +from .configuration_deepseek_vl_hybrid import DeepseekVLHybridConfig + + +@dataclass +@auto_docstring +class BaseModelOutputWithHighResVisionEncodings(BaseModelOutputWithPooling): + r""" + high_res_vision_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 high resolution vision model. + high_res_vision_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 high resolution vision model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the high resolution vision model at the output of each layer plus the optional initial embedding outputs. + high_res_vision_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)` from the high resolution vision model. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + high_res_vision_last_hidden_state: torch.FloatTensor | None = None + high_res_vision_hidden_states: tuple[torch.FloatTensor] | None = None + high_res_vision_attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for DeepseekVLHybrid model's outputs that may also contain a past key/values (to speed up sequential decoding). + """ +) +class DeepseekVLHybridBaseModelOutputWithPast(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 DeepseekVLHybrid causal language model (or autoregressive) outputs. + """ +) +class DeepseekVLHybridCausalLMOutputWithPast(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 DeepseekVLHybridLayerNorm(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 DeepseekVLSamVisionNeck(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False) + self.layer_norm1 = DeepseekVLHybridLayerNorm(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 = DeepseekVLHybridLayerNorm(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 DeepseekVLSamVisionProj(nn.Module): + def __init__(self, config, output_size: int = 24): + super().__init__() + self.config = config + self.output_size = output_size + + self.conv1 = nn.Conv2d( + config.output_channels, config.output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False + ) + self.conv2 = nn.Conv2d( + config.output_channels * 2, config.output_channels * 4, kernel_size=3, stride=2, padding=1, bias=False + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + # interpolate Sam encodings to match Siglip encodings + features = torch.nn.functional.interpolate( + features, + size=(4 * self.output_size, 4 * self.output_size), + mode="bilinear", + align_corners=False, + ) + features = self.conv1(features) + features = self.conv2(features) + return features + + +class DeepseekVLHybridAligner(nn.Module): + def __init__(self, config: DeepseekVLHybridConfig): + super().__init__() + + in_channels = config.vision_config.hidden_size + high_res_in_channels = config.high_res_vision_config.output_channels * 4 + out_channels = config.text_config.hidden_size + + self.vision_proj = nn.Linear(in_channels, out_channels // 2) + self.high_res_vision_proj = nn.Linear(high_res_in_channels, out_channels // 2) + + self.act = nn.GELU() + self.proj = nn.Linear(out_channels, out_channels) + + def forward( + self, + vision_encodings: torch.Tensor, + high_res_vision_encodings: torch.Tensor, + ) -> torch.Tensor: + vision_encodings = self.vision_proj(vision_encodings) + high_res_vision_encodings = self.high_res_vision_proj(high_res_vision_encodings) + + encodings = torch.concat([high_res_vision_encodings, vision_encodings], dim=-1) + encodings = self.act(encodings) + encodings = self.proj(encodings) + + return encodings + + +@auto_docstring +class DeepseekVLHybridPreTrainedModel(PreTrainedModel): + config: DeepseekVLHybridConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values", "causal_mask"] + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.text_config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, DeepseekVLHybridLayerNorm): + init.ones_(module.weight) + init.zeros_(module.bias) + elif isinstance(module, DeepseekVLHybridModel): + init.zeros_(module.high_res_vision_alpha) + + +DEEPSEEK_VL_COMMON_CUSTOM_ARGS = r""" + high_res_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size), *optional*): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. +""" + + +@auto_docstring +class DeepseekVLHybridModel(DeepseekVLHybridPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.output_size = config.vision_config.image_size // config.vision_config.patch_size + self.global_attn_index = config.high_res_vision_config.global_attn_indexes[0] + + self.high_res_vision_model = AutoModel.from_config(config.high_res_vision_config) + self.high_res_vision_neck = DeepseekVLSamVisionNeck(config.high_res_vision_config) + self.high_res_vision_proj = DeepseekVLSamVisionProj( + config.high_res_vision_config, output_size=self.output_size + ) + self.high_res_vision_alpha = nn.Parameter(torch.zeros(1)) + self.config = config + + self.vision_model = AutoModel.from_config(config.vision_config) + self.aligner = DeepseekVLHybridAligner(config) + + self.language_model = AutoModel.from_config(config=config.text_config) + + self.gradient_checkpointing = False + # 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) + + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + high_res_pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithHighResVisionEncodings: + low_res_outputs = self.get_low_res_image_features(pixel_values, **kwargs) + high_res_outputs = self.get_high_res_image_features(high_res_pixel_values, **kwargs) + image_features = self.aligner(low_res_outputs.last_hidden_state, high_res_outputs.last_hidden_state) + + return BaseModelOutputWithHighResVisionEncodings( + last_hidden_state=low_res_outputs.last_hidden_state, + pooler_output=image_features, + hidden_states=low_res_outputs.hidden_states, + attentions=low_res_outputs.attentions, + high_res_vision_last_hidden_state=high_res_outputs.last_hidden_state, + high_res_vision_hidden_states=high_res_outputs.hidden_states, + high_res_vision_attentions=high_res_outputs.attentions, + ) + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + + n_image_tokens = special_image_mask.sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + torch_compilable_check( + inputs_embeds[special_image_mask].numel() == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", + ) + return special_image_mask + + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + high_res_pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> DeepseekVLHybridBaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if pixel_values is not None and high_res_pixel_values is None: + raise ValueError("Both pixel_values and high_res_pixel_values should be specified at the same time") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + if input_ids is None: + image_attention_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + image_attention_mask = image_attention_mask.all(-1) + else: + image_attention_mask = input_ids == self.config.image_token_id + + image_attention_mask = image_attention_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + image_embeds = self.get_image_features(pixel_values, high_res_pixel_values, return_dict=True).pooler_output + image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_attention_mask, image_features) + + lm_output = self.language_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + logits_to_keep=logits_to_keep, + **kwargs, + ) + + return DeepseekVLHybridBaseModelOutputWithPast( + last_hidden_state=lm_output.last_hidden_state, + past_key_values=lm_output.past_key_values, + hidden_states=lm_output.hidden_states, + attentions=lm_output.attentions, + image_hidden_states=image_embeds if pixel_values is not None else None, + ) + + def get_low_res_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.vision_model(pixel_values, return_dict=True, **kwargs) + + def get_high_res_image_features( + self, + pixel_values: torch.FloatTensor, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + high_res_outputs = self.high_res_vision_model( + pixel_values=pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) + last_hidden_state = high_res_outputs.last_hidden_state + last_hidden_state = self.high_res_vision_proj(last_hidden_state) + + hidden_states = high_res_outputs.hidden_states + global_hidden_state = hidden_states[self.global_attn_index + 1] # +1 for embedding layer + global_hidden_state = self.high_res_vision_neck(global_hidden_state) + global_hidden_state = self.high_res_vision_proj(global_hidden_state) + + output = last_hidden_state + global_hidden_state * self.high_res_vision_alpha + + # batch_size, hidden_size, height, width -> batch_size, seq_len, hidden_size + output = output.permute(0, 2, 3, 1) + output = output.reshape(output.shape[0], -1, output.shape[-1]) + high_res_outputs.last_hidden_state = output + + return high_res_outputs + + +class DeepseekVLHybridForConditionalGeneration(DeepseekVLHybridPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + output_modalities = ("text",) + _can_compile_fullgraph = True + + def __init__(self, config: DeepseekVLHybridConfig): + super().__init__(config) + self.config = config + self.model = DeepseekVLHybridModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + + # Initialize weights and apply final processing. + self.post_init() + + def get_input_embeddings(self): + return self.model.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.language_model.set_input_embeddings(value) + + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + high_res_pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> DeepseekVLHybridCausalLMOutputWithPast: + 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]`. + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + high_res_pixel_values=high_res_pixel_values, + 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.text_config.vocab_size, **kwargs + ) + + return DeepseekVLHybridCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + high_res_pixel_values=None, + attention_mask=None, + logits_to_keep=None, + is_first_iteration=False, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + logits_to_keep=logits_to_keep, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if is_first_iteration or not kwargs.get("use_cache", True): + # Pixel values are used only in the first iteration if available + # In subsequent iterations, they are already merged with text and cached + # NOTE: first iteration doesn't have to be prefill, it can be the first + # iteration with a question and cached system prompt (continue generate from cache) + model_inputs["pixel_values"] = pixel_values + model_inputs["high_res_pixel_values"] = high_res_pixel_values + + return model_inputs + + +__all__ = ["DeepseekVLHybridPreTrainedModel", "DeepseekVLHybridModel", "DeepseekVLHybridForConditionalGeneration"] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..99d24c163562f85d7c302503cdb52c9876e261fc --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py @@ -0,0 +1,787 @@ +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Union + +import numpy as np +import torch +import torch.nn as nn +from huggingface_hub.dataclasses import strict + +from ... import initialization as init +from ...cache_utils import Cache +from ...configuration_utils import PreTrainedConfig +from ...image_processing_backends import PilBackend, TorchvisionBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import ImagesKwargs, Unpack +from ...tokenization_utils_base import ( + PreTokenizedInput, + TextInput, +) +from ...utils import ( + TensorType, + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, +) +from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel +from ..deepseek_vl.configuration_deepseek_vl import DeepseekVLConfig +from ..deepseek_vl.image_processing_deepseek_vl import DeepseekVLImageProcessor +from ..deepseek_vl.image_processing_pil_deepseek_vl import DeepseekVLImageProcessorPil +from ..deepseek_vl.modeling_deepseek_vl import ( + DeepseekVLForConditionalGeneration, + DeepseekVLModel, + DeepseekVLPreTrainedModel, +) +from ..deepseek_vl.processing_deepseek_vl import DeepseekVLProcessor, DeepseekVLProcessorKwargs +from ..idefics.modeling_idefics import IdeficsBaseModelOutputWithPast, IdeficsCausalLMOutputWithPast +from ..sam.modeling_sam import SamLayerNorm, SamVisionNeck + + +logger = logging.get_logger(__name__) + + +DEEPSEEK_VL_COMMON_CUSTOM_ARGS = r""" + high_res_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size), *optional*): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. +""" + + +@auto_docstring(checkpoint="deepseek-community/deepseek-vl-7b-chat") +@strict +class DeepseekVLHybridConfig(DeepseekVLConfig): + r""" + high_res_vision_config (`Union[AutoConfig, dict]`, *optional*, defaults to `SamVisionConfig`): + The config object or dictionary of the high resolution vision backbone. + + Example: + + ```python + >>> from transformers import DeepseekVLHybridConfig, DeepseekVLHybridModel + + >>> # Initializing a DeepseekVLHybrid deepseek-community/deepseek-vl-7b-chat style configuration + >>> configuration = DeepseekVLHybridConfig() + + >>> # Initializing a model (with random weights) from the deepseek-community/deepseek-vl-7b-chat style configuration + >>> model = DeepseekVLHybridModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_vl_hybrid" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig, "high_res_vision_config": AutoConfig} + + high_res_vision_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.high_res_vision_config is None: + self.high_res_vision_config = {} + logger.info("`high_res_vision_config` is `None`. Initializing the `SamVisionConfig` with default values.") + + if isinstance(self.high_res_vision_config, dict): + self.high_res_vision_config["model_type"] = self.high_res_vision_config.get( + "model_type", "sam_vision_model" + ) + self.high_res_vision_config = CONFIG_MAPPING[self.high_res_vision_config["model_type"]]( + **self.high_res_vision_config + ) + + super().__post_init__(**kwargs) + + +@dataclass +@auto_docstring +class BaseModelOutputWithHighResVisionEncodings(BaseModelOutputWithPooling): + r""" + high_res_vision_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 high resolution vision model. + high_res_vision_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 high resolution vision model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the high resolution vision model at the output of each layer plus the optional initial embedding outputs. + high_res_vision_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)` from the high resolution vision model. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + high_res_vision_last_hidden_state: torch.FloatTensor | None = None + high_res_vision_hidden_states: tuple[torch.FloatTensor] | None = None + high_res_vision_attentions: tuple[torch.FloatTensor] | None = None + + +class DeepseekVLHybridBaseModelOutputWithPast(IdeficsBaseModelOutputWithPast): + pass + + +class DeepseekVLHybridCausalLMOutputWithPast(IdeficsCausalLMOutputWithPast): + pass + + +class DeepseekVLHybridLayerNorm(SamLayerNorm): + pass + + +class DeepseekVLSamVisionNeck(SamVisionNeck): + def __init__(self, config): + super().__init__(config) + + +class DeepseekVLSamVisionProj(nn.Module): + def __init__(self, config, output_size: int = 24): + super().__init__() + self.config = config + self.output_size = output_size + + self.conv1 = nn.Conv2d( + config.output_channels, config.output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False + ) + self.conv2 = nn.Conv2d( + config.output_channels * 2, config.output_channels * 4, kernel_size=3, stride=2, padding=1, bias=False + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + # interpolate Sam encodings to match Siglip encodings + features = torch.nn.functional.interpolate( + features, + size=(4 * self.output_size, 4 * self.output_size), + mode="bilinear", + align_corners=False, + ) + features = self.conv1(features) + features = self.conv2(features) + return features + + +class DeepseekVLHybridAligner(nn.Module): + def __init__(self, config: DeepseekVLHybridConfig): + super().__init__() + + in_channels = config.vision_config.hidden_size + high_res_in_channels = config.high_res_vision_config.output_channels * 4 + out_channels = config.text_config.hidden_size + + self.vision_proj = nn.Linear(in_channels, out_channels // 2) + self.high_res_vision_proj = nn.Linear(high_res_in_channels, out_channels // 2) + + self.act = nn.GELU() + self.proj = nn.Linear(out_channels, out_channels) + + def forward( + self, + vision_encodings: torch.Tensor, + high_res_vision_encodings: torch.Tensor, + ) -> torch.Tensor: + vision_encodings = self.vision_proj(vision_encodings) + high_res_vision_encodings = self.high_res_vision_proj(high_res_vision_encodings) + + encodings = torch.concat([high_res_vision_encodings, vision_encodings], dim=-1) + encodings = self.act(encodings) + encodings = self.proj(encodings) + + return encodings + + +class DeepseekVLHybridPreTrainedModel(DeepseekVLPreTrainedModel): + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + init.normal_(module.weight, mean=0.0, std=self.config.text_config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, DeepseekVLHybridLayerNorm): + init.ones_(module.weight) + init.zeros_(module.bias) + elif isinstance(module, DeepseekVLHybridModel): + init.zeros_(module.high_res_vision_alpha) + + +class DeepseekVLHybridModel(DeepseekVLModel): + def __init__(self, config): + self.output_size = config.vision_config.image_size // config.vision_config.patch_size + self.global_attn_index = config.high_res_vision_config.global_attn_indexes[0] + + self.high_res_vision_model = AutoModel.from_config(config.high_res_vision_config) + self.high_res_vision_neck = DeepseekVLSamVisionNeck(config.high_res_vision_config) + self.high_res_vision_proj = DeepseekVLSamVisionProj( + config.high_res_vision_config, output_size=self.output_size + ) + self.high_res_vision_alpha = nn.Parameter(torch.zeros(1)) + + super().__init__(config) + + def get_low_res_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.vision_model(pixel_values, return_dict=True, **kwargs) + + def get_high_res_image_features( + self, + pixel_values: torch.FloatTensor, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + high_res_outputs = self.high_res_vision_model( + pixel_values=pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) + last_hidden_state = high_res_outputs.last_hidden_state + last_hidden_state = self.high_res_vision_proj(last_hidden_state) + + hidden_states = high_res_outputs.hidden_states + global_hidden_state = hidden_states[self.global_attn_index + 1] # +1 for embedding layer + global_hidden_state = self.high_res_vision_neck(global_hidden_state) + global_hidden_state = self.high_res_vision_proj(global_hidden_state) + + output = last_hidden_state + global_hidden_state * self.high_res_vision_alpha + + # batch_size, hidden_size, height, width -> batch_size, seq_len, hidden_size + output = output.permute(0, 2, 3, 1) + output = output.reshape(output.shape[0], -1, output.shape[-1]) + high_res_outputs.last_hidden_state = output + + return high_res_outputs + + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + high_res_pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithHighResVisionEncodings: + low_res_outputs = self.get_low_res_image_features(pixel_values, **kwargs) + high_res_outputs = self.get_high_res_image_features(high_res_pixel_values, **kwargs) + image_features = self.aligner(low_res_outputs.last_hidden_state, high_res_outputs.last_hidden_state) + + return BaseModelOutputWithHighResVisionEncodings( + last_hidden_state=low_res_outputs.last_hidden_state, + pooler_output=image_features, + hidden_states=low_res_outputs.hidden_states, + attentions=low_res_outputs.attentions, + high_res_vision_last_hidden_state=high_res_outputs.last_hidden_state, + high_res_vision_hidden_states=high_res_outputs.hidden_states, + high_res_vision_attentions=high_res_outputs.attentions, + ) + + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + high_res_pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs, + ) -> DeepseekVLHybridBaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if pixel_values is not None and high_res_pixel_values is None: + raise ValueError("Both pixel_values and high_res_pixel_values should be specified at the same time") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + if input_ids is None: + image_attention_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + image_attention_mask = image_attention_mask.all(-1) + else: + image_attention_mask = input_ids == self.config.image_token_id + + image_attention_mask = image_attention_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + image_embeds = self.get_image_features(pixel_values, high_res_pixel_values, return_dict=True).pooler_output + image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_attention_mask, image_features) + + lm_output = self.language_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + logits_to_keep=logits_to_keep, + **kwargs, + ) + + return DeepseekVLHybridBaseModelOutputWithPast( + last_hidden_state=lm_output.last_hidden_state, + past_key_values=lm_output.past_key_values, + hidden_states=lm_output.hidden_states, + attentions=lm_output.attentions, + image_hidden_states=image_embeds if pixel_values is not None else None, + ) + + +class DeepseekVLHybridForConditionalGeneration(DeepseekVLForConditionalGeneration): + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + high_res_pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> DeepseekVLHybridCausalLMOutputWithPast: + 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]`. + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + high_res_pixel_values=high_res_pixel_values, + 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.text_config.vocab_size, **kwargs + ) + + return DeepseekVLHybridCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + high_res_pixel_values=None, + attention_mask=None, + logits_to_keep=None, + is_first_iteration=False, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + logits_to_keep=logits_to_keep, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if is_first_iteration or not kwargs.get("use_cache", True): + # Pixel values are used only in the first iteration if available + # In subsequent iterations, they are already merged with text and cached + # NOTE: first iteration doesn't have to be prefill, it can be the first + # iteration with a question and cached system prompt (continue generate from cache) + model_inputs["pixel_values"] = pixel_values + model_inputs["high_res_pixel_values"] = high_res_pixel_values + + return model_inputs + + +class DeepseekVLHybridImageProcessorKwargs(ImagesKwargs, total=False): + r""" + min_size (`int`, *optional*, defaults to 14): + The minimum allowed size for the resized image. Ensures that neither the height nor width + falls below this value after resizing. + high_res_size (`dict`, *optional*, defaults to `{"height": 1024, "width": 1024}`): + Size of the high resolution output image after resizing. Can be overridden by the `high_res_size` parameter in the `preprocess` + method. + high_res_resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `high_res_resample` parameter in the `preprocess` method. + high_res_image_mean (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_MEAN`): + Mean to use if normalizing the high resolution image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `high_res_image_mean` parameter in the `preprocess` method. + high_res_image_std (`float` or `list[float]`, *optional*, defaults to `OPENAI_CLIP_STD`): + Standard deviation to use if normalizing the high resolution image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `high_res_image_std` parameter in the `preprocess` method. + """ + + min_size: int + high_res_size: dict + high_res_resample: Union["PILImageResampling", int] + high_res_image_mean: float | list[float] | tuple[float, ...] + high_res_image_std: float | list[float] | tuple[float, ...] + + +class DeepseekVLHybridImageProcessorPil(DeepseekVLImageProcessorPil): + high_res_image_mean = OPENAI_CLIP_MEAN + high_res_image_std = OPENAI_CLIP_STD + high_res_size = {"height": 1024, "width": 1024} + high_res_resample = PILImageResampling.BICUBIC + model_input_names = ["pixel_values", "high_res_pixel_values"] + + def __init__(self, **kwargs: Unpack[DeepseekVLHybridImageProcessorKwargs]): + if kwargs.get("image_mean") is None: + background_color = (127, 127, 127) + else: + background_color = tuple(int(x * 255) for x in kwargs.get("image_mean")) + if kwargs.get("high_res_image_mean") is None: + high_res_background_color = (127, 127, 127) + else: + high_res_background_color = tuple(int(x * 255) for x in kwargs.get("high_res_image_mean")) + PilBackend.__init__(self, **kwargs) + self.background_color = tuple(background_color) + self.high_res_background_color = tuple(high_res_background_color) + + def _standardize_kwargs( + self, + size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + high_res_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + default_to_square: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + high_res_image_mean: float | list[float] | None = None, + high_res_image_std: float | list[float] | None = None, + **kwargs, + ) -> dict: + """ + Update kwargs that need further processing before being validated + Can be overridden by subclasses to customize the processing of kwargs. + """ + if kwargs is None: + kwargs = {} + if size is not None and not isinstance(size, SizeDict): + size = SizeDict(**get_size_dict(size=size, default_to_square=default_to_square)) + if high_res_size is not None and not isinstance(high_res_size, SizeDict): + high_res_size = SizeDict(**get_size_dict(size=high_res_size, default_to_square=default_to_square)) + if isinstance(image_mean, list): + image_mean = tuple(image_mean) + if isinstance(image_std, list): + image_std = tuple(image_std) + if isinstance(high_res_image_mean, list): + high_res_image_mean = tuple(high_res_image_mean) + if isinstance(high_res_image_std, list): + high_res_image_std = tuple(high_res_image_std) + + kwargs["size"] = size + kwargs["high_res_size"] = high_res_size + kwargs["image_mean"] = image_mean + kwargs["image_std"] = image_std + kwargs["high_res_image_mean"] = high_res_image_mean + kwargs["high_res_image_std"] = high_res_image_std + + return kwargs + + def _preprocess( + self, + images: list[np.ndarray], + do_resize: bool, + size: SizeDict, + high_res_size: SizeDict, + min_size: int, + resample: "PILImageResampling | None", + high_res_resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + high_res_image_mean: float | list[float] | None, + high_res_image_std: float | list[float] | None, + return_tensors: str | TensorType | None, + do_pad: bool = True, + **kwargs, + ) -> BatchFeature: + high_res_processed_images = [] + processed_images = [] + for image in images: + # high_res_image: resize (high) -> rescale -> normalize (high) + # low_res_image: resize (high) -> rescale -> resize (low) -> normalize (low) + high_res_image = image + if do_resize: + high_res_image = self.resize( + image=high_res_image, size=high_res_size, min_size=min_size, resample=high_res_resample + ) + if do_pad: + high_res_image = self.pad_to_square( + high_res_image, background_color=self.high_res_background_color + ) + image = self.resize(image=high_res_image, size=size, min_size=min_size, resample=resample) + if do_pad: + image = self.pad_to_square(image, background_color=self.background_color) + if do_rescale: + high_res_image = self.rescale(high_res_image, rescale_factor) + image = self.rescale(image, rescale_factor) + if do_normalize: + high_res_image = self.normalize(high_res_image, high_res_image_mean, high_res_image_std) + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + high_res_processed_images.append(high_res_image) + + return BatchFeature( + data={"pixel_values": processed_images, "high_res_pixel_values": high_res_processed_images}, + tensor_type=return_tensors, + ) + + +class DeepseekVLHybridImageProcessor(DeepseekVLImageProcessor): + high_res_image_mean = OPENAI_CLIP_MEAN + high_res_image_std = OPENAI_CLIP_STD + high_res_size = {"height": 1024, "width": 1024} + high_res_resample = PILImageResampling.BICUBIC + model_input_names = ["pixel_values", "high_res_pixel_values"] + + def __init__(self, **kwargs: Unpack[DeepseekVLHybridImageProcessorKwargs]): + if kwargs.get("image_mean") is None: + background_color = (127, 127, 127) + else: + background_color = tuple(int(x * 255) for x in kwargs.get("image_mean")) + if kwargs.get("high_res_image_mean") is None: + high_res_background_color = (127, 127, 127) + else: + high_res_background_color = tuple(int(x * 255) for x in kwargs.get("high_res_image_mean")) + TorchvisionBackend.__init__(self, **kwargs) + self.background_color = tuple(background_color) + self.high_res_background_color = tuple(high_res_background_color) + + def _standardize_kwargs( + self, + size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + high_res_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None, + default_to_square: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + high_res_image_mean: float | list[float] | None = None, + high_res_image_std: float | list[float] | None = None, + **kwargs, + ) -> dict: + """ + Update kwargs that need further processing before being validated + Can be overridden by subclasses to customize the processing of kwargs. + """ + if kwargs is None: + kwargs = {} + if size is not None and not isinstance(size, SizeDict): + size = SizeDict(**get_size_dict(size=size, default_to_square=default_to_square)) + if high_res_size is not None and not isinstance(high_res_size, SizeDict): + high_res_size = SizeDict(**get_size_dict(size=high_res_size, default_to_square=default_to_square)) + if isinstance(image_mean, list): + image_mean = tuple(image_mean) + if isinstance(image_std, list): + image_std = tuple(image_std) + if isinstance(high_res_image_mean, list): + high_res_image_mean = tuple(high_res_image_mean) + if isinstance(high_res_image_std, list): + high_res_image_std = tuple(high_res_image_std) + + kwargs["size"] = size + kwargs["high_res_size"] = high_res_size + kwargs["image_mean"] = image_mean + kwargs["image_std"] = image_std + kwargs["high_res_image_mean"] = high_res_image_mean + kwargs["high_res_image_std"] = high_res_image_std + + return kwargs + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + high_res_size: SizeDict, + min_size: int, + resample: "PILImageResampling | None", + high_res_resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + high_res_image_mean: float | list[float] | None, + high_res_image_std: float | list[float] | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + do_pad: bool = True, + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + high_res_resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_high_res_images = self.resize( + image=stacked_images, size=high_res_size, min_size=min_size, resample=high_res_resample + ) + high_res_resized_images_grouped[shape] = stacked_high_res_images + high_res_resized_images = reorder_images(high_res_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_high_res_images, grouped_high_res_images_index = group_images_by_shape( + high_res_resized_images, disable_grouping=disable_grouping + ) + high_res_padded_images = {} + high_res_processed_images_grouped = {} + for shape, stacked_high_res_images in grouped_high_res_images.items(): + if do_pad: + stacked_high_res_images = self.pad_to_square( + stacked_high_res_images, background_color=self.high_res_background_color + ) + high_res_padded_images[shape] = stacked_high_res_images + # Fused rescale and normalize + stacked_high_res_images = self.rescale_and_normalize( + stacked_high_res_images, + do_rescale, + rescale_factor, + do_normalize, + high_res_image_mean, + high_res_image_std, + ) + high_res_processed_images_grouped[shape] = stacked_high_res_images + high_res_processed_images = reorder_images(high_res_processed_images_grouped, grouped_high_res_images_index) + + resized_images_grouped = {} + for shape, stacked_high_res_padded_images in high_res_padded_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_high_res_padded_images, size=size, min_size=min_size, resample=resample + ) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_high_res_images_index) + + grouped_resized_images, grouped_resized_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + for shape, stacked_images in grouped_resized_images.items(): + if do_pad: + stacked_images = self.pad_to_square(stacked_images, background_color=self.background_color) + # 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_resized_images_index) + + return BatchFeature( + data={"pixel_values": processed_images, "high_res_pixel_values": high_res_processed_images}, + tensor_type=return_tensors, + ) + + +class DeepseekVLHybridProcessorKwargs(DeepseekVLProcessorKwargs): + pass + + +class DeepseekVLHybridProcessor(DeepseekVLProcessor): + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + images: ImageInput | None = None, + **kwargs: Unpack[DeepseekVLHybridProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + output_kwargs = self._merge_kwargs( + DeepseekVLHybridProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs + ) + if text is None and images is None: + raise ValueError("You must specify either text or images.") + + if text is not None: + if isinstance(text, str): + text = [text] + elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): + raise ValueError("Invalid input text. Please provide a string, or a list of strings") + + prompt_strings = [] + one_img_tokens = self.image_token * self.num_image_tokens + for prompt in text: + prompt = prompt.replace(self.image_token, one_img_tokens) + prompt_strings.append(prompt) + + data = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"]) + + # process images if pixel_values are provided + if images is not None: + inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + data["pixel_values"] = inputs["pixel_values"] + data["high_res_pixel_values"] = inputs["high_res_pixel_values"] + + return BatchFeature(data=data) + + +__all__ = [ + "DeepseekVLHybridConfig", + "DeepseekVLHybridPreTrainedModel", + "DeepseekVLHybridModel", + "DeepseekVLHybridForConditionalGeneration", + "DeepseekVLHybridImageProcessor", + "DeepseekVLHybridImageProcessorPil", + "DeepseekVLHybridProcessor", +] diff --git a/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/processing_deepseek_vl_hybrid.py b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/processing_deepseek_vl_hybrid.py new file mode 100644 index 0000000000000000000000000000000000000000..7948b954b6d75b4a7faa37067724e7071c6de4b0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/deepseek_vl_hybrid/processing_deepseek_vl_hybrid.py @@ -0,0 +1,119 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.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_deepseek_vl_hybrid.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 Deepseek AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...image_processing_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +class DeepseekVLHybridProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": {"padding": False}, + "common_kwargs": {"return_tensors": "pt"}, + } + + +@auto_docstring +class DeepseekVLHybridProcessor(ProcessorMixin): + def __init__( + self, + image_processor, + tokenizer, + chat_template=None, + num_image_tokens=576, + ): + r""" + num_image_tokens (`int`, *optional*, defaults to 576): + The number of special image tokens used as placeholders for visual content in text sequences. + """ + self.image_token = tokenizer.image_token + self.num_image_tokens = num_image_tokens + + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + images: ImageInput | None = None, + **kwargs: Unpack[DeepseekVLHybridProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + output_kwargs = self._merge_kwargs( + DeepseekVLHybridProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs + ) + if text is None and images is None: + raise ValueError("You must specify either text or images.") + + if text is not None: + if isinstance(text, str): + text = [text] + elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): + raise ValueError("Invalid input text. Please provide a string, or a list of strings") + + prompt_strings = [] + one_img_tokens = self.image_token * self.num_image_tokens + for prompt in text: + prompt = prompt.replace(self.image_token, one_img_tokens) + prompt_strings.append(prompt) + + data = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"]) + + # process images if pixel_values are provided + if images is not None: + inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + data["pixel_values"] = inputs["pixel_values"] + data["high_res_pixel_values"] = inputs["high_res_pixel_values"] + + return BatchFeature(data=data) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + + +__all__ = ["DeepseekVLHybridProcessor"] diff --git a/third_party/transformers/src/transformers/models/diffllama/__init__.py b/third_party/transformers/src/transformers/models/diffllama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c162fce0a48bd164bd0e0a615b942ee4805a12aa --- /dev/null +++ b/third_party/transformers/src/transformers/models/diffllama/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_diffllama import * + from .modeling_diffllama 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/diffllama/configuration_diffllama.py b/third_party/transformers/src/transformers/models/diffllama/configuration_diffllama.py new file mode 100644 index 0000000000000000000000000000000000000000..18fbcf6c347eecf6701e7b6591359ff4414d378f --- /dev/null +++ b/third_party/transformers/src/transformers/models/diffllama/configuration_diffllama.py @@ -0,0 +1,80 @@ +# Copyright 2024 weak-kajuma and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on Llama implementations in this library and Microsoft's +# Differential Transformer implementations. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DiffLlama 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="kajuma/DiffLlama-0.3B-handcut") +@strict +class DiffLlamaConfig(PreTrainedConfig): + r""" + lambda_std_dev (`float`, *optional*, defaults to 0.1): + The standard deviation for initialization of parameter lambda in attention layer. + + ```python + >>> from transformers import DiffLlamaModel, DiffLlamaConfig + + >>> # Initializing a DiffLlama diffllama-7b style configuration + >>> configuration = DiffLlamaConfig() + + >>> # Initializing a model from the diffllama-7b style configuration + >>> model = DiffLlamaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "diffllama" + keys_to_ignore_at_inference = ["past_key_values"] + + vocab_size: int = 32000 + hidden_size: int = 2048 + intermediate_size: int = 8192 + num_hidden_layers: int = 16 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int | None = 0.0 + lambda_std_dev: float | None = 0.1 + head_dim: int | None = None + + def __post_init__(self, **kwargs): + # for backward compatibility + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads + super().__post_init__(**kwargs) + + +__all__ = ["DiffLlamaConfig"] diff --git a/third_party/transformers/src/transformers/models/diffllama/modeling_diffllama.py b/third_party/transformers/src/transformers/models/diffllama/modeling_diffllama.py new file mode 100644 index 0000000000000000000000000000000000000000..d80ccd572dc35d0e0c511f1d0bd2ca87e1b66105 --- /dev/null +++ b/third_party/transformers/src/transformers/models/diffllama/modeling_diffllama.py @@ -0,0 +1,753 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/diffllama/modular_diffllama.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_diffllama.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 weak-kajuma and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on Llama implementations in this library and Microsoft's +# Differential Transformer implementations. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, StaticCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask +from ...modeling_layers import ( + GenericForQuestionAnswering, + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_diffllama import DiffLlamaConfig + + +logger = logging.get_logger(__name__) + + +class DiffLlamaMLP(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 DiffLlamaRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: DiffLlamaConfig, 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: DiffLlamaConfig | 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 lambda_init_fn(layer_idx): + return 0.8 - 0.6 * math.exp(-0.3 * layer_idx) + + +class DiffLlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: DiffLlamaConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + # under this are not used + self.max_position_embeddings = config.max_position_embeddings + self.is_causal = True + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + + self.lambda_init = lambda_init_fn(layer_idx) + self.lambda_q1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_k1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_q2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_k2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.groupnorm = nn.RMSNorm(2 * self.head_dim, eps=config.rms_norm_eps, elementwise_affine=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, target_len, _ = hidden_states.size() + q_len = target_len + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1) + value_states = value_states.repeat(1, 2, 1, 1) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = torch.matmul(attn_weights, value_states) + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1) + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class DiffLlamaFlashAttention2(DiffLlamaAttention): + """ + DiffLlama flash attention module. This module inherits from `DiffLlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask() + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, None]: + if isinstance(past_key_values, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (DiffLlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + device_type = query_states.device.type if query_states.device.type != "mps" else "cpu" + if input_dtype == torch.float32: + if torch.is_autocast_enabled(device_type): + target_dtype = torch.get_autocast_dtype(device_type) + # Handle the case where the model is quantized + elif hasattr(self.config, "_is_quantized"): + target_dtype = self.config.dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + value_states1, value_states2 = torch.chunk(value_states, 2, dim=2) + value_states1 = value_states1.repeat(1, 1, 2, 1) + value_states2 = value_states2.repeat(1, 1, 2, 1) + + attn_output1 = _flash_attention_forward( + query_states, + key_states, + value_states1, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output2 = _flash_attention_forward( + query_states, + key_states, + value_states2, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output = torch.cat([attn_output1, attn_output2], dim=-1) + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=2) + + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, None + + +class DiffLlamaSdpaAttention(DiffLlamaAttention): + """ + DiffLlama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `DiffLlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from DiffLlamaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1) + value_states = value_states.repeat(1, 2, 1, 1) + + causal_mask = attention_mask + if attention_mask is not None: + causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + is_causal = causal_mask is None and q_len > 1 + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1) + + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, None + + +@use_kernel_forward_from_hub("RMSNorm") +class DiffLlamaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + DiffLlamaRMSNorm 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}" + + +DIFFLLAMA_ATTENTION_CLASSES = { + "eager": DiffLlamaAttention, + "flash_attention_2": DiffLlamaFlashAttention2, + "sdpa": DiffLlamaSdpaAttention, +} + + +class DiffLlamaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: DiffLlamaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = DIFFLLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + + self.mlp = DiffLlamaMLP(config) + self.input_layernorm = DiffLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DiffLlamaRMSNorm(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 DiffLlamaPreTrainedModel(PreTrainedModel): + config: DiffLlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["DiffLlamaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = False + + _can_compile_fullgraph = True + _supports_attention_backend = False + _can_record_outputs = { + "hidden_states": DiffLlamaDecoderLayer, + "attentions": DiffLlamaAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, DiffLlamaAttention): + init.normal_(module.lambda_q1, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_k1, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_q2, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_k2, 0, self.config.lambda_std_dev) + + +@auto_docstring +class DiffLlamaModel(DiffLlamaPreTrainedModel): + def __init__(self, config: DiffLlamaConfig): + 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( + [DiffLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = DiffLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = DiffLlamaRotaryEmbedding(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 DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, 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 = DiffLlamaModel(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, DiffLlamaForCausalLM + + >>> model = DiffLlamaForCausalLM.from_pretrained("google/diffllama-7b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/diffllama-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 DiffLlamaForSequenceClassification(GenericForSequenceClassification, DiffLlamaPreTrainedModel): + pass + + +class DiffLlamaForQuestionAnswering(GenericForQuestionAnswering, DiffLlamaPreTrainedModel): + base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model` + + +class DiffLlamaForTokenClassification(GenericForTokenClassification, DiffLlamaPreTrainedModel): + pass + + +__all__ = [ + "DiffLlamaPreTrainedModel", + "DiffLlamaModel", + "DiffLlamaForCausalLM", + "DiffLlamaForSequenceClassification", + "DiffLlamaForQuestionAnswering", + "DiffLlamaForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/diffllama/modular_diffllama.py b/third_party/transformers/src/transformers/models/diffllama/modular_diffllama.py new file mode 100644 index 0000000000000000000000000000000000000000..71087303c97edb10d653606893b5f3cd5746aee1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/diffllama/modular_diffllama.py @@ -0,0 +1,417 @@ +# Copyright 2024 weak-kajuma and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on Llama implementations in this library and Microsoft's +# Differential Transformer implementations. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math + +import torch +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, StaticCache +from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask +from ...modeling_utils import PreTrainedModel +from ...utils import logging +from ..gemma.modeling_gemma import GemmaForCausalLM +from ..llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaForQuestionAnswering, + LlamaForSequenceClassification, + LlamaForTokenClassification, + LlamaModel, + LlamaPreTrainedModel, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + repeat_kv, +) +from ..mistral.modeling_mistral import MistralMLP +from .configuration_diffllama import DiffLlamaConfig + + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "kajuma/DiffLlama-0.3B-handcut" +_CONFIG_FOR_DOC = "DiffLlamaConfig" + + +class DiffLlamaMLP(MistralMLP): + pass + + +def lambda_init_fn(layer_idx): + return 0.8 - 0.6 * math.exp(-0.3 * layer_idx) + + +class DiffLlamaRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class DiffLlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: DiffLlamaConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + # under this are not used + self.max_position_embeddings = config.max_position_embeddings + self.is_causal = True + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + + self.lambda_init = lambda_init_fn(layer_idx) + self.lambda_q1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_k1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_q2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.lambda_k2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,))) + self.groupnorm = nn.RMSNorm(2 * self.head_dim, eps=config.rms_norm_eps, elementwise_affine=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, target_len, _ = hidden_states.size() + q_len = target_len + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1) + value_states = value_states.repeat(1, 2, 1, 1) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = torch.matmul(attn_weights, value_states) + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1) + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class DiffLlamaFlashAttention2(DiffLlamaAttention): + """ + DiffLlama flash attention module. This module inherits from `DiffLlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask() + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, None]: + if isinstance(past_key_values, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (DiffLlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + device_type = query_states.device.type if query_states.device.type != "mps" else "cpu" + if input_dtype == torch.float32: + if torch.is_autocast_enabled(device_type): + target_dtype = torch.get_autocast_dtype(device_type) + # Handle the case where the model is quantized + elif hasattr(self.config, "_is_quantized"): + target_dtype = self.config.dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + value_states1, value_states2 = torch.chunk(value_states, 2, dim=2) + value_states1 = value_states1.repeat(1, 1, 2, 1) + value_states2 = value_states2.repeat(1, 1, 2, 1) + + attn_output1 = _flash_attention_forward( + query_states, + key_states, + value_states1, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output2 = _flash_attention_forward( + query_states, + key_states, + value_states2, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output = torch.cat([attn_output1, attn_output2], dim=-1) + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=2) + + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, None + + +class DiffLlamaSdpaAttention(DiffLlamaAttention): + """ + DiffLlama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `DiffLlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from DiffLlamaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1) + value_states = value_states.repeat(1, 2, 1, 1) + + causal_mask = attention_mask + if attention_mask is not None: + causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + is_causal = causal_mask is None and q_len > 1 + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1) + + lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to( + query_states.dtype + ) + lambda_full = lambda_1 - lambda_2 + self.lambda_init + + attn_output = attn_output1 - lambda_full * attn_output2 + attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, None + + +DIFFLLAMA_ATTENTION_CLASSES = { + "eager": DiffLlamaAttention, + "flash_attention_2": DiffLlamaFlashAttention2, + "sdpa": DiffLlamaSdpaAttention, +} + + +class DiffLlamaDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: DiffLlamaConfig, layer_idx: int): + super().__init__(config, layer_idx) + + self.self_attn = DIFFLLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + + +class DiffLlamaPreTrainedModel(LlamaPreTrainedModel): + _supports_flex_attn = False + _supports_attention_backend = False + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, DiffLlamaAttention): + init.normal_(module.lambda_q1, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_k1, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_q2, 0, self.config.lambda_std_dev) + init.normal_(module.lambda_k2, 0, self.config.lambda_std_dev) + + +class DiffLlamaModel(LlamaModel): + pass + + +class DiffLlamaForCausalLM(GemmaForCausalLM): + pass + + +class DiffLlamaForSequenceClassification(LlamaForSequenceClassification): + pass + + +class DiffLlamaForQuestionAnswering(LlamaForQuestionAnswering): + pass + + +class DiffLlamaForTokenClassification(LlamaForTokenClassification): + pass + + +__all__ = [ + "DiffLlamaPreTrainedModel", + "DiffLlamaModel", + "DiffLlamaForCausalLM", + "DiffLlamaForSequenceClassification", + "DiffLlamaForQuestionAnswering", + "DiffLlamaForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/encodec/__init__.py b/third_party/transformers/src/transformers/models/encodec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3adeea056604d1d31f946a5cd0bf53ea590ea3aa --- /dev/null +++ b/third_party/transformers/src/transformers/models/encodec/__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_encodec import * + from .feature_extraction_encodec import * + from .modeling_encodec 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/encodec/configuration_encodec.py b/third_party/transformers/src/transformers/models/encodec/configuration_encodec.py new file mode 100644 index 0000000000000000000000000000000000000000..5eff8811f4578d0c8ddfe81e01d7b986f70cb164 --- /dev/null +++ b/third_party/transformers/src/transformers/models/encodec/configuration_encodec.py @@ -0,0 +1,155 @@ +# Copyright 2023 Meta Platforms, Inc. and affiliates, 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. +"""EnCodec model configuration""" + +import math + +import numpy as np +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="facebook/encodec_24khz") +@strict +class EncodecConfig(PreTrainedConfig): + r""" + target_bandwidths (`list[float]`, *optional*, defaults to `[1.5, 3.0, 6.0, 12.0, 24.0]`): + The range of different bandwidths the model can encode audio with. + normalize (`bool`, *optional*, defaults to `False`): + Whether the audio shall be normalized when passed. + chunk_length_s (`float`, *optional*): + If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded. + overlap (`float`, *optional*): + Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following + formulae : `int((1.0 - self.overlap) * self.chunk_length)`. + num_filters (`int`, *optional*, defaults to 32): + Number of convolution kernels of first `EncodecConv1d` down sampling layer. + num_residual_layers (`int`, *optional*, defaults to 1): + Number of residual layers. + upsampling_ratios (`Sequence[int]` , *optional*, defaults to `[8, 5, 4, 2]`): + Kernel size and stride ratios. The encoder uses downsampling ratios instead of upsampling ratios, hence it + will use the ratios in the reverse order to the ones specified here that must match the decoder order. + norm_type (`str`, *optional*, defaults to `"weight_norm"`): + Normalization method. Should be in `["weight_norm", "time_group_norm"]` + kernel_size (`int`, *optional*, defaults to 7): + Kernel size for the initial convolution. + last_kernel_size (`int`, *optional*, defaults to 7): + Kernel size for the last convolution layer. + residual_kernel_size (`int`, *optional*, defaults to 3): + Kernel size for the residual layers. + dilation_growth_rate (`int`, *optional*, defaults to 2): + How much to increase the dilation with each layer. + use_causal_conv (`bool`, *optional*, defaults to `True`): + Whether to use fully causal convolution. + pad_mode (`str`, *optional*, defaults to `"reflect"`): + Padding mode for the convolutions. + compress (`int`, *optional*, defaults to 2): + Reduced dimensionality in residual branches (from Demucs v3). + num_lstm_layers (`int`, *optional*, defaults to 2): + Number of LSTM layers at the end of the encoder. + trim_right_ratio (`float`, *optional*, defaults to 1.0): + Ratio for trimming at the right of the transposed convolution under the `use_causal_conv = True` setup. If + equal to 1.0, it means that all the trimming is done at the right. + use_conv_shortcut (`bool`, *optional*, defaults to `True`): + Whether to use a convolutional layer as the 'skip' connection in the `EncodecResnetBlock` block. If False, + an identity function will be used, giving a generic residual connection. + + Example: + + ```python + >>> from transformers import EncodecModel, EncodecConfig + + >>> # Initializing a "facebook/encodec_24khz" style configuration + >>> configuration = EncodecConfig() + + >>> # Initializing a model (with random weights) from the "facebook/encodec_24khz" style configuration + >>> model = EncodecModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "encodec" + + target_bandwidths: list[float] | tuple[float, ...] = (1.5, 3.0, 6.0, 12.0, 24.0) + sampling_rate: int = 24_000 + audio_channels: int = 1 + normalize: bool = False + chunk_length_s: int | float | None = None + overlap: float | None = None + hidden_size: int = 128 + num_filters: int = 32 + num_residual_layers: int = 1 + upsampling_ratios: list[int] | tuple[int, ...] = (8, 5, 4, 2) + norm_type: str = "weight_norm" + kernel_size: int = 7 + last_kernel_size: int = 7 + residual_kernel_size: int = 3 + dilation_growth_rate: int = 2 + use_causal_conv: bool = True + pad_mode: str = "reflect" + compress: int = 2 + num_lstm_layers: int = 2 + trim_right_ratio: float = 1.0 + codebook_size: int = 1024 + codebook_dim: int | None = None + use_conv_shortcut: bool = True + + def __post_init__(self, **kwargs): + self.codebook_dim = self.codebook_dim if self.codebook_dim is not None else self.hidden_size + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_length(self) -> int | None: + if self.chunk_length_s is None: + return None + else: + return int(self.chunk_length_s * self.sampling_rate) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_stride(self) -> int | None: + if self.chunk_length_s is None or self.overlap is None: + return None + else: + return max(1, int((1.0 - self.overlap) * self.chunk_length)) + + @property + def hop_length(self) -> int: + return int(np.prod(self.upsampling_ratios)) + + @property + def codebook_nbits(self) -> int: + return math.ceil(math.log2(self.codebook_size)) + + @property + def frame_rate(self) -> int: + return math.ceil(self.sampling_rate / self.hop_length) + + @property + def num_quantizers(self) -> int: + return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * self.codebook_nbits)) + + +__all__ = ["EncodecConfig"] diff --git a/third_party/transformers/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..6b2348f7785758af37105858e90a98880638c240 --- /dev/null +++ b/third_party/transformers/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py @@ -0,0 +1,364 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert EnCodec checkpoints.""" + +import argparse + +import torch + +from transformers import ( + EncodecConfig, + EncodecFeatureExtractor, + EncodecModel, + logging, +) + + +# checkpoints downloaded from: +# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th +# https://huggingface.co/facebook/musicgen-small/resolve/main/compression_state_dict.bin +# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th + + +logging.set_verbosity_info() +logger = logging.get_logger("transformers.models.encodec") + +MAPPING_QUANTIZER = { + "quantizer.vq.layers.*._codebook.inited": "quantizer.layers.*.codebook.inited", + "quantizer.vq.layers.*._codebook.cluster_size": "quantizer.layers.*.codebook.cluster_size", + "quantizer.vq.layers.*._codebook.embed": "quantizer.layers.*.codebook.embed", + "quantizer.vq.layers.*._codebook.embed_avg": "quantizer.layers.*.codebook.embed_avg", +} +MAPPING_ENCODER = { + "encoder.model.0.conv.conv": "encoder.layers.0.conv", + "encoder.model.1.block.1.conv.conv": "encoder.layers.1.block.1.conv", + "encoder.model.1.block.3.conv.conv": "encoder.layers.1.block.3.conv", + "encoder.model.1.shortcut.conv.conv": "encoder.layers.1.shortcut.conv", + "encoder.model.3.conv.conv": "encoder.layers.3.conv", + "encoder.model.4.block.1.conv.conv": "encoder.layers.4.block.1.conv", + "encoder.model.4.block.3.conv.conv": "encoder.layers.4.block.3.conv", + "encoder.model.4.shortcut.conv.conv": "encoder.layers.4.shortcut.conv", + "encoder.model.6.conv.conv": "encoder.layers.6.conv", + "encoder.model.7.block.1.conv.conv": "encoder.layers.7.block.1.conv", + "encoder.model.7.block.3.conv.conv": "encoder.layers.7.block.3.conv", + "encoder.model.7.shortcut.conv.conv": "encoder.layers.7.shortcut.conv", + "encoder.model.9.conv.conv": "encoder.layers.9.conv", + "encoder.model.10.block.1.conv.conv": "encoder.layers.10.block.1.conv", + "encoder.model.10.block.3.conv.conv": "encoder.layers.10.block.3.conv", + "encoder.model.10.shortcut.conv.conv": "encoder.layers.10.shortcut.conv", + "encoder.model.12.conv.conv": "encoder.layers.12.conv", + "encoder.model.13.lstm": "encoder.layers.13.lstm", + "encoder.model.15.conv.conv": "encoder.layers.15.conv", +} +MAPPING_ENCODER_48K = { + "encoder.model.0.conv.norm": "encoder.layers.0.norm", + "encoder.model.1.block.1.conv.norm": "encoder.layers.1.block.1.norm", + "encoder.model.1.block.3.conv.norm": "encoder.layers.1.block.3.norm", + "encoder.model.1.shortcut.conv.norm": "encoder.layers.1.shortcut.norm", + "encoder.model.3.conv.norm": "encoder.layers.3.norm", + "encoder.model.4.block.1.conv.norm": "encoder.layers.4.block.1.norm", + "encoder.model.4.block.3.conv.norm": "encoder.layers.4.block.3.norm", + "encoder.model.4.shortcut.conv.norm": "encoder.layers.4.shortcut.norm", + "encoder.model.6.conv.norm": "encoder.layers.6.norm", + "encoder.model.7.block.1.conv.norm": "encoder.layers.7.block.1.norm", + "encoder.model.7.block.3.conv.norm": "encoder.layers.7.block.3.norm", + "encoder.model.7.shortcut.conv.norm": "encoder.layers.7.shortcut.norm", + "encoder.model.9.conv.norm": "encoder.layers.9.norm", + "encoder.model.10.block.1.conv.norm": "encoder.layers.10.block.1.norm", + "encoder.model.10.block.3.conv.norm": "encoder.layers.10.block.3.norm", + "encoder.model.10.shortcut.conv.norm": "encoder.layers.10.shortcut.norm", + "encoder.model.12.conv.norm": "encoder.layers.12.norm", + "encoder.model.15.conv.norm": "encoder.layers.15.norm", +} +MAPPING_DECODER = { + "decoder.model.0.conv.conv": "decoder.layers.0.conv", + "decoder.model.1.lstm": "decoder.layers.1.lstm", + "decoder.model.3.convtr.convtr": "decoder.layers.3.conv", + "decoder.model.4.block.1.conv.conv": "decoder.layers.4.block.1.conv", + "decoder.model.4.block.3.conv.conv": "decoder.layers.4.block.3.conv", + "decoder.model.4.shortcut.conv.conv": "decoder.layers.4.shortcut.conv", + "decoder.model.6.convtr.convtr": "decoder.layers.6.conv", + "decoder.model.7.block.1.conv.conv": "decoder.layers.7.block.1.conv", + "decoder.model.7.block.3.conv.conv": "decoder.layers.7.block.3.conv", + "decoder.model.7.shortcut.conv.conv": "decoder.layers.7.shortcut.conv", + "decoder.model.9.convtr.convtr": "decoder.layers.9.conv", + "decoder.model.10.block.1.conv.conv": "decoder.layers.10.block.1.conv", + "decoder.model.10.block.3.conv.conv": "decoder.layers.10.block.3.conv", + "decoder.model.10.shortcut.conv.conv": "decoder.layers.10.shortcut.conv", + "decoder.model.12.convtr.convtr": "decoder.layers.12.conv", + "decoder.model.13.block.1.conv.conv": "decoder.layers.13.block.1.conv", + "decoder.model.13.block.3.conv.conv": "decoder.layers.13.block.3.conv", + "decoder.model.13.shortcut.conv.conv": "decoder.layers.13.shortcut.conv", + "decoder.model.15.conv.conv": "decoder.layers.15.conv", +} +MAPPING_DECODER_48K = { + "decoder.model.0.conv.norm": "decoder.layers.0.norm", + "decoder.model.3.convtr.norm": "decoder.layers.3.norm", + "decoder.model.4.block.1.conv.norm": "decoder.layers.4.block.1.norm", + "decoder.model.4.block.3.conv.norm": "decoder.layers.4.block.3.norm", + "decoder.model.4.shortcut.conv.norm": "decoder.layers.4.shortcut.norm", + "decoder.model.6.convtr.norm": "decoder.layers.6.norm", + "decoder.model.7.block.1.conv.norm": "decoder.layers.7.block.1.norm", + "decoder.model.7.block.3.conv.norm": "decoder.layers.7.block.3.norm", + "decoder.model.7.shortcut.conv.norm": "decoder.layers.7.shortcut.norm", + "decoder.model.9.convtr.norm": "decoder.layers.9.norm", + "decoder.model.10.block.1.conv.norm": "decoder.layers.10.block.1.norm", + "decoder.model.10.block.3.conv.norm": "decoder.layers.10.block.3.norm", + "decoder.model.10.shortcut.conv.norm": "decoder.layers.10.shortcut.norm", + "decoder.model.12.convtr.norm": "decoder.layers.12.norm", + "decoder.model.13.block.1.conv.norm": "decoder.layers.13.block.1.norm", + "decoder.model.13.block.3.conv.norm": "decoder.layers.13.block.3.norm", + "decoder.model.13.shortcut.conv.norm": "decoder.layers.13.shortcut.norm", + "decoder.model.15.conv.norm": "decoder.layers.15.norm", +} +MAPPING_24K = { + **MAPPING_QUANTIZER, + **MAPPING_ENCODER, + **MAPPING_DECODER, +} +MAPPING_48K = { + **MAPPING_QUANTIZER, + **MAPPING_ENCODER, + **MAPPING_ENCODER_48K, + **MAPPING_DECODER, + **MAPPING_DECODER_48K, +} +TOP_LEVEL_KEYS = [] +IGNORE_KEYS = [] + + +def set_recursively(hf_pointer, key, value, full_name, weight_type): + for attribute in key.split("."): + hf_pointer = getattr(hf_pointer, attribute) + + if weight_type is not None: + hf_shape = getattr(hf_pointer, weight_type).shape + else: + hf_shape = hf_pointer.shape + + if hf_shape != value.shape: + raise ValueError( + f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be" + f" {value.shape} for {full_name}" + ) + + if weight_type == "weight": + hf_pointer.weight.data = value + elif weight_type == "weight_g": + hf_pointer.weight_g.data = value + elif weight_type == "weight_v": + hf_pointer.weight_v.data = value + elif weight_type == "bias": + hf_pointer.bias.data = value + elif weight_type == "running_mean": + hf_pointer.running_mean.data = value + elif weight_type == "running_var": + hf_pointer.running_var.data = value + elif weight_type == "num_batches_tracked": + hf_pointer.num_batches_tracked.data = value + elif weight_type == "weight_ih_l0": + hf_pointer.weight_ih_l0.data = value + elif weight_type == "weight_hh_l0": + hf_pointer.weight_hh_l0.data = value + elif weight_type == "bias_ih_l0": + hf_pointer.bias_ih_l0.data = value + elif weight_type == "bias_hh_l0": + hf_pointer.bias_hh_l0.data = value + elif weight_type == "weight_ih_l1": + hf_pointer.weight_ih_l1.data = value + elif weight_type == "weight_hh_l1": + hf_pointer.weight_hh_l1.data = value + elif weight_type == "bias_ih_l1": + hf_pointer.bias_ih_l1.data = value + elif weight_type == "bias_hh_l1": + hf_pointer.bias_hh_l1.data = value + else: + hf_pointer.data = value + + logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.") + + +def should_ignore(name, ignore_keys): + for key in ignore_keys: + if key.endswith(".*"): + if name.startswith(key[:-1]): + return True + elif ".*." in key: + prefix, suffix = key.split(".*.") + if prefix in name and suffix in name: + return True + elif key in name: + return True + return False + + +def recursively_load_weights(orig_dict, hf_model, model_name): + unused_weights = [] + + if model_name in ["encodec_24khz", "encodec_32khz"]: + MAPPING = MAPPING_24K + elif model_name == "encodec_48khz": + MAPPING = MAPPING_48K + else: + raise ValueError(f"Unsupported model: {model_name}") + + for name, value in orig_dict.items(): + if should_ignore(name, IGNORE_KEYS): + logger.info(f"{name} was ignored") + continue + + is_used = False + for key, mapped_key in MAPPING.items(): + if "*" in key: + prefix, suffix = key.split(".*.") + if prefix in name and suffix in name: + key = suffix + + if key in name: + # HACK otherwise .embed gets initialized with .embed_avg too + if key.endswith("embed") and name.endswith("embed_avg"): + continue + + is_used = True + if "*" in mapped_key: + layer_index = name.split(key)[0].split(".")[-2] + mapped_key = mapped_key.replace("*", layer_index) + if "weight_g" in name: + weight_type = "weight_g" + elif "weight_v" in name: + weight_type = "weight_v" + elif "weight_ih_l0" in name: + weight_type = "weight_ih_l0" + elif "weight_hh_l0" in name: + weight_type = "weight_hh_l0" + elif "bias_ih_l0" in name: + weight_type = "bias_ih_l0" + elif "bias_hh_l0" in name: + weight_type = "bias_hh_l0" + elif "weight_ih_l1" in name: + weight_type = "weight_ih_l1" + elif "weight_hh_l1" in name: + weight_type = "weight_hh_l1" + elif "bias_ih_l1" in name: + weight_type = "bias_ih_l1" + elif "bias_hh_l1" in name: + weight_type = "bias_hh_l1" + elif "bias" in name: + weight_type = "bias" + elif "weight" in name: + weight_type = "weight" + elif "running_mean" in name: + weight_type = "running_mean" + elif "running_var" in name: + weight_type = "running_var" + elif "num_batches_tracked" in name: + weight_type = "num_batches_tracked" + else: + weight_type = None + set_recursively(hf_model, mapped_key, value, name, weight_type) + continue + if not is_used: + unused_weights.append(name) + + logger.warning(f"Unused weights: {unused_weights}") + + +@torch.no_grad() +def convert_checkpoint( + model_name, + checkpoint_path, + pytorch_dump_folder_path, + config_path=None, + repo_id=None, +): + """ + Copy/paste/tweak model's weights to transformers design. + """ + if config_path is not None: + config = EncodecConfig.from_pretrained(config_path) + else: + config = EncodecConfig() + + if model_name == "encodec_24khz": + pass # config is already correct + elif model_name == "encodec_32khz": + config.upsampling_ratios = [8, 5, 4, 4] + config.target_bandwidths = [2.2] + config.num_filters = 64 + config.sampling_rate = 32_000 + config.codebook_size = 2048 + config.use_causal_conv = False + config.normalize = False + config.use_conv_shortcut = False + elif model_name == "encodec_48khz": + config.upsampling_ratios = [8, 5, 4, 2] + config.target_bandwidths = [3.0, 6.0, 12.0, 24.0] + config.sampling_rate = 48_000 + config.audio_channels = 2 + config.use_causal_conv = False + config.norm_type = "time_group_norm" + config.normalize = True + config.chunk_length_s = 1.0 + config.overlap = 0.01 + else: + raise ValueError(f"Unknown model name: {model_name}") + + model = EncodecModel(config) + + feature_extractor = EncodecFeatureExtractor( + feature_size=config.audio_channels, + sampling_rate=config.sampling_rate, + chunk_length_s=config.chunk_length_s, + overlap=config.overlap, + ) + feature_extractor.save_pretrained(pytorch_dump_folder_path) + + original_checkpoint = torch.load(checkpoint_path, weights_only=True) + if "best_state" in original_checkpoint: + # we might have a training state saved, in which case discard the yaml results and just retain the weights + original_checkpoint = original_checkpoint["best_state"] + recursively_load_weights(original_checkpoint, model, model_name) + model.save_pretrained(pytorch_dump_folder_path) + + if repo_id: + print("Pushing to the hub...") + feature_extractor.push_to_hub(repo_id) + model.push_to_hub(repo_id) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + default="encodec_24khz", + type=str, + help="The model to convert. Should be one of 'encodec_24khz', 'encodec_32khz', 'encodec_48khz'.", + ) + parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint") + parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") + parser.add_argument( + "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model." + ) + parser.add_argument( + "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the Hugging Face hub." + ) + + args = parser.parse_args() + convert_checkpoint( + args.model, + args.checkpoint_path, + args.pytorch_dump_folder_path, + args.config_path, + args.push_to_hub, + ) diff --git a/third_party/transformers/src/transformers/models/encodec/feature_extraction_encodec.py b/third_party/transformers/src/transformers/models/encodec/feature_extraction_encodec.py new file mode 100644 index 0000000000000000000000000000000000000000..383936000243c86c00bd031e99846cf63af51634 --- /dev/null +++ b/third_party/transformers/src/transformers/models/encodec/feature_extraction_encodec.py @@ -0,0 +1,205 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Feature extractor class for EnCodec.""" + +import numpy as np + +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import PaddingStrategy, TensorType, logging + + +logger = logging.get_logger(__name__) + + +class EncodecFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs an EnCodec feature extractor. + + This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains + most of the main methods. Users should refer to this superclass for more information regarding those methods. + + Instantiating a feature extractor with the defaults will yield a similar configuration to that of the + [facebook/encodec_24khz](https://huggingface.co/facebook/encodec_24khz) architecture. + + Args: + feature_size (`int`, *optional*, defaults to 1): + The feature dimension of the extracted features. Use 1 for mono, 2 for stereo. + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values. + chunk_length_s (`float`, *optional*): + If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded. + overlap (`float`, *optional*): + Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following + formulae : `int((1.0 - self.overlap) * self.chunk_length)`. + """ + + model_input_names = ["input_values", "padding_mask"] + + def __init__( + self, + feature_size: int = 1, + sampling_rate: int = 24000, + padding_value: float = 0.0, + chunk_length_s: float | None = None, + overlap: float | None = None, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + self.chunk_length_s = chunk_length_s + self.overlap = overlap + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_length(self) -> int | None: + if self.chunk_length_s is None: + return None + else: + return int(self.chunk_length_s * self.sampling_rate) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_stride(self) -> int | None: + if self.chunk_length_s is None or self.overlap is None: + return None + else: + return max(1, int((1.0 - self.overlap) * self.chunk_length)) + + def __call__( + self, + raw_audio: np.ndarray | list[float] | list[np.ndarray] | list[list[float]], + padding: bool | str | PaddingStrategy | None = None, + truncation: bool | None = False, + max_length: int | None = None, + return_tensors: str | TensorType | None = None, + sampling_rate: int | None = None, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). + + Args: + raw_audio (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`): + The sequence or batch of sequences to be processed. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. The numpy array must be of shape + `(num_samples,)` for mono audio (`feature_size = 1`), or `(2, num_samples)` for stereo audio + (`feature_size = 2`). + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, *optional*, defaults to `False`): + Activates truncation to cut input sequences longer than `max_length` to `max_length`. + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + sampling_rate (`int`, *optional*): + The sampling rate at which the `audio` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors. + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of" + f" {self.sampling_rate}. Please make sure that the provided audio input was sampled with" + f" {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + if padding and truncation: + raise ValueError("Both padding and truncation were set. Make sure you only set one.") + elif padding is None: + # by default let's pad the inputs + padding = True + + is_batched = bool( + isinstance(raw_audio, (list, tuple)) and (isinstance(raw_audio[0], (np.ndarray, tuple, list))) + ) + + if is_batched: + raw_audio = [np.asarray(audio, dtype=np.float32).T for audio in raw_audio] + elif not is_batched and not isinstance(raw_audio, np.ndarray): + raw_audio = np.asarray(raw_audio, dtype=np.float32) + elif isinstance(raw_audio, np.ndarray) and raw_audio.dtype is np.dtype(np.float64): + raw_audio = raw_audio.astype(np.float32) + + # always return batch + if not is_batched: + raw_audio = [np.asarray(raw_audio).T] + + # verify inputs are valid + for idx, example in enumerate(raw_audio): + if example.ndim > 2: + raise ValueError(f"Expected input shape (channels, length) but got shape {example.shape}") + if self.feature_size == 1 and example.ndim != 1: + raise ValueError(f"Expected mono audio but example has {example.shape[-1]} channels") + if self.feature_size == 2 and example.shape[-1] != 2: + raise ValueError(f"Expected stereo audio but example has {example.shape[-1]} channels") + + padded_inputs = None + input_values = BatchFeature({"input_values": raw_audio}) + if self.chunk_stride is not None and self.chunk_length is not None and max_length is None: + if truncation: + max_length = min(array.shape[0] for array in raw_audio) + nb_step = int(np.floor(max_length / self.chunk_stride)) + max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length + elif padding: + max_length = max(array.shape[0] for array in raw_audio) + nb_step = int(np.ceil(max_length / self.chunk_stride)) + max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length + padding = "max_length" + else: + padded_inputs = input_values + + # normal padding on batch + if padded_inputs is None: + padded_inputs = self.pad( + input_values, + max_length=max_length, + truncation=truncation, + padding=padding, + return_attention_mask=padding, + ) + if padding: + padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask") + + input_values = [] + for example in padded_inputs.pop("input_values"): + if self.feature_size == 1: + example = example[..., None] + input_values.append(example.T) + + padded_inputs["input_values"] = input_values + if return_tensors is not None: + padded_inputs = padded_inputs.convert_to_tensors(return_tensors) + + return padded_inputs + + +__all__ = ["EncodecFeatureExtractor"] diff --git a/third_party/transformers/src/transformers/models/encodec/modeling_encodec.py b/third_party/transformers/src/transformers/models/encodec/modeling_encodec.py new file mode 100644 index 0000000000000000000000000000000000000000..352a1e94006ccc357ecedb1a532ff0c001c90f3b --- /dev/null +++ b/third_party/transformers/src/transformers/models/encodec/modeling_encodec.py @@ -0,0 +1,822 @@ +# Copyright 2023 Meta Platforms, Inc. and affiliates, 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 EnCodec model.""" + +import math +from dataclasses import dataclass + +import torch +from torch import nn + +from ... import initialization as init +from ...modeling_utils import PreTrainedAudioTokenizerBase +from ...utils import ( + ModelOutput, + auto_docstring, + logging, +) +from .configuration_encodec import EncodecConfig + + +logger = logging.get_logger(__name__) + + +# General docstring + + +@dataclass +@auto_docstring +class EncodecOutput(ModelOutput): + r""" + audio_codes (`torch.LongTensor` of shape `(nb_frames, batch_size, nb_quantizers, frame_len)`, *optional*): + Discrete code embeddings computed using `model.encode`. + audio_values (`torch.FloatTensor` of shape `(batch_size, segment_length)`, *optional*): + Decoded audio values, obtained using the decoder part of Encodec. + """ + + audio_codes: torch.LongTensor | None = None + audio_values: torch.FloatTensor | None = None + + +@dataclass +@auto_docstring +class EncodecEncoderOutput(ModelOutput): + r""" + audio_codes (`torch.LongTensor` of shape `(nb_frames, batch_size, nb_quantizers, frame_len)`, *optional*): + Discrete code embeddings computed using `model.encode`. + audio_scales (list of length `nb_frames` of `torch.Tensor` of shape `(batch_size, 1)`, *optional*): + Scaling factor for each `audio_codes` input. This is used to unscale each chunk of audio when decoding. + last_frame_pad_length (`int`, *optional*): + The length of the padding in the last frame, if any. This is used to ensure that the encoded frames can be + outputted as a tensor. This value should be passed during decoding to ensure padding is removed from the + encoded frames. + """ + + audio_codes: torch.LongTensor | None = None + audio_scales: torch.FloatTensor | None = None + last_frame_pad_length: int | None = None + + +@dataclass +@auto_docstring +class EncodecDecoderOutput(ModelOutput): + r""" + audio_values (`torch.FloatTensor` of shape `(batch_size, segment_length)`, *optional*): + Decoded audio values, obtained using the decoder part of Encodec. + """ + + audio_values: torch.FloatTensor | None = None + + +class EncodecConv1d(nn.Module): + """Conv1d with asymmetric or causal padding and normalization.""" + + def __init__( + self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1 + ): + super().__init__() + self.causal = config.use_causal_conv + self.pad_mode = config.pad_mode + self.norm_type = config.norm_type + + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + # warn user on unusual setup between dilation and stride + if stride > 1 and dilation > 1: + logger.warning( + "EncodecConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation) + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + if self.norm_type == "weight_norm": + self.conv = weight_norm(self.conv) + elif self.norm_type == "time_group_norm": + self.norm = nn.GroupNorm(1, out_channels) + + kernel_size = self.conv.kernel_size[0] + stride = torch.tensor(self.conv.stride[0], dtype=torch.int64) + dilation = self.conv.dilation[0] + + # Effective kernel size with dilations. + kernel_size = torch.tensor((kernel_size - 1) * dilation + 1, dtype=torch.int64) + + self.register_buffer("stride", stride, persistent=False) + self.register_buffer("kernel_size", kernel_size, persistent=False) + self.register_buffer("padding_total", kernel_size - stride, persistent=False) + + def _get_extra_padding_for_conv1d( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """See `pad_for_conv1d`.""" + length = hidden_states.shape[-1] + n_frames = (length - self.kernel_size + self.padding_total) / self.stride + 1 + n_frames = torch.ceil(n_frames).to(torch.int64) - 1 + ideal_length = n_frames * self.stride + self.kernel_size - self.padding_total + + return ideal_length - length + + @staticmethod + def _pad1d(hidden_states: torch.Tensor, paddings: tuple[int, int], mode: str = "zero", value: float = 0.0): + """Tiny wrapper around torch.nn.functional.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happens. + """ + length = hidden_states.shape[-1] + padding_left, padding_right = paddings + if mode != "reflect": + return nn.functional.pad(hidden_states, paddings, mode, value) + + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + hidden_states = nn.functional.pad(hidden_states, (0, extra_pad)) + padded = nn.functional.pad(hidden_states, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + + def forward(self, hidden_states): + extra_padding = self._get_extra_padding_for_conv1d(hidden_states) + + if self.causal: + # Left padding for causal + hidden_states = self._pad1d(hidden_states, (self.padding_total, extra_padding), mode=self.pad_mode) + else: + # Asymmetric padding required for odd strides + padding_right = self.padding_total // 2 + padding_left = self.padding_total - padding_right + hidden_states = self._pad1d( + hidden_states, (padding_left, padding_right + extra_padding), mode=self.pad_mode + ) + + hidden_states = self.conv(hidden_states) + + if self.norm_type == "time_group_norm": + hidden_states = self.norm(hidden_states) + + return hidden_states + + +class EncodecConvTranspose1d(nn.Module): + """ConvTranspose1d with asymmetric or causal padding and normalization.""" + + def __init__(self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1): + super().__init__() + self.causal = config.use_causal_conv + self.trim_right_ratio = config.trim_right_ratio + self.norm_type = config.norm_type + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride) + + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + if config.norm_type == "weight_norm": + self.conv = weight_norm(self.conv) + elif config.norm_type == "time_group_norm": + self.norm = nn.GroupNorm(1, out_channels) + + if not (self.causal or self.trim_right_ratio == 1.0): + raise ValueError("`trim_right_ratio` != 1.0 only makes sense for causal convolutions") + + def forward(self, hidden_states): + kernel_size = self.conv.kernel_size[0] + stride = self.conv.stride[0] + padding_total = kernel_size - stride + + hidden_states = self.conv(hidden_states) + + if self.norm_type == "time_group_norm": + hidden_states = self.norm(hidden_states) + + # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be + # removed at the very end, when keeping only the right length for the output, + # as removing it here would require also passing the length at the matching layer + # in the encoder. + if self.causal: + # Trim the padding on the right according to the specified ratio + # if trim_right_ratio = 1.0, trim everything from right + padding_right = math.ceil(padding_total * self.trim_right_ratio) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + + padding_left = padding_total - padding_right + + # unpad + end = hidden_states.shape[-1] - padding_right + hidden_states = hidden_states[..., padding_left:end] + return hidden_states + + +class EncodecLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. Expects input as convolutional layout. + """ + + def __init__(self, config: EncodecConfig, dimension: int): + super().__init__() + self.lstm = nn.LSTM(dimension, dimension, config.num_lstm_layers) + + def forward(self, hidden_states): + hidden_states = hidden_states.permute(2, 0, 1) + hidden_states = self.lstm(hidden_states)[0] + hidden_states + hidden_states = hidden_states.permute(1, 2, 0) + return hidden_states + + +class EncodecResnetBlock(nn.Module): + """ + Residual block from SEANet model as used by EnCodec. + """ + + def __init__(self, config: EncodecConfig, dim: int, dilations: list[int]): + super().__init__() + kernel_sizes = (config.residual_kernel_size, 1) + if len(kernel_sizes) != len(dilations): + raise ValueError("Number of kernel sizes should match number of dilations") + + hidden = dim // config.compress + block = [] + for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)): + in_chs = dim if i == 0 else hidden + out_chs = dim if i == len(kernel_sizes) - 1 else hidden + block += [nn.ELU()] + block += [EncodecConv1d(config, in_chs, out_chs, kernel_size, dilation=dilation)] + self.block = nn.ModuleList(block) + + if config.use_conv_shortcut: + self.shortcut = EncodecConv1d(config, dim, dim, kernel_size=1) + else: + self.shortcut = nn.Identity() + + def forward(self, hidden_states): + residual = hidden_states + for layer in self.block: + hidden_states = layer(hidden_states) + + return self.shortcut(residual) + hidden_states + + +class EncodecEncoder(nn.Module): + """SEANet encoder as used by EnCodec.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + model = [EncodecConv1d(config, config.audio_channels, config.num_filters, config.kernel_size)] + scaling = 1 + + # Downsample to raw audio scale + for ratio in reversed(config.upsampling_ratios): + current_scale = scaling * config.num_filters + # Add residual layers + for j in range(config.num_residual_layers): + model += [EncodecResnetBlock(config, current_scale, [config.dilation_growth_rate**j, 1])] + # Add downsampling layers + model += [nn.ELU()] + model += [EncodecConv1d(config, current_scale, current_scale * 2, kernel_size=ratio * 2, stride=ratio)] + scaling *= 2 + + model += [EncodecLSTM(config, scaling * config.num_filters)] + model += [nn.ELU()] + model += [EncodecConv1d(config, scaling * config.num_filters, config.hidden_size, config.last_kernel_size)] + + self.layers = nn.ModuleList(model) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class EncodecDecoder(nn.Module): + """SEANet decoder as used by EnCodec.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + scaling = int(2 ** len(config.upsampling_ratios)) + model = [EncodecConv1d(config, config.hidden_size, scaling * config.num_filters, config.kernel_size)] + + model += [EncodecLSTM(config, scaling * config.num_filters)] + + # Upsample to raw audio scale + for ratio in config.upsampling_ratios: + current_scale = scaling * config.num_filters + # Add upsampling layers + model += [nn.ELU()] + model += [ + EncodecConvTranspose1d(config, current_scale, current_scale // 2, kernel_size=ratio * 2, stride=ratio) + ] + # Add residual layers + for j in range(config.num_residual_layers): + model += [EncodecResnetBlock(config, current_scale // 2, (config.dilation_growth_rate**j, 1))] + scaling //= 2 + + # Add final layers + model += [nn.ELU()] + model += [EncodecConv1d(config, config.num_filters, config.audio_channels, config.last_kernel_size)] + self.layers = nn.ModuleList(model) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class EncodecEuclideanCodebook(nn.Module): + """Codebook with Euclidean distance.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + embed = torch.zeros(config.codebook_size, config.codebook_dim) + + self.codebook_size = config.codebook_size + + self.register_buffer("inited", torch.Tensor([True])) + self.register_buffer("cluster_size", torch.zeros(config.codebook_size)) + self.register_buffer("embed", embed) + self.register_buffer("embed_avg", embed.clone()) + + def quantize(self, hidden_states): + embed = self.embed.t() + scaled_states = hidden_states.pow(2).sum(1, keepdim=True) + dist = -(scaled_states - 2 * hidden_states @ embed + embed.pow(2).sum(0, keepdim=True)) + embed_ind = dist.max(dim=-1).indices + return embed_ind + + def encode(self, hidden_states): + shape = hidden_states.shape + # pre-process + hidden_states = hidden_states.reshape((-1, shape[-1])) + # quantize + embed_ind = self.quantize(hidden_states) + # post-process + embed_ind = embed_ind.view(*shape[:-1]) + return embed_ind + + def decode(self, embed_ind): + quantize = nn.functional.embedding(embed_ind, self.embed) + return quantize + + +class EncodecVectorQuantization(nn.Module): + """ + Vector quantization implementation. Currently supports only euclidean distance. + """ + + def __init__(self, config: EncodecConfig): + super().__init__() + self.codebook = EncodecEuclideanCodebook(config) + + def encode(self, hidden_states): + hidden_states = hidden_states.permute(0, 2, 1) + embed_in = self.codebook.encode(hidden_states) + return embed_in + + def decode(self, embed_ind): + quantize = self.codebook.decode(embed_ind) + quantize = quantize.permute(0, 2, 1) + return quantize + + +class EncodecResidualVectorQuantizer(nn.Module): + """Residual Vector Quantizer.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + self.codebook_size = config.codebook_size + self.frame_rate = config.frame_rate + self.num_quantizers = config.num_quantizers + self.layers = nn.ModuleList([EncodecVectorQuantization(config) for _ in range(config.num_quantizers)]) + + def get_num_quantizers_for_bandwidth(self, bandwidth: float | None = None) -> int: + """Return num_quantizers based on specified target bandwidth.""" + bw_per_q = math.log2(self.codebook_size) * self.frame_rate + num_quantizers = self.num_quantizers + if bandwidth is not None and bandwidth > 0.0: + num_quantizers = int(max(1, math.floor(bandwidth * 1000 / bw_per_q))) + return num_quantizers + + def encode(self, embeddings: torch.Tensor, bandwidth: float | None = None) -> torch.Tensor: + """ + Encode a given input tensor with the specified frame rate at the given bandwidth. The RVQ encode method sets + the appropriate number of quantizers to use and returns indices for each quantizer. + """ + num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth) + residual = embeddings + all_indices = [] + for layer in self.layers[:num_quantizers]: + indices = layer.encode(residual) + quantized = layer.decode(indices) + residual = residual - quantized + all_indices.append(indices) + out_indices = torch.stack(all_indices) + return out_indices + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation.""" + quantized_out = torch.tensor(0.0, device=codes.device) + for i, indices in enumerate(codes): + layer = self.layers[i] + quantized = layer.decode(indices) + quantized_out = quantized_out + quantized + return quantized_out + + +@auto_docstring +class EncodecPreTrainedModel(PreTrainedAudioTokenizerBase): + config: EncodecConfig + base_model_prefix = "encodec" + main_input_name = "input_values" + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.GroupNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + elif isinstance(module, nn.Conv1d): + init.kaiming_normal_(module.weight) + if module.bias is not None: + k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) + init.uniform_(module.bias, a=-k, b=k) + elif isinstance(module, nn.ConvTranspose1d): + module.reset_parameters() + elif isinstance(module, nn.LSTM): + for name, param in module.named_parameters(): + if "weight" in name: + init.xavier_uniform_(param) + elif "bias" in name: + init.constant_(param, 0.0) + elif isinstance(module, EncodecConv1d): + kernel_size = module.conv.kernel_size[0] + stride = torch.tensor(module.conv.stride[0], dtype=torch.int64) + dilation = module.conv.dilation[0] + # Effective kernel size with dilations. + kernel_size = torch.tensor((kernel_size - 1) * dilation + 1, dtype=torch.int64) + init.copy_(module.stride, stride) + init.copy_(module.kernel_size, kernel_size) + init.copy_(module.padding_total, kernel_size - stride) + elif isinstance(module, EncodecEuclideanCodebook): + init.copy_(module.inited, torch.Tensor([True])) + init.zeros_(module.cluster_size) + init.zeros_(module.embed) + init.zeros_(module.embed_avg) + + +@auto_docstring( + custom_intro=""" + The EnCodec neural audio codec model. + """ +) +class EncodecModel(EncodecPreTrainedModel): + def __init__(self, config: EncodecConfig): + super().__init__(config) + self.config = config + + self.encoder = EncodecEncoder(config) + self.decoder = EncodecDecoder(config) + + self.quantizer = EncodecResidualVectorQuantizer(config) + + self.bits_per_codebook = int(math.log2(self.config.codebook_size)) + if 2**self.bits_per_codebook != self.config.codebook_size: + raise ValueError("The codebook_size must be a power of 2.") + + # Initialize weights and apply final processing + self.post_init() + + def _encode_frame(self, input_values: torch.Tensor, bandwidth: float) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Encodes the given input using the underlying VQVAE. If `config.normalize` is set to `True` the input is first + normalized. The padding mask is required to compute the correct scale. + """ + length = input_values.shape[-1] + duration = length / self.config.sampling_rate + + if self.config.chunk_length_s is not None and duration > 1e-5 + self.config.chunk_length_s: + raise RuntimeError(f"Duration of frame ({duration}) is longer than chunk {self.config.chunk_length_s}") + + scale = None + if self.config.normalize: + mono = torch.sum(input_values, 1, keepdim=True) / input_values.shape[1] + scale = mono.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-8 + input_values = input_values / scale + scale = scale.view(-1, 1) + + embeddings = self.encoder(input_values) + codes = self.quantizer.encode(embeddings, bandwidth) + codes = codes.transpose(0, 1) + return codes, scale + + def encode( + self, + input_values: torch.Tensor, + padding_mask: torch.Tensor | None = None, + bandwidth: float | None = None, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, int] | EncodecEncoderOutput: + """ + Encodes the input audio waveform into discrete codes of shape + `(nb_frames, batch_size, nb_quantizers, frame_len)`. + + - `nb_frames=1` if `self.config.chunk_length=None` (as the encoder is applied on the full audio), which is the + case for the 24kHz model. Otherwise, `nb_frames=ceil(input_length/self.config.chunk_stride)`, which is the case + for the 48kHz model. + - `frame_len` is the length of each frame, which is equal to `ceil(input_length/self.config.hop_length)` if + `self.config.chunk_length=None` (e.g., for the 24kHz model). Otherwise, if `self.config.chunk_length` is + defined, `frame_len=self.config.chunk_length/self.config.hop_length`, e.g., the case for the 48kHz model with + `frame_len=150`. + + Args: + input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Float values of the input audio waveform. + padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Padding mask used to pad the `input_values`. + bandwidth (`float`, *optional*): + The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible + bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented + as bandwidth == 6.0 + + Returns: + EncodecEncoderOutput dict or a tuple containing: + - audio_codes (`torch.LongTensor` of shape `(nb_frames, batch_size, nb_quantizers, frame_len)`, *optional*), + - audio_scales (list of length `nb_frames` of `torch.Tensor` of shape `(batch_size, 1)`, *optional*), + - last_frame_pad_length (`int`, *optional*). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if bandwidth is None: + bandwidth = self.config.target_bandwidths[0] + if bandwidth not in self.config.target_bandwidths: + raise ValueError( + f"This model doesn't support the bandwidth {bandwidth}. Select one of {self.config.target_bandwidths}." + ) + + _, channels, input_length = input_values.shape + + if channels < 1 or channels > 2: + raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}") + + chunk_length = self.config.chunk_length + if chunk_length is None: + chunk_length = input_length + stride = input_length + else: + stride = self.config.chunk_stride + + if padding_mask is None: + padding_mask = torch.ones_like(input_values).bool() + else: + padding_mask = padding_mask.view(padding_mask.shape[0], -1, padding_mask.shape[-1]) + + encoded_frames = [] + scales = [] + for offset in range(0, input_length, stride): + mask = padding_mask[..., offset : offset + chunk_length].bool() + frame = mask * input_values[..., offset : offset + chunk_length] + encoded_frame, scale = self._encode_frame(frame, bandwidth) + encoded_frames.append(encoded_frame) + scales.append(scale) + + # pad last frame (if necessary) to be able to apply `torch.stack` + last_frame_pad_length = encoded_frames[0].shape[-1] - encoded_frames[-1].shape[-1] + if last_frame_pad_length > 0: + last_frame = nn.functional.pad(encoded_frames[-1], (0, last_frame_pad_length), value=0) + encoded_frames[-1] = last_frame + encoded_frames = torch.stack(encoded_frames) + + if not return_dict: + return (encoded_frames, scales, last_frame_pad_length) + return EncodecEncoderOutput(encoded_frames, scales, last_frame_pad_length) + + @staticmethod + def _linear_overlap_add(frames: list[torch.Tensor], stride: int): + # Generic overlap add, with linear fade-in/fade-out, supporting complex scenario + # e.g., more than 2 frames per position. + # The core idea is to use a weight function that is a triangle, + # with a maximum value at the middle of the chunk. + # We use this weighting when summing the frames, and divide by the sum of weights + # for each positions at the end. Thus: + # - if a frame is the only one to cover a position, the weighting is a no-op. + # - if 2 frames cover a position: + # ... ... + # / \/ \ + # / /\ \ + # S T , i.e. S offset of second frame starts, T end of first frame. + # Then the weight function for each one is: (t - S), (T - t), with `t` a given offset. + # After the final normalization, the weight of the second frame at position `t` is + # (t - S) / (t - S + (T - t)) = (t - S) / (T - S), which is exactly what we want. + # + # - if more than 2 frames overlap at a given point, we hope that by induction + # something sensible happens. + if len(frames) == 0: + raise ValueError("`frames` cannot be an empty list.") + + device = frames[0].device + dtype = frames[0].dtype + shape = frames[0].shape[:-1] + total_size = stride * (len(frames) - 1) + frames[-1].shape[-1] + + frame_length = frames[0].shape[-1] + time_vec = torch.linspace(0, 1, frame_length + 2, device=device, dtype=dtype)[1:-1] + weight = 0.5 - (time_vec - 0.5).abs() + + sum_weight = torch.zeros(total_size, device=device, dtype=dtype) + out = torch.zeros(*shape, total_size, device=device, dtype=dtype) + offset: int = 0 + + for frame in frames: + frame_length = frame.shape[-1] + out[..., offset : offset + frame_length] += weight[:frame_length] * frame + sum_weight[offset : offset + frame_length] += weight[:frame_length] + offset += stride + + if sum_weight.min() == 0: + raise ValueError(f"`sum_weight` minimum element must be bigger than zero: {sum_weight}`") + + return out / sum_weight + + def _decode_frame(self, codes: torch.Tensor, scale: torch.Tensor | None = None) -> torch.Tensor: + codes = codes.transpose(0, 1) + embeddings = self.quantizer.decode(codes) + outputs = self.decoder(embeddings) + if scale is not None: + outputs = outputs * scale.view(-1, 1, 1) + return outputs + + def decode( + self, + audio_codes: torch.LongTensor, + audio_scales: torch.Tensor, + padding_mask: torch.Tensor | None = None, + return_dict: bool | None = None, + last_frame_pad_length: int | None = 0, + ) -> tuple[torch.Tensor, torch.Tensor] | EncodecDecoderOutput: + """ + Decodes the given frames into an output audio waveform. + + Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be + trimmed. + + Args: + audio_codes (`torch.LongTensor` of shape `(nb_frames, batch_size, nb_quantizers, frame_len)`, *optional*): + Discrete code embeddings computed using `model.encode`. + audio_scales (list of length `nb_frames` of `torch.Tensor` of shape `(batch_size, 1)`, *optional*): + Scaling factor for each `audio_codes` input. + padding_mask (`torch.Tensor` of shape `(channels, sequence_length)`): + Padding mask used to pad the `input_values`. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + last_frame_pad_length (`int`, *optional*): + Integer representing the length of the padding in the last frame, which is removed during decoding. + + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + chunk_length = self.config.chunk_length + if chunk_length is None: + if len(audio_codes) != 1: + raise ValueError(f"Expected one frame, got {len(audio_codes)}") + frame = audio_codes[0] + if last_frame_pad_length > 0: + frame = frame[..., :-last_frame_pad_length] + audio_values = self._decode_frame(frame, audio_scales[0]) + else: + decoded_frames = [] + for i, (frame, scale) in enumerate(zip(audio_codes, audio_scales)): + if i == len(audio_codes) - 1 and last_frame_pad_length > 0: + frame = frame[..., :-last_frame_pad_length] + frames = self._decode_frame(frame, scale) + decoded_frames.append(frames) + + audio_values = self._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1) + + # truncate based on padding mask + if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]: + audio_values = audio_values[..., : padding_mask.shape[-1]] + + if not return_dict: + return (audio_values,) + return EncodecDecoderOutput(audio_values) + + @auto_docstring + def forward( + self, + input_values: torch.FloatTensor, + padding_mask: torch.BoolTensor | None = None, + bandwidth: float | None = None, + audio_codes: torch.LongTensor | None = None, + audio_scales: torch.Tensor | None = None, + return_dict: bool | None = None, + last_frame_pad_length: int | None = 0, + ) -> tuple[torch.Tensor, torch.Tensor] | EncodecOutput: + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Raw audio input converted to Float and padded to the appropriate length in order to be encoded using chunks + of length self.chunk_length and a stride of `config.chunk_stride`. + padding_mask (`torch.BoolTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Mask to avoid computing scaling factors on padding token indices (can we avoid computing conv on these+). + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + + + `padding_mask` should always be passed, unless the input was truncated or not padded. This is because in + order to process tensors effectively, the input audio should be padded so that `input_length % stride = + step` with `step = chunk_length-stride`. This ensures that all chunks are of the same shape + + + bandwidth (`float`, *optional*): + The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible + bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented as + `bandwidth == 6.0` + audio_codes (`torch.LongTensor` of shape `(nb_frames, batch_size, nb_quantizers, frame_len)`, *optional*): + Discrete code embeddings computed using `model.encode`. + audio_scales (list of length `nb_frames` of `torch.Tensor` of shape `(batch_size, 1)`, *optional*): + Scaling factor for each `audio_codes` input. + return_dict (`bool`, *optional*): + Whether to return outputs as a dict. + last_frame_pad_length (`int`, *optional*): + The length of the padding in the last frame, if any. This is used to ensure that the encoded frames can be + outputted as a tensor. This value should be passed during decoding to ensure padding is removed from the + encoded frames. + + Examples: + + ```python + >>> from datasets import load_dataset + >>> from transformers import AutoProcessor, EncodecModel + + >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example") + >>> audio_sample = dataset["train"]["audio"][0]["array"] + + >>> model_id = "facebook/encodec_24khz" + >>> model = EncodecModel.from_pretrained(model_id) + >>> processor = AutoProcessor.from_pretrained(model_id) + + >>> inputs = processor(raw_audio=audio_sample, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> audio_codes = outputs.audio_codes + >>> audio_values = outputs.audio_values + ```""" + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if padding_mask is None: + padding_mask = torch.ones_like(input_values).bool() + else: + # ensure that channel dimension is present + padding_mask = padding_mask.view(padding_mask.shape[0], -1, padding_mask.shape[-1]) + + if audio_codes is not None and audio_scales is None: + raise ValueError("You specified `audio_codes` but did not specify the `audio_scales`") + + if audio_scales is not None and audio_codes is None: + raise ValueError("You specified `audio_scales` but did not specify the `audio_codes`") + + if audio_scales is None and audio_codes is None: + audio_codes, audio_scales, last_frame_pad_length = self.encode( + input_values, padding_mask, bandwidth, False + ) + + audio_values = self.decode( + audio_codes, + audio_scales, + padding_mask, + return_dict=return_dict, + last_frame_pad_length=last_frame_pad_length, + )[0] + if not return_dict: + return (audio_codes, audio_values) + + return EncodecOutput(audio_codes=audio_codes, audio_values=audio_values) + + +__all__ = ["EncodecModel", "EncodecPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/herbert/__init__.py b/third_party/transformers/src/transformers/models/herbert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..77c7ec616f1f0f5c85812bba8fe1c696dc0c3b9a --- /dev/null +++ b/third_party/transformers/src/transformers/models/herbert/__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_herbert 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/herbert/tokenization_herbert.py b/third_party/transformers/src/transformers/models/herbert/tokenization_herbert.py new file mode 100644 index 0000000000000000000000000000000000000000..eb0543187b62851a4cee148169cb0048f8466739 --- /dev/null +++ b/third_party/transformers/src/transformers/models/herbert/tokenization_herbert.py @@ -0,0 +1,111 @@ +# Copyright 2020 The Google AI Language Team Authors, Allegro.pl, 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. + + +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import BPE + +from ...tokenization_utils_tokenizers import TokenizersBackend +from ...utils import logging + + +logger = logging.get_logger(__name__) + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt"} + + +class HerbertTokenizer(TokenizersBackend): + """ + Construct a BPE tokenizer for HerBERT (backed by HuggingFace's tokenizers library). + + Peculiarities: + + - uses BERT's pre-tokenizer: BertPreTokenizer splits tokens on spaces, and also on punctuation. Each occurrence of + a punctuation character will be treated separately. + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the methods. Users should refer to the + superclass for more information regarding methods. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + merges_file (`str`): + Path to the merges file. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. + pad_token (`str`, *optional*, defaults to `""`): + The padding token. + mask_token (`str`, *optional*, defaults to `""`): + The mask token. + sep_token (`str`, *optional*, defaults to `""`): + The separator token. + vocab (`str`, `dict` or `list`, *optional*): + Custom vocabulary dictionary. + merges (`str` or `list[str]`, *optional*): + Custom merges list. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = BPE + + def __init__( + self, + vocab: str | dict[str, int] | None = None, + merges: str | list[str] | None = None, + cls_token: str = "", + unk_token: str = "", + pad_token: str = "", + mask_token: str = "", + sep_token: str = "", + vocab_file: str | None = None, + merges_file: str | None = None, + **kwargs, + ): + self._vocab = vocab if vocab is not None else {str(unk_token): 0} + self._merges = merges or [] + self._tokenizer = Tokenizer( + BPE( + vocab=self._vocab, + merges=self._merges, + dropout=None, + unk_token=str(unk_token), + end_of_word_suffix="", + ) + ) + + self._tokenizer.normalizer = normalizers.BertNormalizer( + lowercase=False, strip_accents=False, clean_text=True, handle_chinese_chars=True + ) + self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + self._tokenizer.decoder = decoders.BPEDecoder(suffix="") + + super().__init__( + cls_token=cls_token, + unk_token=unk_token, + pad_token=pad_token, + mask_token=mask_token, + sep_token=sep_token, + **kwargs, + ) + + self._tokenizer.post_processor = processors.BertProcessing( + sep=(self.sep_token, 2), + cls=(self.cls_token, 0), + ) + + +__all__ = ["HerbertTokenizer"] diff --git a/third_party/transformers/src/transformers/models/maskformer/__init__.py b/third_party/transformers/src/transformers/models/maskformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..096e81b0bfa4a21c99aa257f40e5b95f15a576ba --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/__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_maskformer import * + from .configuration_maskformer_swin import * + from .feature_extraction_maskformer import * + from .image_processing_maskformer import * + from .image_processing_pil_maskformer import * + from .modeling_maskformer import * + from .modeling_maskformer_swin 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/maskformer/configuration_maskformer.py b/third_party/transformers/src/transformers/models/maskformer/configuration_maskformer.py new file mode 100644 index 0000000000000000000000000000000000000000..abfe6ae0154ed34f6babc93755591180fe2640de --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/configuration_maskformer.py @@ -0,0 +1,226 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/maskformer/modular_maskformer.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_maskformer.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2022 Meta Platforms, Inc.s and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring, logging +from ..auto import CONFIG_MAPPING, AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") +@strict +class MaskFormerDetrConfig(PreTrainedConfig): + r""" + num_queries (`int`, *optional*, defaults to 100): + Number of object queries, i.e. detection slots. This is the maximal number of objects + [`ConditionalDetrModel`] can detect in a single image. For COCO, we recommend 100 queries. + position_embedding_type (`str`, *optional*, defaults to `"sine"`): + Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`. + dilation (`bool`, *optional*, defaults to `False`): + Whether to replace stride with dilation in the last convolutional block (DC5). Only supported when + `use_timm_backbone` = `True`. + + Examples: + + ```python + >>> from transformers import MaskFormerDetrConfig, MaskFormerDetrModel + + >>> # Initializing a MASK_FORMER_DETR facebook/mask_former_detr-resnet-50 style configuration + >>> configuration = MaskFormerDetrConfig() + + >>> # Initializing a model (with random weights) from the facebook/mask_former_detr-resnet-50 style configuration + >>> model = MaskFormerDetrModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "detr" + sub_configs = {"backbone_config": AutoConfig} + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + "num_hidden_layers": "encoder_layers", + } + + backbone_config: dict | PreTrainedConfig | None = None + num_channels: int = 3 + num_queries: int = 100 + encoder_layers: int = 6 + encoder_ffn_dim: int = 2048 + encoder_attention_heads: int = 8 + decoder_layers: int = 6 + decoder_ffn_dim: int = 2048 + decoder_attention_heads: int = 8 + encoder_layerdrop: float | int = 0.0 + decoder_layerdrop: float | int = 0.0 + is_encoder_decoder: bool = True + activation_function: str = "relu" + d_model: int = 256 + dropout: float | int = 0.1 + attention_dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + init_std: float = 0.02 + init_xavier_std: float = 1.0 + auxiliary_loss: bool = False + position_embedding_type: str = "sine" + dilation: bool = False + class_cost: int = 1 + bbox_cost: int = 5 + giou_cost: int = 2 + mask_loss_coefficient: int = 1 + dice_loss_coefficient: int = 1 + bbox_loss_coefficient: int = 5 + giou_loss_coefficient: int = 2 + eos_coefficient: float = 0.1 + + def __post_init__(self, **kwargs): + backbone_kwargs = kwargs.get("backbone_kwargs", {}) + timm_default_kwargs = { + "num_channels": backbone_kwargs.get("num_channels", self.num_channels), + "features_only": True, + "use_pretrained_backbone": False, + "out_indices": backbone_kwargs.get("out_indices", [1, 2, 3, 4]), + } + if self.dilation: + timm_default_kwargs["output_stride"] = backbone_kwargs.get("output_stride", 16) + + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_backbone="resnet50", + default_config_type="resnet", + default_config_kwargs={"out_features": ["stage4"]}, + timm_default_kwargs=timm_default_kwargs, + **kwargs, + ) + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") +@strict +class MaskFormerConfig(PreTrainedConfig): + r""" + fpn_feature_size (`int`, *optional*, defaults to 256): + The Feature Pyramid Network's features size. + mask_feature_size (`int`, *optional*, defaults to 256): + The masks' features size, this value will also be used to specify the Feature Pyramid Network features' + size. + decoder_config (`Dict`, *optional*): + The configuration passed to the transformer decoder model, if unset the base config for `detr-resnet-50` + will be used. + cross_entropy_weight (`float`, *optional*, defaults to 1.0): + The weight for the cross entropy loss. + output_auxiliary_logits (`bool`, *optional*): + Should the model output its `auxiliary_logits` or not. + + Raises: + `ValueError`: + Raised if the backbone model type selected is not in `["swin"]` or the decoder model type selected is not + in `["detr"]` + + Examples: + + ```python + >>> from transformers import MaskFormerConfig, MaskFormerModel + + >>> # Initializing a MaskFormer facebook/maskformer-swin-base-ade configuration + >>> configuration = MaskFormerConfig() + + >>> # Initializing a model (with random weights) from the facebook/maskformer-swin-base-ade style configuration + >>> model = MaskFormerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + + """ + + model_type = "maskformer" + sub_configs = {"backbone_config": AutoConfig, "decoder_config": AutoConfig} + attribute_map = {"hidden_size": "mask_feature_size"} + backbones_supported = ["resnet", "swin"] + decoders_supported = ["detr"] + + fpn_feature_size: int = 256 + mask_feature_size: int = 256 + no_object_weight: float = 0.1 + use_auxiliary_loss: bool = False + backbone_config: dict | PreTrainedConfig | None = None + decoder_config: dict | PreTrainedConfig | None = None + init_std: float = 0.02 + init_xavier_std: float = 1.0 + dice_weight: float = 1.0 + cross_entropy_weight: float = 1.0 + mask_weight: float = 20.0 + output_auxiliary_logits: bool | None = None + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="swin", + default_config_kwargs={ + "depths": [2, 2, 18, 2], + "drop_path_rate": 0.3, + "image_size": 384, + "embed_dim": 128, + "num_heads": [4, 8, 16, 32], + "window_size": 12, + "out_features": ["stage1", "stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + # verify that the backbone is supported + if self.backbone_config is not None and self.backbone_config.model_type not in self.backbones_supported: + logger.warning_once( + f"Backbone {self.backbone_config.model_type} is not a supported model and may not be compatible with MaskFormer. " + f"Supported model types: {','.join(self.backbones_supported)}" + ) + + if self.decoder_config is None: + # fall back to https://huggingface.co/facebook/detr-resnet-50 + self.decoder_config = MaskFormerDetrConfig() + else: + # verify that the decoder is supported + decoder_type = ( + self.decoder_config.pop("model_type") + if isinstance(self.decoder_config, dict) + else self.decoder_config.model_type + ) + if decoder_type not in self.decoders_supported: + raise ValueError( + f"Transformer Decoder {decoder_type} not supported, please use one of" + f" {','.join(self.decoders_supported)}" + ) + if isinstance(self.decoder_config, dict): + config_class = CONFIG_MAPPING[decoder_type] + self.decoder_config = config_class.from_dict(self.decoder_config) + + self.num_attention_heads = self.decoder_config.encoder_attention_heads + self.num_hidden_layers = self.decoder_config.num_hidden_layers + super().__post_init__(**kwargs) + + +__all__ = ["MaskFormerConfig", "MaskFormerDetrConfig"] diff --git a/third_party/transformers/src/transformers/models/maskformer/configuration_maskformer_swin.py b/third_party/transformers/src/transformers/models/maskformer/configuration_maskformer_swin.py new file mode 100644 index 0000000000000000000000000000000000000000..917e0d5cd7b54e5a02755dfdcf536f86c703a708 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/configuration_maskformer_swin.py @@ -0,0 +1,82 @@ +# Copyright 2022 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MaskFormer Swin Transformer model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import BackboneConfigMixin +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/swin-tiny-patch4-window7-224") +@strict +class MaskFormerSwinConfig(BackboneConfigMixin, PreTrainedConfig): + r""" + window_size (`int`, *optional*, defaults to 7): + Size of windows. + + Example: + + ```python + >>> from transformers import MaskFormerSwinConfig, MaskFormerSwinModel + + >>> # Initializing a microsoft/swin-tiny-patch4-window7-224 style configuration + >>> configuration = MaskFormerSwinConfig() + + >>> # Initializing a model (with random weights) from the microsoft/swin-tiny-patch4-window7-224 style configuration + >>> model = MaskFormerSwinModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "maskformer-swin" + + attribute_map = { + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + } + + image_size: int | list[int] | tuple[int, int] = 224 + patch_size: int | list[int] | tuple[int, int] = 4 + num_channels: int = 3 + embed_dim: int = 96 + depths: list[int] | tuple[int, ...] = (2, 2, 6, 2) + num_heads: list[int] | tuple[int, ...] = (3, 6, 12, 24) + window_size: int = 7 + mlp_ratio: float = 4.0 + qkv_bias: bool = True + hidden_dropout_prob: float | int = 0.0 + attention_probs_dropout_prob: float | int = 0.0 + drop_path_rate: float | int = 0.1 + hidden_act: str = "gelu" + use_absolute_embeddings: bool = False + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + _out_features: list[str] | None = None + _out_indices: list[int] | None = None + + def __post_init__(self, **kwargs): + # we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel + # this indicates the channel dimension after the last stage of the model + self.hidden_size = int(self.embed_dim * 2 ** (len(self.depths) - 1)) + self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)] + self.set_output_features_output_indices( + out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None) + ) + super().__post_init__(**kwargs) + + +__all__ = ["MaskFormerSwinConfig"] diff --git a/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..0fa3c564b6ea174be4f16e60384149d1ec6fe2a8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,724 @@ +# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import sys +from argparse import ArgumentParser +from collections.abc import Iterator +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from pprint import pformat +from typing import Any + +import httpx +import torch +import torchvision.transforms as T +from detectron2.checkpoint import DetectionCheckpointer +from detectron2.config import get_cfg +from detectron2.data import MetadataCatalog +from detectron2.projects.deeplab import add_deeplab_config +from PIL import Image +from torch import Tensor, nn + +from transformers.models.maskformer.feature_extraction_maskformer import MaskFormerImageProcessor +from transformers.models.maskformer.modeling_maskformer import ( + MaskFormerConfig, + MaskFormerForInstanceSegmentation, + MaskFormerForInstanceSegmentationOutput, + MaskFormerModel, + MaskFormerModelOutput, +) +from transformers.utils import logging + + +StateDict = dict[str, Tensor] + +logging.set_verbosity_info() +logger = logging.get_logger() + +torch.manual_seed(0) + + +class TrackedStateDict: + def __init__(self, to_track: dict): + """This class "tracks" a python dictionary by keeping track of which item is accessed. + + Args: + to_track (Dict): The dictionary we wish to track + """ + self.to_track = to_track + self._seen: set[str] = set() + + def __getitem__(self, key: str) -> Any: + return self.to_track[key] + + def __setitem__(self, key: str, item: Any): + self._seen.add(key) + self.to_track[key] = item + + def diff(self) -> list[str]: + """This method returns a set difference between the keys in the tracked state dict and the one we have access so far. + This is an effective method to check if we have update all the keys + + Returns: + list[str]: List of keys not yet updated + """ + return set(self.to_track.keys()) - self._seen + + def copy(self) -> dict: + # proxy the call to the internal dictionary + return self.to_track.copy() + + +# We will verify our results on an image of cute cats +def prepare_img(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + with httpx.stream("GET", url) as response: + image = Image.open(BytesIO(response.read())) + return image + + +@dataclass +class Args: + """Fake command line arguments needed by maskformer/detectron implementation""" + + config_file: str + + +def setup_cfg(args: Args): + # load config from file and command-line arguments + cfg = get_cfg() + add_deeplab_config(cfg) + add_mask_former_config(cfg) + cfg.merge_from_file(args.config_file) + cfg.freeze() + return cfg + + +class OriginalMaskFormerConfigToOursConverter: + def __call__(self, original_config: object) -> MaskFormerConfig: + model = original_config.MODEL + mask_former = model.MASK_FORMER + swin = model.SWIN + + dataset_catalog = MetadataCatalog.get(original_config.DATASETS.TEST[0]) + id2label = dict(enumerate(dataset_catalog.stuff_classes)) + label2id = {label: idx for idx, label in id2label.items()} + + config: MaskFormerConfig = MaskFormerConfig( + fpn_feature_size=model.SEM_SEG_HEAD.CONVS_DIM, + mask_feature_size=model.SEM_SEG_HEAD.MASK_DIM, + num_labels=model.SEM_SEG_HEAD.NUM_CLASSES, + no_object_weight=mask_former.NO_OBJECT_WEIGHT, + num_queries=mask_former.NUM_OBJECT_QUERIES, + backbone_config={ + "pretrain_img_size": swin.PRETRAIN_IMG_SIZE, + "image_size": swin.PRETRAIN_IMG_SIZE, + "in_channels": 3, + "patch_size": swin.PATCH_SIZE, + "embed_dim": swin.EMBED_DIM, + "depths": swin.DEPTHS, + "num_heads": swin.NUM_HEADS, + "window_size": swin.WINDOW_SIZE, + "drop_path_rate": swin.DROP_PATH_RATE, + "model_type": "swin", + }, + dice_weight=mask_former.DICE_WEIGHT, + ce_weight=1.0, + mask_weight=mask_former.MASK_WEIGHT, + decoder_config={ + "model_type": "detr", + "max_position_embeddings": 1024, + "encoder_layers": 6, + "encoder_ffn_dim": 2048, + "encoder_attention_heads": 8, + "decoder_layers": mask_former.DEC_LAYERS, + "decoder_ffn_dim": mask_former.DIM_FEEDFORWARD, + "decoder_attention_heads": mask_former.NHEADS, + "encoder_layerdrop": 0.0, + "decoder_layerdrop": 0.0, + "d_model": mask_former.HIDDEN_DIM, + "dropout": mask_former.DROPOUT, + "attention_dropout": 0.0, + "activation_dropout": 0.0, + "init_std": 0.02, + "init_xavier_std": 1.0, + "scale_embedding": False, + "auxiliary_loss": False, + "dilation": False, + # default pretrained config values + }, + id2label=id2label, + label2id=label2id, + ) + + return config + + +class OriginalMaskFormerConfigToImageProcessorConverter: + def __call__(self, original_config: object) -> MaskFormerImageProcessor: + model = original_config.MODEL + model_input = original_config.INPUT + dataset_catalog = MetadataCatalog.get(original_config.DATASETS.TEST[0]) + + return MaskFormerImageProcessor( + 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, + size_divisibility=32, # 32 is required by swin + ) + + +class OriginalMaskFormerCheckpointToOursConverter: + def __init__(self, original_model: nn.Module, config: MaskFormerConfig): + self.original_model = original_model + self.config = config + + def pop_all(self, renamed_keys: list[tuple[str, str]], dst_state_dict: StateDict, src_state_dict: StateDict): + for src_key, dst_key in renamed_keys: + dst_state_dict[dst_key] = src_state_dict.pop(src_key) + + def replace_backbone(self, dst_state_dict: StateDict, src_state_dict: StateDict, config: MaskFormerConfig): + dst_prefix: str = "pixel_level_module.encoder" + src_prefix: str = "backbone" + + renamed_keys = [ + ( + f"{src_prefix}.patch_embed.proj.weight", + f"{dst_prefix}.model.embeddings.patch_embeddings.projection.weight", + ), + (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.model.embeddings.patch_embeddings.projection.bias"), + (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.model.embeddings.norm.weight"), + (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.model.embeddings.norm.bias"), + ] + num_layers = len(config.backbone_config.depths) + for layer_idx in range(num_layers): + for block_idx in range(config.backbone_config.depths[layer_idx]): + renamed_keys.extend( + [ # src, dst + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table", + ), + ] + ) + # now we need to handle the attentions + # read in weights + bias of input projection layer of cross-attention + + src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"] + src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"] + + size = src_att_weight.shape[0] + offset = size // 3 + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight" + ] = src_att_weight[:offset, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias" + ] = src_att_bias[:offset] + + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight" + ] = src_att_weight[offset : offset * 2, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias" + ] = src_att_bias[offset : offset * 2] + + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight" + ] = src_att_weight[-offset:, :] + dst_state_dict[ + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias" + ] = src_att_bias[-offset:] + + # let's pop them + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight") + src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias") + # proj + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias", + ), + ] + ) + + # second norm + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias", + ), + ] + ) + + # mlp + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias", + ), + ] + ) + + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index", + ) + ] + ) + + if layer_idx < num_layers - 1: + # patch merging + renamed_keys.extend( + [ + ( + f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.reduction.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.weight", + ), + ( + f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias", + f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.bias", + ), + ] + ) + + # hidden states norms + renamed_keys.extend( + [ + ( + f"{src_prefix}.norm{layer_idx}.weight", + f"{dst_prefix}.hidden_states_norms.{layer_idx}.weight", + ), + ( + f"{src_prefix}.norm{layer_idx}.bias", + f"{dst_prefix}.hidden_states_norms.{layer_idx}.bias", + ), + ] + ) + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_pixel_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "pixel_level_module.decoder" + src_prefix: str = "sem_seg_head.pixel_decoder" + + self.replace_backbone(dst_state_dict, src_state_dict, self.config) + + def rename_keys_for_conv(detectron_conv: str, mine_conv: str): + return [ + (f"{detectron_conv}.weight", f"{mine_conv}.0.weight"), + # 2 cuz the have act in the middle -> rename it + (f"{detectron_conv}.norm.weight", f"{mine_conv}.1.weight"), + (f"{detectron_conv}.norm.bias", f"{mine_conv}.1.bias"), + ] + + renamed_keys = [ + (f"{src_prefix}.mask_features.weight", f"{dst_prefix}.mask_projection.weight"), + (f"{src_prefix}.mask_features.bias", f"{dst_prefix}.mask_projection.bias"), + # the layers in the original one are in reverse order, stem is the last one! + ] + + renamed_keys.extend(rename_keys_for_conv(f"{src_prefix}.layer_4", f"{dst_prefix}.fpn.stem")) + + # add all the fpn layers (here we need some config parameters to know the size in advance) + for src_i, dst_i in zip(range(3, 0, -1), range(0, 3)): + renamed_keys.extend( + rename_keys_for_conv(f"{src_prefix}.adapter_{src_i}", f"{dst_prefix}.fpn.layers.{dst_i}.proj") + ) + renamed_keys.extend( + rename_keys_for_conv(f"{src_prefix}.layer_{src_i}", f"{dst_prefix}.fpn.layers.{dst_i}.block") + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def rename_keys_in_detr_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder" + src_prefix: str = "sem_seg_head.predictor.transformer.decoder" + # not sure why we are not popping direcetly here! + # here we list all keys to be renamed (original name on the left, our name on the right) + rename_keys = [] + for i in range(self.config.decoder_config.decoder_layers): + # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms + rename_keys.append( + ( + f"{src_prefix}.layers.{i}.self_attn.out_proj.weight", + f"{dst_prefix}.layers.{i}.self_attn.out_proj.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.layers.{i}.self_attn.out_proj.bias", + f"{dst_prefix}.layers.{i}.self_attn.out_proj.bias", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.layers.{i}.multihead_attn.out_proj.weight", + f"{dst_prefix}.layers.{i}.encoder_attn.out_proj.weight", + ) + ) + rename_keys.append( + ( + f"{src_prefix}.layers.{i}.multihead_attn.out_proj.bias", + f"{dst_prefix}.layers.{i}.encoder_attn.out_proj.bias", + ) + ) + rename_keys.append((f"{src_prefix}.layers.{i}.linear1.weight", f"{dst_prefix}.layers.{i}.fc1.weight")) + rename_keys.append((f"{src_prefix}.layers.{i}.linear1.bias", f"{dst_prefix}.layers.{i}.fc1.bias")) + rename_keys.append((f"{src_prefix}.layers.{i}.linear2.weight", f"{dst_prefix}.layers.{i}.fc2.weight")) + rename_keys.append((f"{src_prefix}.layers.{i}.linear2.bias", f"{dst_prefix}.layers.{i}.fc2.bias")) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm1.weight", f"{dst_prefix}.layers.{i}.self_attn_layer_norm.weight") + ) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm1.bias", f"{dst_prefix}.layers.{i}.self_attn_layer_norm.bias") + ) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm2.weight", f"{dst_prefix}.layers.{i}.encoder_attn_layer_norm.weight") + ) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm2.bias", f"{dst_prefix}.layers.{i}.encoder_attn_layer_norm.bias") + ) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm3.weight", f"{dst_prefix}.layers.{i}.final_layer_norm.weight") + ) + rename_keys.append( + (f"{src_prefix}.layers.{i}.norm3.bias", f"{dst_prefix}.layers.{i}.final_layer_norm.bias") + ) + + return rename_keys + + def replace_q_k_v_in_detr_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder" + src_prefix: str = "sem_seg_head.predictor.transformer.decoder" + for i in range(self.config.decoder_config.decoder_layers): + # read in weights + bias of input projection layer of self-attention + in_proj_weight = src_state_dict.pop(f"{src_prefix}.layers.{i}.self_attn.in_proj_weight") + in_proj_bias = src_state_dict.pop(f"{src_prefix}.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}.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :] + dst_state_dict[f"{dst_prefix}.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256] + dst_state_dict[f"{dst_prefix}.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :] + dst_state_dict[f"{dst_prefix}.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512] + dst_state_dict[f"{dst_prefix}.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :] + dst_state_dict[f"{dst_prefix}.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:] + # read in weights + bias of input projection layer of cross-attention + in_proj_weight_cross_attn = src_state_dict.pop(f"{src_prefix}.layers.{i}.multihead_attn.in_proj_weight") + in_proj_bias_cross_attn = src_state_dict.pop(f"{src_prefix}.layers.{i}.multihead_attn.in_proj_bias") + # next, add query, keys and values (in that order) of cross-attention to the state dict + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.q_proj.weight"] = in_proj_weight_cross_attn[:256, :] + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.q_proj.bias"] = in_proj_bias_cross_attn[:256] + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.k_proj.weight"] = in_proj_weight_cross_attn[ + 256:512, : + ] + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.k_proj.bias"] = in_proj_bias_cross_attn[256:512] + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.v_proj.weight"] = in_proj_weight_cross_attn[-256:, :] + dst_state_dict[f"{dst_prefix}.layers.{i}.encoder_attn.v_proj.bias"] = in_proj_bias_cross_attn[-256:] + + def replace_detr_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module.decoder" + src_prefix: str = "sem_seg_head.predictor.transformer.decoder" + renamed_keys = self.rename_keys_in_detr_decoder(dst_state_dict, src_state_dict) + # add more + renamed_keys.extend( + [ + (f"{src_prefix}.norm.weight", f"{dst_prefix}.layernorm.weight"), + (f"{src_prefix}.norm.bias", f"{dst_prefix}.layernorm.bias"), + ] + ) + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + self.replace_q_k_v_in_detr_decoder(dst_state_dict, src_state_dict) + + def replace_transformer_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + dst_prefix: str = "transformer_module" + src_prefix: str = "sem_seg_head.predictor" + + self.replace_detr_decoder(dst_state_dict, src_state_dict) + + renamed_keys = [ + (f"{src_prefix}.query_embed.weight", f"{dst_prefix}.queries_embedder.weight"), + (f"{src_prefix}.input_proj.weight", f"{dst_prefix}.input_projection.weight"), + (f"{src_prefix}.input_proj.bias", f"{dst_prefix}.input_projection.bias"), + ] + + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def replace_instance_segmentation_module(self, dst_state_dict: StateDict, src_state_dict: StateDict): + # NOTE in our case we don't have a prefix, thus we removed the "." from the keys later on! + dst_prefix: str = "" + src_prefix: str = "sem_seg_head.predictor" + + renamed_keys = [ + (f"{src_prefix}.class_embed.weight", f"{dst_prefix}class_predictor.weight"), + (f"{src_prefix}.class_embed.bias", f"{dst_prefix}class_predictor.bias"), + ] + + mlp_len = 3 + for i in range(mlp_len): + renamed_keys.extend( + [ + (f"{src_prefix}.mask_embed.layers.{i}.weight", f"{dst_prefix}mask_embedder.{i}.0.weight"), + (f"{src_prefix}.mask_embed.layers.{i}.bias", f"{dst_prefix}mask_embedder.{i}.0.bias"), + ] + ) + logger.info(f"Replacing keys {pformat(renamed_keys)}") + self.pop_all(renamed_keys, dst_state_dict, src_state_dict) + + def convert(self, mask_former: MaskFormerModel) -> MaskFormerModel: + dst_state_dict = TrackedStateDict(mask_former.state_dict()) + src_state_dict = self.original_model.state_dict() + + self.replace_pixel_module(dst_state_dict, src_state_dict) + self.replace_transformer_module(dst_state_dict, src_state_dict) + + logger.info(f"Missed keys are {pformat(dst_state_dict.diff())}") + logger.info(f"Not copied keys are {pformat(src_state_dict.keys())}") + logger.info("🙌 Done") + + mask_former.load_state_dict(dst_state_dict) + + return mask_former + + def convert_instance_segmentation( + self, mask_former: MaskFormerForInstanceSegmentation + ) -> MaskFormerForInstanceSegmentation: + dst_state_dict = TrackedStateDict(mask_former.state_dict()) + src_state_dict = self.original_model.state_dict() + + self.replace_instance_segmentation_module(dst_state_dict, src_state_dict) + + mask_former.load_state_dict(dst_state_dict) + + return mask_former + + @staticmethod + def using_dirs(checkpoints_dir: Path, config_dir: Path) -> Iterator[tuple[object, Path, Path]]: + checkpoints: list[Path] = checkpoints_dir.glob("**/*.pkl") + + for checkpoint in checkpoints: + logger.info(f"Converting {checkpoint.stem}") + # find associated config file + config: Path = config_dir / checkpoint.parents[0].stem / "swin" / f"{checkpoint.stem}.yaml" + + yield config, checkpoint + + +def test(original_model, our_model: MaskFormerForInstanceSegmentation, image_processor: MaskFormerImageProcessor): + with torch.no_grad(): + original_model = original_model.eval() + our_model = our_model.eval() + + im = prepare_img() + + tr = T.Compose( + [ + T.Resize((384, 384)), + 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) + + original_model_backbone_features = original_model.backbone(x.clone()) + + our_model_output: MaskFormerModelOutput = our_model.model(x.clone(), 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=1e-3), ( + "The backbone features are not the same." + ) + + original_model_pixel_out = original_model.sem_seg_head.pixel_decoder.forward_features( + original_model_backbone_features + ) + + assert torch.allclose( + original_model_pixel_out[0], our_model_output.pixel_decoder_last_hidden_state, atol=1e-4 + ), "The pixel decoder feature are not the same" + + # let's test the full model + original_model_out = original_model([{"image": x.squeeze(0)}]) + + original_segmentation = original_model_out[0]["sem_seg"] + + our_model_out: MaskFormerForInstanceSegmentationOutput = our_model(x) + + our_segmentation = image_processor.post_process_segmentation(our_model_out, target_size=(384, 384)) + + 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 + # model_name_raw is something like maskformer_panoptic_swin_base_IN21k_384_bs64_554k + parent_name: str = checkpoint_file.parents[0].stem + backbone = "swin" + dataset = "" + if "coco" in parent_name: + dataset = "coco" + elif "ade" in parent_name: + dataset = "ade" + else: + raise ValueError(f"{parent_name} must be wrong since we didn't find 'coco' or 'ade' in it ") + + backbone_types = ["tiny", "small", "base", "large"] + + backbone_type = list(filter(lambda x: x in model_name_raw, backbone_types))[0] + + model_name = f"maskformer-{backbone}-{backbone_type}-{dataset}" + + return model_name + + +if __name__ == "__main__": + parser = ArgumentParser( + description="Command line to convert the original maskformers (with swin backbone) to our implementations." + ) + + parser.add_argument( + "--checkpoints_dir", + type=Path, + help=( + "A directory containing the model's checkpoints. The directory has to have the following structure:" + " //.pkl\n" + "Given the files are in the pickle format, please be wary of passing it files you trust." + ), + ) + parser.add_argument( + "--configs_dir", + type=Path, + help=( + "A directory containing the model's configs, see detectron2 doc. The directory has to have the following" + " structure: //.yaml" + ), + ) + parser.add_argument( + "--pytorch_dump_folder_path", + required=True, + type=Path, + help="Path to the folder to output PyTorch models.", + ) + parser.add_argument( + "--maskformer_dir", + required=True, + type=Path, + help=( + "A path to MaskFormer's original implementation directory. You can download from here:" + " https://github.com/facebookresearch/MaskFormer" + ), + ) + + args = parser.parse_args() + + checkpoints_dir: Path = args.checkpoints_dir + config_dir: Path = args.configs_dir + save_directory: Path = args.pytorch_dump_folder_path + maskformer_dir: Path = args.maskformer_dir + # append the path to the parents to maskformer dir + sys.path.append(str(maskformer_dir.parent)) + # and import what's needed + from MaskFormer.mask_former import add_mask_former_config + from MaskFormer.mask_former.mask_former_model import MaskFormer as OriginalMaskFormer + + if not save_directory.exists(): + save_directory.mkdir(parents=True) + + for config_file, checkpoint_file in OriginalMaskFormerCheckpointToOursConverter.using_dirs( + checkpoints_dir, config_dir + ): + image_processor = OriginalMaskFormerConfigToImageProcessorConverter()(setup_cfg(Args(config_file=config_file))) + + original_config = setup_cfg(Args(config_file=config_file)) + mask_former_kwargs = OriginalMaskFormer.from_config(original_config) + + original_model = OriginalMaskFormer(**mask_former_kwargs).eval() + + DetectionCheckpointer(original_model).load(str(checkpoint_file)) + + config: MaskFormerConfig = OriginalMaskFormerConfigToOursConverter()(original_config) + + mask_former = MaskFormerModel(config=config).eval() + + converter = OriginalMaskFormerCheckpointToOursConverter(original_model, config) + + maskformer = converter.convert(mask_former) + + mask_former_for_instance_segmentation = MaskFormerForInstanceSegmentation(config=config).eval() + + mask_former_for_instance_segmentation.model = mask_former + mask_former_for_instance_segmentation = converter.convert_instance_segmentation( + mask_former_for_instance_segmentation + ) + + test(original_model, mask_former_for_instance_segmentation, image_processor) + + model_name = get_name(checkpoint_file) + logger.info(f"Saving {model_name}") + + image_processor.save_pretrained(save_directory / model_name) + mask_former_for_instance_segmentation.save_pretrained(save_directory / model_name) + + image_processor.push_to_hub(repo_id=model_name) + mask_former_for_instance_segmentation.push_to_hub(repo_id=model_name) diff --git a/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_resnet_to_pytorch.py b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_resnet_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..89eb07cdc9c6b8aa9ee3e338f498dd344688c1e3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_resnet_to_pytorch.py @@ -0,0 +1,403 @@ +# 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 MaskFormer checkpoints with ResNet backbone from the original repository. URL: +https://github.com/facebookresearch/MaskFormer""" + +import argparse +import json +import os +import pickle +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 MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, ResNetConfig +from transformers.utils import logging + +from ...utils import strtobool + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_maskformer_config(model_name: str): + if "resnet101c" in model_name: + # TODO add support for ResNet-C backbone, which uses a "deeplab" stem + raise NotImplementedError("To do") + elif "resnet101" in model_name: + backbone_config = ResNetConfig.from_pretrained( + "microsoft/resnet-101", out_features=["stage1", "stage2", "stage3", "stage4"] + ) + else: + backbone_config = ResNetConfig.from_pretrained( + "microsoft/resnet-50", out_features=["stage1", "stage2", "stage3", "stage4"] + ) + config = MaskFormerConfig(backbone_config=backbone_config) + + repo_id = "huggingface/label-files" + if "ade20k-full" in model_name: + config.num_labels = 847 + filename = "maskformer-ade20k-full-id2label.json" + elif "ade" in model_name: + config.num_labels = 150 + filename = "ade20k-id2label.json" + elif "coco-stuff" in model_name: + config.num_labels = 171 + filename = "maskformer-coco-stuff-id2label.json" + elif "coco" in model_name: + # TODO + config.num_labels = 133 + filename = "coco-panoptic-id2label.json" + elif "cityscapes" in model_name: + config.num_labels = 19 + filename = "cityscapes-id2label.json" + elif "vistas" in model_name: + config.num_labels = 65 + filename = "mapillary-vistas-id2label.json" + + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + config.id2label = id2label + config.label2id = {v: k for k, v in id2label.items()} + + return config + + +def create_rename_keys(config): + rename_keys = [] + # stem + # fmt: off + rename_keys.append(("backbone.stem.conv1.weight", "model.pixel_level_module.encoder.embedder.embedder.convolution.weight")) + rename_keys.append(("backbone.stem.conv1.norm.weight", "model.pixel_level_module.encoder.embedder.embedder.normalization.weight")) + rename_keys.append(("backbone.stem.conv1.norm.bias", "model.pixel_level_module.encoder.embedder.embedder.normalization.bias")) + rename_keys.append(("backbone.stem.conv1.norm.running_mean", "model.pixel_level_module.encoder.embedder.embedder.normalization.running_mean")) + rename_keys.append(("backbone.stem.conv1.norm.running_var", "model.pixel_level_module.encoder.embedder.embedder.normalization.running_var")) + # fmt: on + # stages + for stage_idx in range(len(config.backbone_config.depths)): + for layer_idx in range(config.backbone_config.depths[stage_idx]): + # shortcut + if layer_idx == 0: + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.shortcut.weight", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.shortcut.norm.weight", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.shortcut.norm.bias", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.shortcut.norm.running_mean", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.shortcut.norm.running_var", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var", + ) + ) + # 3 convs + for i in range(3): + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.conv{i + 1}.weight", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.conv{i + 1}.norm.weight", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.conv{i + 1}.norm.bias", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.conv{i + 1}.norm.running_mean", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean", + ) + ) + rename_keys.append( + ( + f"backbone.res{stage_idx + 2}.{layer_idx}.conv{i + 1}.norm.running_var", + f"model.pixel_level_module.encoder.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var", + ) + ) + + # FPN + # fmt: off + rename_keys.append(("sem_seg_head.layer_4.weight", "model.pixel_level_module.decoder.fpn.stem.0.weight")) + rename_keys.append(("sem_seg_head.layer_4.norm.weight", "model.pixel_level_module.decoder.fpn.stem.1.weight")) + rename_keys.append(("sem_seg_head.layer_4.norm.bias", "model.pixel_level_module.decoder.fpn.stem.1.bias")) + for source_index, target_index in zip(range(3, 0, -1), range(0, 3)): + rename_keys.append((f"sem_seg_head.adapter_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight")) + rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight")) + rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias")) + rename_keys.append(("sem_seg_head.mask_features.weight", "model.pixel_level_module.decoder.mask_projection.weight")) + rename_keys.append(("sem_seg_head.mask_features.bias", "model.pixel_level_module.decoder.mask_projection.bias")) + # fmt: on + + # Transformer decoder + # fmt: off + for idx in range(config.decoder_config.decoder_layers): + # self-attention out projection + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias")) + # cross-attention out projection + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias")) + # MLP 1 + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight", f"model.transformer_module.decoder.layers.{idx}.fc1.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias", f"model.transformer_module.decoder.layers.{idx}.fc1.bias")) + # MLP 2 + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight", f"model.transformer_module.decoder.layers.{idx}.fc2.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias", f"model.transformer_module.decoder.layers.{idx}.fc2.bias")) + # layernorm 1 (self-attention layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias")) + # layernorm 2 (cross-attention layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias")) + # layernorm 3 (final layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias")) + + rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.weight", "model.transformer_module.decoder.layernorm.weight")) + rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.bias", "model.transformer_module.decoder.layernorm.bias")) + # fmt: on + + # heads on top + # fmt: off + rename_keys.append(("sem_seg_head.predictor.query_embed.weight", "model.transformer_module.queries_embedder.weight")) + + rename_keys.append(("sem_seg_head.predictor.input_proj.weight", "model.transformer_module.input_projection.weight")) + rename_keys.append(("sem_seg_head.predictor.input_proj.bias", "model.transformer_module.input_projection.bias")) + + rename_keys.append(("sem_seg_head.predictor.class_embed.weight", "class_predictor.weight")) + rename_keys.append(("sem_seg_head.predictor.class_embed.bias", "class_predictor.bias")) + + for i in range(3): + rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.weight", f"mask_embedder.{i}.0.weight")) + rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.bias", f"mask_embedder.{i}.0.bias")) + # fmt: on + + return rename_keys + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_decoder_q_k_v(state_dict, config): + # fmt: off + hidden_size = config.decoder_config.hidden_size + for idx in range(config.decoder_config.decoder_layers): + # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight") + in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.bias"] = in_proj_bias[-hidden_size :] + # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight") + in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.bias"] = in_proj_bias[-hidden_size :] + # fmt: on + + +# We will verify our results on an image of cute cats +def prepare_img() -> torch.Tensor: + 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_maskformer_checkpoint( + model_name: str, checkpoint_path: str, pytorch_dump_folder_path: str, push_to_hub: bool = False +): + """ + Copy/paste/tweak model's weights to our MaskFormer structure. + """ + config = get_maskformer_config(model_name) + + if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")): + raise ValueError( + "This part uses `pickle.load` which is insecure and will execute arbitrary code that is potentially " + "malicious. It's recommended to never unpickle data that could have come from an untrusted source, or " + "that could have been tampered with. If you already verified the pickle data and decided to use it, " + "you can set the environment variable `TRUST_REMOTE_CODE` to `True` to allow it." + ) + # load original state_dict + with open(checkpoint_path, "rb") as f: + data = pickle.load(f) + state_dict = data["model"] + + # rename keys + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + read_in_decoder_q_k_v(state_dict, config) + + # update to torch tensors + for key, value in state_dict.items(): + state_dict[key] = torch.from_numpy(value) + + # load 🤗 model + model = MaskFormerForInstanceSegmentation(config) + model.eval() + + model.load_state_dict(state_dict) + + # verify results + image = prepare_img() + if "vistas" in model_name: + ignore_index = 65 + elif "cityscapes" in model_name: + ignore_index = 65535 + else: + ignore_index = 255 + do_reduce_labels = "ade" in model_name + image_processor = MaskFormerImageProcessor(ignore_index=ignore_index, do_reduce_labels=do_reduce_labels) + + inputs = image_processor(image, return_tensors="pt") + + outputs = model(**inputs) + + if model_name == "maskformer-resnet50-ade": + expected_logits = torch.tensor( + [[6.7710, -0.1452, -3.5687], [1.9165, -1.0010, -1.8614], [3.6209, -0.2950, -1.3813]] + ) + elif model_name == "maskformer-resnet101-ade": + expected_logits = torch.tensor( + [[4.0381, -1.1483, -1.9688], [2.7083, -1.9147, -2.2555], [3.4367, -1.3711, -2.1609]] + ) + elif model_name == "maskformer-resnet50-coco-stuff": + expected_logits = torch.tensor( + [[3.2309, -3.0481, -2.8695], [5.4986, -5.4242, -2.4211], [6.2100, -5.2279, -2.7786]] + ) + elif model_name == "maskformer-resnet101-coco-stuff": + expected_logits = torch.tensor( + [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]] + ) + elif model_name == "maskformer-resnet101-cityscapes": + expected_logits = torch.tensor( + [[-1.8861, -1.5465, 0.6749], [-2.3677, -1.6707, -0.0867], [-2.2314, -1.9530, -0.9132]] + ) + elif model_name == "maskformer-resnet50-vistas": + expected_logits = torch.tensor( + [[-6.3917, -1.5216, -1.1392], [-5.5335, -4.5318, -1.8339], [-4.3576, -4.0301, 0.2162]] + ) + elif model_name == "maskformer-resnet50-ade20k-full": + expected_logits = torch.tensor( + [[3.6146, -1.9367, -3.2534], [4.0099, 0.2027, -2.7576], [3.3913, -2.3644, -3.9519]] + ) + elif model_name == "maskformer-resnet101-ade20k-full": + expected_logits = torch.tensor( + [[3.2211, -1.6550, -2.7605], [2.8559, -2.4512, -2.9574], [2.6331, -2.6775, -2.1844]] + ) + + assert torch.allclose(outputs.class_queries_logits[0, :3, :3], expected_logits, atol=1e-4) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + print(f"Saving model and image processor of {model_name} to {pytorch_dump_folder_path}") + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + model.save_pretrained(pytorch_dump_folder_path) + image_processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print(f"Pushing model and image processor of {model_name} to the hub...") + model.push_to_hub(f"facebook/{model_name}") + image_processor.push_to_hub(f"facebook/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="maskformer-resnet50-ade", + type=str, + required=True, + choices=[ + "maskformer-resnet50-ade", + "maskformer-resnet101-ade", + "maskformer-resnet50-coco-stuff", + "maskformer-resnet101-coco-stuff", + "maskformer-resnet101-cityscapes", + "maskformer-resnet50-vistas", + "maskformer-resnet50-ade20k-full", + "maskformer-resnet101-ade20k-full", + ], + help=("Name of the MaskFormer model you'd like to convert",), + ) + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Path to the original pickle file (.pkl) of the original checkpoint.\n" + "Given the files are in the pickle format, please be wary of passing it files you trust.", + ) + 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_maskformer_checkpoint( + args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub + ) diff --git a/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_swin_to_pytorch.py b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_swin_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..9ddec753c83dc6c88b8ff2f55030d46712f8c567 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/convert_maskformer_swin_to_pytorch.py @@ -0,0 +1,346 @@ +# 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 MaskFormer checkpoints with Swin backbone from the original repository. URL: +https://github.com/facebookresearch/MaskFormer""" + +import argparse +import json +import os +import pickle +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 MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig +from transformers.utils import logging + +from ...utils import strtobool + + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +def get_maskformer_config(model_name: str): + backbone_config = SwinConfig.from_pretrained( + "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"] + ) + config = MaskFormerConfig(backbone_config=backbone_config) + + repo_id = "huggingface/label-files" + if "ade20k-full" in model_name: + # this should be ok + config.num_labels = 847 + filename = "maskformer-ade20k-full-id2label.json" + elif "ade" in model_name: + # this should be ok + config.num_labels = 150 + filename = "ade20k-id2label.json" + elif "coco-stuff" in model_name: + # this should be ok + config.num_labels = 171 + filename = "maskformer-coco-stuff-id2label.json" + elif "coco" in model_name: + # TODO + config.num_labels = 133 + filename = "coco-panoptic-id2label.json" + elif "cityscapes" in model_name: + # this should be ok + config.num_labels = 19 + filename = "cityscapes-id2label.json" + elif "vistas" in model_name: + # this should be ok + config.num_labels = 65 + filename = "mapillary-vistas-id2label.json" + + id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r")) + id2label = {int(k): v for k, v in id2label.items()} + + return config + + +def create_rename_keys(config): + rename_keys = [] + # stem + # fmt: off + rename_keys.append(("backbone.patch_embed.proj.weight", "model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight")) + rename_keys.append(("backbone.patch_embed.proj.bias", "model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias")) + rename_keys.append(("backbone.patch_embed.norm.weight", "model.pixel_level_module.encoder.model.embeddings.norm.weight")) + rename_keys.append(("backbone.patch_embed.norm.bias", "model.pixel_level_module.encoder.model.embeddings.norm.bias")) + # stages + for i in range(len(config.backbone_config.depths)): + for j in range(config.backbone_config.depths[i]): + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_index", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight")) + rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias")) + + if i < 3: + rename_keys.append((f"backbone.layers.{i}.downsample.reduction.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight")) + rename_keys.append((f"backbone.layers.{i}.downsample.norm.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight")) + rename_keys.append((f"backbone.layers.{i}.downsample.norm.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias")) + rename_keys.append((f"backbone.norm{i}.weight", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.weight")) + rename_keys.append((f"backbone.norm{i}.bias", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.bias")) + + # FPN + rename_keys.append(("sem_seg_head.layer_4.weight", "model.pixel_level_module.decoder.fpn.stem.0.weight")) + rename_keys.append(("sem_seg_head.layer_4.norm.weight", "model.pixel_level_module.decoder.fpn.stem.1.weight")) + rename_keys.append(("sem_seg_head.layer_4.norm.bias", "model.pixel_level_module.decoder.fpn.stem.1.bias")) + for source_index, target_index in zip(range(3, 0, -1), range(0, 3)): + rename_keys.append((f"sem_seg_head.adapter_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight")) + rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight")) + rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight")) + rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias")) + rename_keys.append(("sem_seg_head.mask_features.weight", "model.pixel_level_module.decoder.mask_projection.weight")) + rename_keys.append(("sem_seg_head.mask_features.bias", "model.pixel_level_module.decoder.mask_projection.bias")) + + # Transformer decoder + for idx in range(config.decoder_config.decoder_layers): + # self-attention out projection + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias")) + # cross-attention out projection + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias")) + # MLP 1 + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight", f"model.transformer_module.decoder.layers.{idx}.fc1.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias", f"model.transformer_module.decoder.layers.{idx}.fc1.bias")) + # MLP 2 + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight", f"model.transformer_module.decoder.layers.{idx}.fc2.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias", f"model.transformer_module.decoder.layers.{idx}.fc2.bias")) + # layernorm 1 (self-attention layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias")) + # layernorm 2 (cross-attention layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias")) + # layernorm 3 (final layernorm) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight")) + rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias")) + + rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.weight", "model.transformer_module.decoder.layernorm.weight")) + rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.bias", "model.transformer_module.decoder.layernorm.bias")) + + # heads on top + rename_keys.append(("sem_seg_head.predictor.query_embed.weight", "model.transformer_module.queries_embedder.weight")) + + rename_keys.append(("sem_seg_head.predictor.input_proj.weight", "model.transformer_module.input_projection.weight")) + rename_keys.append(("sem_seg_head.predictor.input_proj.bias", "model.transformer_module.input_projection.bias")) + + rename_keys.append(("sem_seg_head.predictor.class_embed.weight", "class_predictor.weight")) + rename_keys.append(("sem_seg_head.predictor.class_embed.bias", "class_predictor.bias")) + + for i in range(3): + rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.weight", f"mask_embedder.{i}.0.weight")) + rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.bias", f"mask_embedder.{i}.0.bias")) + # fmt: on + + return rename_keys + + +def rename_key(dct, old, new): + val = dct.pop(old) + dct[new] = val + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_swin_q_k_v(state_dict, backbone_config): + num_features = [int(backbone_config.embed_dim * 2**i) for i in range(len(backbone_config.depths))] + for i in range(len(backbone_config.depths)): + dim = num_features[i] + for j in range(backbone_config.depths[i]): + # fmt: off + # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.weight") + in_proj_bias = state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.query.weight"] = in_proj_weight[:dim, :] + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.query.bias"] = in_proj_bias[: dim] + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.key.weight"] = in_proj_weight[ + dim : dim * 2, : + ] + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.key.bias"] = in_proj_bias[ + dim : dim * 2 + ] + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.value.weight"] = in_proj_weight[ + -dim :, : + ] + state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.value.bias"] = in_proj_bias[-dim :] + # fmt: on + + +# we split up the matrix of each encoder layer into queries, keys and values +def read_in_decoder_q_k_v(state_dict, config): + # fmt: off + hidden_size = config.decoder_config.hidden_size + for idx in range(config.decoder_config.decoder_layers): + # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight") + in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.bias"] = in_proj_bias[-hidden_size :] + # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight") + in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :] + state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.bias"] = in_proj_bias[-hidden_size :] + # fmt: on + + +# We will verify our results on an image of cute cats +def prepare_img() -> torch.Tensor: + 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_maskformer_checkpoint( + model_name: str, checkpoint_path: str, pytorch_dump_folder_path: str, push_to_hub: bool = False +): + """ + Copy/paste/tweak model's weights to our MaskFormer structure. + """ + config = get_maskformer_config(model_name) + + if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")): + raise ValueError( + "This part uses `pickle.load` which is insecure and will execute arbitrary code that is potentially " + "malicious. It's recommended to never unpickle data that could have come from an untrusted source, or " + "that could have been tampered with. If you already verified the pickle data and decided to use it, " + "you can set the environment variable `TRUST_REMOTE_CODE` to `True` to allow it." + ) + # load original state_dict + with open(checkpoint_path, "rb") as f: + data = pickle.load(f) + state_dict = data["model"] + + # for name, param in state_dict.items(): + # print(name, param.shape) + + # rename keys + rename_keys = create_rename_keys(config) + for src, dest in rename_keys: + rename_key(state_dict, src, dest) + read_in_swin_q_k_v(state_dict, config.backbone_config) + read_in_decoder_q_k_v(state_dict, config) + + # update to torch tensors + for key, value in state_dict.items(): + state_dict[key] = torch.from_numpy(value) + + # load 🤗 model + model = MaskFormerForInstanceSegmentation(config) + model.eval() + + for name, param in model.named_parameters(): + print(name, param.shape) + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + assert missing_keys == [ + "model.pixel_level_module.encoder.model.layernorm.weight", + "model.pixel_level_module.encoder.model.layernorm.bias", + ] + assert len(unexpected_keys) == 0, f"Unexpected keys: {unexpected_keys}" + + # verify results + image = prepare_img() + if "vistas" in model_name: + ignore_index = 65 + elif "cityscapes" in model_name: + ignore_index = 65535 + else: + ignore_index = 255 + do_reduce_labels = "ade" in model_name + image_processor = MaskFormerImageProcessor(ignore_index=ignore_index, do_reduce_labels=do_reduce_labels) + + inputs = image_processor(image, return_tensors="pt") + + outputs = model(**inputs) + + print("Logits:", outputs.class_queries_logits[0, :3, :3]) + + if model_name == "maskformer-swin-tiny-ade": + expected_logits = torch.tensor( + [[3.6353, -4.4770, -2.6065], [0.5081, -4.2394, -3.5343], [2.1909, -5.0353, -1.9323]] + ) + assert torch.allclose(outputs.class_queries_logits[0, :3, :3], expected_logits, atol=1e-4) + print("Looks ok!") + + if pytorch_dump_folder_path is not None: + print(f"Saving model and image processor to {pytorch_dump_folder_path}") + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + model.save_pretrained(pytorch_dump_folder_path) + image_processor.save_pretrained(pytorch_dump_folder_path) + + if push_to_hub: + print("Pushing model and image processor to the hub...") + model.push_to_hub(f"nielsr/{model_name}") + image_processor.push_to_hub(f"nielsr/{model_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="maskformer-swin-tiny-ade", + type=str, + help=("Name of the MaskFormer model you'd like to convert",), + ) + parser.add_argument( + "--checkpoint_path", + default="/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl", + type=str, + help="Path to the original state dict (.pth file).\n" + "Given the files are in the pickle format, please be wary of passing it files you trust.", + ) + 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_maskformer_checkpoint( + args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub + ) diff --git a/third_party/transformers/src/transformers/models/maskformer/image_processing_maskformer.py b/third_party/transformers/src/transformers/models/maskformer/image_processing_maskformer.py new file mode 100644 index 0000000000000000000000000000000000000000..9415bcf60e0b657cfef546669054fa131e2885c0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/image_processing_maskformer.py @@ -0,0 +1,806 @@ +# 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 MaskFormer.""" + +import math +from typing import Any, Optional, Union + +import numpy as np +import torch +from torch import nn +from torchvision.transforms.v2 import functional as tvF + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import get_size_with_aspect_ratio, group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size_for_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +# Helper functions for post-processing (PyTorch-based) +def binary_mask_to_rle(mask: "torch.Tensor | np.ndarray") -> list[int]: + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `np.ndarray`): + A binary mask of shape `(height, width)` where 0 denotes background and 1 denotes the target + segment_id or class_id. + Returns: + `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE + format. + """ + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask) + + pixels = mask.flatten() + zero = torch.zeros(1, device=pixels.device, dtype=pixels.dtype) + pixels = torch.cat([zero, pixels, zero]) + runs = torch.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return runs.tolist() + + +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `list[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels): + """ + Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and + `labels`. + + Args: + masks (`torch.Tensor`): + A tensor of shape `(num_queries, height, width)`. + scores (`torch.Tensor`): + A tensor of shape `(num_queries)`. + labels (`torch.Tensor`): + A tensor of shape `(num_queries)`. + object_mask_threshold (`float`): + A number between 0 and 1 used to binarize the masks. + Raises: + `ValueError`: Raised when the first dimension doesn't match in all input tensors. + Returns: + `tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region + < `object_mask_threshold`. + """ + if not (masks.shape[0] == scores.shape[0] == labels.shape[0]): + raise ValueError("mask, scores and labels must have the same shape!") + + to_keep = labels.ne(num_labels) & (scores > object_mask_threshold) + + return masks[to_keep], scores[to_keep], labels[to_keep] + + +def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8): + # Get the mask associated with the k class + mask_k = mask_labels == k + mask_k_area = mask_k.sum() + + # Compute the area of all the stuff in query k + original_area = (mask_probs[k] >= mask_threshold).sum() + mask_exists = mask_k_area > 0 and original_area > 0 + + # Eliminate disconnected tiny segments + if mask_exists: + area_ratio = mask_k_area / original_area + if not area_ratio.item() > overlap_mask_area_threshold: + mask_exists = False + + return mask_exists, mask_k + + +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_size: tuple[int, int] | None = None, +): + height = mask_probs.shape[1] if target_size is None else target_size[0] + width = mask_probs.shape[2] if target_size is None else target_size[1] + + segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = nn.functional.interpolate( + mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False + )[0] + + current_segment_id = 0 + + # Weigh each mask by its prediction score + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + # Keep track of instances of each class + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + # Check if mask exists and large enough to be a segment + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + # Add current object segment to final segmentation map + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +def convert_segmentation_map_to_binary_masks_fast( + segmentation_map: "torch.Tensor", + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, +): + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = torch.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = torch.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] # drop background label if applicable + + binary_masks = [(segmentation_map == i) for i in all_labels] + if binary_masks: + binary_masks = torch.stack(binary_masks, dim=0) + else: + binary_masks = torch.zeros((0, *segmentation_map.shape), device=segmentation_map.device) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = torch.zeros(all_labels.shape[0], device=segmentation_map.device) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(label.item() + 1 if do_reduce_labels else label.item())] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels + return binary_masks.float(), labels.long() + + +class MaskFormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ignore_index (`int`, *optional*): + Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels + denoted with 0 (background) will be replaced with `ignore_index`. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). + The background label will be replaced by `ignore_index`. + num_labels (`int`, *optional*): + The number of labels in the segmentation map. + size_divisor (`int`, *optional*, defaults to `32`): + Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in + Swin Transformer. + pad_size (`SizeDict`, *optional*): + The size to pad the images to. Must be larger than any image size provided for preprocessing. If `pad_size` + is not provided, images will be padded to the largest height and width in the batch. + """ + + ignore_index: int | None + do_reduce_labels: bool + num_labels: int | None + size_divisor: int + pad_size: SizeDict | None + + +@auto_docstring +class MaskFormerImageProcessor(TorchvisionBackend): + valid_kwargs = MaskFormerImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_pad = True + model_input_names = ["pixel_values", "pixel_mask"] + size_divisor = 32 + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[MaskFormerImageProcessorKwargs]) -> None: + size = kwargs.pop("size", None) + max_size = kwargs.pop("max_size", None) + + if size is None and max_size is not None: + size = self.size.copy() + size["longest_edge"] = max_size + elif size is None: + size = self.size + + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + super().__init__(**kwargs) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the + `_max_size` attribute from the dictionary. + """ + image_processor_dict = super().to_dict() + image_processor_dict.pop("_max_size", None) + return image_processor_dict + + def reduce_label(self, labels: list["torch.Tensor"]): + for idx in range(len(labels)): + label = labels[idx] + label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype), label) + label = label - 1 + label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype), label) + labels[idx] = label + + def resize( + self, + image: torch.Tensor, + size: SizeDict, + size_divisor: int = 0, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + **kwargs, + ) -> torch.Tensor: + """ + Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an + int, smaller edge of the image will be matched to this number. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Size of the image's `(height, width)` dimensions after resizing. + size_divisor (`int`, *optional*, defaults to 0): + If `size_divisor` is given, the output image size will be divisible by the number. + resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*): + Resampling filter to use if resizing the image. + """ + + if size.shortest_edge and size.longest_edge: + # Resize the image so that the shortest edge or the longest edge is of the given size + # while maintaining the aspect ratio of the original image. + new_size = get_size_with_aspect_ratio( + image.size()[-2:], + size.shortest_edge, + size.longest_edge, + ) + elif size.max_height and size.max_width: + new_size = get_image_size_for_max_height_width(image.size()[-2:], size.max_height, size.max_width) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError( + f"Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got {size}." + ) + if size_divisor > 0: + height, width = new_size + height = int(math.ceil(height / size_divisor) * size_divisor) + width = int(math.ceil(width / size_divisor) * size_divisor) + new_size = (height, width) + + image = super().resize( + image, size=SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs + ) + return image + + def pad( + self, + images: torch.Tensor, + padded_size: tuple[int, int], + segmentation_maps: torch.Tensor | None = None, + fill: int = 0, + ignore_index: int = 255, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + original_size = images.size()[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + if original_size != padded_size: + padding = [0, 0, padding_right, padding_bottom] + images = tvF.pad(images, padding, fill=fill) + if segmentation_maps is not None: + segmentation_maps = [tvF.pad(mask, padding, fill=ignore_index) for mask in segmentation_maps] + + # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + pixel_mask = torch.zeros((images.shape[0], *padded_size), dtype=torch.int64, device=images.device) + pixel_mask[:, : original_size[0], : original_size[1]] = 1 + + return images, pixel_mask, segmentation_maps + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[MaskFormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + device: Union[str, "torch.device"] | None = None, + **kwargs: Unpack[MaskFormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + segmentation_maps: Optional["torch.Tensor"], + instance_id_to_semantic_id: dict[int, int] | None, + do_resize: bool | None, + size: SizeDict | None, + pad_size: SizeDict | None, + size_divisor: int | None, + resample: Union["PILImageResampling", "tvF.InterpolationMode"] | None, + do_rescale: bool | None, + rescale_factor: float | None, + do_normalize: bool | None, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + from ...image_utils import get_max_height_width + + if segmentation_maps is not None and len(images) != len(segmentation_maps): + raise ValueError("Images and segmentation maps must have the same length.") + + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + if segmentation_maps is not None: + grouped_segmentation_maps, grouped_segmentation_maps_index = group_images_by_shape( + segmentation_maps, disable_grouping=disable_grouping + ) + resized_segmentation_maps_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize( + image=stacked_images, size=size, size_divisor=size_divisor, resample=resample + ) + if segmentation_maps is not None: + stacked_segmentation_maps = grouped_segmentation_maps[shape] + if do_resize: + stacked_segmentation_maps = self.resize( + image=stacked_segmentation_maps, + size=size, + size_divisor=size_divisor, + resample=tvF.InterpolationMode.NEAREST_EXACT, + ) + resized_images_grouped[shape] = stacked_images + if segmentation_maps is not None: + resized_segmentation_maps_grouped[shape] = stacked_segmentation_maps + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + if segmentation_maps is not None: + resized_segmentation_maps = reorder_images( + resized_segmentation_maps_grouped, grouped_segmentation_maps_index + ) + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(resized_images) + + if segmentation_maps is not None: + mask_labels = [] + class_labels = [] + # Convert to list of binary masks and labels + for idx, segmentation_map in enumerate(resized_segmentation_maps): + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + # Use instance2class_id mapping per image + masks, classes = convert_segmentation_map_to_binary_masks_fast( + segmentation_map.squeeze(0), + instance_id, + ignore_index=ignore_index, + do_reduce_labels=do_reduce_labels, + ) + mask_labels.append(masks) + class_labels.append(classes) + + if segmentation_maps is not None: + # group mask_labels as paired inputs and not images so as not to stack them + grouped_images, grouped_segmentation_maps, grouped_images_index = group_images_by_shape( + resized_images, mask_labels, disable_grouping=disable_grouping + ) + processed_segmentation_maps_grouped = {} + else: + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_pixel_masks_grouped = {} + for shape, stacked_images in grouped_images.items(): + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + padded_images, pixel_masks, padded_segmentation_maps = self.pad( + images=stacked_images, + segmentation_maps=grouped_segmentation_maps[shape] if segmentation_maps is not None else None, + padded_size=padded_size, + ignore_index=ignore_index, + ) + processed_images_grouped[shape] = padded_images + processed_pixel_masks_grouped[shape] = pixel_masks + if segmentation_maps is not None: + processed_segmentation_maps_grouped[shape] = padded_segmentation_maps + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + processed_pixel_masks = reorder_images(processed_pixel_masks_grouped, grouped_images_index) + encoded_inputs = BatchFeature( + data={"pixel_values": processed_images, "pixel_mask": processed_pixel_masks}, + tensor_type=return_tensors, + ) + if segmentation_maps is not None: + mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_images_index) + # we cannot batch them since they don't share a common class size + encoded_inputs["mask_labels"] = mask_labels + encoded_inputs["class_labels"] = class_labels + + return encoded_inputs + + # Copied from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.post_process_semantic_segmentation + 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 + + # Copied from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.post_process_instance_segmentation + def post_process_instance_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into instance segmentation predictions. Only + supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`MaskFormerForInstanceSegmentation`]): + Raw outputs of the model. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If left to None, predictions will not be resized. + return_coco_annotation (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE) format. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id`, or + `list[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to + `True`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps is set to `True`. + Set to `None` if no mask if found above `threshold`. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- An integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + if return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros(masks_queries_logits.shape[2:]) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + # 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__ = ["MaskFormerImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/maskformer/image_processing_pil_maskformer.py b/third_party/transformers/src/transformers/models/maskformer/image_processing_pil_maskformer.py new file mode 100644 index 0000000000000000000000000000000000000000..e1217b78d53a6b8f712e3c33cf349f314e27dd49 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/image_processing_pil_maskformer.py @@ -0,0 +1,845 @@ +# 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 MaskFormer.""" + +import math +from typing import Any + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import PaddingMode, get_size_with_aspect_ratio +from ...image_transforms import pad as np_pad +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, + get_image_size_for_max_height_width, + get_max_height_width, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available, logging, requires_backends +from ...utils.import_utils import requires + + +if is_torch_available(): + import torch + from torch import nn + +logger = logging.get_logger(__name__) + + +def convert_segmentation_map_to_binary_masks( + segmentation_map: np.ndarray, + instance_id_to_semantic_id: dict[int, int] | None = None, + ignore_index: int | None = None, + do_reduce_labels: bool = False, +): + """Convert segmentation map to binary masks using NumPy operations.""" + if do_reduce_labels and ignore_index is None: + raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.") + + if do_reduce_labels: + segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1) + + all_labels = np.unique(segmentation_map) + + if ignore_index is not None: + all_labels = all_labels[all_labels != ignore_index] + + binary_masks = [(segmentation_map == i) for i in all_labels] + if binary_masks: + binary_masks = np.stack(binary_masks, axis=0) + else: + binary_masks = np.zeros((0, *segmentation_map.shape), dtype=np.float32) + + # Convert instance ids to class ids + if instance_id_to_semantic_id is not None: + labels = np.zeros(all_labels.shape[0], dtype=np.int64) + + for i, label in enumerate(all_labels): + class_id = instance_id_to_semantic_id[(int(label) + 1 if do_reduce_labels else int(label))] + labels[i] = class_id - 1 if do_reduce_labels else class_id + else: + labels = all_labels.astype(np.int64) + return binary_masks.astype(np.float32), labels + + +# Adapted from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessorKwargs +class MaskFormerImageProcessorKwargs(ImagesKwargs, total=False): + r""" + ignore_index (`int`, *optional*): + Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels + denoted with 0 (background) will be replaced with `ignore_index`. + do_reduce_labels (`bool`, *optional*, defaults to `False`): + Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0 + is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). + The background label will be replaced by `ignore_index`. + num_labels (`int`, *optional*): + The number of labels in the segmentation map. + size_divisor (`int`, *optional*, defaults to `32`): + Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in + Swin Transformer. + pad_size (`SizeDict`, *optional*): + The size to pad the images to. Must be larger than any image size provided for preprocessing. If `pad_size` + is not provided, images will be padded to the largest height and width in the batch. + """ + + ignore_index: int | None + do_reduce_labels: bool + num_labels: int | None + size_divisor: int + pad_size: SizeDict | None + + +# Adapted from transformers.models.maskformer.image_processing_maskformer.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.maskformer.image_processing_maskformer.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.maskformer.image_processing_maskformer.compute_segments +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_size: tuple[int, int] | None = None, +): + height = mask_probs.shape[1] if target_size is None else target_size[0] + width = mask_probs.shape[2] if target_size is None else target_size[1] + + segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = nn.functional.interpolate( + mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False + )[0] + + current_segment_id = 0 + + # Weigh each mask by its prediction score + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + # Keep track of instances of each class + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + # Check if mask exists and large enough to be a segment + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + # Add current object segment to final segmentation map + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +# Adapted from transformers.models.maskformer.image_processing_maskformer.convert_segmentation_to_rle +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `list[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +# Adapted from transformers.models.maskformer.image_processing_maskformer.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 MaskFormerImageProcessorPil(PilBackend): + valid_kwargs = MaskFormerImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + do_resize = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_pad = True + model_input_names = ["pixel_values", "pixel_mask"] + size_divisor = 32 + do_reduce_labels = False + + def __init__(self, **kwargs: Unpack[MaskFormerImageProcessorKwargs]) -> None: + size = kwargs.pop("size", None) + max_size = kwargs.pop("max_size", None) + + # Store max_size as private attribute for backward compatibility + self._max_size = max_size if max_size is not None else 1333 + + if size is None and max_size is not None: + size = self.size.copy() + size["longest_edge"] = max_size + elif size is None: + size = self.size + + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + super().__init__(**kwargs) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the + `_max_size` attribute from the dictionary. + """ + image_processor_dict = super().to_dict() + image_processor_dict.pop("_max_size", None) + return image_processor_dict + + def reduce_label(self, labels: list[np.ndarray]): + """Reduce label values by 1, replacing 0 with 255.""" + for idx in range(len(labels)): + label = labels[idx].copy() + label[label == 0] = 255 + label = label - 1 + label[label == 254] = 255 + labels[idx] = label + + def resize( + self, + image: np.ndarray, + size: SizeDict, + size_divisor: int = 0, + resample: PILImageResampling | None = None, + **kwargs, + ) -> np.ndarray: + """ + Resize the image to the given size with optional size_divisor. + + Args: + image (`np.ndarray`): + Image to resize. + size (`SizeDict`): + Size of the image's `(height, width)` dimensions after resizing. + size_divisor (`int`, *optional*, defaults to 0): + If `size_divisor` is given, the output image size will be divisible by the number. + resample (`PILImageResampling | int | None`, *optional*): + Resampling filter to use if resizing the image. + """ + + if size.shortest_edge and size.longest_edge: + # Get current image size + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + new_size = get_size_with_aspect_ratio((height, width), size.shortest_edge, size.longest_edge) + elif size.max_height and size.max_width: + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + new_size = get_image_size_for_max_height_width((height, width), size.max_height, size.max_width) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError( + f"Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got {size}." + ) + if size_divisor > 0: + height, width = new_size + height = int(math.ceil(height / size_divisor) * size_divisor) + width = int(math.ceil(width / size_divisor) * size_divisor) + new_size = (height, width) + + return super().resize(image, size=SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def pad( + self, + images: list[np.ndarray], + padded_size: tuple[int, int], + segmentation_maps: list[np.ndarray] | None = None, + fill: int = 0, + ignore_index: int = 255, + ) -> tuple[list[np.ndarray], list[np.ndarray], list[np.ndarray] | None]: + """ + Pad images and optionally segmentation maps to the given size. + + Args: + images (`list[np.ndarray]`): + Images to pad. + padded_size (`tuple[int, int]`): + Target size (height, width) to pad to. + segmentation_maps (`list[np.ndarray]`, *optional*): + Segmentation maps to pad. + fill (`int`, *optional*, defaults to 0): + Fill value for images. + ignore_index (`int`, *optional*, defaults to 255): + Fill value for segmentation maps. + + Returns: + `tuple`: (padded_images, pixel_masks, padded_segmentation_maps) + """ + padded_images = [] + pixel_masks = [] + + for image in images: + original_size = image.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + image = np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=fill, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + padded_images.append(image) + + # Make a pixel mask for the image + pixel_mask = np.zeros(padded_size, dtype=np.int64) + pixel_mask[: original_size[0], : original_size[1]] = 1 + pixel_masks.append(pixel_mask) + + padded_segmentation_maps = None + if segmentation_maps is not None: + padded_segmentation_maps = [] + for mask in segmentation_maps: + original_size = mask.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + mask = np_pad( + mask, + padding, + mode=PaddingMode.CONSTANT, + constant_values=ignore_index, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + padded_segmentation_maps.append(mask) + + return padded_images, pixel_masks, padded_segmentation_maps + + @auto_docstring + def preprocess( + self, + images: ImageInput, + segmentation_maps: ImageInput | None = None, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None = None, + **kwargs: Unpack[MaskFormerImageProcessorKwargs], + ) -> BatchFeature: + r""" + segmentation_maps (`ImageInput`, *optional*): + The segmentation maps. + instance_id_to_semantic_id (`Union[list[dict[int, int]], dict[int, int]]`, *optional*): + A mapping from instance IDs to semantic IDs. + """ + return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess_image_like_inputs( + self, + images: ImageInput, + segmentation_maps: ImageInput, + instance_id_to_semantic_id: list[dict[int, int]] | dict[int, int] | None, + do_convert_rgb: bool, + input_data_format: ChannelDimension, + **kwargs: Unpack[MaskFormerImageProcessorKwargs], + ) -> BatchFeature: + """ + Preprocess image-like inputs. + To be overridden by subclasses when image-like inputs other than images should be processed. + It can be used for segmentation maps, depth maps, etc. + """ + # Prepare input images + images = self._prepare_image_like_inputs( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format + ) + if segmentation_maps is not None: + segmentation_maps = self._prepare_image_like_inputs( + images=segmentation_maps, + expected_ndims=2, + do_convert_rgb=False, + input_data_format=ChannelDimension.FIRST, + ) + return self._preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + segmentation_maps: list[np.ndarray] | None, + instance_id_to_semantic_id: dict[int, int] | None, + do_resize: bool | None, + size: SizeDict | None, + pad_size: SizeDict | None, + size_divisor: int | None, + resample: PILImageResampling | None, + do_rescale: bool | None, + rescale_factor: float | None, + do_normalize: bool | None, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + ignore_index: int | None, + do_reduce_labels: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + if segmentation_maps is not None and len(images) != len(segmentation_maps): + raise ValueError("Images and segmentation maps must have the same length.") + + # Process images one by one (no batching in PIL backend) + resized_images = [] + resized_segmentation_maps = None + if segmentation_maps is not None: + resized_segmentation_maps = [] + + for idx, image in enumerate(images): + if do_resize: + image = self.resize(image=image, size=size, size_divisor=size_divisor, resample=resample) + resized_images.append(image) + + if segmentation_maps is not None: + seg_map = segmentation_maps[idx] + if do_resize: + seg_map = self.resize( + image=seg_map, size=size, size_divisor=size_divisor, resample=PILImageResampling.NEAREST + ) + resized_segmentation_maps.append(seg_map) + + # Determine padded size + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(resized_images, input_data_format=ChannelDimension.FIRST) + + # Convert segmentation maps to binary masks if provided + mask_labels = None + class_labels = None + if segmentation_maps is not None: + mask_labels = [] + class_labels = [] + for idx, segmentation_map in enumerate(resized_segmentation_maps): + if isinstance(instance_id_to_semantic_id, list): + instance_id = instance_id_to_semantic_id[idx] + else: + instance_id = instance_id_to_semantic_id + # Squeeze channel dimension if present + if segmentation_map.ndim == 3 and segmentation_map.shape[0] == 1: + segmentation_map = segmentation_map.squeeze(0) + masks, classes = convert_segmentation_map_to_binary_masks( + segmentation_map, instance_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels + ) + mask_labels.append(masks) + class_labels.append(classes) + + # Process images: rescale, normalize, pad + processed_images = [] + for image in resized_images: + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + processed_images.append(image) + + # Pad images and create pixel masks (also pad mask_labels to match padded image size) + padded_images, pixel_masks, padded_mask_labels = self.pad( + images=processed_images, + padded_size=padded_size, + segmentation_maps=mask_labels, + fill=0, + ignore_index=ignore_index, # Match Torchvision backend for cross-backend equivalence + ) + + encoded_inputs = BatchFeature( + data={"pixel_values": padded_images, "pixel_mask": pixel_masks}, tensor_type=return_tensors + ) + # we cannot batch them since they don't share a common class size + if segmentation_maps is not None: + encoded_inputs["mask_labels"] = [ + torch.from_numpy(mask_label) if return_tensors == "pt" else mask_label + for mask_label in padded_mask_labels + ] + encoded_inputs["class_labels"] = [ + torch.from_numpy(class_label) if return_tensors == "pt" else class_label + for class_label in class_labels + ] + + return encoded_inputs + + def post_process_semantic_segmentation( + self, outputs, target_sizes: list[tuple[int, int]] | None = None + ) -> "torch.Tensor": + """ + Converts the output of [`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. + """ + requires_backends(self, "torch") + + class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width] + + # Remove the null class `[..., :-1]` + masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1] + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = torch.nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + def post_process_instance_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + return_binary_maps: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into instance segmentation predictions. Only + supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps + to `True` to get the correct segmentation result. + + Args: + outputs ([`MaskFormerForInstanceSegmentation`]): + Raw outputs of the model. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If left to None, predictions will not be resized. + return_coco_annotation (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE) format. + return_binary_maps (`bool`, *optional*, defaults to `False`): + If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps + (one per detected instance). + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id`, or + `list[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to + `True`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps is set to `True`. + Set to `None` if no mask if found above `threshold`. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- An integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + requires_backends(self, "torch") + + if return_coco_annotation and return_binary_maps: + raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.") + + # [batch_size, num_queries, num_classes+1] + class_queries_logits = outputs.class_queries_logits + # [batch_size, num_queries, height, width] + masks_queries_logits = outputs.masks_queries_logits + + device = masks_queries_logits.device + num_classes = class_queries_logits.shape[-1] - 1 + num_queries = class_queries_logits.shape[-2] + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(class_queries_logits.shape[0]): + mask_pred = masks_queries_logits[i] + mask_cls = class_queries_logits[i] + + scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1] + labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1) + + scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False) + labels_per_image = labels[topk_indices] + + topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor") + mask_pred = mask_pred[topk_indices] + pred_masks = (mask_pred > 0).float() + + # Calculate average mask prob + mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / ( + pred_masks.flatten(1).sum(1) + 1e-6 + ) + pred_scores = scores_per_image * mask_scores_per_image + pred_classes = labels_per_image + + segmentation = torch.zeros(masks_queries_logits.shape[2:]) - 1 + if target_sizes is not None: + segmentation = torch.zeros(target_sizes[i]) - 1 + pred_masks = torch.nn.functional.interpolate( + pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest" + )[0] + + instance_maps, segments = [], [] + current_segment_id = 0 + for j in range(num_queries): + score = pred_scores[j].item() + + if not torch.all(pred_masks[j] == 0) and score >= threshold: + segmentation[pred_masks[j] == 1] = current_segment_id + segments.append( + { + "id": current_segment_id, + "label_id": pred_classes[j].item(), + "was_fused": False, + "score": round(score, 6), + } + ) + current_segment_id += 1 + instance_maps.append(pred_masks[j]) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + # Return a concatenated tensor of binary instance maps + if return_binary_maps and len(instance_maps) != 0: + segmentation = torch.stack(instance_maps, dim=0) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + # Adapted from transformers.models.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__ = ["MaskFormerImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer.py b/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer.py new file mode 100644 index 0000000000000000000000000000000000000000..788775a52fcb089a51f7742dde30474ca029b8e2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer.py @@ -0,0 +1,2071 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/maskformer/modular_maskformer.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_maskformer.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2022 Meta Platforms, Inc.s and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Callable +from dataclasses import dataclass +from numbers import Number + +import numpy as np +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithCrossAttentions +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + is_accelerate_available, + is_scipy_available, + requires_backends, +) +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..auto import AutoBackbone +from .configuration_maskformer import MaskFormerConfig, MaskFormerDetrConfig +from .configuration_maskformer_swin import MaskFormerSwinConfig + + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import reduce + +if is_scipy_available(): + from scipy.optimize import linear_sum_assignment + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the MASK_FORMER_DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions, + namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them + gone through a layernorm. This is useful when training the model with auxiliary decoding losses. + """ +) +class DetrDecoderOutput(BaseModelOutputWithCrossAttentions): + r""" + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, + used to compute the weighted average in the cross-attention heads. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + """ + + intermediate_hidden_states: torch.FloatTensor | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + MaskFormer'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 MaskFormerSwin Transformer and the `decoder` is a Feature + Pyramid Network (FPN). + + The `encoder_last_hidden_state` are referred on the paper as **images features**, while `decoder_last_hidden_state` + as **pixel embeddings** + """ +) +class MaskFormerPixelLevelModuleOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder. + decoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the 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 model at + the output of each stage. + 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 model at + the output of each stage. + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + MaskFormer's pixel decoder module output, practically a Feature Pyramid Network. It returns the last hidden state + and (optionally) the hidden states. + """ +) +class MaskFormerPixelDecoderOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the model. + """ + + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`MaskFormerModel`]. This class returns all the needed hidden states to compute the logits. + """ +) +class MaskFormerModelOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). + transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Last hidden states (final feature map) of the last stage of the transformer decoder model. + 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. + 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and + `decoder_hidden_states` + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`MaskFormerForInstanceSegmentation`]. + + This output can be directly passed to [`~MaskFormerImageProcessor.post_process_semantic_segmentation`] or + [`~MaskFormerImageProcessor.post_process_instance_segmentation`] or + [`~MaskFormerImageProcessor.post_process_panoptic_segmentation`] depending on the task. Please, see + [`~MaskFormerImageProcessor] for details regarding usage. + """ +) +class MaskFormerForInstanceSegmentationOutput(ModelOutput): + r""" + loss (`torch.Tensor`, *optional*): + The computed loss, returned when labels are present. + class_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each + query. Note the `+ 1` is needed because we incorporate the null class. + masks_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each + query. + auxiliary_logits (`Dict[str, torch.FloatTensor]`, *optional*, returned when `output_auxiliary_logits=True`): + Dictionary containing auxiliary predictions for each decoder layer when auxiliary losses are enabled. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). + transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Last hidden states (final feature map) of the last stage of the transformer decoder model. + 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 of the transformer decoder at the output + of each stage. + 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and + `decoder_hidden_states`. + """ + + loss: torch.FloatTensor | None = None + class_queries_logits: torch.FloatTensor | None = None + masks_queries_logits: torch.FloatTensor | None = None + auxiliary_logits: torch.FloatTensor | None = None + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the MASK_FORMER_DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions, + namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them + gone through a layernorm. This is useful when training the model with auxiliary decoding losses. + """ +) +class MaskFormerDetrDecoderOutput(BaseModelOutputWithCrossAttentions): + r""" + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, + used to compute the weighted average in the cross-attention heads. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + """ + + intermediate_hidden_states: torch.FloatTensor | None = None + + +class MaskFormerDetrLearnedPositionEmbedding(nn.Module): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, embedding_dim=256): + super().__init__() + self.row_embeddings = nn.Embedding(50, embedding_dim) + self.column_embeddings = nn.Embedding(50, embedding_dim) + + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: torch.Tensor | None = None, + ): + height, width = shape[-2:] + width_values = torch.arange(width, device=device) + height_values = torch.arange(height, device=device) + x_emb = self.column_embeddings(width_values) + y_emb = self.row_embeddings(height_values) + pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1) + pos = pos.permute(2, 0, 1) + pos = pos.unsqueeze(0) + pos = pos.repeat(shape[0], 1, 1, 1) + # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format + # expected by the encoder + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float | None = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class MaskFormerDetrSelfAttention(nn.Module): + """ + Multi-headed self-attention from 'Attention Is All You Need' paper. + + In MASK_FORMER_DETR, position embeddings are added to both queries and keys (but not values) in self-attention. + """ + + def __init__( + self, + config: MaskFormerDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.config = config + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Position embeddings are added to both queries and keys (but not values). + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_key_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states + + query_states = self.q_proj(query_key_input).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(query_key_input).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class MaskFormerDetrCrossAttention(nn.Module): + """ + Multi-headed cross-attention from 'Attention Is All You Need' paper. + + In MASK_FORMER_DETR, queries get their own position embeddings, while keys get encoder position embeddings. + Values don't get any position embeddings. + """ + + def __init__( + self, + config: MaskFormerDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.config = config + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + encoder_position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Position embeddings logic: + - Queries get position_embeddings + - Keys get encoder_position_embeddings + - Values don't get any position embeddings + """ + query_input_shape = hidden_states.shape[:-1] + query_hidden_shape = (*query_input_shape, -1, self.head_dim) + + kv_input_shape = key_value_states.shape[:-1] + kv_hidden_shape = (*kv_input_shape, -1, self.head_dim) + + query_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states + key_input = ( + key_value_states + encoder_position_embeddings + if encoder_position_embeddings is not None + else key_value_states + ) + + query_states = self.q_proj(query_input).view(query_hidden_shape).transpose(1, 2) + key_states = self.k_proj(key_input).view(kv_hidden_shape).transpose(1, 2) + value_states = self.v_proj(key_value_states).view(kv_hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*query_input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class MaskFormerDetrMLP(nn.Module): + def __init__(self, config: MaskFormerDetrConfig, hidden_size: int, intermediate_size: int): + super().__init__() + self.fc1 = nn.Linear(hidden_size, intermediate_size) + self.fc2 = nn.Linear(intermediate_size, hidden_size) + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.dropout = config.dropout + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + return hidden_states + + +class MaskFormerDetrDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MaskFormerDetrConfig): + super().__init__() + self.hidden_size = config.d_model + + self.self_attn = MaskFormerDetrSelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + self.dropout = config.dropout + + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.encoder_attn = MaskFormerDetrCrossAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.mlp = MaskFormerDetrMLP(config, self.hidden_size, config.decoder_ffn_dim) + self.final_layer_norm = nn.LayerNorm(self.hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + spatial_position_embeddings: torch.Tensor | None = None, + object_queries_position_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + spatial_position_embeddings (`torch.FloatTensor`, *optional*): + Spatial position embeddings (2D positional encodings from encoder) that are added to the keys only + in the cross-attention layer (not to values). + object_queries_position_embeddings (`torch.FloatTensor`, *optional*): + Position embeddings for the object query slots. In self-attention, these are added to both queries + and keys (not values). In cross-attention, these are added to queries only (not to keys or values). + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(batch, seq_len, hidden_size)` + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + """ + residual = hidden_states + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=object_queries_position_embeddings, + attention_mask=attention_mask, + **kwargs, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Cross-Attention Block + if encoder_hidden_states is not None: + residual = hidden_states + + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_embeddings=object_queries_position_embeddings, + encoder_position_embeddings=spatial_position_embeddings, + **kwargs, + ) + + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + # Fully Connected + residual = hidden_states + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states + + +class MaskFormerDetrConvBlock(nn.Module): + """Basic conv block: Conv3x3 -> GroupNorm -> Activation.""" + + def __init__(self, in_channels: int, out_channels: int, activation: str = "relu"): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm = nn.GroupNorm(min(8, out_channels), out_channels) + self.activation = ACT2FN[activation] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.activation(self.norm(self.conv(x))) + + +class MaskFormerDetrFPNFusionStage(nn.Module): + """Single FPN fusion stage combining low-resolution features with high-resolution FPN features.""" + + def __init__(self, fpn_channels: int, current_channels: int, output_channels: int, activation: str = "relu"): + super().__init__() + self.fpn_adapter = nn.Conv2d(fpn_channels, current_channels, kernel_size=1) + self.refine = MaskFormerDetrConvBlock(current_channels, output_channels, activation) + + def forward(self, features: torch.Tensor, fpn_features: torch.Tensor) -> torch.Tensor: + """ + Args: + features: Current features to upsample, shape (B*Q, current_channels, H_in, W_in) + fpn_features: FPN features at target resolution, shape (B*Q, fpn_channels, H_out, W_out) + + Returns: + Fused and refined features, shape (B*Q, output_channels, H_out, W_out) + """ + fpn_features = self.fpn_adapter(fpn_features) + features = nn.functional.interpolate(features, size=fpn_features.shape[-2:], mode="nearest") + return self.refine(fpn_features + features) + + +class MaskFormerDetrMaskHeadSmallConv(nn.Module): + """ + Segmentation mask head that generates per-query masks using FPN-based progressive upsampling. + + Combines attention maps (spatial localization) with encoder features (semantics) and progressively + upsamples through multiple scales, fusing with FPN features for high-resolution detail. + """ + + def __init__( + self, + input_channels: int, + fpn_channels: list[int], + hidden_size: int, + activation_function: str = "relu", + ): + super().__init__() + if input_channels % 8 != 0: + raise ValueError(f"input_channels must be divisible by 8, got {input_channels}") + + self.conv1 = MaskFormerDetrConvBlock(input_channels, input_channels, activation_function) + self.conv2 = MaskFormerDetrConvBlock(input_channels, hidden_size // 2, activation_function) + + # Progressive channel reduction: /2 -> /4 -> /8 -> /16 + self.fpn_stages = nn.ModuleList( + [ + MaskFormerDetrFPNFusionStage(fpn_channels[0], hidden_size // 2, hidden_size // 4, activation_function), + MaskFormerDetrFPNFusionStage(fpn_channels[1], hidden_size // 4, hidden_size // 8, activation_function), + MaskFormerDetrFPNFusionStage( + fpn_channels[2], hidden_size // 8, hidden_size // 16, activation_function + ), + ] + ) + + self.output_conv = nn.Conv2d(hidden_size // 16, 1, kernel_size=3, padding=1) + + def forward( + self, + features: torch.Tensor, + attention_masks: torch.Tensor, + fpn_features: list[torch.Tensor], + ) -> torch.Tensor: + """ + Args: + features: Encoder output features, shape (batch_size, hidden_size, H, W) + attention_masks: Cross-attention maps from decoder, shape (batch_size, num_queries, num_heads, H, W) + fpn_features: List of 3 FPN features from low to high resolution, each (batch_size, C, H, W) + + Returns: + Predicted masks, shape (batch_size * num_queries, 1, output_H, output_W) + """ + num_queries = attention_masks.shape[1] + + # Expand to (batch_size * num_queries) dimension + features = features.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) + attention_masks = attention_masks.flatten(0, 1) + fpn_features = [ + fpn_feat.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) for fpn_feat in fpn_features + ] + + hidden_states = torch.cat([features, attention_masks], dim=1) + hidden_states = self.conv1(hidden_states) + hidden_states = self.conv2(hidden_states) + + for fpn_stage, fpn_feat in zip(self.fpn_stages, fpn_features): + hidden_states = fpn_stage(hidden_states, fpn_feat) + + return self.output_conv(hidden_states) + + +class MaskFormerDetrMHAttentionMap(nn.Module): + """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, query_states: torch.Tensor, key_states: torch.Tensor, attention_mask: torch.Tensor | None = None + ): + query_hidden_shape = (*query_states.shape[:-1], -1, self.head_dim) + key_hidden_shape = (key_states.shape[0], -1, self.head_dim, *key_states.shape[-2:]) + + query_states = self.q_proj(query_states).view(query_hidden_shape) + key_states = nn.functional.conv2d( + key_states, self.k_proj.weight.unsqueeze(-1).unsqueeze(-1), self.k_proj.bias + ).view(key_hidden_shape) + + batch_size, num_queries, num_heads, head_dim = query_states.shape + _, _, _, height, width = key_states.shape + query_shape = (batch_size * num_heads, num_queries, head_dim) + key_shape = (batch_size * num_heads, height * width, head_dim) + attn_weights_shape = (batch_size, num_heads, num_queries, height, width) + + query = query_states.transpose(1, 2).contiguous().view(query_shape) + key = key_states.permute(0, 1, 3, 4, 2).contiguous().view(key_shape) + + attn_weights = ( + (torch.matmul(query * self.scaling, key.transpose(1, 2))).view(attn_weights_shape).transpose(1, 2) + ) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights.flatten(2), dim=-1).view(attn_weights.size()) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + + return attn_weights + + +@auto_docstring +class MaskFormerDetrPreTrainedModel(PreTrainedModel): + config: MaskFormerDetrConfig + base_model_prefix = "model" + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = [r"MaskFormerDetrConvEncoder", r"MaskFormerDetrEncoderLayer", r"MaskFormerDetrDecoderLayer"] + supports_gradient_checkpointing = True + _supports_sdpa = True + _supports_flash_attn = True + _supports_attention_backend = True + _supports_flex_attn = True # Uses create_bidirectional_masks for attention masking + _keys_to_ignore_on_load_unexpected = [ + r"mask_former_detr\.model\.backbone\.model\.layer\d+\.0\.downsample\.1\.num_batches_tracked" + ] + + @torch.no_grad() + def _init_weights(self, module): + std = self.config.init_std + xavier_std = self.config.init_xavier_std + + if isinstance(module, MaskFormerDetrMaskHeadSmallConv): + # MaskFormerDetrMaskHeadSmallConv uses kaiming initialization for all its Conv2d layers + for m in module.modules(): + if isinstance(m, nn.Conv2d): + init.kaiming_uniform_(m.weight, a=1) + if m.bias is not None: + init.constant_(m.bias, 0) + elif isinstance(module, MaskFormerDetrMHAttentionMap): + init.zeros_(module.k_proj.bias) + init.zeros_(module.q_proj.bias) + init.xavier_uniform_(module.k_proj.weight, gain=xavier_std) + init.xavier_uniform_(module.q_proj.weight, gain=xavier_std) + elif isinstance(module, MaskFormerDetrLearnedPositionEmbedding): + init.uniform_(module.row_embeddings.weight) + init.uniform_(module.column_embeddings.weight) + elif isinstance(module, (nn.Linear, nn.Conv2d)): + init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, mean=0.0, std=std) + # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag + if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): + init.zeros_(module.weight[module.padding_idx]) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.ones_(module.weight) + init.zeros_(module.bias) + + +class MaskFormerDetrDecoder(MaskFormerDetrPreTrainedModel): + """ + Transformer decoder that refines a set of object queries. It is composed of a stack of [`MaskFormerDetrDecoderLayer`] modules, + which apply self-attention to the queries and cross-attention to the encoder's outputs. + + Args: + config (`MaskFormerDetrConfig`): Model configuration object. + """ + + _can_record_outputs = { + "hidden_states": MaskFormerDetrDecoderLayer, + "attentions": MaskFormerDetrSelfAttention, + "cross_attentions": MaskFormerDetrCrossAttention, + } + + def __init__(self, config: MaskFormerDetrConfig): + super().__init__(config) + self.dropout = config.dropout + + self.layers = nn.ModuleList([MaskFormerDetrDecoderLayer(config) for _ in range(config.decoder_layers)]) + # in MASK_FORMER_DETR, the decoder uses layernorm after the last decoder layer output + self.layernorm = nn.LayerNorm(config.d_model) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + spatial_position_embeddings=None, + object_queries_position_embeddings=None, + **kwargs: Unpack[TransformersKwargs], + ) -> MaskFormerDetrDecoderOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + The query embeddings that are passed into the decoder. + + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`: + + - 1 for queries that are **not masked**, + - 0 for queries that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention + of the decoder. + encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): + Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected + in `[0, 1]`: + + - 1 for pixels that are real (i.e. **not masked**), + - 0 for pixels that are padding (i.e. **masked**). + + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Spatial position embeddings (2D positional encodings from encoder) that are added to the keys in each cross-attention layer. + object_queries_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Position embeddings for the object query slots that are added to the queries and keys in each self-attention layer. + """ + + if inputs_embeds is not None: + hidden_states = inputs_embeds + + if attention_mask is not None: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + # expand encoder attention mask (for cross-attention on encoder outputs) + if encoder_hidden_states is not None and encoder_attention_mask is not None: + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + # optional intermediate hidden states + intermediate = () if self.config.auxiliary_loss else None + + # decoder layers + + for idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask, + spatial_position_embeddings, + object_queries_position_embeddings, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + **kwargs, + ) + + if self.config.auxiliary_loss: + hidden_states = self.layernorm(hidden_states) + intermediate += (hidden_states,) + + # finally, apply layernorm + hidden_states = self.layernorm(hidden_states) + + # stack intermediate decoder activations + if self.config.auxiliary_loss: + intermediate = torch.stack(intermediate) + + return MaskFormerDetrDecoderOutput(last_hidden_state=hidden_states, intermediate_hidden_states=intermediate) + + +# refactored from original implementation +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 + + +# refactored from original implementation +def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float = 0.25, gamma: float = 2.0) -> Tensor: + r""" + A pair wise version of the focal loss, see `sigmoid_focal_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). + alpha (float, *optional*, defaults to 0.25): + Weighting factor in range (0,1) to balance positive vs negative examples. + gamma (float, *optional*, defaults to 2.0): + Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. + + Returns: + `torch.Tensor`: The computed loss between each pairs. + """ + if alpha < 0: + raise ValueError("alpha must be positive") + + height_and_width = inputs.shape[1] + + criterion = nn.BCEWithLogitsLoss(reduction="none") + prob = inputs.sigmoid() + cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) + focal_pos = ((1 - prob) ** gamma) * cross_entropy_loss_pos + focal_pos *= alpha + + cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) + + focal_neg = (prob**gamma) * cross_entropy_loss_neg + focal_neg *= 1 - alpha + + loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T) + + return loss / height_and_width + + +# refactored from original implementation +class MaskFormerHungarianMatcher(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): + """Creates the matcher + + 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 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 + """ + 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 + + @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_target_boxes). + """ + indices: list[tuple[np.array]] = [] + + 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): + # downsample the target mask, save memory + target_mask = nn.functional.interpolate(target_mask[:, None], size=pred_mask.shape[-2:], mode="nearest") + 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] + # flatten spatial dimension "q h w -> q (h w)" + pred_mask_flat = pred_mask.flatten(1) # [num_queries, height*width] + # same for target_mask "c h w -> c (h w)" + target_mask_flat = target_mask[:, 0].flatten(1) # [num_total_labels, height*width] + # compute the focal loss between each mask pairs -> shape (num_queries, num_labels) + cost_mask = pair_wise_sigmoid_focal_loss(pred_mask_flat, target_mask_flat) + # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels) + cost_dice = pair_wise_dice_loss(pred_mask_flat, target_mask_flat) + # final cost matrix + cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice + # 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 __repr__(self): + head = "Matcher " + self.__class__.__name__ + body = [ + f"cost_class: {self.cost_class}", + f"cost_mask: {self.cost_mask}", + f"cost_dice: {self.cost_dice}", + ] + _repr_indent = 4 + lines = [head] + [" " * _repr_indent + line for line in body] + return "\n".join(lines) + + +# refactored from original implementation +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 + + +# refactored from original implementation +def sigmoid_focal_loss( + inputs: Tensor, labels: Tensor, num_masks: int, alpha: float = 0.25, gamma: float = 2 +) -> Tensor: + r""" + Focal loss proposed in [Focal Loss for Dense Object Detection](https://huggingface.co/papers/1708.02002) originally used in + RetinaNet. The loss is computed as follows: + + $$ \mathcal{L}_{\text{focal loss} = -(1 - p_t)^{\gamma}\log{(p_t)} $$ + + where \\(CE(p_t) = -\log{(p_t)}}\\), CE is the standard Cross Entropy Loss + + Please refer to equation (1,2,3) of the paper for a better understanding. + + 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). + num_masks (`int`): + The number of masks present in the current batch, used for normalization. + alpha (float, *optional*, defaults to 0.25): + Weighting factor in range (0,1) to balance positive vs negative examples. + gamma (float, *optional*, defaults to 2.0): + Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. + + Returns: + `torch.Tensor`: The computed loss. + """ + criterion = nn.BCEWithLogitsLoss(reduction="none") + probs = inputs.sigmoid() + cross_entropy_loss = criterion(inputs, labels) + p_t = probs * labels + (1 - probs) * (1 - labels) + loss = cross_entropy_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * labels + (1 - alpha) * (1 - labels) + loss = alpha_t * loss + + loss = loss.mean(1).sum() / num_masks + return loss + + +# copied and adapted from original implementation +class MaskFormerLoss(nn.Module): + def __init__( + self, + num_labels: int, + matcher: MaskFormerHungarianMatcher, + weight_dict: dict[str, float], + eos_coef: float, + ): + """ + The MaskFormer 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: + num_labels (`int`): + The number of classes. + matcher (`MaskFormerHungarianMatcher`): + 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. + """ + + super().__init__() + requires_backends(self, ["scipy"]) + self.num_labels = num_labels + self.matcher = matcher + self.weight_dict = weight_dict + self.eos_coef = eos_coef + empty_weight = torch.ones(self.num_labels + 1) + empty_weight[-1] = self.eos_coef + self.register_buffer("empty_weight", empty_weight) + + 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_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_labels, dtype=torch.int64, device=pred_logits.device + ) + target_classes[idx] = target_classes_o + # target_classes is a (batch_size, num_labels, num_queries), we need to permute pred_logits "b q c -> b c q" + 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 focal 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] + # upsample predictions to the target size, we have to add one dim to use interpolate + pred_masks = nn.functional.interpolate( + pred_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False + ) + pred_masks = pred_masks[:, 0].flatten(1) + + target_masks = target_masks.flatten(1) + losses = { + "loss_mask": sigmoid_focal_loss(pred_masks, target_masks, num_masks), + "loss_dice": dice_loss(pred_masks, target_masks, num_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 forward( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: list[Tensor], + class_labels: list[Tensor], + auxiliary_predictions: dict[str, Tensor] | None = None, + ) -> 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` + 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 [`MaskFormerConfig`], then it contains the logits from the + inner layers of the Detr's Decoder. + + 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 focal 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 [`MaskFormerConfig`], 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: Number = self.get_num_masks(class_labels, device=class_labels[0].device) + # get all the losses + losses: dict[str, Tensor] = { + **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), + **self.loss_labels(class_queries_logits, class_labels, indices), + } + # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if auxiliary_predictions is not None: + for idx, aux_outputs in enumerate(auxiliary_predictions): + masks_queries_logits = aux_outputs["masks_queries_logits"] + class_queries_logits = aux_outputs["class_queries_logits"] + loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} + losses.update(loss_dict) + + return losses + + def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Computes the average number of target masks across the batch, for normalization purposes. + """ + num_masks = sum(len(classes) for classes in class_labels) + num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device) + world_size = 1 + if is_accelerate_available(): + if PartialState._shared_state != {}: + num_masks = reduce(num_masks) + world_size = PartialState().num_processes + + num_masks = torch.clamp(num_masks / world_size, min=1) + return num_masks + + +class MaskFormerFPNConvLayer(nn.Module): + def __init__(self, in_features: int, out_features: int, kernel_size: int = 3, padding: int = 1): + """ + A basic module that executes conv - norm - in sequence used in MaskFormer. + + Args: + in_features (`int`): + The number of input features (channels). + out_features (`int`): + The number of outputs features (channels). + """ + super().__init__() + self.layers = [ + nn.Conv2d(in_features, out_features, kernel_size=kernel_size, padding=padding, bias=False), + nn.GroupNorm(32, out_features), + nn.ReLU(inplace=True), + ] + for i, layer in enumerate(self.layers): + # Provide backwards compatibility from when the class inherited from nn.Sequential + # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. + # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. + # self.my_layer_name = Layer() + # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register + # explicitly + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class MaskFormerFPNLayer(nn.Module): + def __init__(self, in_features: int, lateral_features: int): + """ + A Feature Pyramid Network Layer (FPN) layer. It creates a feature map by aggregating features from the previous + and backbone layer. Due to the spatial mismatch, the tensor coming from the previous layer is upsampled. + + Args: + in_features (`int`): + The number of input features (channels). + lateral_features (`int`): + The number of lateral features (channels). + """ + super().__init__() + self.proj = nn.Sequential( + nn.Conv2d(lateral_features, in_features, kernel_size=1, padding=0, bias=False), + nn.GroupNorm(32, in_features), + ) + + self.block = MaskFormerFPNConvLayer(in_features, in_features) + + def forward(self, down: Tensor, left: Tensor) -> Tensor: + left = self.proj(left) + down = nn.functional.interpolate(down, size=left.shape[-2:], mode="nearest") + down += left + down = self.block(down) + return down + + +class MaskFormerFPNModel(nn.Module): + def __init__(self, in_features: int, lateral_widths: list[int], feature_size: int = 256): + """ + Feature Pyramid Network, given an input tensor and a set of feature map of different feature/spatial size, it + creates a list of feature maps with the same feature size. + + Args: + in_features (`int`): + The number of input features (channels). + lateral_widths (`list[int]`): + A list with the features (channels) size of each lateral connection. + feature_size (int, *optional*, defaults to 256): + The features (channels) of the resulting feature maps. + """ + super().__init__() + self.stem = MaskFormerFPNConvLayer(in_features, feature_size) + self.layers = nn.Sequential( + *[MaskFormerFPNLayer(feature_size, lateral_width) for lateral_width in lateral_widths[::-1]] + ) + + def forward(self, features: list[Tensor]) -> list[Tensor]: + fpn_features = [] + last_feature = features[-1] + other_features = features[:-1] + output = self.stem(last_feature) + for layer, left in zip(self.layers, other_features[::-1]): + output = layer(output, left) + fpn_features.append(output) + return fpn_features + + +class MaskFormerPixelDecoder(nn.Module): + def __init__(self, *args, feature_size: int = 256, mask_feature_size: int = 256, **kwargs): + r""" + Pixel Decoder Module proposed in [Per-Pixel Classification is Not All You Need for Semantic + Segmentation](https://huggingface.co/papers/2107.06278). It first runs the backbone's features into a Feature Pyramid + Network creating a list of feature maps. Then, it projects the last one to the correct `mask_size`. + + Args: + feature_size (`int`, *optional*, defaults to 256): + The feature size (channel dimension) of the FPN feature maps. + mask_feature_size (`int`, *optional*, defaults to 256): + The features (channels) of the target masks size \\(C_{\epsilon}\\) in the paper. + """ + super().__init__() + + self.fpn = MaskFormerFPNModel(*args, feature_size=feature_size, **kwargs) + self.mask_projection = nn.Conv2d(feature_size, mask_feature_size, kernel_size=3, padding=1) + + def forward( + self, features: list[Tensor], output_hidden_states: bool = False, return_dict: bool = True + ) -> MaskFormerPixelDecoderOutput: + fpn_features = self.fpn(features) + # we use the last feature map + last_feature_projected = self.mask_projection(fpn_features[-1]) + + if not return_dict: + return (last_feature_projected, tuple(fpn_features)) if output_hidden_states else (last_feature_projected,) + + return MaskFormerPixelDecoderOutput( + last_hidden_state=last_feature_projected, hidden_states=tuple(fpn_features) if output_hidden_states else () + ) + + +# copied and adapted from original implementation, also practically equal to DetrSinePositionEmbedding +class MaskFormerSinePositionEmbedding(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 + + +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 MaskformerMLPPredictionHead(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3): + """ + A classic Multi Layer Perceptron (MLP). + + Args: + input_dim (`int`): + The input dimensions. + hidden_dim (`int`): + The hidden dimensions. + output_dim (`int`): + The output dimensions. + num_layers (int, *optional*, defaults to 3): + The number of layers. + """ + super().__init__() + in_dims = [input_dim] + [hidden_dim] * (num_layers - 1) + out_dims = [hidden_dim] * (num_layers - 1) + [output_dim] + + self.layers = [] + for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)): + activation = nn.ReLU() if i < num_layers - 1 else nn.Identity() + layer = PredictionBlock(in_dim, out_dim, activation=activation) + self.layers.append(layer) + # Provide backwards compatibility from when the class inherited from nn.Sequential + # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. + # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. + # self.my_layer_name = Layer() + # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register + # explicitly + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class MaskFormerPixelLevelModule(nn.Module): + def __init__(self, config: MaskFormerConfig): + """ + Pixel Level Module proposed in [Per-Pixel Classification is Not All You Need for Semantic + Segmentation](https://huggingface.co/papers/2107.06278). It runs the input image through a backbone and a pixel + decoder, generating an image feature map and pixel embeddings. + + Args: + config ([`MaskFormerConfig`]): + The configuration used to instantiate this model. + """ + super().__init__() + if getattr(config, "backbone_config") is not None and config.backbone_config.model_type == "swin": + # for backwards compatibility + backbone_config = config.backbone_config + backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict()) + backbone_config.out_features = ["stage1", "stage2", "stage3", "stage4"] + config.backbone_config = backbone_config + self.encoder = AutoBackbone.from_config(config=config.backbone_config) + + feature_channels = self.encoder.channels + self.decoder = MaskFormerPixelDecoder( + in_features=feature_channels[-1], + feature_size=config.fpn_feature_size, + mask_feature_size=config.mask_feature_size, + lateral_widths=feature_channels[:-1], + ) + + def forward( + self, pixel_values: Tensor, output_hidden_states: bool = False, return_dict: bool = True + ) -> MaskFormerPixelLevelModuleOutput: + features = self.encoder(pixel_values).feature_maps + decoder_output = self.decoder(features, output_hidden_states, return_dict=return_dict) + + if not return_dict: + last_hidden_state = decoder_output[0] + outputs = (features[-1], last_hidden_state) + if output_hidden_states: + hidden_states = decoder_output[1] + outputs = outputs + (tuple(features),) + (hidden_states,) + return outputs + + return MaskFormerPixelLevelModuleOutput( + # the last feature is actually the output from the last layer + encoder_last_hidden_state=features[-1], + decoder_last_hidden_state=decoder_output.last_hidden_state, + encoder_hidden_states=tuple(features) if output_hidden_states else (), + decoder_hidden_states=decoder_output.hidden_states if output_hidden_states else (), + ) + + +class MaskFormerTransformerModule(nn.Module): + """ + The MaskFormer's transformer module. + """ + + def __init__(self, in_features: int, config: MaskFormerConfig): + super().__init__() + hidden_size = config.decoder_config.hidden_size + should_project = in_features != hidden_size + self.position_embedder = MaskFormerSinePositionEmbedding(num_pos_feats=hidden_size // 2, normalize=True) + self.queries_embedder = nn.Embedding(config.decoder_config.num_queries, hidden_size) + self.input_projection = nn.Conv2d(in_features, hidden_size, kernel_size=1) if should_project else None + self.decoder = MaskFormerDetrDecoder(config=config.decoder_config) + + def forward( + self, + image_features: Tensor, + output_hidden_states: bool = False, + output_attentions: bool = False, + return_dict: bool | None = None, + ) -> DetrDecoderOutput: + if self.input_projection is not None: + image_features = self.input_projection(image_features) + object_queries = self.position_embedder(image_features.shape, image_features.device, image_features.dtype) + # repeat the queries "q c -> b q c" + batch_size = image_features.shape[0] + queries_embeddings = self.queries_embedder.weight.unsqueeze(0).repeat(batch_size, 1, 1) + inputs_embeds = torch.zeros_like(queries_embeddings, requires_grad=self.training) + + # torch.export.export does no support requires_grad + if self.training: + inputs_embeds.requires_grad_(True) + + batch_size, num_channels, height, width = image_features.shape + # rearrange both image_features and object_queries "b c h w -> b (h w) c" + image_features = image_features.view(batch_size, num_channels, height * width).permute(0, 2, 1) + object_queries = object_queries.view(batch_size, num_channels, height * width).permute(0, 2, 1) + + decoder_output: DetrDecoderOutput = self.decoder( + inputs_embeds=inputs_embeds, + attention_mask=None, + encoder_hidden_states=image_features, + encoder_attention_mask=None, + spatial_position_embeddings=object_queries, + object_queries_position_embeddings=queries_embeddings, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + return decoder_output + + +@auto_docstring +class MaskFormerPreTrainedModel(PreTrainedModel): + config: MaskFormerConfig + 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, MaskFormerTransformerModule): + if module.input_projection is not None: + init.xavier_uniform_(module.input_projection.weight, gain=xavier_std) + init.constant_(module.input_projection.bias, 0) + # FPN + elif isinstance(module, MaskFormerFPNModel): + init.xavier_uniform_(module.stem.get_submodule("0").weight, gain=xavier_std) + + elif isinstance(module, MaskFormerFPNLayer): + init.xavier_uniform_(module.proj[0].weight, gain=xavier_std) + + elif isinstance(module, MaskFormerFPNConvLayer): + init.xavier_uniform_(module.get_submodule("0").weight, gain=xavier_std) + # The MLP head + elif isinstance(module, MaskformerMLPPredictionHead): + # I was not able to find the correct initializer in the original implementation + # we'll use xavier + 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.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + # copied from DETR + if 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.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, MaskFormerLoss): + empty_weight = torch.ones(module.num_labels + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + + +@auto_docstring +class MaskFormerModel(MaskFormerPreTrainedModel): + def __init__(self, config: MaskFormerConfig): + super().__init__(config) + self.pixel_level_module = MaskFormerPixelLevelModule(config) + self.transformer_module = MaskFormerTransformerModule( + in_features=self.pixel_level_module.encoder.channels[-1], config=config + ) + + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + pixel_mask: Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskFormerModelOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerModel + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") + >>> model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-base-ade") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = image_processor(image, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # the decoder of MaskFormer outputs hidden states of shape (batch_size, num_queries, hidden_size) + >>> transformer_decoder_last_hidden_state = outputs.transformer_decoder_last_hidden_state + >>> list(transformer_decoder_last_hidden_state.shape) + [1, 100, 256] + ```""" + + 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, return_dict=return_dict + ) + image_features = pixel_level_module_output[0] + pixel_embeddings = pixel_level_module_output[1] + + transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions) + queries = transformer_module_output.last_hidden_state + + encoder_hidden_states = None + pixel_decoder_hidden_states = None + transformer_decoder_hidden_states = None + hidden_states = None + + if output_hidden_states: + encoder_hidden_states = pixel_level_module_output[2] + pixel_decoder_hidden_states = pixel_level_module_output[3] + transformer_decoder_hidden_states = transformer_module_output[1] + hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states + + output = MaskFormerModelOutput( + encoder_last_hidden_state=image_features, + pixel_decoder_last_hidden_state=pixel_embeddings, + transformer_decoder_last_hidden_state=queries, + encoder_hidden_states=encoder_hidden_states, + pixel_decoder_hidden_states=pixel_decoder_hidden_states, + transformer_decoder_hidden_states=transformer_decoder_hidden_states, + hidden_states=hidden_states, + attentions=transformer_module_output.attentions, + ) + + if not return_dict: + output = tuple(v for v in output.values()) + + return output + + +class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel): + def __init__(self, config: MaskFormerConfig): + super().__init__(config) + self.model = MaskFormerModel(config) + hidden_size = config.decoder_config.hidden_size + # + 1 because we add the "null" class + self.class_predictor = nn.Linear(hidden_size, config.num_labels + 1) + self.mask_embedder = MaskformerMLPPredictionHead(hidden_size, hidden_size, config.mask_feature_size) + + self.matcher = MaskFormerHungarianMatcher( + cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight + ) + + self.weight_dict: dict[str, float] = { + "loss_cross_entropy": config.cross_entropy_weight, + "loss_mask": config.mask_weight, + "loss_dice": config.dice_weight, + } + + self.criterion = MaskFormerLoss( + config.num_labels, + matcher=self.matcher, + weight_dict=self.weight_dict, + eos_coef=config.no_object_weight, + ) + + self.post_init() + + def get_loss_dict( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: Tensor, + class_labels: Tensor, + auxiliary_logits: dict[str, Tensor], + ) -> dict[str, Tensor]: + loss_dict: dict[str, Tensor] = self.criterion( + masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits + ) + # weight each loss by `self.weight_dict[]` including auxiliary losses + for key, weight in self.weight_dict.items(): + for loss_key, loss in loss_dict.items(): + if key in loss_key: + loss *= weight + + return loss_dict + + def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: + return sum(loss_dict.values()) + + def get_logits(self, outputs: MaskFormerModelOutput) -> tuple[Tensor, Tensor, dict[str, Tensor]]: + pixel_embeddings = outputs.pixel_decoder_last_hidden_state + # get the auxiliary predictions (one for each decoder's layer) + auxiliary_logits: list[str, Tensor] = [] + + # This code is a little bit cumbersome, an improvement can be to return a list of predictions. If we have auxiliary loss then we are going to return more than one element in the list + if self.config.use_auxiliary_loss: + stacked_transformer_decoder_outputs = torch.stack(outputs.transformer_decoder_hidden_states) + classes = self.class_predictor(stacked_transformer_decoder_outputs) + class_queries_logits = classes[-1] + # get the masks + mask_embeddings = self.mask_embedder(stacked_transformer_decoder_outputs) + binaries_masks = torch.einsum("lbqc, bchw -> lbqhw", mask_embeddings, pixel_embeddings) + + masks_queries_logits = binaries_masks[-1] + # go til [:-1] because the last one is always used + for aux_binary_masks, aux_classes in zip(binaries_masks[:-1], classes[:-1]): + auxiliary_logits.append( + {"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes} + ) + + else: + transformer_decoder_hidden_states = outputs.transformer_decoder_last_hidden_state + classes = self.class_predictor(transformer_decoder_hidden_states) + class_queries_logits = classes + # get the masks + mask_embeddings = self.mask_embedder(transformer_decoder_hidden_states) + # sum up over the channels + masks_queries_logits = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings) + + return class_queries_logits, masks_queries_logits, auxiliary_logits + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + pixel_mask: Tensor | None = None, + output_auxiliary_logits: bool | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskFormerForInstanceSegmentationOutput: + r""" + mask_labels (`list[torch.Tensor]`, *optional*): + List of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + output_auxiliary_logits (`bool`, *optional*): + Whether or not to output auxiliary logits. + + Examples: + + Semantic segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") + >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-ade") + + >>> url = ( + ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg" + ... ) + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> 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 image_processor for postprocessing + >>> predicted_semantic_map = image_processor.post_process_semantic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0] + + >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) + >>> list(predicted_semantic_map.shape) + [512, 683] + ``` + + Panoptic segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on COCO panoptic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-coco") + >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-coco") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> 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 image_processor for postprocessing + >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(image.height, image.width)])[0] + + >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) + >>> predicted_panoptic_map = result["segmentation"] + >>> list(predicted_panoptic_map.shape) + [480, 640] + ``` + """ + + 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 + + raw_outputs = self.model( + pixel_values, + pixel_mask, + output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, + return_dict=return_dict, + output_attentions=output_attentions, + ) + # We need to have raw_outputs optionally be returned as a dict to use torch.compile. For backwards + # compatibility we convert to a dataclass for the rest of the model logic + outputs = MaskFormerModelOutput( + encoder_last_hidden_state=raw_outputs[0], + pixel_decoder_last_hidden_state=raw_outputs[1], + transformer_decoder_last_hidden_state=raw_outputs[2], + encoder_hidden_states=raw_outputs[3] if output_hidden_states else None, + pixel_decoder_hidden_states=raw_outputs[4] if output_hidden_states else None, + transformer_decoder_hidden_states=raw_outputs[5] if output_hidden_states else None, + hidden_states=raw_outputs[6] if output_hidden_states else None, + attentions=raw_outputs[-1] if output_attentions else None, + ) + + loss, loss_dict, auxiliary_logits = None, None, None + + class_queries_logits, masks_queries_logits, auxiliary_logits = self.get_logits(outputs) + + if mask_labels is not None and class_labels is not None: + loss_dict: dict[str, Tensor] = self.get_loss_dict( + masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits + ) + 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_logits = None + + if not return_dict: + output = tuple( + v + for v in (loss, class_queries_logits, masks_queries_logits, auxiliary_logits, *outputs.values()) + if v is not None + ) + return output + + return MaskFormerForInstanceSegmentationOutput( + loss=loss, + **outputs, + class_queries_logits=class_queries_logits, + masks_queries_logits=masks_queries_logits, + auxiliary_logits=auxiliary_logits, + ) + + +__all__ = [ + "MaskFormerForInstanceSegmentation", + "MaskFormerModel", + "MaskFormerPreTrainedModel", + "MaskFormerDetrPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer_swin.py b/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer_swin.py new file mode 100644 index 0000000000000000000000000000000000000000..fc30dd865dc03a55d249647344755164c28ff3ac --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/modeling_maskformer_swin.py @@ -0,0 +1,872 @@ +# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MaskFormer Swin Transformer. The reason Swin Transformer is implemented here is because MaskFormer uses the hidden +states before downsampling, which is different from the default Swin Transformer.""" + +import collections.abc +import math +from dataclasses import dataclass + +import torch +from torch import Tensor, nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import BackboneMixin, filter_output_hidden_states +from ...file_utils import ModelOutput +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BackboneOutput +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, torch_int +from ...utils.generic import can_return_tuple +from .configuration_maskformer_swin import MaskFormerSwinConfig + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for MaskFormerSwinModel's outputs that also contains the spatial dimensions of the hidden states. + """ +) +class MaskFormerSwinModelOutputWithPooling(ModelOutput): + r""" + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state after a mean pooling operation. + hidden_states_spatial_dimensions (`tuple(tuple(int, int))`, *optional*): + A tuple containing the spatial dimension of each `hidden_state` needed to reshape the `hidden_states` to + `batch, channels, height, width`. Due to padding, their spatial size cannot be inferred before the + `forward` method. + """ + + last_hidden_state: torch.FloatTensor | None = None + pooler_output: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states_spatial_dimensions: tuple[tuple[int, int]] = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for SwinEncoder's outputs. + """ +) +class MaskFormerSwinBaseModelOutput(ModelOutput): + r""" + hidden_states_spatial_dimensions (`tuple(tuple(int, int))`, *optional*): + A tuple containing the spatial dimension of each `hidden_state` needed to reshape the `hidden_states` to + `batch, channels, height, width`. Due to padding, their spatial size cannot inferred before the `forward` + method. + """ + + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states_spatial_dimensions: tuple[tuple[int, int]] = None + attentions: tuple[torch.FloatTensor] | None = None + + +# Copied from transformers.models.swin.modeling_swin.window_partition +def window_partition(input_feature, window_size): + """ + Partitions the given input into windows. + """ + batch_size, height, width, num_channels = input_feature.shape + input_feature = input_feature.view( + batch_size, height // window_size, window_size, width // window_size, window_size, num_channels + ) + windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels) + return windows + + +# Copied from transformers.models.swin.modeling_swin.window_reverse +def window_reverse(windows, window_size, height, width): + """ + Merges windows to produce higher resolution features. + """ + num_channels = windows.shape[-1] + windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels) + windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels) + return windows + + +# Copied from transformers.models.swin.modeling_swin.drop_path +def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +class MaskFormerSwinEmbeddings(nn.Module): + """ + Construct the patch and position embeddings. + """ + + def __init__(self, config): + super().__init__() + + self.patch_embeddings = MaskFormerSwinPatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.patch_grid = self.patch_embeddings.grid_size + + if config.use_absolute_embeddings: + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim)) + else: + self.position_embeddings = None + + self.norm = nn.LayerNorm(config.embed_dim) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.patch_size = config.patch_size + + # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embeddings + + class_pos_embed = self.position_embeddings[:, :1] + patch_pos_embed = self.position_embeddings[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward(self, pixel_values, interpolate_pos_encoding): + _, num_channels, height, width = pixel_values.shape + embeddings, output_dimensions = self.patch_embeddings(pixel_values) + embeddings = self.norm(embeddings) + + if self.position_embeddings is not None: + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embeddings + + embeddings = self.dropout(embeddings) + + return embeddings, output_dimensions + + +# Copied from transformers.models.swin.modeling_swin.SwinPatchEmbeddings with Swin->MaskFormerSwin +class MaskFormerSwinPatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.embed_dim + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1]) + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def maybe_pad(self, pixel_values, height, width): + if width % self.patch_size[1] != 0: + pad_values = (0, self.patch_size[1] - width % self.patch_size[1]) + pixel_values = nn.functional.pad(pixel_values, pad_values) + if height % self.patch_size[0] != 0: + pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0]) + pixel_values = nn.functional.pad(pixel_values, pad_values) + return pixel_values + + def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]: + _, num_channels, height, width = pixel_values.shape + # pad the input to be divisible by self.patch_size, if needed + pixel_values = self.maybe_pad(pixel_values, height, width) + embeddings = self.projection(pixel_values) + _, _, height, width = embeddings.shape + output_dimensions = (height, width) + embeddings = embeddings.flatten(2).transpose(1, 2) + + return embeddings, output_dimensions + + +# Copied from transformers.models.swin.modeling_swin.SwinPatchMerging +class MaskFormerSwinPatchMerging(nn.Module): + """ + Patch Merging Layer. + + Args: + input_resolution (`tuple[int]`): + Resolution of input feature. + dim (`int`): + Number of input channels. + norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`): + Normalization layer class. + """ + + def __init__(self, input_resolution: tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None: + super().__init__() + self.input_resolution = input_resolution + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def maybe_pad(self, input_feature, height, width): + should_pad = (height % 2 == 1) or (width % 2 == 1) + if should_pad: + pad_values = (0, 0, 0, width % 2, 0, height % 2) + input_feature = nn.functional.pad(input_feature, pad_values) + + return input_feature + + def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor: + height, width = input_dimensions + # `dim` is height * width + batch_size, dim, num_channels = input_feature.shape + + input_feature = input_feature.view(batch_size, height, width, num_channels) + # pad input to be divisible by width and height, if needed + input_feature = self.maybe_pad(input_feature, height, width) + # [batch_size, height/2, width/2, num_channels] + input_feature_0 = input_feature[:, 0::2, 0::2, :] + # [batch_size, height/2, width/2, num_channels] + input_feature_1 = input_feature[:, 1::2, 0::2, :] + # [batch_size, height/2, width/2, num_channels] + input_feature_2 = input_feature[:, 0::2, 1::2, :] + # [batch_size, height/2, width/2, num_channels] + input_feature_3 = input_feature[:, 1::2, 1::2, :] + # batch_size height/2 width/2 4*num_channels + input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1) + input_feature = input_feature.view(batch_size, -1, 4 * num_channels) # batch_size height/2*width/2 4*C + + input_feature = self.norm(input_feature) + input_feature = self.reduction(input_feature) + + return input_feature + + +# Copied from transformers.models.swin.modeling_swin.SwinDropPath with Swin->MaskFormerSwin +class MaskFormerSwinDropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float | None = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return f"p={self.drop_prob}" + + +# Copied from transformers.models.swin.modeling_swin.SwinSelfAttention with Swin->MaskFormerSwin +class MaskFormerSwinSelfAttention(nn.Module): + def __init__(self, config, dim, num_heads, window_size): + super().__init__() + if dim % num_heads != 0: + raise ValueError( + f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})" + ) + + self.num_attention_heads = num_heads + self.attention_head_size = int(dim / num_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.window_size = ( + window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size) + ) + + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads) + ) + + self.register_buffer("relative_position_index", self.create_relative_position_index()) + + self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias) + self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias) + self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor]: + batch_size, dim, num_channels = hidden_states.shape + hidden_shape = (batch_size, dim, -1, self.attention_head_size) + + query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2) + key_layer = self.key(hidden_states).view(hidden_shape).transpose(1, 2) + value_layer = self.value(hidden_states).view(hidden_shape).transpose(1, 2) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)] + relative_position_bias = relative_position_bias.view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + attention_scores = attention_scores + relative_position_bias.unsqueeze(0) + + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in MaskFormerSwinModel forward() function) + mask_shape = attention_mask.shape[0] + attention_scores = attention_scores.view( + batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim + ) + attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0) + attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim) + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + context_layer = torch.matmul(attention_probs, value_layer) + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + + return outputs + + def create_relative_position_index(self): + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += self.window_size[0] - 1 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) + return relative_position_index + + +# Copied from transformers.models.swin.modeling_swin.SwinSelfOutput with Swin->MaskFormerSwin +class MaskFormerSwinSelfOutput(nn.Module): + def __init__(self, config, dim): + super().__init__() + self.dense = nn.Linear(dim, dim) + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +# Copied from transformers.models.swin.modeling_swin.SwinAttention with Swin->MaskFormerSwin +class MaskFormerSwinAttention(nn.Module): + def __init__(self, config, dim, num_heads, window_size): + super().__init__() + self.self = MaskFormerSwinSelfAttention(config, dim, num_heads, window_size) + self.output = MaskFormerSwinSelfOutput(config, dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor]: + self_outputs = self.self(hidden_states, attention_mask, output_attentions) + attention_output = self.output(self_outputs[0], hidden_states) + outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.swin.modeling_swin.SwinIntermediate with Swin->MaskFormerSwin +class MaskFormerSwinIntermediate(nn.Module): + def __init__(self, config, dim): + super().__init__() + self.dense = nn.Linear(dim, int(config.mlp_ratio * dim)) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +# Copied from transformers.models.swin.modeling_swin.SwinOutput with Swin->MaskFormerSwin +class MaskFormerSwinOutput(nn.Module): + def __init__(self, config, dim): + super().__init__() + self.dense = nn.Linear(int(config.mlp_ratio * dim), dim) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +class MaskFormerSwinLayer(nn.Module): + def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0): + super().__init__() + self.shift_size = shift_size + self.window_size = config.window_size + self.input_resolution = input_resolution + self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps) + self.attention = MaskFormerSwinAttention(config, dim, num_heads, self.window_size) + self.drop_path = MaskFormerSwinDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps) + self.intermediate = MaskFormerSwinIntermediate(config, dim) + self.output = MaskFormerSwinOutput(config, dim) + + def get_attn_mask(self, input_resolution): + if self.shift_size > 0: + # calculate attention mask for SW-MSA + height, width = input_resolution + img_mask = torch.zeros((1, height, width, 1)) + height_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + width_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + count = 0 + for height_slice in height_slices: + for width_slice in width_slices: + img_mask[:, height_slice, width_slice, :] = count + count += 1 + + mask_windows = window_partition(img_mask, self.window_size) + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0) + else: + attn_mask = None + return attn_mask + + def maybe_pad(self, hidden_states, height, width): + pad_left = pad_top = 0 + pad_right = (self.window_size - width % self.window_size) % self.window_size + pad_bottom = (self.window_size - height % self.window_size) % self.window_size + pad_values = (0, 0, pad_left, pad_right, pad_top, pad_bottom) + hidden_states = nn.functional.pad(hidden_states, pad_values) + return hidden_states, pad_values + + def forward(self, hidden_states, input_dimensions, output_attentions=False): + height, width = input_dimensions + batch_size, dim, channels = hidden_states.size() + shortcut = hidden_states + + hidden_states = self.layernorm_before(hidden_states) + hidden_states = hidden_states.view(batch_size, height, width, channels) + # pad hidden_states to multiples of window size + hidden_states, pad_values = self.maybe_pad(hidden_states, height, width) + + _, height_pad, width_pad, _ = hidden_states.shape + # cyclic shift + if self.shift_size > 0: + shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + else: + shifted_hidden_states = hidden_states + + # partition windows + hidden_states_windows = window_partition(shifted_hidden_states, self.window_size) + hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels) + attn_mask = self.get_attn_mask((height_pad, width_pad)) + if attn_mask is not None: + attn_mask = attn_mask.to(hidden_states_windows.device) + + self_attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions) + + attention_output = self_attention_outputs[0] + + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels) + shifted_windows = window_reverse( + attention_windows, self.window_size, height_pad, width_pad + ) # B height' width' C + + # reverse cyclic shift + if self.shift_size > 0: + attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + attention_windows = shifted_windows + + was_padded = pad_values[3] > 0 or pad_values[5] > 0 + if was_padded: + attention_windows = attention_windows[:, :height, :width, :].contiguous() + + attention_windows = attention_windows.view(batch_size, height * width, channels) + + hidden_states = shortcut + self.drop_path(attention_windows) + + layer_output = self.layernorm_after(hidden_states) + layer_output = self.intermediate(layer_output) + layer_output = hidden_states + self.output(layer_output) + + outputs = (layer_output,) + outputs + + return outputs + + +class MaskFormerSwinStage(GradientCheckpointingLayer): + # Copied from transformers.models.swin.modeling_swin.SwinStage.__init__ with Swin->MaskFormerSwin + def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample): + super().__init__() + self.config = config + self.dim = dim + self.blocks = nn.ModuleList( + [ + MaskFormerSwinLayer( + config=config, + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + drop_path_rate=drop_path[i], + shift_size=0 if (i % 2 == 0) else config.window_size // 2, + ) + for i in range(depth) + ] + ) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm) + else: + self.downsample = None + + self.pointing = False + + def forward(self, hidden_states, input_dimensions, output_attentions=False, output_hidden_states=False): + all_hidden_states = () if output_hidden_states else None + + height, width = input_dimensions + for i, block_module in enumerate(self.blocks): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + block_hidden_states = block_module(hidden_states, input_dimensions, output_attentions) + + hidden_states = block_hidden_states[0] + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.downsample is not None: + height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2 + output_dimensions = (height, width, height_downsampled, width_downsampled) + hidden_states = self.downsample(hidden_states, input_dimensions) + else: + output_dimensions = (height, width, height, width) + + return hidden_states, output_dimensions, all_hidden_states + + +class MaskFormerSwinEncoder(nn.Module): + # Copied from transformers.models.swin.modeling_swin.SwinEncoder.__init__ with Swin->MaskFormerSwin + def __init__(self, config, grid_size): + super().__init__() + self.num_layers = len(config.depths) + self.config = config + dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")] + self.layers = nn.ModuleList( + [ + MaskFormerSwinStage( + config=config, + dim=int(config.embed_dim * 2**i_layer), + input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)), + depth=config.depths[i_layer], + num_heads=config.num_heads[i_layer], + drop_path=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])], + downsample=MaskFormerSwinPatchMerging if (i_layer < self.num_layers - 1) else None, + ) + for i_layer in range(self.num_layers) + ] + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + input_dimensions, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ) -> tuple | MaskFormerSwinBaseModelOutput: + all_hidden_states = () if output_hidden_states else None + all_input_dimensions = () + all_self_attentions = () if output_attentions else None + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + for i, layer_module in enumerate(self.layers): + layer_hidden_states, output_dimensions, layer_all_hidden_states = layer_module( + hidden_states, + input_dimensions, + output_attentions, + output_hidden_states, + ) + + input_dimensions = (output_dimensions[-2], output_dimensions[-1]) + all_input_dimensions += (input_dimensions,) + if output_hidden_states: + all_hidden_states += (layer_all_hidden_states,) + + hidden_states = layer_hidden_states + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_all_hidden_states[1],) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + + return MaskFormerSwinBaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + hidden_states_spatial_dimensions=all_input_dimensions, + attentions=all_self_attentions, + ) + + +@auto_docstring +class MaskFormerSwinPreTrainedModel(PreTrainedModel): + config: MaskFormerSwinConfig + base_model_prefix = "model" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MaskFormerSwinStage"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, MaskFormerSwinEmbeddings): + if module.position_embeddings is not None: + init.zeros_(module.position_embeddings) + elif isinstance(module, MaskFormerSwinSelfAttention): + init.zeros_(module.relative_position_bias_table) + init.copy_(module.relative_position_index, module.create_relative_position_index()) + + +class MaskFormerSwinModel(MaskFormerSwinPreTrainedModel): + def __init__(self, config, add_pooling_layer=True): + super().__init__(config) + self.config = config + self.num_layers = len(config.depths) + self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1)) + + self.embeddings = MaskFormerSwinEmbeddings(config) + self.encoder = MaskFormerSwinEncoder(config, self.embeddings.patch_grid) + + self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps) + self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.patch_embeddings + + def forward( + self, + pixel_values=None, + output_attentions=None, + output_hidden_states=None, + interpolate_pos_encoding=False, + return_dict=None, + **kwargs, + ) -> tuple | MaskFormerSwinModelOutputWithPooling: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + embedding_output, input_dimensions = self.embeddings( + pixel_values, interpolate_pos_encoding=interpolate_pos_encoding + ) + + encoder_outputs = self.encoder( + embedding_output, + input_dimensions, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = encoder_outputs.last_hidden_state if return_dict else encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + + pooled_output = None + if self.pooler is not None: + pooled_output = self.pooler(sequence_output.transpose(1, 2)) + pooled_output = torch.flatten(pooled_output, 1) + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + hidden_states_spatial_dimensions = (input_dimensions,) + encoder_outputs.hidden_states_spatial_dimensions + + return MaskFormerSwinModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + hidden_states_spatial_dimensions=hidden_states_spatial_dimensions, + attentions=encoder_outputs.attentions, + ) + + +class MaskFormerSwinBackbone(BackboneMixin, MaskFormerSwinPreTrainedModel): + """ + MaskFormerSwin backbone, designed especially for the MaskFormer framework. + + This classes reshapes `hidden_states` from (`batch_size, sequence_length, hidden_size)` to (`batch_size, + num_channels, height, width)`). It also adds additional layernorms after each stage. + + Args: + config (`MaskFormerSwinConfig`): + The configuration used by [`MaskFormerSwinModel`]. + """ + + def __init__(self, config: MaskFormerSwinConfig): + super().__init__(config) + + self.model = MaskFormerSwinModel(config) + if "stem" in self.out_features: + raise ValueError("This backbone does not support 'stem' in the `out_features`.") + self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))] + self.hidden_states_norms = nn.ModuleList( + [nn.LayerNorm(num_channels) for num_channels in self.num_features[1:]] + ) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @filter_output_hidden_states + def forward( + self, + pixel_values: Tensor, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BackboneOutput: + return_dict = return_dict if return_dict is not None else self.config.return_dict + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + + outputs = self.model( + pixel_values, output_hidden_states=True, output_attentions=output_attentions, return_dict=True + ) + + # we skip the stem + hidden_states = outputs.hidden_states[1:] + + # we need to reshape the hidden states to their original spatial dimensions + # spatial dimensions contains all the heights and widths of each stage, including after the embeddings + spatial_dimensions: tuple[tuple[int, int]] = outputs.hidden_states_spatial_dimensions + feature_maps = () + for i, (hidden_state, stage, (height, width)) in enumerate( + zip(hidden_states, self.stage_names[1:], spatial_dimensions) + ): + norm = self.hidden_states_norms[i] + # the last element correspond to the layer's last block output but before patch merging + hidden_state_unpolled = hidden_state[-1] + hidden_state_norm = norm(hidden_state_unpolled) + # the pixel decoder (FPN) expects 3D tensors (features) + batch_size, _, hidden_size = hidden_state_norm.shape + # reshape "b (h w) d -> b d h w" + hidden_state_permuted = ( + hidden_state_norm.permute(0, 2, 1).view((batch_size, hidden_size, height, width)).contiguous() + ) + if stage in self.out_features: + feature_maps += (hidden_state_permuted,) + + if not return_dict: + output = (feature_maps,) + if output_hidden_states: + output += (outputs.hidden_states,) + if output_attentions: + output += (outputs.attentions,) + return output + + return BackboneOutput( + feature_maps=feature_maps, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=outputs.attentions, + ) + + +__all__ = ["MaskFormerSwinBackbone", "MaskFormerSwinModel", "MaskFormerSwinPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/maskformer/modular_maskformer.py b/third_party/transformers/src/transformers/models/maskformer/modular_maskformer.py new file mode 100644 index 0000000000000000000000000000000000000000..06705906c89184ffd593c9f39b8ccd03d166698e --- /dev/null +++ b/third_party/transformers/src/transformers/models/maskformer/modular_maskformer.py @@ -0,0 +1,1526 @@ +# Copyright 2022 Meta Platforms, Inc.s 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 MaskFormer model.""" + +import math +from dataclasses import dataclass +from numbers import Number + +import numpy as np +import torch +from huggingface_hub.dataclasses import strict +from torch import Tensor, nn + +from ... import initialization as init +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +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, +) +from ..auto import CONFIG_MAPPING, AutoBackbone, AutoConfig +from ..detr.configuration_detr import DetrConfig +from ..detr.modeling_detr import DetrDecoder, DetrDecoderOutput +from .configuration_maskformer_swin import MaskFormerSwinConfig + + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import reduce + +if is_scipy_available(): + from scipy.optimize import linear_sum_assignment + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") +@strict +class MaskFormerDetrConfig(DetrConfig): + model_type = "detr" + + +@auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") +@strict +class MaskFormerConfig(PreTrainedConfig): + r""" + fpn_feature_size (`int`, *optional*, defaults to 256): + The Feature Pyramid Network's features size. + mask_feature_size (`int`, *optional*, defaults to 256): + The masks' features size, this value will also be used to specify the Feature Pyramid Network features' + size. + decoder_config (`Dict`, *optional*): + The configuration passed to the transformer decoder model, if unset the base config for `detr-resnet-50` + will be used. + cross_entropy_weight (`float`, *optional*, defaults to 1.0): + The weight for the cross entropy loss. + output_auxiliary_logits (`bool`, *optional*): + Should the model output its `auxiliary_logits` or not. + + Raises: + `ValueError`: + Raised if the backbone model type selected is not in `["swin"]` or the decoder model type selected is not + in `["detr"]` + + Examples: + + ```python + >>> from transformers import MaskFormerConfig, MaskFormerModel + + >>> # Initializing a MaskFormer facebook/maskformer-swin-base-ade configuration + >>> configuration = MaskFormerConfig() + + >>> # Initializing a model (with random weights) from the facebook/maskformer-swin-base-ade style configuration + >>> model = MaskFormerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + + """ + + model_type = "maskformer" + sub_configs = {"backbone_config": AutoConfig, "decoder_config": AutoConfig} + attribute_map = {"hidden_size": "mask_feature_size"} + backbones_supported = ["resnet", "swin"] + decoders_supported = ["detr"] + + fpn_feature_size: int = 256 + mask_feature_size: int = 256 + no_object_weight: float = 0.1 + use_auxiliary_loss: bool = False + backbone_config: dict | PreTrainedConfig | None = None + decoder_config: dict | PreTrainedConfig | None = None + init_std: float = 0.02 + init_xavier_std: float = 1.0 + dice_weight: float = 1.0 + cross_entropy_weight: float = 1.0 + mask_weight: float = 20.0 + output_auxiliary_logits: bool | None = None + + def __post_init__(self, **kwargs): + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="swin", + default_config_kwargs={ + "depths": [2, 2, 18, 2], + "drop_path_rate": 0.3, + "image_size": 384, + "embed_dim": 128, + "num_heads": [4, 8, 16, 32], + "window_size": 12, + "out_features": ["stage1", "stage2", "stage3", "stage4"], + }, + **kwargs, + ) + + # verify that the backbone is supported + if self.backbone_config is not None and self.backbone_config.model_type not in self.backbones_supported: + logger.warning_once( + f"Backbone {self.backbone_config.model_type} is not a supported model and may not be compatible with MaskFormer. " + f"Supported model types: {','.join(self.backbones_supported)}" + ) + + if self.decoder_config is None: + # fall back to https://huggingface.co/facebook/detr-resnet-50 + self.decoder_config = MaskFormerDetrConfig() + else: + # verify that the decoder is supported + decoder_type = ( + self.decoder_config.pop("model_type") + if isinstance(self.decoder_config, dict) + else self.decoder_config.model_type + ) + if decoder_type not in self.decoders_supported: + raise ValueError( + f"Transformer Decoder {decoder_type} not supported, please use one of" + f" {','.join(self.decoders_supported)}" + ) + if isinstance(self.decoder_config, dict): + config_class = CONFIG_MAPPING[decoder_type] + self.decoder_config = config_class.from_dict(self.decoder_config) + + self.num_attention_heads = self.decoder_config.encoder_attention_heads + self.num_hidden_layers = self.decoder_config.num_hidden_layers + super().__post_init__(**kwargs) + + +class DetrDecoderOutput(DetrDecoderOutput): + pass + + +@dataclass +@auto_docstring( + custom_intro=""" + MaskFormer'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 MaskFormerSwin Transformer and the `decoder` is a Feature + Pyramid Network (FPN). + + The `encoder_last_hidden_state` are referred on the paper as **images features**, while `decoder_last_hidden_state` + as **pixel embeddings** + """ +) +class MaskFormerPixelLevelModuleOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder. + decoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the 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 model at + the output of each stage. + 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 model at + the output of each stage. + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + MaskFormer's pixel decoder module output, practically a Feature Pyramid Network. It returns the last hidden state + and (optionally) the hidden states. + """ +) +class MaskFormerPixelDecoderOutput(ModelOutput): + r""" + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the model. + """ + + last_hidden_state: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`MaskFormerModel`]. This class returns all the needed hidden states to compute the logits. + """ +) +class MaskFormerModelOutput(ModelOutput): + r""" + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). + transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Last hidden states (final feature map) of the last stage of the transformer decoder model. + 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. + 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and + `decoder_hidden_states` + """ + + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Class for outputs of [`MaskFormerForInstanceSegmentation`]. + + This output can be directly passed to [`~MaskFormerImageProcessor.post_process_semantic_segmentation`] or + [`~MaskFormerImageProcessor.post_process_instance_segmentation`] or + [`~MaskFormerImageProcessor.post_process_panoptic_segmentation`] depending on the task. Please, see + [`~MaskFormerImageProcessor] for details regarding usage. + """ +) +class MaskFormerForInstanceSegmentationOutput(ModelOutput): + r""" + loss (`torch.Tensor`, *optional*): + The computed loss, returned when labels are present. + class_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each + query. Note the `+ 1` is needed because we incorporate the null class. + masks_queries_logits (`torch.FloatTensor`): + A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each + query. + auxiliary_logits (`Dict[str, torch.FloatTensor]`, *optional*, returned when `output_auxiliary_logits=True`): + Dictionary containing auxiliary predictions for each decoder layer when auxiliary losses are enabled. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the encoder model (backbone). + pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). + transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Last hidden states (final feature map) of the last stage of the transformer decoder model. + 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 of the transformer decoder at the output + of each stage. + 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and + `decoder_hidden_states`. + """ + + loss: torch.FloatTensor | None = None + class_queries_logits: torch.FloatTensor | None = None + masks_queries_logits: torch.FloatTensor | None = None + auxiliary_logits: torch.FloatTensor | None = None + encoder_last_hidden_state: torch.FloatTensor | None = None + pixel_decoder_last_hidden_state: torch.FloatTensor | None = None + transformer_decoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + + +def upsample_like(pixel_values: Tensor, like: Tensor, mode: str = "bilinear") -> Tensor: + """ + An utility function that upsamples `pixel_values` to match the dimension of `like`. + + Args: + pixel_values (`torch.Tensor`): + The tensor we wish to upsample. + like (`torch.Tensor`): + The tensor we wish to use as size target. + mode (str, *optional*, defaults to `"bilinear"`): + The interpolation mode. + + Returns: + `torch.Tensor`: The upsampled tensor + """ + _, _, height, width = like.shape + upsampled = nn.functional.interpolate(pixel_values, size=(height, width), mode=mode, align_corners=False) + return upsampled + + +# refactored from original implementation +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 + + +# refactored from original implementation +def sigmoid_focal_loss( + inputs: Tensor, labels: Tensor, num_masks: int, alpha: float = 0.25, gamma: float = 2 +) -> Tensor: + r""" + Focal loss proposed in [Focal Loss for Dense Object Detection](https://huggingface.co/papers/1708.02002) originally used in + RetinaNet. The loss is computed as follows: + + $$ \mathcal{L}_{\text{focal loss} = -(1 - p_t)^{\gamma}\log{(p_t)} $$ + + where \\(CE(p_t) = -\log{(p_t)}}\\), CE is the standard Cross Entropy Loss + + Please refer to equation (1,2,3) of the paper for a better understanding. + + 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). + num_masks (`int`): + The number of masks present in the current batch, used for normalization. + alpha (float, *optional*, defaults to 0.25): + Weighting factor in range (0,1) to balance positive vs negative examples. + gamma (float, *optional*, defaults to 2.0): + Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. + + Returns: + `torch.Tensor`: The computed loss. + """ + criterion = nn.BCEWithLogitsLoss(reduction="none") + probs = inputs.sigmoid() + cross_entropy_loss = criterion(inputs, labels) + p_t = probs * labels + (1 - probs) * (1 - labels) + loss = cross_entropy_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * labels + (1 - alpha) * (1 - labels) + loss = alpha_t * loss + + loss = loss.mean(1).sum() / num_masks + return loss + + +# refactored from original implementation +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 + + +# refactored from original implementation +def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float = 0.25, gamma: float = 2.0) -> Tensor: + r""" + A pair wise version of the focal loss, see `sigmoid_focal_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). + alpha (float, *optional*, defaults to 0.25): + Weighting factor in range (0,1) to balance positive vs negative examples. + gamma (float, *optional*, defaults to 2.0): + Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. + + Returns: + `torch.Tensor`: The computed loss between each pairs. + """ + if alpha < 0: + raise ValueError("alpha must be positive") + + height_and_width = inputs.shape[1] + + criterion = nn.BCEWithLogitsLoss(reduction="none") + prob = inputs.sigmoid() + cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) + focal_pos = ((1 - prob) ** gamma) * cross_entropy_loss_pos + focal_pos *= alpha + + cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) + + focal_neg = (prob**gamma) * cross_entropy_loss_neg + focal_neg *= 1 - alpha + + loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T) + + return loss / height_and_width + + +class MaskFormerDetrDecoder(DetrDecoder): + pass + + +# refactored from original implementation +class MaskFormerHungarianMatcher(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): + """Creates the matcher + + 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 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 + """ + 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 + + @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_target_boxes). + """ + indices: list[tuple[np.array]] = [] + + 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): + # downsample the target mask, save memory + target_mask = nn.functional.interpolate(target_mask[:, None], size=pred_mask.shape[-2:], mode="nearest") + 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] + # flatten spatial dimension "q h w -> q (h w)" + pred_mask_flat = pred_mask.flatten(1) # [num_queries, height*width] + # same for target_mask "c h w -> c (h w)" + target_mask_flat = target_mask[:, 0].flatten(1) # [num_total_labels, height*width] + # compute the focal loss between each mask pairs -> shape (num_queries, num_labels) + cost_mask = pair_wise_sigmoid_focal_loss(pred_mask_flat, target_mask_flat) + # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels) + cost_dice = pair_wise_dice_loss(pred_mask_flat, target_mask_flat) + # final cost matrix + cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice + # 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 __repr__(self): + head = "Matcher " + self.__class__.__name__ + body = [ + f"cost_class: {self.cost_class}", + f"cost_mask: {self.cost_mask}", + f"cost_dice: {self.cost_dice}", + ] + _repr_indent = 4 + lines = [head] + [" " * _repr_indent + line for line in body] + return "\n".join(lines) + + +# copied and adapted from original implementation +class MaskFormerLoss(nn.Module): + def __init__( + self, + num_labels: int, + matcher: MaskFormerHungarianMatcher, + weight_dict: dict[str, float], + eos_coef: float, + ): + """ + The MaskFormer 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: + num_labels (`int`): + The number of classes. + matcher (`MaskFormerHungarianMatcher`): + 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. + """ + + super().__init__() + requires_backends(self, ["scipy"]) + self.num_labels = num_labels + self.matcher = matcher + self.weight_dict = weight_dict + self.eos_coef = eos_coef + empty_weight = torch.ones(self.num_labels + 1) + empty_weight[-1] = self.eos_coef + self.register_buffer("empty_weight", empty_weight) + + 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_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_labels, dtype=torch.int64, device=pred_logits.device + ) + target_classes[idx] = target_classes_o + # target_classes is a (batch_size, num_labels, num_queries), we need to permute pred_logits "b q c -> b c q" + 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 focal 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] + # upsample predictions to the target size, we have to add one dim to use interpolate + pred_masks = nn.functional.interpolate( + pred_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False + ) + pred_masks = pred_masks[:, 0].flatten(1) + + target_masks = target_masks.flatten(1) + losses = { + "loss_mask": sigmoid_focal_loss(pred_masks, target_masks, num_masks), + "loss_dice": dice_loss(pred_masks, target_masks, num_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 forward( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: list[Tensor], + class_labels: list[Tensor], + auxiliary_predictions: dict[str, Tensor] | None = None, + ) -> 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` + 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 [`MaskFormerConfig`], then it contains the logits from the + inner layers of the Detr's Decoder. + + 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 focal 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 [`MaskFormerConfig`], 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: Number = self.get_num_masks(class_labels, device=class_labels[0].device) + # get all the losses + losses: dict[str, Tensor] = { + **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), + **self.loss_labels(class_queries_logits, class_labels, indices), + } + # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. + if auxiliary_predictions is not None: + for idx, aux_outputs in enumerate(auxiliary_predictions): + masks_queries_logits = aux_outputs["masks_queries_logits"] + class_queries_logits = aux_outputs["class_queries_logits"] + loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels) + loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} + losses.update(loss_dict) + + return losses + + def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Computes the average number of target masks across the batch, for normalization purposes. + """ + num_masks = sum(len(classes) for classes in class_labels) + num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device) + world_size = 1 + if is_accelerate_available(): + if PartialState._shared_state != {}: + num_masks = reduce(num_masks) + world_size = PartialState().num_processes + + num_masks = torch.clamp(num_masks / world_size, min=1) + return num_masks + + +class MaskFormerFPNConvLayer(nn.Module): + def __init__(self, in_features: int, out_features: int, kernel_size: int = 3, padding: int = 1): + """ + A basic module that executes conv - norm - in sequence used in MaskFormer. + + Args: + in_features (`int`): + The number of input features (channels). + out_features (`int`): + The number of outputs features (channels). + """ + super().__init__() + self.layers = [ + nn.Conv2d(in_features, out_features, kernel_size=kernel_size, padding=padding, bias=False), + nn.GroupNorm(32, out_features), + nn.ReLU(inplace=True), + ] + for i, layer in enumerate(self.layers): + # Provide backwards compatibility from when the class inherited from nn.Sequential + # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. + # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. + # self.my_layer_name = Layer() + # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register + # explicitly + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class MaskFormerFPNLayer(nn.Module): + def __init__(self, in_features: int, lateral_features: int): + """ + A Feature Pyramid Network Layer (FPN) layer. It creates a feature map by aggregating features from the previous + and backbone layer. Due to the spatial mismatch, the tensor coming from the previous layer is upsampled. + + Args: + in_features (`int`): + The number of input features (channels). + lateral_features (`int`): + The number of lateral features (channels). + """ + super().__init__() + self.proj = nn.Sequential( + nn.Conv2d(lateral_features, in_features, kernel_size=1, padding=0, bias=False), + nn.GroupNorm(32, in_features), + ) + + self.block = MaskFormerFPNConvLayer(in_features, in_features) + + def forward(self, down: Tensor, left: Tensor) -> Tensor: + left = self.proj(left) + down = nn.functional.interpolate(down, size=left.shape[-2:], mode="nearest") + down += left + down = self.block(down) + return down + + +class MaskFormerFPNModel(nn.Module): + def __init__(self, in_features: int, lateral_widths: list[int], feature_size: int = 256): + """ + Feature Pyramid Network, given an input tensor and a set of feature map of different feature/spatial size, it + creates a list of feature maps with the same feature size. + + Args: + in_features (`int`): + The number of input features (channels). + lateral_widths (`list[int]`): + A list with the features (channels) size of each lateral connection. + feature_size (int, *optional*, defaults to 256): + The features (channels) of the resulting feature maps. + """ + super().__init__() + self.stem = MaskFormerFPNConvLayer(in_features, feature_size) + self.layers = nn.Sequential( + *[MaskFormerFPNLayer(feature_size, lateral_width) for lateral_width in lateral_widths[::-1]] + ) + + def forward(self, features: list[Tensor]) -> list[Tensor]: + fpn_features = [] + last_feature = features[-1] + other_features = features[:-1] + output = self.stem(last_feature) + for layer, left in zip(self.layers, other_features[::-1]): + output = layer(output, left) + fpn_features.append(output) + return fpn_features + + +class MaskFormerPixelDecoder(nn.Module): + def __init__(self, *args, feature_size: int = 256, mask_feature_size: int = 256, **kwargs): + r""" + Pixel Decoder Module proposed in [Per-Pixel Classification is Not All You Need for Semantic + Segmentation](https://huggingface.co/papers/2107.06278). It first runs the backbone's features into a Feature Pyramid + Network creating a list of feature maps. Then, it projects the last one to the correct `mask_size`. + + Args: + feature_size (`int`, *optional*, defaults to 256): + The feature size (channel dimension) of the FPN feature maps. + mask_feature_size (`int`, *optional*, defaults to 256): + The features (channels) of the target masks size \\(C_{\epsilon}\\) in the paper. + """ + super().__init__() + + self.fpn = MaskFormerFPNModel(*args, feature_size=feature_size, **kwargs) + self.mask_projection = nn.Conv2d(feature_size, mask_feature_size, kernel_size=3, padding=1) + + def forward( + self, features: list[Tensor], output_hidden_states: bool = False, return_dict: bool = True + ) -> MaskFormerPixelDecoderOutput: + fpn_features = self.fpn(features) + # we use the last feature map + last_feature_projected = self.mask_projection(fpn_features[-1]) + + if not return_dict: + return (last_feature_projected, tuple(fpn_features)) if output_hidden_states else (last_feature_projected,) + + return MaskFormerPixelDecoderOutput( + last_hidden_state=last_feature_projected, hidden_states=tuple(fpn_features) if output_hidden_states else () + ) + + +# copied and adapted from original implementation, also practically equal to DetrSinePositionEmbedding +class MaskFormerSinePositionEmbedding(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 + + +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 MaskformerMLPPredictionHead(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3): + """ + A classic Multi Layer Perceptron (MLP). + + Args: + input_dim (`int`): + The input dimensions. + hidden_dim (`int`): + The hidden dimensions. + output_dim (`int`): + The output dimensions. + num_layers (int, *optional*, defaults to 3): + The number of layers. + """ + super().__init__() + in_dims = [input_dim] + [hidden_dim] * (num_layers - 1) + out_dims = [hidden_dim] * (num_layers - 1) + [output_dim] + + self.layers = [] + for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)): + activation = nn.ReLU() if i < num_layers - 1 else nn.Identity() + layer = PredictionBlock(in_dim, out_dim, activation=activation) + self.layers.append(layer) + # Provide backwards compatibility from when the class inherited from nn.Sequential + # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. + # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. + # self.my_layer_name = Layer() + # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register + # explicitly + self.add_module(str(i), layer) + + def forward(self, input: Tensor) -> Tensor: + hidden_state = input + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class MaskFormerPixelLevelModule(nn.Module): + def __init__(self, config: MaskFormerConfig): + """ + Pixel Level Module proposed in [Per-Pixel Classification is Not All You Need for Semantic + Segmentation](https://huggingface.co/papers/2107.06278). It runs the input image through a backbone and a pixel + decoder, generating an image feature map and pixel embeddings. + + Args: + config ([`MaskFormerConfig`]): + The configuration used to instantiate this model. + """ + super().__init__() + if getattr(config, "backbone_config") is not None and config.backbone_config.model_type == "swin": + # for backwards compatibility + backbone_config = config.backbone_config + backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict()) + backbone_config.out_features = ["stage1", "stage2", "stage3", "stage4"] + config.backbone_config = backbone_config + self.encoder = AutoBackbone.from_config(config=config.backbone_config) + + feature_channels = self.encoder.channels + self.decoder = MaskFormerPixelDecoder( + in_features=feature_channels[-1], + feature_size=config.fpn_feature_size, + mask_feature_size=config.mask_feature_size, + lateral_widths=feature_channels[:-1], + ) + + def forward( + self, pixel_values: Tensor, output_hidden_states: bool = False, return_dict: bool = True + ) -> MaskFormerPixelLevelModuleOutput: + features = self.encoder(pixel_values).feature_maps + decoder_output = self.decoder(features, output_hidden_states, return_dict=return_dict) + + if not return_dict: + last_hidden_state = decoder_output[0] + outputs = (features[-1], last_hidden_state) + if output_hidden_states: + hidden_states = decoder_output[1] + outputs = outputs + (tuple(features),) + (hidden_states,) + return outputs + + return MaskFormerPixelLevelModuleOutput( + # the last feature is actually the output from the last layer + encoder_last_hidden_state=features[-1], + decoder_last_hidden_state=decoder_output.last_hidden_state, + encoder_hidden_states=tuple(features) if output_hidden_states else (), + decoder_hidden_states=decoder_output.hidden_states if output_hidden_states else (), + ) + + +class MaskFormerTransformerModule(nn.Module): + """ + The MaskFormer's transformer module. + """ + + def __init__(self, in_features: int, config: MaskFormerConfig): + super().__init__() + hidden_size = config.decoder_config.hidden_size + should_project = in_features != hidden_size + self.position_embedder = MaskFormerSinePositionEmbedding(num_pos_feats=hidden_size // 2, normalize=True) + self.queries_embedder = nn.Embedding(config.decoder_config.num_queries, hidden_size) + self.input_projection = nn.Conv2d(in_features, hidden_size, kernel_size=1) if should_project else None + self.decoder = MaskFormerDetrDecoder(config=config.decoder_config) + + def forward( + self, + image_features: Tensor, + output_hidden_states: bool = False, + output_attentions: bool = False, + return_dict: bool | None = None, + ) -> DetrDecoderOutput: + if self.input_projection is not None: + image_features = self.input_projection(image_features) + object_queries = self.position_embedder(image_features.shape, image_features.device, image_features.dtype) + # repeat the queries "q c -> b q c" + batch_size = image_features.shape[0] + queries_embeddings = self.queries_embedder.weight.unsqueeze(0).repeat(batch_size, 1, 1) + inputs_embeds = torch.zeros_like(queries_embeddings, requires_grad=self.training) + + # torch.export.export does no support requires_grad + if self.training: + inputs_embeds.requires_grad_(True) + + batch_size, num_channels, height, width = image_features.shape + # rearrange both image_features and object_queries "b c h w -> b (h w) c" + image_features = image_features.view(batch_size, num_channels, height * width).permute(0, 2, 1) + object_queries = object_queries.view(batch_size, num_channels, height * width).permute(0, 2, 1) + + decoder_output: DetrDecoderOutput = self.decoder( + inputs_embeds=inputs_embeds, + attention_mask=None, + encoder_hidden_states=image_features, + encoder_attention_mask=None, + spatial_position_embeddings=object_queries, + object_queries_position_embeddings=queries_embeddings, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + return decoder_output + + +@auto_docstring +class MaskFormerPreTrainedModel(PreTrainedModel): + config: MaskFormerConfig + 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, MaskFormerTransformerModule): + if module.input_projection is not None: + init.xavier_uniform_(module.input_projection.weight, gain=xavier_std) + init.constant_(module.input_projection.bias, 0) + # FPN + elif isinstance(module, MaskFormerFPNModel): + init.xavier_uniform_(module.stem.get_submodule("0").weight, gain=xavier_std) + + elif isinstance(module, MaskFormerFPNLayer): + init.xavier_uniform_(module.proj[0].weight, gain=xavier_std) + + elif isinstance(module, MaskFormerFPNConvLayer): + init.xavier_uniform_(module.get_submodule("0").weight, gain=xavier_std) + # The MLP head + elif isinstance(module, MaskformerMLPPredictionHead): + # I was not able to find the correct initializer in the original implementation + # we'll use xavier + 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.LayerNorm): + init.zeros_(module.bias) + init.ones_(module.weight) + # copied from DETR + if 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.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, MaskFormerLoss): + empty_weight = torch.ones(module.num_labels + 1) + empty_weight[-1] = module.eos_coef + init.copy_(module.empty_weight, empty_weight) + + +@auto_docstring +class MaskFormerModel(MaskFormerPreTrainedModel): + def __init__(self, config: MaskFormerConfig): + super().__init__(config) + self.pixel_level_module = MaskFormerPixelLevelModule(config) + self.transformer_module = MaskFormerTransformerModule( + in_features=self.pixel_level_module.encoder.channels[-1], config=config + ) + + self.post_init() + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + pixel_mask: Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskFormerModelOutput: + r""" + Examples: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerModel + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") + >>> model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-base-ade") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = image_processor(image, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # the decoder of MaskFormer outputs hidden states of shape (batch_size, num_queries, hidden_size) + >>> transformer_decoder_last_hidden_state = outputs.transformer_decoder_last_hidden_state + >>> list(transformer_decoder_last_hidden_state.shape) + [1, 100, 256] + ```""" + + 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, return_dict=return_dict + ) + image_features = pixel_level_module_output[0] + pixel_embeddings = pixel_level_module_output[1] + + transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions) + queries = transformer_module_output.last_hidden_state + + encoder_hidden_states = None + pixel_decoder_hidden_states = None + transformer_decoder_hidden_states = None + hidden_states = None + + if output_hidden_states: + encoder_hidden_states = pixel_level_module_output[2] + pixel_decoder_hidden_states = pixel_level_module_output[3] + transformer_decoder_hidden_states = transformer_module_output[1] + hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states + + output = MaskFormerModelOutput( + encoder_last_hidden_state=image_features, + pixel_decoder_last_hidden_state=pixel_embeddings, + transformer_decoder_last_hidden_state=queries, + encoder_hidden_states=encoder_hidden_states, + pixel_decoder_hidden_states=pixel_decoder_hidden_states, + transformer_decoder_hidden_states=transformer_decoder_hidden_states, + hidden_states=hidden_states, + attentions=transformer_module_output.attentions, + ) + + if not return_dict: + output = tuple(v for v in output.values()) + + return output + + +class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel): + def __init__(self, config: MaskFormerConfig): + super().__init__(config) + self.model = MaskFormerModel(config) + hidden_size = config.decoder_config.hidden_size + # + 1 because we add the "null" class + self.class_predictor = nn.Linear(hidden_size, config.num_labels + 1) + self.mask_embedder = MaskformerMLPPredictionHead(hidden_size, hidden_size, config.mask_feature_size) + + self.matcher = MaskFormerHungarianMatcher( + cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight + ) + + self.weight_dict: dict[str, float] = { + "loss_cross_entropy": config.cross_entropy_weight, + "loss_mask": config.mask_weight, + "loss_dice": config.dice_weight, + } + + self.criterion = MaskFormerLoss( + config.num_labels, + matcher=self.matcher, + weight_dict=self.weight_dict, + eos_coef=config.no_object_weight, + ) + + self.post_init() + + def get_loss_dict( + self, + masks_queries_logits: Tensor, + class_queries_logits: Tensor, + mask_labels: Tensor, + class_labels: Tensor, + auxiliary_logits: dict[str, Tensor], + ) -> dict[str, Tensor]: + loss_dict: dict[str, Tensor] = self.criterion( + masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits + ) + # weight each loss by `self.weight_dict[]` including auxiliary losses + for key, weight in self.weight_dict.items(): + for loss_key, loss in loss_dict.items(): + if key in loss_key: + loss *= weight + + return loss_dict + + def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: + return sum(loss_dict.values()) + + def get_logits(self, outputs: MaskFormerModelOutput) -> tuple[Tensor, Tensor, dict[str, Tensor]]: + pixel_embeddings = outputs.pixel_decoder_last_hidden_state + # get the auxiliary predictions (one for each decoder's layer) + auxiliary_logits: list[str, Tensor] = [] + + # This code is a little bit cumbersome, an improvement can be to return a list of predictions. If we have auxiliary loss then we are going to return more than one element in the list + if self.config.use_auxiliary_loss: + stacked_transformer_decoder_outputs = torch.stack(outputs.transformer_decoder_hidden_states) + classes = self.class_predictor(stacked_transformer_decoder_outputs) + class_queries_logits = classes[-1] + # get the masks + mask_embeddings = self.mask_embedder(stacked_transformer_decoder_outputs) + binaries_masks = torch.einsum("lbqc, bchw -> lbqhw", mask_embeddings, pixel_embeddings) + + masks_queries_logits = binaries_masks[-1] + # go til [:-1] because the last one is always used + for aux_binary_masks, aux_classes in zip(binaries_masks[:-1], classes[:-1]): + auxiliary_logits.append( + {"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes} + ) + + else: + transformer_decoder_hidden_states = outputs.transformer_decoder_last_hidden_state + classes = self.class_predictor(transformer_decoder_hidden_states) + class_queries_logits = classes + # get the masks + mask_embeddings = self.mask_embedder(transformer_decoder_hidden_states) + # sum up over the channels + masks_queries_logits = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings) + + return class_queries_logits, masks_queries_logits, auxiliary_logits + + @auto_docstring + def forward( + self, + pixel_values: Tensor, + mask_labels: list[Tensor] | None = None, + class_labels: list[Tensor] | None = None, + pixel_mask: Tensor | None = None, + output_auxiliary_logits: bool | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskFormerForInstanceSegmentationOutput: + r""" + mask_labels (`list[torch.Tensor]`, *optional*): + List of mask labels of shape `(num_labels, height, width)` to be fed to a model + class_labels (`list[torch.LongTensor]`, *optional*): + list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the + labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. + output_auxiliary_logits (`bool`, *optional*): + Whether or not to output auxiliary logits. + + Examples: + + Semantic segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") + >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-ade") + + >>> url = ( + ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg" + ... ) + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> 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 image_processor for postprocessing + >>> predicted_semantic_map = image_processor.post_process_semantic_segmentation( + ... outputs, target_sizes=[(image.height, image.width)] + ... )[0] + + >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) + >>> list(predicted_semantic_map.shape) + [512, 683] + ``` + + Panoptic segmentation example: + + ```python + >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + + >>> # load MaskFormer fine-tuned on COCO panoptic segmentation + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-coco") + >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-coco") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> 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 image_processor for postprocessing + >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(image.height, image.width)])[0] + + >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) + >>> predicted_panoptic_map = result["segmentation"] + >>> list(predicted_panoptic_map.shape) + [480, 640] + ``` + """ + + 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 + + raw_outputs = self.model( + pixel_values, + pixel_mask, + output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, + return_dict=return_dict, + output_attentions=output_attentions, + ) + # We need to have raw_outputs optionally be returned as a dict to use torch.compile. For backwards + # compatibility we convert to a dataclass for the rest of the model logic + outputs = MaskFormerModelOutput( + encoder_last_hidden_state=raw_outputs[0], + pixel_decoder_last_hidden_state=raw_outputs[1], + transformer_decoder_last_hidden_state=raw_outputs[2], + encoder_hidden_states=raw_outputs[3] if output_hidden_states else None, + pixel_decoder_hidden_states=raw_outputs[4] if output_hidden_states else None, + transformer_decoder_hidden_states=raw_outputs[5] if output_hidden_states else None, + hidden_states=raw_outputs[6] if output_hidden_states else None, + attentions=raw_outputs[-1] if output_attentions else None, + ) + + loss, loss_dict, auxiliary_logits = None, None, None + + class_queries_logits, masks_queries_logits, auxiliary_logits = self.get_logits(outputs) + + if mask_labels is not None and class_labels is not None: + loss_dict: dict[str, Tensor] = self.get_loss_dict( + masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits + ) + 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_logits = None + + if not return_dict: + output = tuple( + v + for v in (loss, class_queries_logits, masks_queries_logits, auxiliary_logits, *outputs.values()) + if v is not None + ) + return output + + return MaskFormerForInstanceSegmentationOutput( + loss=loss, + **outputs, + class_queries_logits=class_queries_logits, + masks_queries_logits=masks_queries_logits, + auxiliary_logits=auxiliary_logits, + ) + + +__all__ = [ + "MaskFormerConfig", + "MaskFormerDetrConfig", + "MaskFormerForInstanceSegmentation", + "MaskFormerModel", + "MaskFormerPreTrainedModel", + "MaskFormerDetrPreTrainedModel", # noqa F821 +] diff --git a/third_party/transformers/src/transformers/models/mbart50/__init__.py b/third_party/transformers/src/transformers/models/mbart50/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e66676802277c6c929556955daac9d681842d375 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mbart50/__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_mbart50 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/mbart50/tokenization_mbart50.py b/third_party/transformers/src/transformers/models/mbart50/tokenization_mbart50.py new file mode 100644 index 0000000000000000000000000000000000000000..a6eba73c73896dc4c2b54a4832df9e1c0b94abe5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mbart50/tokenization_mbart50.py @@ -0,0 +1,312 @@ +# Copyright 2021 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import Unigram + +from ...tokenization_python import AddedToken, BatchEncoding +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"} + + +FAIRSEQ_LANGUAGE_CODES = ["ar_AR", "cs_CZ", "de_DE", "en_XX", "es_XX", "et_EE", "fi_FI", "fr_XX", "gu_IN", "hi_IN", "it_IT", "ja_XX", "kk_KZ", "ko_KR", "lt_LT", "lv_LV", "my_MM", "ne_NP", "nl_XX", "ro_RO", "ru_RU", "si_LK", "tr_TR", "vi_VN", "zh_CN", "af_ZA", "az_AZ", "bn_IN", "fa_IR", "he_IL", "hr_HR", "id_ID", "ka_GE", "km_KH", "mk_MK", "ml_IN", "mn_MN", "mr_IN", "pl_PL", "ps_AF", "pt_XX", "sv_SE", "sw_KE", "ta_IN", "te_IN", "th_TH", "tl_XX", "uk_UA", "ur_PK", "xh_ZA", "gl_ES", "sl_SI"] # fmt: skip + + +class MBart50Tokenizer(TokenizersBackend): + """ + Construct a MBart50 tokenizer (backed by HuggingFace's *tokenizers* library). Based on + [Unigram](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=unigram#models). + + This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + Path to the vocabulary file. + src_lang (`str`, *optional*): + A string representing the source language. + tgt_lang (`str`, *optional*): + A string representing the target language. + eos_token (`str`, *optional*, defaults to `""`): + The end of sequence token. + sep_token (`str`, *optional*, defaults to `""`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + cls_token (`str`, *optional*, defaults to `""`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + unk_token (`str`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + pad_token (`str`, *optional*, defaults to `""`): + The token used for padding, for example when batching sequences of different lengths. + mask_token (`str`, *optional*, defaults to `""`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + + Examples: + + ```python + >>> from transformers import MBart50Tokenizer + + >>> tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50", src_lang="en_XX", tgt_lang="ro_RO") + >>> src_text = " UN Chief Says There Is No Military Solution in Syria" + >>> tgt_text = "Şeful ONU declară că nu există o soluţie militară în Siria" + >>> model_inputs = tokenizer(src_text, text_target=tgt_text, return_tensors="pt") + >>> # model(**model_inputs) should work + ```""" + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + model = Unigram + + prefix_tokens: list[int] = [] + suffix_tokens: list[int] = [] + + def __init__( + self, + vocab: str | dict | list | None = None, + _spm_precompiled_charsmap: str | None = None, + src_lang=None, + tgt_lang=None, + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + **kwargs, + ): + mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token + + # Do not pass language codes via extra_special_tokens to super().__init__. + # We will mark them as special AFTER backend construction to avoid re-adding tokens + # when loading from pretrained files. + + # Always construct a tokenizer_object without referencing external tokenizer files + if isinstance(vocab, list): + # MBart50 uses fairseq vocab alignment matching MBart50Converter: + # =0, =1, =2, =3, then tokens, lang codes, + + vocab = [(str(item[0]), float(item[1])) for item in vocab] + + vocab_tokens = [item[0] for item in vocab] + has_language_codes = any(lang_code in vocab_tokens for lang_code in FAIRSEQ_LANGUAGE_CODES) + + if has_language_codes: + self._vocab_scores = vocab + else: + # Vocab from SentencePieceExtractor is in sentencepiece format: + # =0, =1, =2, then tokens + # We need to reorder to fairseq format: =0, =1, =2, =3, then tokens + + # Reorder: fairseq expects , , , , then rest of vocab starting from index 3 + vocab_list = [ + (str(cls_token), 0.0), # 0: + (str(pad_token), 0.0), # 1: + (str(eos_token), 0.0), # 2: + (str(unk_token), 0.0), # 3: + ] + # Add remaining tokens from position 3 onwards (skip , , from sentencepiece) + vocab_list.extend(vocab[3:]) + + # Add language codes + for lang_code in FAIRSEQ_LANGUAGE_CODES: + vocab_list.append((str(lang_code), 0.0)) + + # Add mask token + vocab_list.append((str(mask_token), 0.0)) + + self._vocab_scores = vocab_list + else: + # Minimal fallback: small vocab with specials and language codes + self._vocab_scores = [ + (str(cls_token), 0.0), + (str(pad_token), 0.0), + (str(eos_token), 0.0), + (str(unk_token), 0.0), + ("▁", -2.0), + ] + for lang_code in FAIRSEQ_LANGUAGE_CODES: + self._vocab_scores.append((lang_code, 0.0)) + self._vocab_scores.append((str(mask_token), 0.0)) + + # Build backend tokenizer from self._vocab_scores (both branches above set it) + self._tokenizer = Tokenizer( + Unigram( + self._vocab_scores, + unk_id=3, + byte_fallback=False, + ) + ) + + normalizers_ = [normalizers.Replace(Regex(r" {2,}"), " ")] + if _spm_precompiled_charsmap is not None: + normalizers_ = [normalizers.Precompiled(_spm_precompiled_charsmap)] + normalizers_ + + self._tokenizer.normalizer = normalizers.Sequence(normalizers_) + self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True) + + self._tokenizer.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True) + additional_special_tokens = kwargs.pop("additional_special_tokens", []) or [] + additional_special_tokens.extend(FAIRSEQ_LANGUAGE_CODES) + super().__init__( + src_lang=src_lang, + tgt_lang=tgt_lang, + eos_token=eos_token, + sep_token=sep_token, + cls_token=cls_token, + unk_token=unk_token, + pad_token=pad_token, + mask_token=mask_token, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + self.fairseq_offset = 1 + + # Mark language codes as extra special tokens without re-adding them to the backend. + # Merge with any pre-existing extra_special_tokens (e.g., restored from config on load). + try: + lang_tokens = [AddedToken(code, special=True) for code in FAIRSEQ_LANGUAGE_CODES] + except Exception: + lang_tokens = list(FAIRSEQ_LANGUAGE_CODES) + existing_extra = getattr(self, "_extra_special_tokens", []) or [] + # Preserve order: keep existing, append missing language codes + existing_strs = {str(t) for t in existing_extra} + merged_extra = list(existing_extra) + [t for t in lang_tokens if str(t) not in existing_strs] + self._extra_special_tokens = merged_extra + + self._src_lang = src_lang if src_lang is not None else "en_XX" + self.tgt_lang = tgt_lang + + # Build language code mappings and fairseq mappings + # This will be called again in _post_init after tokenizer.json is loaded + self._build_language_code_mappings() + + self.cur_lang_code_id = self.lang_code_to_id[self._src_lang] + self.set_src_lang_special_tokens(self._src_lang) + + def _build_language_code_mappings(self): + """Build language code to ID mappings and fairseq compatibility mappings.""" + self.lang_code_to_id = { + lang_code: self.convert_tokens_to_ids(lang_code) for lang_code in FAIRSEQ_LANGUAGE_CODES + } + self.id_to_lang_code = {v: k for k, v in self.lang_code_to_id.items()} + + # Build fairseq token mappings for backward compatibility + self.fairseq_tokens_to_ids = { + "": 0, + "": 1, + "": 2, + "": 3, + } + self.fairseq_tokens_to_ids.update(self.lang_code_to_id) + mask_token = getattr(self, "mask_token", "") + self.fairseq_tokens_to_ids[""] = self.convert_tokens_to_ids(str(mask_token)) + self.fairseq_ids_to_tokens = {v: k for k, v in self.fairseq_tokens_to_ids.items()} + + def _post_init(self): + """Called after tokenizer.json is loaded in from_pretrained.""" + # Rebuild language code mappings with the loaded tokenizer + self._build_language_code_mappings() + # Update cur_lang_code_id with the correct ID + if hasattr(self, "_src_lang"): + self.cur_lang_code_id = self.lang_code_to_id[self._src_lang] + self.set_src_lang_special_tokens(self._src_lang) + + @property + def src_lang(self) -> str: + return self._src_lang + + @src_lang.setter + def src_lang(self, new_src_lang: str) -> None: + self._src_lang = new_src_lang + self.set_src_lang_special_tokens(self._src_lang) + + def prepare_seq2seq_batch( + self, + src_texts: list[str], + src_lang: str = "en_XX", + tgt_texts: list[str] | None = None, + tgt_lang: str = "ro_RO", + **kwargs, + ) -> BatchEncoding: + self.src_lang = src_lang + self.tgt_lang = tgt_lang + return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs) + + def _switch_to_input_mode(self): + return self.set_src_lang_special_tokens(self.src_lang) + + def _switch_to_target_mode(self): + if self.tgt_lang is None: + self.tgt_lang = self._src_lang + return self.set_tgt_lang_special_tokens(self.tgt_lang) + + def set_src_lang_special_tokens(self, src_lang: str) -> None: + """Reset the special tokens to the source lang setting. prefix=[src_lang_code] and suffix=[eos].""" + self.cur_lang_code_id = self.convert_tokens_to_ids(src_lang) + self.prefix_tokens = [self.cur_lang_code_id] + self.suffix_tokens = [self.eos_token_id] + + prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens) + suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=prefix_tokens_str + ["$A"] + suffix_tokens_str, + pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str, + special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str, self.prefix_tokens + self.suffix_tokens)), + ) + + def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None: + """Reset the special tokens to the target language setting. prefix=[tgt_lang_code] and suffix=[eos].""" + self.cur_lang_code_id = self.convert_tokens_to_ids(tgt_lang) + self.prefix_tokens = [self.cur_lang_code_id] + self.suffix_tokens = [self.eos_token_id] + + prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens) + suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens) + + self._tokenizer.post_processor = processors.TemplateProcessing( + single=prefix_tokens_str + ["$A"] + suffix_tokens_str, + pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str, + special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str, self.prefix_tokens + self.suffix_tokens)), + ) + + def _build_translation_inputs( + self, raw_inputs, return_tensors: str, src_lang: str | None, tgt_lang: str | None, **extra_kwargs + ): + """Used by translation pipeline, to prepare inputs for the generate function""" + if src_lang is None or tgt_lang is None: + raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model") + self.src_lang = src_lang + inputs = self(raw_inputs, add_special_tokens=True, return_tensors=return_tensors, **extra_kwargs) + tgt_lang_id = self.convert_tokens_to_ids(tgt_lang) + inputs["forced_bos_token_id"] = tgt_lang_id + return inputs + + +__all__ = ["MBart50Tokenizer"] + +# Backward alias +MBart50TokenizerFast = MBart50Tokenizer diff --git a/third_party/transformers/src/transformers/models/olmoe/__init__.py b/third_party/transformers/src/transformers/models/olmoe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e18d77cbad11bf218b6cd304f703b4d7935eef --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmoe/__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_olmoe import * + from .modeling_olmoe 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/olmoe/configuration_olmoe.py b/third_party/transformers/src/transformers/models/olmoe/configuration_olmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..16bedbe698f8cf52299963c504421cd6f3e75d26 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmoe/configuration_olmoe.py @@ -0,0 +1,89 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OLMoE model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="allenai/OLMoE-1B-7B-0924") +@strict +class OlmoeConfig(PreTrainedConfig): + r""" + clip_qkv (`float`, *optional*): + If not `None`, elements of query, key and value attention states are clipped so that their + absolute value does not exceed this value. + + ```python + >>> from transformers import OlmoeModel, OlmoeConfig + + >>> # Initializing a OLMoE 7B A1B style configuration + >>> configuration = OlmoeConfig() + + >>> # Initializing a model from the OLMoE 7B A1B style configuration + >>> model = OlmoeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "olmoe" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_local_experts": "num_experts"} + + # Default tensor parallel plan for base model `Olmoe` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise_gather_output", # due to the norm, we have to gather + "layers.*.self_attn.k_proj": "colwise_gather_output", # due to the norm, we have to gather + "layers.*.self_attn.v_proj": "colwise_gather_output", # due to the norm, we have to gather + "layers.*.self_attn.o_proj": "rowwise_split_input", # due to the norm, we have to gather + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + + vocab_size: int = 50304 + hidden_size: int = 2048 + intermediate_size: int = 2048 + num_hidden_layers: int = 16 + num_attention_heads: int = 16 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 4096 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-05 + use_cache: bool = True + pad_token_id: int | None = 1 + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = 50279 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + clip_qkv: float | None = None + num_experts_per_tok: int = 8 + num_experts: int = 64 + output_router_logits: bool = False + router_aux_loss_coef: float = 0.01 + norm_topk_prob: bool = False + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + super().__post_init__(**kwargs) + + +__all__ = ["OlmoeConfig"] diff --git a/third_party/transformers/src/transformers/models/olmoe/convert_olmoe_weights_to_hf.py b/third_party/transformers/src/transformers/models/olmoe/convert_olmoe_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..35f223acf461b9ee4b27ba5691a49fbad7dfeee4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmoe/convert_olmoe_weights_to_hf.py @@ -0,0 +1,277 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Example for running: +0. Cp ckpts to local +aws s3 cp --recursive s3://ai2-llm/checkpoints/OLMoE/olmoe-8x1b-newhp-newds-final-annealFrom1200000/step23842 /data/niklas/llm/checkpoints/olmoe-8x1b-newhp-newds-final-annealFrom1200000_step23842 +1. Unshard your OLMoE checkpoint using https://github.com/allenai/OLMo/blob/7d63fe09d23cf23714da5aa633a44a90180195da/scripts/unshard.py +python OLMo/scripts/unshard.py /data/niklas/llm/checkpoints/23485/step954000 /data/niklas/llm/checkpoints/1b-954000-unsharded --model-only +python OLMo/scripts/unshard.py /data/niklas/llm/checkpoints/23485/step954000 /data/niklas/llm/checkpoints/1b-954000-unsharded --model-only +python OLMo/scripts/unshard.py /data/niklas/llm/checkpoints/olmoe-8x1b-newhp-newds-final-annealFrom1200000_step23842 /data/niklas/llm/checkpoints/olmoe-8x1b-newhp-newds-final-annealFrom1200000_step23842-unsharded --model-only +2. Convert to transformers +rm -rf olmoe; mkdir olmoe; python /data/niklas/transformers/src/transformers/models/olmoe/convert_olmoe_weights_to_hf.py --input_dir /data/niklas/llm/checkpoints/olmoe-8x1b-newhp-newds-final-annealFrom1200000_step23842-unsharded --tokenizer_json_path /data/niklas/llm/checkpoints/olmoe-step1200000-unsharded/tokenizer.json --output_dir olmoe +3. Load model via: +``` +from transformers import OlmoeForCausalLM, AutoTokenizer +import torch +model = OlmoeForCausalLM.from_pretrained("../transformers/olmoe", dtype=torch.bfloat16).cuda() +model = OlmoeForCausalLM.from_pretrained("../transformers/olmoe").cuda() +tokenizer = AutoTokenizer.from_pretrained("../transformers/olmoe") +inputs = tokenizer("Bitcoin is", return_tensors="pt") +inputs = {k: v.cuda() for k, v in inputs.items()} +out = model.generate(**inputs, max_length=64) +print(tokenizer.decode(out[0])) +# > # Bitcoin is a digital currency that is created and held electronically. No one controls it. Bitcoins aren’t printed, like dollars or euros – they’re produced by people and businesses running computers all around the world, using software that solves mathematical +# Or quick sanity check: +o = model(torch.tensor([[0, 1]]).cuda()) +# If the checkpoint is not converted to BF16 but kept in FP32: +# > # Bitcoin is a digital currency that is not controlled by any central authority. It is a peer-to-peer payment system that allows users to send and receive payments from anywhere in the world. Bitcoin is also known as a cryptocurrency because it uses cryptography to secure transactions and prevent fraud. +``` + +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). + +Compare with OLMo codebase: +``` +from olmo.model import OLMo +import torch +model = OLMo.from_checkpoint("/data/niklas/llm/checkpoints/olmoe-step1200000-unsharded-pt") +model = model.cuda() +model = model.to(torch.bfloat16) +from transformers import AutoTokenizer +tokenizer = AutoTokenizer.from_pretrained("../transformers/olmoe") +inputs = tokenizer("Bitcoin is", return_tensors="pt") +inputs = {k: v.cuda() for k, v in inputs.items()} +out = model.generate(**inputs) +print(tokenizer.decode(out[0][0][0])) +# Bitcoin is a digital currency that is created and held electronically. No one controls it. Bitcoins aren’t printed, like dollars or euros – they’re produced by people and businesses running computers all around the world, using software that solves mathematical problems. It’s the first example of a growing category of money +# Or quick sanity check: +o = model(torch.tensor([[0, 1]]).cuda()) +``` +""" + +import argparse +import gc +import json +import os +import shutil +from pathlib import Path + +import torch +import yaml +from tokenizers import Tokenizer + +from transformers import OlmoeConfig, OlmoeForCausalLM +from transformers.models.gpt_neox.tokenization_gpt_neox_fast import GPTNeoXTokenizerFast + + +def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256): + return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of) + + +def read_json(path): + with open(path, "r") as f: + return json.load(f) + + +def write_json(text, path): + with open(path, "w") as f: + json.dump(text, f) + + +def write_model(model_path, input_base_path, tokenizer_path=None, fix_eos_token_id=True): + os.makedirs(model_path, exist_ok=True) + tmp_model_path = os.path.join(model_path, "tmp") + os.makedirs(tmp_model_path, exist_ok=True) + + config_path = Path(input_base_path) / "config.yaml" + olmoe_config = yaml.safe_load(config_path.read_text())["model"] + + if fix_eos_token_id: + olmoe_config["eos_token_id"] = 50279 + + n_layers = olmoe_config["n_layers"] + n_heads = olmoe_config["n_heads"] + dim = olmoe_config["d_model"] + dims_per_head = dim // n_heads + base = 10000.0 + inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)) + max_position_embeddings = olmoe_config["max_sequence_length"] + + vocab_size = olmoe_config.get("embedding_size", olmoe_config["vocab_size"]) + + if olmoe_config.get("n_kv_heads", None) is not None: + num_key_value_heads = olmoe_config["n_kv_heads"] # for GQA / MQA + elif olmoe_config["multi_query_attention"]: # compatibility with other checkpoints + num_key_value_heads = 1 + else: + num_key_value_heads = n_heads + + print(f"Fetching all parameters from the checkpoint at {input_base_path}.") + + # Not sharded + loaded = torch.load(os.path.join(input_base_path, "model.pt"), map_location="cpu", weights_only=True) + + param_count = 0 + index_dict = {"weight_map": {}} + for layer_i in range(n_layers): + filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin" + fused_dims = [dim, dims_per_head * num_key_value_heads, dims_per_head * num_key_value_heads] + q_proj_weight, k_proj_weight, v_proj_weight = torch.split( + loaded[f"transformer.blocks.{layer_i}.att_proj.weight"], fused_dims, dim=0 + ) + state_dict = { + f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj_weight, + f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj_weight, + f"model.layers.{layer_i}.self_attn.v_proj.weight": v_proj_weight, + f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"transformer.blocks.{layer_i}.attn_out.weight"], + f"model.layers.{layer_i}.self_attn.q_norm.weight": loaded[f"transformer.blocks.{layer_i}.q_norm.weight"], + f"model.layers.{layer_i}.self_attn.k_norm.weight": loaded[f"transformer.blocks.{layer_i}.k_norm.weight"], + f"model.layers.{layer_i}.mlp.gate.weight": loaded[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"], + f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"transformer.blocks.{layer_i}.attn_norm.weight"], + f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[ + f"transformer.blocks.{layer_i}.ff_norm.weight" + ], + } + + num_experts = loaded[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"].shape[0] + dim_per_expert = loaded[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"].shape[0] // num_experts + for expert_i in range(num_experts): + state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.gate_proj.weight"] = loaded[ + f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1" + ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :] + state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.up_proj.weight"] = loaded[ + f"transformer.blocks.{layer_i}.ffn.experts.mlp.v1" + ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :] + state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.down_proj.weight"] = loaded[ + f"transformer.blocks.{layer_i}.ffn.experts.mlp.w2" + ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :].T.contiguous() + + state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq + + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin" + + # Unsharded + state_dict = { + "model.embed_tokens.weight": loaded["transformer.wte.weight"], + "lm_head.weight": loaded["transformer.ff_out.weight"], + "model.norm.weight": loaded["transformer.ln_f.weight"], + } + + for k, v in state_dict.items(): + index_dict["weight_map"][k] = filename + param_count += v.numel() + torch.save(state_dict, os.path.join(tmp_model_path, filename)) + + # Write configs + index_dict["metadata"] = {"total_size": param_count * 2} + write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json")) + + config = OlmoeConfig( + vocab_size=vocab_size, + hidden_size=dim, + intermediate_size=dim_per_expert, + num_hidden_layers=n_layers, + num_attention_heads=n_heads, + num_key_value_heads=num_key_value_heads, + max_position_embeddings=max_position_embeddings, + pad_token_id=olmoe_config["pad_token_id"], + bos_token_id=None, + eos_token_id=olmoe_config["eos_token_id"], + tie_word_embeddings=olmoe_config["weight_tying"], + rope_theta=base, + clip_qkv=olmoe_config.get("clip_qkv"), + ) + config.save_pretrained(tmp_model_path) + + # Make space so we can load the model properly now. + del state_dict + del loaded + gc.collect() + + if tokenizer_path is not None: + _write_tokenizer(model_path, config, tokenizer_path, fix_eos_token_id) + + print("Loading the checkpoint in a OLMoE model.") + model = OlmoeForCausalLM.from_pretrained(tmp_model_path, dtype=torch.bfloat16) + # Avoid saving this as part of the config. + del model.config._name_or_path + print("Saving in the Transformers format.") + model.save_pretrained(model_path) + shutil.rmtree(tmp_model_path) + + +def _write_tokenizer( + output_path: Path, config: OlmoeConfig, input_tokenizer_path: Path, fix_eos_token_id: bool = True +) -> None: + print(f"Saving a {GPTNeoXTokenizerFast.__name__} to {output_path}.") + + base_tokenizer = Tokenizer.from_file(str(input_tokenizer_path)) + + eos_token_id = config.eos_token_id if config.eos_token_id is not None else base_tokenizer.get_vocab_size() - 1 + pad_token_id = config.pad_token_id if config.pad_token_id is not None else eos_token_id + + if fix_eos_token_id and eos_token_id == 0: + # Fixing a bug in OLMo where eos token id was incorrectly set + print("Changing eos_token_id from 0 to 50279.") + eos_token_id = 50279 + + tokenizer = GPTNeoXTokenizerFast( + tokenizer_object=base_tokenizer, + eos_token=base_tokenizer.decode([eos_token_id], skip_special_tokens=False), + pad_token=base_tokenizer.decode([pad_token_id], skip_special_tokens=False), + unk_token=None, + bos_token=None, + ) + + tokenizer.save_pretrained(output_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + required=True, + help="Location of OLMoE weights, which contains config.yaml and model.pt.", + ) + parser.add_argument( + "--tokenizer_json_path", + default=None, + help="Location of OLMoE tokenizer json file.", + ) + parser.add_argument( + "--output_dir", + required=True, + help="Location to write HF model and tokenizer", + ) + parser.add_argument( + "--no_fix_eos_token_id", + action="store_false", + dest="fix_eos_token_id", + help="If set, does not change eos token id from 0 to 50279 if it is 0. Changing 0 to 50279 is a bug fix, so use this option with care.", + ) + args = parser.parse_args() + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + tokenizer_path=args.tokenizer_json_path, + fix_eos_token_id=args.fix_eos_token_id, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/olmoe/modeling_olmoe.py b/third_party/transformers/src/transformers/models/olmoe/modeling_olmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..8a83315a5820b2639337730b0a8a0937d398084a --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmoe/modeling_olmoe.py @@ -0,0 +1,708 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/olmoe/modular_olmoe.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_olmoe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_olmoe import OlmoeConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class OlmoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-5) -> None: + """ + OlmoeRMSNorm 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 OlmoeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: OlmoeConfig, 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: OlmoeConfig | 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 OlmoeMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def 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 OlmoeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: OlmoeConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.k_norm = OlmoeRMSNorm( + (config.hidden_size // config.num_attention_heads) * config.num_key_value_heads, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states)) + key_states = self.k_norm(self.k_proj(hidden_states)) + value_states = self.v_proj(hidden_states) + + if self.config.clip_qkv is not None: # Diff with llama + query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + + query_states = query_states.view(*hidden_shape).transpose(1, 2) + key_states = key_states.view(*hidden_shape).transpose(1, 2) + value_states = value_states.view(*hidden_shape).transpose(1, 2) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +@use_experts_implementation +class OlmoeExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config: OlmoeConfig): + super().__init__() + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.intermediate_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) + self.act_fn = ACT2FN[config.hidden_act] + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states + + +class OlmoeTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.norm_topk_prob = config.norm_topk_prob + 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) + if self.norm_topk_prob: + 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 OlmoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.gate = OlmoeTopKRouter(config) + self.experts = OlmoeExperts(config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + _, top_k_weights, top_k_index = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states, top_k_index, top_k_weights).reshape( + batch_size, sequence_length, hidden_dim + ) + return final_hidden_states + + +class OlmoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: OlmoeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = OlmoeAttention(config=config, layer_idx=layer_idx) + self.mlp = OlmoeSparseMoeBlock(config) + self.input_layernorm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = OlmoeRMSNorm(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 OlmoePreTrainedModel(PreTrainedModel): + config: OlmoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["OlmoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _can_record_outputs = { + "router_logits": OutputRecorder(OlmoeTopKRouter, index=0), + "hidden_states": OlmoeDecoderLayer, + "attentions": OlmoeAttention, + } + + _supports_attention_backend = True + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, OlmoeExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, OlmoeTopKRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class OlmoeModel(OlmoePreTrainedModel): + def __init__(self, config: OlmoeConfig): + 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( + [OlmoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = OlmoeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( # diff with mixtral: no sliding + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +def load_balancing_loss_func( + gate_logits: torch.Tensor | tuple[torch.Tensor] | None, + num_experts: int | None = None, + top_k=2, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor | int: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +@auto_docstring +class OlmoeForCausalLM(OlmoePreTrainedModel, 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 = OlmoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.num_experts + self.num_experts_per_tok = config.num_experts_per_tok + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_router_logits: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeCausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, OlmoeForCausalLM + + >>> model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924") + >>> tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924") + + >>> 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 sure if you’re conscious of this, but I’m' + ``` + """ + + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs: MoeModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_router_logits=output_router_logits, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) + + +__all__ = ["OlmoeForCausalLM", "OlmoeModel", "OlmoePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/olmoe/modular_olmoe.py b/third_party/transformers/src/transformers/models/olmoe/modular_olmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..9fee40493496aa4a476ab2621473c522d0cb3d40 --- /dev/null +++ b/third_party/transformers/src/transformers/models/olmoe/modular_olmoe.py @@ -0,0 +1,279 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 OLMoE model.""" + +from collections.abc import Callable + +import torch +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_outputs import MoeModelOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.output_capturing import OutputRecorder +from ..gemma.modeling_gemma import GemmaMLP +from ..llama.modeling_llama import ( + LlamaAttention, + LlamaDecoderLayer, + LlamaRMSNorm, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from ..mixtral.modeling_mixtral import MixtralExperts, MixtralForCausalLM, MixtralModel +from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeTopKRouter +from .configuration_olmoe import OlmoeConfig + + +logger = logging.get_logger(__name__) + + +class OlmoeRMSNorm(LlamaRMSNorm): + def __init__(self, hidden_size, eps=1e-5): + super().__init__(hidden_size, eps) + + +class OlmoeRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class OlmoeMLP(GemmaMLP): + pass + + +class OlmoeAttention(LlamaAttention): + def __init__(self, config: OlmoeConfig, layer_idx: int | None = None): + super().__init__(config, layer_idx) + self.q_norm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.k_norm = OlmoeRMSNorm( + (config.hidden_size // config.num_attention_heads) * config.num_key_value_heads, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states)) + key_states = self.k_norm(self.k_proj(hidden_states)) + value_states = self.v_proj(hidden_states) + + if self.config.clip_qkv is not None: # Diff with llama + query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + + query_states = query_states.view(*hidden_shape).transpose(1, 2) + key_states = key_states.view(*hidden_shape).transpose(1, 2) + value_states = value_states.view(*hidden_shape).transpose(1, 2) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class OlmoeExperts(MixtralExperts): + pass + + +class OlmoeTopKRouter(Qwen2MoeTopKRouter): + pass + + +class OlmoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.gate = OlmoeTopKRouter(config) + self.experts = OlmoeExperts(config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + _, top_k_weights, top_k_index = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states, top_k_index, top_k_weights).reshape( + batch_size, sequence_length, hidden_dim + ) + return final_hidden_states + + +class OlmoeDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: OlmoeConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.hidden_size = config.hidden_size + self.self_attn = OlmoeAttention(config=config, layer_idx=layer_idx) + self.mlp = OlmoeSparseMoeBlock(config) + self.input_layernorm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +@auto_docstring +class OlmoePreTrainedModel(PreTrainedModel): + config: OlmoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["OlmoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _can_record_outputs = { + "router_logits": OutputRecorder(OlmoeTopKRouter, index=0), + "hidden_states": OlmoeDecoderLayer, + "attentions": OlmoeAttention, + } + + _supports_attention_backend = True + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, OlmoeExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, OlmoeTopKRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class OlmoeModel(MixtralModel): + def __init__(self, config: OlmoeConfig): + super().__init__(config) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [OlmoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = OlmoeRotaryEmbedding(config=config) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> MoeModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( # diff with mixtral: no sliding + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + + return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class OlmoeForCausalLM(MixtralForCausalLM, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = OlmoeModel(config) + self.num_experts = config.num_experts + + def forward(self, **super_kwargs): + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, OlmoeForCausalLM + + >>> model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924") + >>> tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924") + + >>> 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 sure if you’re conscious of this, but I’m' + ``` + """ + return super().forward(**super_kwargs) + + +__all__ = ["OlmoeForCausalLM", "OlmoeModel", "OlmoePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/vipllava/__init__.py b/third_party/transformers/src/transformers/models/vipllava/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2ff5c6838dfaff4eb96663c6e2a808e5f082f0de --- /dev/null +++ b/third_party/transformers/src/transformers/models/vipllava/__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_vipllava import * + from .modeling_vipllava 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/vipllava/configuration_vipllava.py b/third_party/transformers/src/transformers/models/vipllava/configuration_vipllava.py new file mode 100644 index 0000000000000000000000000000000000000000..891a2c32f27a22bf9215fa99c240f404305d5f94 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vipllava/configuration_vipllava.py @@ -0,0 +1,92 @@ +# Copyright 2023 Microsoft Research & University of Wisconsin-Madison and the HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VipLlava model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import CONFIG_MAPPING, AutoConfig + + +@auto_docstring(checkpoint="llava-hf/vip-llava-7b-hf") +@strict +class VipLlavaConfig(PreTrainedConfig): + r""" + projector_layernorm_eps (`float`, *optional*, defaults to 1e-05): + The layer norm epsilon of the projector layernorm + vision_feature_layers (`Union[int, list[int]]`, *optional*, defaults to `[-2, -5, -8, -11, 6]`): + The vision feature layer, or list of layers to select the vision features from. + + Example: + + ```python + >>> from transformers import VipLlavaForConditionalGeneration, VipLlavaConfig, CLIPVisionConfig, LlamaConfig + + >>> # Initializing a CLIP-vision config + >>> vision_config = CLIPVisionConfig() + + >>> # Initializing a Llama config + >>> text_config = LlamaConfig() + + >>> # Initializing a VipLlava vipllava-7b style configuration + >>> configuration = VipLlavaConfig(vision_config, text_config) + + >>> # Initializing a model from the vipllava-7b style configuration + >>> model = VipLlavaForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "vipllava" + attribute_map = { + "image_token_id": "image_token_index", + } + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_token_index: int = 32000 + projector_hidden_act: str = "gelu" + projector_layernorm_eps: float = 1e-5 + vision_feature_layers: int | list[int] | tuple[int, ...] = (-2, -5, -8, -11, 6) + image_seq_length: int = 576 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "clip_vision_model") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config) + elif self.vision_config is None: + self.vision_config = CONFIG_MAPPING["clip_vision_model"]( + intermediate_size=4096, + hidden_size=1024, + patch_size=14, + image_size=336, + num_hidden_layers=24, + num_attention_heads=16, + vocab_size=32000, + projection_dim=768, + ) + + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "llama") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["llama"]() + + super().__post_init__(**kwargs) + + +__all__ = ["VipLlavaConfig"] diff --git a/third_party/transformers/src/transformers/models/vipllava/convert_vipllava_weights_to_hf.py b/third_party/transformers/src/transformers/models/vipllava/convert_vipllava_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..47f58cc6e10a5d3ce30e0fb5663dd934d0e116d4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vipllava/convert_vipllava_weights_to_hf.py @@ -0,0 +1,132 @@ +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse + +import torch +from huggingface_hub import hf_hub_download + +from transformers import ( + AddedToken, + AutoConfig, + AutoTokenizer, + CLIPImageProcessor, + LlavaProcessor, + VipLlavaConfig, + VipLlavaForConditionalGeneration, +) + + +KEYS_TO_MODIFY_MAPPING = { + "model.vision_tower.": "", + "model.mm_projector": "multi_modal_projector", + "model": "model.model", + "vision_model.model": "vision_model", + "lm_head": "language_model.lm_head", + "model.model": "language_model.model", + "multi_modal_projector.0": "multi_modal_projector.linear_1", + "multi_modal_projector.2": "multi_modal_projector.linear_2", + "final_linear.0": "linear_1", + "final_linear.2": "linear_2", + "multi_modal_projector.clip_layernorm": "multi_modal_projector.projector_layernorm", +} + + +# Copied from transformers.models.llava.convert_llava_weights_to_hf.convert_state_dict_to_hf +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 convert_vipllava_llama_to_hf(text_model_id, vision_model_id, output_hub_path, old_state_dict_id): + torch.set_default_dtype(torch.float16) + text_config = AutoConfig.from_pretrained(text_model_id) + + tokenizer = AutoTokenizer.from_pretrained(text_model_id) + tokenizer.add_tokens(AddedToken("", special=True, normalized=False), special_tokens=True) + tokenizer.add_special_tokens({"pad_token": ""}) + + image_processor = CLIPImageProcessor.from_pretrained(vision_model_id) + + processor = LlavaProcessor(tokenizer=tokenizer, image_processor=image_processor) + + config = VipLlavaConfig(text_config=text_config) + config.pad_token_id = 32001 + + with torch.device("meta"): + model = VipLlavaForConditionalGeneration(config) + + # Pad to 64 for performance reasons + pad_shape = 64 + + state_dict_path = hf_hub_download(old_state_dict_id, "model_state_dict_7b.bin") + + state_dict = torch.load(state_dict_path, map_location="cpu", weights_only=True) + state_dict = convert_state_dict_to_hf(state_dict) + + model.load_state_dict(state_dict, strict=True, assign=True) + + pre_expansion_embeddings = model.language_model.model.embed_tokens.weight.data + mu = torch.mean(pre_expansion_embeddings, dim=0).float() + n = pre_expansion_embeddings.size()[0] + sigma = ((pre_expansion_embeddings - mu).T @ (pre_expansion_embeddings - mu)) / n + dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, covariance_matrix=1e-5 * sigma) + + # We add an image token so we resize the model + model.resize_token_embeddings(config.text_config.vocab_size + 2, pad_shape) + model.language_model.model.embed_tokens.weight.data[32000:] = torch.stack( + tuple(dist.sample() for _ in range(model.language_model.model.embed_tokens.weight.data[32000:].shape[0])), + dim=0, + ) + model.language_model.lm_head.weight.data[32000:] = torch.stack( + tuple(dist.sample() for _ in range(model.language_model.lm_head.weight.data[32000:].shape[0])), + dim=0, + ) + + model.push_to_hub(output_hub_path) + processor.push_to_hub(output_hub_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--text_model_id", + help="Hub location of the text model", + ) + parser.add_argument( + "--vision_model_id", + help="Hub location of the vision model", + ) + parser.add_argument( + "--output_hub_path", + help="Location on the hub of the converted model", + ) + parser.add_argument( + "--old_state_dict_id", + help="Location on the hub of the raw state dict of the original model. The filename needs to be `model_state_dict.bin`", + ) + args = parser.parse_args() + convert_vipllava_llama_to_hf( + args.text_model_id, args.vision_model_id, args.output_hub_path, args.old_state_dict_id + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/vipllava/modular_vipllava.py b/third_party/transformers/src/transformers/models/vipllava/modular_vipllava.py new file mode 100644 index 0000000000000000000000000000000000000000..0b3352e54c868f55609709eb6443050b39f35015 --- /dev/null +++ b/third_party/transformers/src/transformers/models/vipllava/modular_vipllava.py @@ -0,0 +1,279 @@ +# Copyright 2023 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from torch import nn + +from transformers.models.llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, + LlavaForConditionalGeneration, + LlavaModel, + LlavaModelOutputWithPast, + LlavaPreTrainedModel, +) + +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import can_return_tuple +from .configuration_vipllava import VipLlavaConfig + + +logger = logging.get_logger(__name__) + + +class VipLlavaModelOutputWithPast(LlavaModelOutputWithPast): + pass + + +class VipLlavaCausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass + + +class VipLlavaMultiModalProjector(nn.Module): + def __init__(self, config: VipLlavaConfig): + super().__init__() + num_feature_layers = 1 if isinstance(config.vision_feature_layers, int) else len(config.vision_feature_layers) + self.projector_layernorm = nn.LayerNorm( + num_feature_layers * config.vision_config.hidden_size, eps=config.projector_layernorm_eps + ) + + self.linear_1 = nn.Linear( + num_feature_layers * config.vision_config.hidden_size, + config.text_config.hidden_size, + bias=True, + ) + self.act = ACT2FN[config.projector_hidden_act] + self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True) + + def forward(self, hidden_states): + hidden_states = self.projector_layernorm(hidden_states) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class VipLlavaPreTrainedModel(LlavaPreTrainedModel): + pass + + +class VipLlavaModel(LlavaModel): + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + """ + vision_feature_layers = ( + vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers + ) + # We need hidden states to select intermediate vision features by layer index below. + kwargs["output_hidden_states"] = True + image_outputs = self.vision_tower( + pixel_values, + **kwargs, + ) + + # If multiple feature layers are provided (which is usually the case) + # then the image features are concatenated after the CLS is removed. + if isinstance(vision_feature_layers, int): + image_features = image_outputs.hidden_states[vision_feature_layers][:, 1:] + else: + # Usually, we select the features from index 1: the layers -2, -5, -8, -11 and 6 + image_features = [image_outputs.hidden_states[index][:, 1:] for index in vision_feature_layers] + image_features = torch.cat(image_features, dim=-1) + image_features = self.multi_modal_projector(image_features) + image_outputs.pooler_output = image_features + + return image_outputs + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layers: int | list[int] | None = None, + use_cache: bool | None = None, + **lm_kwargs: Unpack[TransformersKwargs], + ) -> tuple | VipLlavaModelOutputWithPast: + r""" + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + """ + vision_feature_layers = ( + vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers + ).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs: BaseModelOutputWithPast = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **lm_kwargs, + ) + + output = VipLlavaModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + return output + + +class VipLlavaForConditionalGeneration(LlavaForConditionalGeneration): + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + """ + return self.model.get_image_features( + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers, **kwargs + ) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layers: int | list[int] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **lm_kwargs: Unpack[TransformersKwargs], + ) -> tuple | VipLlavaCausalLMOutputWithPast: + r""" + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + 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 + >>> import torch + >>> from PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, VipLlavaForConditionalGeneration + + >>> model = VipLlavaForConditionalGeneration.from_pretrained("llava-hf/vip-llava-7b-hf", device_map="auto", dtype=torch.float16) + >>> processor = AutoProcessor.from_pretrained("llava-hf/vip-llava-7b-hf") + + >>> prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.###Human: \n{}###Assistant:" + >>> question = "Can you please describe this image?" + >>> prompt = prompt.format(question) + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/compel-neg.png" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> inputs = processor(text=text, images=image, return_tensors="pt").to(0, torch.float16) + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_new_tokens=20) + >>> processor.decode(generate_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True) + The image features a brown and white cat sitting on a green surface, with a red ball in its + ```""" + + vision_feature_layers = ( + vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers + ) + + outputs: VipLlavaModelOutputWithPast = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + vision_feature_layers=vision_feature_layers, + **lm_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.text_config.vocab_size) + + return VipLlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + +__all__ = ["VipLlavaModel", "VipLlavaForConditionalGeneration", "VipLlavaPreTrainedModel"] diff --git a/third_party/transformers/src/transformers/optimization.py b/third_party/transformers/src/transformers/optimization.py new file mode 100644 index 0000000000000000000000000000000000000000..64559c9b591059431a6274f7f2d970c768b94dea --- /dev/null +++ b/third_party/transformers/src/transformers/optimization.py @@ -0,0 +1,1342 @@ +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch optimization for BERT model.""" + +from __future__ import annotations + +import math +import warnings +from functools import partial +from typing import Any + +import torch +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau + +from .trainer_pt_utils import LayerWiseDummyOptimizer, LayerWiseDummyScheduler +from .trainer_utils import SchedulerType +from .utils import logging + + +logger = logging.get_logger(__name__) + + +def _get_constant_lambda(_=None): + return 1 + + +def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1): + """ + Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch) + + +def get_reduce_on_plateau_schedule(optimizer: Optimizer, **kwargs): + """ + Create a schedule with a constant learning rate that decreases when a metric has stopped improving. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + kwargs (`dict`, *optional*): + Extra parameters to be passed to the scheduler. See `torch.optim.lr_scheduler.ReduceLROnPlateau` + for possible parameters. + + Return: + `torch.optim.lr_scheduler.ReduceLROnPlateau` with the appropriate schedule. + """ + + return ReduceLROnPlateau(optimizer, **kwargs) + + +def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1.0, num_warmup_steps)) + return 1.0 + + +def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1): + """ + Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate + increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps) + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))) + + +def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1): + """ + Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after + a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_linear_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_cosine_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) + + +def get_cosine_schedule_with_warmup( + optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1 +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_cosine_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + if progress >= 1.0: + return 0.0 + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0)))) + + +def get_cosine_with_hard_restarts_schedule_with_warmup( + optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1 +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases + linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`int`, *optional*, defaults to 1): + The number of hard restarts to use. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_polynomial_decay_schedule_with_warmup_lr_lambda( + current_step: int, + *, + num_warmup_steps: int, + num_training_steps: int, + lr_end: float, + power: float, + lr_init: int, +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + elif current_step > num_training_steps: + return lr_end / lr_init # as LambdaLR multiplies by lr_init + else: + lr_range = lr_init - lr_end + decay_steps = num_training_steps - num_warmup_steps + pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps + decay = lr_range * pct_remaining**power + lr_end + return decay / lr_init # as LambdaLR multiplies by lr_init + + +def get_polynomial_decay_schedule_with_warmup( + optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1 +): + """ + Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the + optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + lr_end (`float`, *optional*, defaults to 1e-7): + The end LR. + power (`float`, *optional*, defaults to 1.0): + Power factor. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT + implementation at + https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37 + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + + """ + + lr_init = optimizer.defaults["lr"] + if not (lr_init > lr_end): + raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})") + + lr_lambda = partial( + _get_polynomial_decay_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + lr_end=lr_end, + power=power, + lr_init=lr_init, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: int | None = None): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + shift = timescale - num_warmup_steps + decay = 1.0 / math.sqrt((current_step + shift) / timescale) + return decay + + +def get_inverse_sqrt_schedule( + optimizer: Optimizer, num_warmup_steps: int, timescale: int | None = None, last_epoch: int = -1 +): + """ + Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a + warmup period which increases lr linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + timescale (`int`, *optional*, defaults to `num_warmup_steps`): + Time scale. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + # Note: this implementation is adapted from + # https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930 + + if timescale is None: + timescale = num_warmup_steps or 10_000 + + lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale) + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def _get_cosine_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float, min_lr_rate: float = 0.0 +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) + factor = factor * (1 - min_lr_rate) + min_lr_rate + return max(0, factor) + + +def get_cosine_with_min_lr_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float = 0.5, + last_epoch: int = -1, + min_lr: float | None = None, + min_lr_rate: float | None = None, +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + min_lr (`float`, *optional*): + The minimum learning rate to reach after the cosine schedule. + min_lr_rate (`float`, *optional*): + The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + if min_lr is not None and min_lr_rate is not None: + raise ValueError("Only one of min_lr or min_lr_rate should be set") + elif min_lr is not None: + min_lr_rate = min_lr / optimizer.defaults["lr"] + elif min_lr_rate is None: + raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`") + + lr_lambda = partial( + _get_cosine_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + min_lr_rate=min_lr_rate, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_cosine_with_min_lr_schedule_with_warmup_lr_rate_lambda( + current_step: int, + *, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float, + min_lr_rate: float = 0.0, + warmup_lr_rate: float | None = None, +): + current_step = float(current_step) + num_warmup_steps = float(num_warmup_steps) + num_training_steps = float(num_training_steps) + + if current_step < num_warmup_steps: + if warmup_lr_rate is None: + return (current_step + 1.0) / max(1.0, num_warmup_steps) + else: + warmup_lr_rate = float(warmup_lr_rate) + return warmup_lr_rate + (1.0 - warmup_lr_rate) * (current_step) / (max(1, num_warmup_steps - 1)) + progress = (current_step - num_warmup_steps + 1.0) / (max(1.0, num_training_steps - num_warmup_steps)) + factor = 0.5 * (1.0 + math.cos(math.pi * num_cycles * 2.0 * progress)) + factor = factor * (1 - min_lr_rate) + min_lr_rate + return max(0, factor) + + +def get_cosine_with_min_lr_schedule_with_warmup_lr_rate( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float = 0.5, + last_epoch: int = -1, + min_lr: float | None = None, + min_lr_rate: float | None = None, + warmup_lr_rate: float | None = None, +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + min_lr (`float`, *optional*): + The minimum learning rate to reach after the cosine schedule. + min_lr_rate (`float`, *optional*): + The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set. + warmup_lr_rate (`float`, *optional*): + The minimum learning rate as a ratio of the start learning rate. If not set, `warmup_lr_rate` will be treated as float(1/num_warmup_steps). + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + if min_lr is not None and min_lr_rate is not None: + raise ValueError("Only one of min_lr or min_lr_rate should be set") + elif min_lr is not None: + min_lr_rate = min_lr / optimizer.defaults["lr"] + elif min_lr_rate is None: + raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`") + + lr_lambda = partial( + _get_cosine_with_min_lr_schedule_with_warmup_lr_rate_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + min_lr_rate=min_lr_rate, + warmup_lr_rate=warmup_lr_rate, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_wsd_scheduler_lambda( + current_step: int, + *, + num_warmup_steps: int, + num_stable_steps: int, + num_decay_steps: int, + warmup_type: str, + decay_type: str, + min_lr_ratio: float, + num_cycles: float, +): + if current_step < num_warmup_steps: + progress = float(current_step) / float(max(1, num_warmup_steps)) + if warmup_type == "linear": + factor = progress + elif warmup_type == "cosine": + factor = 0.5 * (1.0 - math.cos(math.pi * progress)) + elif warmup_type == "1-sqrt": + factor = 1.0 - math.sqrt(1.0 - progress) + factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio + return max(0.0, factor) + + if current_step < num_warmup_steps + num_stable_steps: + return 1.0 + + if current_step < num_warmup_steps + num_stable_steps + num_decay_steps: + progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps)) + if decay_type == "linear": + factor = 1.0 - progress + elif decay_type == "cosine": + factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) + elif decay_type == "1-sqrt": + factor = 1.0 - math.sqrt(progress) + factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio + return max(0.0, factor) + return min_lr_ratio + + +def get_wsd_schedule( + optimizer: Optimizer, + num_warmup_steps: int, + num_decay_steps: int, + num_training_steps: int | None = None, + num_stable_steps: int | None = None, + warmup_type: str = "linear", + decay_type: str = "cosine", + min_lr_ratio: float = 0, + num_cycles: float = 0.5, + last_epoch: int = -1, +): + """ + Create a schedule with a learning rate that has three stages: + 1. warmup: increase from min_lr_ratio times the initial learning rate to the initial learning rate following a warmup_type. + 2. stable: constant learning rate. + 3. decay: decrease from the initial learning rate to min_lr_ratio times the initial learning rate following a decay_type. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_decay_steps (`int`): + The number of steps for the decay phase. + num_training_steps (`int`, *optional*): + The total number of training steps. This is the sum of the warmup, stable and decay steps. If `num_stable_steps` is not provided, the stable phase will be `num_training_steps - num_warmup_steps - num_decay_steps`. + num_stable_steps (`int`, *optional*): + The number of steps for the stable phase. Please ensure that `num_warmup_steps + num_stable_steps + num_decay_steps` equals `num_training_steps`, otherwise the other steps will default to the minimum learning rate. + warmup_type (`str`, *optional*, defaults to "linear"): + The type of warmup to use. Can be 'linear', 'cosine' or '1-sqrt'. + decay_type (`str`, *optional*, defaults to "cosine"): + The type of decay to use. Can be 'linear', 'cosine' or '1-sqrt'. + min_lr_ratio (`float`, *optional*, defaults to 0): + The minimum learning rate as a ratio of the initial learning rate. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + if num_training_steps is None and num_stable_steps is None: + raise ValueError("Either num_training_steps or num_stable_steps must be specified.") + + if num_training_steps is not None and num_stable_steps is not None: + warnings.warn("Both num_training_steps and num_stable_steps are specified. num_stable_steps will be used.") + + if warmup_type not in ["linear", "cosine", "1-sqrt"]: + raise ValueError(f"Unknown warmup type: {warmup_type}, expected 'linear', 'cosine' or '1-sqrt'") + + if decay_type not in ["linear", "cosine", "1-sqrt"]: + raise ValueError(f"Unknown decay type: {decay_type}, expected 'linear', 'cosine' or '1-sqrt'") + + if num_stable_steps is None: + num_stable_steps = num_training_steps - num_warmup_steps - num_decay_steps + + lr_lambda = partial( + _get_wsd_scheduler_lambda, + num_warmup_steps=num_warmup_steps, + num_stable_steps=num_stable_steps, + num_decay_steps=num_decay_steps, + warmup_type=warmup_type, + decay_type=decay_type, + min_lr_ratio=min_lr_ratio, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +class StreamingAverage: + """Rolling window average for smoothing metric values. + + Maintains a sliding window of values and computes their average, + useful for smoothing noisy metric values before making learning rate decisions. + + Args: + window_size (`int`): + The maximum number of values to keep in the rolling window. + """ + + def __init__(self, window_size: int) -> None: + self.window_size: int = window_size + self.values: list[float] = [] + self.sum: float = 0.0 + + def streamavg(self, value: float) -> float: + """Add a value and return the current rolling average.""" + self.values.append(value) + self.sum += value + + if len(self.values) > self.window_size: + removed = self.values.pop(0) + self.sum -= removed + + return self.sum / len(self.values) + + def state_dict(self) -> dict[str, Any]: + return { + "window_size": self.window_size, + "values": self.values.copy(), + "sum": self.sum, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.window_size = state_dict.get("window_size", self.window_size) + self.values = state_dict.get("values", []).copy() + self.sum = state_dict.get("sum", 0.0) + + +class GreedyLR: + """Adaptive learning rate scheduler that responds to training metrics. + + GreedyLR dynamically adjusts the learning rate based on training performance: + - Increases LR when metrics improve consistently (divides by factor) + - Decreases LR when metrics plateau (multiplies by factor) + + This differs from traditional schedulers like cosine annealing by responding + to actual training dynamics rather than following a predetermined schedule. + + Reference: `GreedyLR: A Novel Adaptive Learning Rate Scheduler `_ + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + mode (`str`, *optional*, defaults to `"min"`): + One of 'min' or 'max'. In 'min' mode, LR will be reduced when the + metric has stopped decreasing; in 'max' mode when it has stopped increasing. + factor (`float`, *optional*, defaults to 0.95): + Factor by which the learning rate will be adjusted. LR is multiplied by + factor on plateau and divided by factor on improvement. Must be < 1.0. + patience (`int`, *optional*, defaults to 10): + Number of epochs with no improvement after which learning rate will be adjusted. + threshold (`float`, *optional*, defaults to 1e-06): + Threshold for measuring the new optimum. + threshold_mode (`str`, *optional*, defaults to `"abs"`): + One of 'rel' or 'abs'. + cooldown (`int`, *optional*, defaults to 0): + Number of epochs to wait before resuming normal operation after LR has been reduced. + warmup (`int`, *optional*, defaults to 0): + Number of epochs to wait before resuming normal operation after LR has been increased. + min_lr (`float` or `list[float]`, *optional*, defaults to 0.001): + A lower bound on the learning rate. + max_lr (`float` or `list[float]`, *optional*, defaults to 1.0): + An upper bound on the learning rate. + eps (`float`, *optional*, defaults to 1e-08): + Minimal decay applied to lr. + verbose (`bool`, *optional*, defaults to `False`): + If True, prints a message to stdout for each update. + smooth (`bool`, *optional*, defaults to `False`): + If True, applies streaming average smoothing to metrics. + window_size (`int`, *optional*, defaults to 50): + The window size for the streaming average when smooth=True. + reset_start (`int`, *optional*, defaults to 500): + Number of steps to wait at min_lr before resetting to initial state. + + Example: + ```python + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + >>> scheduler = GreedyLR(optimizer, mode="min", patience=10) + >>> for epoch in range(100): + ... train(...) + ... val_loss = validate(...) + ... scheduler.step(val_loss) + ``` + """ + + def __init__( + self, + optimizer: Optimizer, + mode: str = "min", + factor: float = 0.95, + patience: int = 10, + threshold: float = 1e-6, + threshold_mode: str = "abs", + cooldown: int = 0, + warmup: int = 0, + min_lr: float | list[float] = 1e-3, + max_lr: float | list[float] = 1.0, + eps: float = 1e-8, + verbose: bool = False, + smooth: bool = False, + window_size: int = 50, + reset_start: int = 500, + ) -> None: + if factor >= 1.0: + raise ValueError("Factor should be < 1.0.") + if not isinstance(optimizer, Optimizer): + raise TypeError(f"{type(optimizer).__name__} is not an Optimizer") + + self.optimizer = optimizer + self.factor = factor + self.patience = patience + self.verbose = verbose + self.cooldown = cooldown + self.warmup = warmup + self.cooldown_counter = 0 + self.warmup_counter = 0 + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + self.eps = eps + self.smooth = smooth + self.window_size = window_size + self.reset_start = reset_start + self.reset_start_original = reset_start + self.last_epoch = 0 + + if isinstance(min_lr, (list, tuple)): + if len(min_lr) != len(optimizer.param_groups): + raise ValueError(f"expected {len(optimizer.param_groups)} min_lrs, got {len(min_lr)}") + self.min_lrs = list(min_lr) + else: + self.min_lrs = [min_lr] * len(optimizer.param_groups) + + if isinstance(max_lr, (list, tuple)): + if len(max_lr) != len(optimizer.param_groups): + raise ValueError(f"expected {len(optimizer.param_groups)} max_lrs, got {len(max_lr)}") + self.max_lrs = list(max_lr) + else: + self.max_lrs = [max_lr] * len(optimizer.param_groups) + + self._init_lrs = [group["lr"] for group in optimizer.param_groups] + self._last_lr = self._init_lrs.copy() + + self.best: float = float("inf") if mode == "min" else float("-inf") + self.num_bad_epochs = 0 + self.num_good_epochs = 0 + + if mode not in ("min", "max"): + raise ValueError(f"mode {mode} is unknown!") + if threshold_mode not in ("rel", "abs"): + raise ValueError(f"threshold mode {threshold_mode} is unknown!") + + self._streaming_avg: StreamingAverage | None = None + if smooth: + self._streaming_avg = StreamingAverage(window_size) + + def step(self, metrics: float, epoch: int | None = None) -> None: + """Perform a scheduler step based on the given metrics. + + Args: + metrics (`float`): + The metric value to use for LR adjustment decisions. + epoch (`int`, *optional*): + The current epoch number. If None, uses internal counter. + """ + current = float(metrics) + + if self.smooth and self._streaming_avg is not None: + current = self._streaming_avg.streamavg(current) + + if epoch is None: + epoch = self.last_epoch + 1 + self.last_epoch = epoch + + if self.cooldown_counter > 0: + self.cooldown_counter -= 1 + self.num_bad_epochs = 0 + self.num_good_epochs = 0 + elif self.warmup_counter > 0: + self.warmup_counter -= 1 + self.num_bad_epochs = 0 + self.num_good_epochs = 0 + else: + if self.is_better(current, self.best): + self.best = current + self.num_bad_epochs = 0 + self.num_good_epochs += 1 + else: + self.num_bad_epochs += 1 + self.num_good_epochs = 0 + + if self.num_good_epochs > self.patience: + self._increase_lr(epoch) + self.warmup_counter = self.warmup + self.num_good_epochs = 0 + elif self.num_bad_epochs > self.patience: + self._reduce_lr(epoch) + self.cooldown_counter = self.cooldown + self.num_bad_epochs = 0 + + self._last_lr = [group["lr"] for group in self.optimizer.param_groups] + + def is_better(self, current: float, best: float) -> bool: + if self.mode == "min": + if self.threshold_mode == "rel": + return current < best * (1.0 - self.threshold) + else: + return current < best - self.threshold + else: + if self.threshold_mode == "rel": + return current > best * (1.0 + self.threshold) + else: + return current > best + self.threshold + + def _reduce_lr(self, epoch: int) -> None: + all_at_min = True + for i, param_group in enumerate(self.optimizer.param_groups): + old_lr = float(param_group["lr"]) + new_lr = max(old_lr * self.factor, self.min_lrs[i]) + + if old_lr - new_lr > self.eps: + param_group["lr"] = new_lr + if self.verbose: + print(f"Epoch {epoch}: reducing learning rate of group {i} to {new_lr:.4e}.") + + if param_group["lr"] > self.min_lrs[i]: + all_at_min = False + + if all_at_min: + self.reset_start -= 1 + if self.reset_start <= 0: + self._reset() + + def _increase_lr(self, epoch: int) -> None: + for i, param_group in enumerate(self.optimizer.param_groups): + old_lr = float(param_group["lr"]) + new_lr = min(old_lr / self.factor, self.max_lrs[i]) + + if new_lr - old_lr > self.eps: + param_group["lr"] = new_lr + if self.verbose: + print(f"Epoch {epoch}: increasing learning rate of group {i} to {new_lr:.4e}.") + + self.reset_start = self.reset_start_original + + def _reset(self) -> None: + for i, param_group in enumerate(self.optimizer.param_groups): + param_group["lr"] = self._init_lrs[i] + + self.best = float("inf") if self.mode == "min" else float("-inf") + self.num_bad_epochs = 0 + self.num_good_epochs = 0 + self.cooldown_counter = 0 + self.warmup_counter = 0 + self.reset_start = self.reset_start_original + + if self.smooth and self._streaming_avg is not None: + self._streaming_avg = StreamingAverage(self.window_size) + + if self.verbose: + print("Scheduler reset to initial state.") + + def get_last_lr(self) -> list[float]: + """Return last computed learning rate by current scheduler.""" + return self._last_lr + + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a dictionary.""" + state = { + "factor": self.factor, + "min_lrs": self.min_lrs, + "max_lrs": self.max_lrs, + "patience": self.patience, + "verbose": self.verbose, + "cooldown": self.cooldown, + "warmup": self.warmup, + "cooldown_counter": self.cooldown_counter, + "warmup_counter": self.warmup_counter, + "mode": self.mode, + "threshold": self.threshold, + "threshold_mode": self.threshold_mode, + "best": self.best, + "num_bad_epochs": self.num_bad_epochs, + "num_good_epochs": self.num_good_epochs, + "eps": self.eps, + "last_epoch": self.last_epoch, + "smooth": self.smooth, + "window_size": self.window_size, + "reset_start": self.reset_start, + "reset_start_original": self.reset_start_original, + "_last_lr": self._last_lr, + "_init_lrs": self._init_lrs, + } + + if self.smooth and self._streaming_avg is not None: + state["_streaming_avg"] = self._streaming_avg.state_dict() + + return state + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load state from a dictionary.""" + self.factor = state_dict.get("factor", self.factor) + self.min_lrs = state_dict.get("min_lrs", self.min_lrs) + self.max_lrs = state_dict.get("max_lrs", self.max_lrs) + self.patience = state_dict.get("patience", self.patience) + self.verbose = state_dict.get("verbose", self.verbose) + self.cooldown = state_dict.get("cooldown", self.cooldown) + self.warmup = state_dict.get("warmup", self.warmup) + self.cooldown_counter = state_dict.get("cooldown_counter", self.cooldown_counter) + self.warmup_counter = state_dict.get("warmup_counter", self.warmup_counter) + self.mode = state_dict.get("mode", self.mode) + self.threshold = state_dict.get("threshold", self.threshold) + self.threshold_mode = state_dict.get("threshold_mode", self.threshold_mode) + self.best = state_dict.get("best", self.best) + self.num_bad_epochs = state_dict.get("num_bad_epochs", self.num_bad_epochs) + self.num_good_epochs = state_dict.get("num_good_epochs", self.num_good_epochs) + self.eps = state_dict.get("eps", self.eps) + self.last_epoch = state_dict.get("last_epoch", self.last_epoch) + self.smooth = state_dict.get("smooth", self.smooth) + self.window_size = state_dict.get("window_size", self.window_size) + self.reset_start = state_dict.get("reset_start", self.reset_start) + self.reset_start_original = state_dict.get("reset_start_original", self.reset_start_original) + self._last_lr = state_dict.get("_last_lr", self._last_lr) + self._init_lrs = state_dict.get("_init_lrs", self._init_lrs) + + if "_streaming_avg" in state_dict: + if self._streaming_avg is None: + self._streaming_avg = StreamingAverage(self.window_size) + self._streaming_avg.load_state_dict(state_dict["_streaming_avg"]) + + if "_last_lr" in state_dict: + for param_group, lr in zip(self.optimizer.param_groups, self._last_lr): + param_group["lr"] = lr + + +def get_greedy_schedule(optimizer: Optimizer, **kwargs): + """ + Create an adaptive learning rate scheduler that adjusts LR based on training metrics. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + kwargs (`dict`, *optional*): + Extra parameters passed to the scheduler. See [`GreedyLR`] for possible parameters. + + Return: + [`GreedyLR`] with the appropriate schedule. + """ + return GreedyLR(optimizer, **kwargs) + + +TYPE_TO_SCHEDULER_FUNCTION = { + SchedulerType.LINEAR: get_linear_schedule_with_warmup, + SchedulerType.COSINE: get_cosine_schedule_with_warmup, + SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup, + SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup, + SchedulerType.CONSTANT: get_constant_schedule, + SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup, + SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule, + SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule, + SchedulerType.COSINE_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup, + SchedulerType.COSINE_WARMUP_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup_lr_rate, + SchedulerType.WARMUP_STABLE_DECAY: get_wsd_schedule, + SchedulerType.GREEDY: get_greedy_schedule, +} + + +def get_scheduler( + name: str | SchedulerType, + optimizer: Optimizer, + num_warmup_steps: int | None = None, + num_training_steps: int | None = None, + scheduler_specific_kwargs: dict | None = None, +): + """ + Unified API to get any scheduler from its name. + + Args: + name (`str` or `SchedulerType`): + The name of the scheduler to use. + optimizer (`torch.optim.Optimizer`): + The optimizer that will be used during training. + num_warmup_steps (`int`, *optional*): + The number of warmup steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_training_steps (`int``, *optional*): + The number of training steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + scheduler_specific_kwargs (`dict`, *optional*): + Extra parameters for schedulers such as cosine with restarts. Mismatched scheduler types and scheduler + parameters will cause the scheduler function to raise a TypeError. + """ + name = SchedulerType(name) + schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] + + # If a `LayerWiseDummyOptimizer` is passed we extract the optimizer dict and + # recursively call `get_scheduler` to get the proper schedulers on each parameter + if optimizer is not None and isinstance(optimizer, LayerWiseDummyOptimizer): + optimizer_dict = optimizer.optimizer_dict + scheduler_dict = {} + + for param in optimizer_dict: + scheduler_dict[param] = get_scheduler( + name, + optimizer=optimizer_dict[param], + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + scheduler_specific_kwargs=scheduler_specific_kwargs, + ) + + def scheduler_hook(param): + # Since the optimizer hook has been already attached we only need to + # attach the scheduler hook, the gradients have been zeroed here + scheduler_dict[param].step() + + for param in optimizer_dict: + if param.requires_grad: + param.register_post_accumulate_grad_hook(scheduler_hook) + + return LayerWiseDummyScheduler(optimizer_dict=optimizer_dict, lr=optimizer.defaults["lr"]) + + if name == SchedulerType.CONSTANT: + return schedule_func(optimizer) + + if scheduler_specific_kwargs is None: + scheduler_specific_kwargs = {} + + if name == SchedulerType.REDUCE_ON_PLATEAU: + return schedule_func(optimizer, **scheduler_specific_kwargs) + + if name == SchedulerType.GREEDY: + return schedule_func(optimizer, **scheduler_specific_kwargs) + + # All other schedulers require `num_warmup_steps` + if num_warmup_steps is None: + raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") + + if name == SchedulerType.CONSTANT_WITH_WARMUP: + return schedule_func(optimizer, num_warmup_steps=num_warmup_steps) + + if name == SchedulerType.INVERSE_SQRT: + return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, **scheduler_specific_kwargs) + + # wsd scheduler requires either num_training_steps or num_stable_steps + if name == SchedulerType.WARMUP_STABLE_DECAY: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + **scheduler_specific_kwargs, + ) + + # All other schedulers require `num_training_steps` + if num_training_steps is None: + raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") + + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + **scheduler_specific_kwargs, + ) + + +class Adafactor(Optimizer): + """ + AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code: + https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py + + Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://huggingface.co/papers/1804.04235 Note that + this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and + `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and + `relative_step=False`. + + Arguments: + params (`Iterable[nn.parameter.Parameter]`): + Iterable of parameters to optimize or dictionaries defining parameter groups. + lr (`float`, *optional*): + The external learning rate. + eps (`tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`): + Regularization constants for square gradient and parameter scale respectively + clip_threshold (`float`, *optional*, defaults to 1.0): + Threshold of root mean square of final gradient update + decay_rate (`float`, *optional*, defaults to -0.8): + Coefficient used to compute running averages of square + beta1 (`float`, *optional*): + Coefficient used for computing running averages of gradient + weight_decay (`float`, *optional*, defaults to 0.0): + Weight decay (L2 penalty) + scale_parameter (`bool`, *optional*, defaults to `True`): + If True, learning rate is scaled by root mean square + relative_step (`bool`, *optional*, defaults to `True`): + If True, time-dependent learning rate is computed instead of external learning rate + warmup_init (`bool`, *optional*, defaults to `False`): + Time-dependent learning rate computation depends on whether warm-up initialization is being used + + This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested. + + Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3): + + - Training without LR warmup or clip_threshold is not recommended. + + - use scheduled LR warm-up to fixed LR + - use clip_threshold=1.0 (https://huggingface.co/papers/1804.04235) + - Disable relative updates + - Use scale_parameter=False + - Additional optimizer operations like gradient clipping should not be used alongside Adafactor + + Example: + + ```python + Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3) + ``` + + Others reported the following combination to work well: + + ```python + Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None) + ``` + + When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`] + scheduler as following: + + ```python + from transformers.optimization import Adafactor, AdafactorSchedule + + optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None) + lr_scheduler = AdafactorSchedule(optimizer) + trainer = Trainer(..., optimizers=(optimizer, lr_scheduler)) + ``` + + Usage: + + ```python + # replace AdamW with Adafactor + optimizer = Adafactor( + model.parameters(), + lr=1e-3, + eps=(1e-30, 1e-3), + clip_threshold=1.0, + decay_rate=-0.8, + beta1=None, + weight_decay=0.0, + relative_step=False, + scale_parameter=False, + warmup_init=False, + ) + ```""" + + def __init__( + self, + params, + lr=None, + eps=(1e-30, 1e-3), + clip_threshold=1.0, + decay_rate=-0.8, + beta1=None, + weight_decay=0.0, + scale_parameter=True, + relative_step=True, + warmup_init=False, + ): + if lr is not None and relative_step: + raise ValueError("Cannot combine manual `lr` and `relative_step=True` options") + if warmup_init and not relative_step: + raise ValueError("`warmup_init=True` requires `relative_step=True`") + + defaults = { + "lr": lr, + "eps": eps, + "clip_threshold": clip_threshold, + "decay_rate": decay_rate, + "beta1": beta1, + "weight_decay": weight_decay, + "scale_parameter": scale_parameter, + "relative_step": relative_step, + "warmup_init": warmup_init, + } + super().__init__(params, defaults) + + @staticmethod + def _get_lr(param_group, param_state): + rel_step_sz = param_group["lr"] + if param_group["relative_step"]: + min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2 + rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"])) + param_scale = 1.0 + if param_group["scale_parameter"]: + param_scale = max(param_group["eps"][1], param_state["RMS"]) + return param_scale * rel_step_sz + + @staticmethod + def _get_options(param_group, param_shape): + factored = len(param_shape) >= 2 + use_first_moment = param_group["beta1"] is not None + return factored, use_first_moment + + @staticmethod + def _rms(tensor): + return tensor.norm(2) / (tensor.numel() ** 0.5) + + @staticmethod + def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col): + # copy from fairseq's adafactor implementation: + # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505 + r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1) + c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt() + return torch.mul(r_factor, c_factor) + + @torch.no_grad() + def step(self, closure=None): + """ + Performs a single optimization step + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if grad.dtype in {torch.float16, torch.bfloat16}: + grad = grad.float() + if grad.is_sparse: + raise RuntimeError("Adafactor does not support sparse gradients.") + + state = self.state[p] + grad_shape = grad.shape + + factored, use_first_moment = self._get_options(group, grad_shape) + # State Initialization + if len(state) == 0: + state["step"] = 0 + + if use_first_moment: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(grad) + if factored: + state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad) + state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad) + else: + state["exp_avg_sq"] = torch.zeros_like(grad) + + state["RMS"] = 0 + else: + if use_first_moment: + state["exp_avg"] = state["exp_avg"].to(grad) + if factored: + state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad) + state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad) + else: + state["exp_avg_sq"] = state["exp_avg_sq"].to(grad) + + p_data_fp32 = p + if p.dtype in {torch.float16, torch.bfloat16}: + p_data_fp32 = p_data_fp32.float() + + state["step"] += 1 + state["RMS"] = self._rms(p_data_fp32) + lr = self._get_lr(group, state) + + beta2t = 1.0 - math.pow(state["step"], group["decay_rate"]) + update = (grad**2) + group["eps"][0] + if factored: + exp_avg_sq_row = state["exp_avg_sq_row"] + exp_avg_sq_col = state["exp_avg_sq_col"] + + exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t)) + exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t)) + + # Approximation of exponential moving average of square of gradient + update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) + update.mul_(grad) + else: + exp_avg_sq = state["exp_avg_sq"] + + exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t)) + update = exp_avg_sq.rsqrt().mul_(grad) + + update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)) + update.mul_(lr) + + if use_first_moment: + exp_avg = state["exp_avg"] + exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"])) + update = exp_avg + + if group["weight_decay"] != 0: + p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr)) + + p_data_fp32.add_(-update) + + if p.dtype in {torch.float16, torch.bfloat16}: + p.copy_(p_data_fp32) + + return loss + + +class AdafactorSchedule(LambdaLR): + """ + Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g., + for logging), this class creates a proxy object that retrieves the current lr values from the optimizer. + + It returns `initial_lr` during startup and the actual `lr` during stepping. + """ + + def __init__(self, optimizer, initial_lr=0.0): + def lr_lambda(_): + return initial_lr + + for group in optimizer.param_groups: + group["initial_lr"] = initial_lr + super().__init__(optimizer, lr_lambda) + for group in optimizer.param_groups: + del group["initial_lr"] + + def get_lr(self): + opt = self.optimizer + lrs = [ + opt._get_lr(group, opt.state[group["params"][0]]) + for group in opt.param_groups + if group["params"][0].grad is not None + ] + if len(lrs) == 0: + lrs = self.base_lrs # if called before stepping + return lrs + + +def get_adafactor_schedule(optimizer, initial_lr=0.0): + """ + Get a proxy schedule for [`~optimization.Adafactor`] + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + initial_lr (`float`, *optional*, defaults to 0.0): + Initial lr + + Return: + [`~optimization.Adafactor`] proxy schedule object. + + + """ + return AdafactorSchedule(optimizer, initial_lr) diff --git a/third_party/transformers/src/transformers/time_series_utils.py b/third_party/transformers/src/transformers/time_series_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ebe1a2cbd3c2819276a0cd0d6105caa5a6397f0d --- /dev/null +++ b/third_party/transformers/src/transformers/time_series_utils.py @@ -0,0 +1,225 @@ +# Copyright 2023 The HuggingFace Inc. team. +# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Time series distributional output classes and utilities. +""" + +from collections.abc import Callable + +import torch +from torch import nn +from torch.distributions import ( + AffineTransform, + Distribution, + Independent, + NegativeBinomial, + Normal, + StudentT, + TransformedDistribution, +) + + +class AffineTransformed(TransformedDistribution): + def __init__(self, base_distribution: Distribution, loc=None, scale=None, event_dim=0): + self.scale = 1.0 if scale is None else scale + self.loc = 0.0 if loc is None else loc + + super().__init__(base_distribution, [AffineTransform(loc=self.loc, scale=self.scale, event_dim=event_dim)]) + + @property + def mean(self): + """ + Returns the mean of the distribution. + """ + return self.base_dist.mean * self.scale + self.loc + + @property + def variance(self): + """ + Returns the variance of the distribution. + """ + return self.base_dist.variance * self.scale**2 + + @property + def stddev(self): + """ + Returns the standard deviation of the distribution. + """ + return self.variance.sqrt() + + +class ParameterProjection(nn.Module): + def __init__( + self, in_features: int, args_dim: dict[str, int], domain_map: Callable[..., tuple[torch.Tensor]], **kwargs + ) -> None: + super().__init__(**kwargs) + self.args_dim = args_dim + self.proj = nn.ModuleList([nn.Linear(in_features, dim) for dim in args_dim.values()]) + self.domain_map = domain_map + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]: + params_unbounded = [proj(x) for proj in self.proj] + + return self.domain_map(*params_unbounded) + + +class LambdaLayer(nn.Module): + def __init__(self, function): + super().__init__() + self.function = function + + def forward(self, x, *args): + return self.function(x, *args) + + +class DistributionOutput: + distribution_class: type + in_features: int + args_dim: dict[str, int] + + def __init__(self, dim: int = 1) -> None: + self.dim = dim + self.args_dim = {k: dim * self.args_dim[k] for k in self.args_dim} + + def _base_distribution(self, distr_args): + if self.dim == 1: + return self.distribution_class(*distr_args) + else: + return Independent(self.distribution_class(*distr_args), 1) + + def distribution( + self, + distr_args, + loc: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + ) -> Distribution: + distr = self._base_distribution(distr_args) + if loc is None and scale is None: + return distr + else: + return AffineTransformed(distr, loc=loc, scale=scale, event_dim=self.event_dim) + + @property + def event_shape(self) -> tuple: + r""" + Shape of each individual event contemplated by the distributions that this object constructs. + """ + return () if self.dim == 1 else (self.dim,) + + @property + def event_dim(self) -> int: + r""" + Number of event dimensions, i.e., length of the `event_shape` tuple, of the distributions that this object + constructs. + """ + return len(self.event_shape) + + @property + def value_in_support(self) -> float: + r""" + A float that will have a valid numeric value when computing the log-loss of the corresponding distribution. By + default 0.0. This value will be used when padding data series. + """ + return 0.0 + + def get_parameter_projection(self, in_features: int) -> nn.Module: + r""" + Return the parameter projection layer that maps the input to the appropriate parameters of the distribution. + """ + return ParameterProjection( + in_features=in_features, + args_dim=self.args_dim, + domain_map=LambdaLayer(self.domain_map), + ) + + def domain_map(self, *args: torch.Tensor): + r""" + Converts arguments to the right shape and domain. The domain depends on the type of distribution, while the + correct shape is obtained by reshaping the trailing axis in such a way that the returned tensors define a + distribution of the right event_shape. + """ + raise NotImplementedError() + + @staticmethod + def squareplus(x: torch.Tensor) -> torch.Tensor: + r""" + Helper to map inputs to the positive orthant by applying the square-plus operation. Reference: + https://twitter.com/jon_barron/status/1387167648669048833 + """ + return (x + torch.sqrt(torch.square(x) + 4.0)) / 2.0 + + +class StudentTOutput(DistributionOutput): + """ + Student-T distribution output class. + """ + + args_dim: dict[str, int] = {"df": 1, "loc": 1, "scale": 1} + distribution_class: type = StudentT + + @classmethod + def domain_map(cls, df: torch.Tensor, loc: torch.Tensor, scale: torch.Tensor): + scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps) + df = 2.0 + cls.squareplus(df) + return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1) + + +class NormalOutput(DistributionOutput): + """ + Normal distribution output class. + """ + + args_dim: dict[str, int] = {"loc": 1, "scale": 1} + distribution_class: type = Normal + + @classmethod + def domain_map(cls, loc: torch.Tensor, scale: torch.Tensor): + scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps) + return loc.squeeze(-1), scale.squeeze(-1) + + +class NegativeBinomialOutput(DistributionOutput): + """ + Negative Binomial distribution output class. + """ + + args_dim: dict[str, int] = {"total_count": 1, "logits": 1} + distribution_class: type = NegativeBinomial + + @classmethod + def domain_map(cls, total_count: torch.Tensor, logits: torch.Tensor): + total_count = cls.squareplus(total_count) + return total_count.squeeze(-1), logits.squeeze(-1) + + def _base_distribution(self, distr_args) -> Distribution: + total_count, logits = distr_args + if self.dim == 1: + return self.distribution_class(total_count=total_count, logits=logits) + else: + return Independent(self.distribution_class(total_count=total_count, logits=logits), 1) + + # Overwrites the parent class method. We cannot scale using the affine + # transformation since negative binomial should return integers. Instead + # we scale the parameters. + def distribution( + self, distr_args, loc: torch.Tensor | None = None, scale: torch.Tensor | None = None + ) -> Distribution: + total_count, logits = distr_args + + if scale is not None: + # See scaling property of Gamma. + logits += scale.log() + + return self._base_distribution((total_count, logits)) diff --git a/third_party/transformers/src/transformers/tokenization_mistral_common.py b/third_party/transformers/src/transformers/tokenization_mistral_common.py new file mode 100644 index 0000000000000000000000000000000000000000..1f218fe408733a3b40699be163950c578d979a55 --- /dev/null +++ b/third_party/transformers/src/transformers/tokenization_mistral_common.py @@ -0,0 +1,1685 @@ +# Copyright 2025 Mistral AI and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import re +import shutil +from collections.abc import Callable, Sequence +from enum import Enum +from pathlib import Path +from typing import Any, Literal, Union, overload + +import numpy as np +from huggingface_hub import create_repo + +from transformers.audio_utils import load_audio_as +from transformers.image_utils import get_image_size +from transformers.tokenization_utils_base import ( + VERY_LARGE_INTEGER, + AddedToken, + BatchEncoding, + EncodedInput, + PreTokenizedInput, + PreTrainedTokenizerBase, + TextInput, + TruncationStrategy, +) +from transformers.utils import PaddingStrategy, TensorType, add_end_docstrings, logging, to_py_obj +from transformers.utils.import_utils import is_mistral_common_available, is_torch_available, requires + + +if is_mistral_common_available(): + from mistral_common.protocol.instruct.request import ChatCompletionRequest, ReasoningEffort + from mistral_common.protocol.instruct.validator import ValidationMode + from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy, SpecialTokens + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer + from mistral_common.tokens.tokenizers.tekken import Tekkenizer + from mistral_common.tokens.tokenizers.utils import ( + download_tokenizer_from_hf_hub, + get_one_valid_tokenizer_file, + ) + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to add special tokens when encoding the sequences. This will use the underlying + `PretrainedTokenizerBase.build_inputs_with_special_tokens` function, which defines which tokens are + automatically added to the input ids. This is useful if you want to add `bos` or `eos` tokens + automatically. When Tokenizer is loading with `finetuning` mode it adds both `bos` and `eos`. Else, for "test" mode it only adds `bos`. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence is 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. + - `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. Requires `padding` to be activated. + This is especially useful to enable the use of Tensor Cores 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_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. +""" + +ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. For `MistralCommonBackend` it returns a list of zeros of the sequence length as only one sequence is supported. + + [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_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. + return_offsets_mapping (`Literal[False]`, *optional*): False, kept to match Transformers' signature. + split_special_tokens (`Literal[False]`, *optional*): False, kept to match Transformers' signature. + **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) + + - **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) + + - **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 MistralTokenizerType(str, Enum): + """Enum for the different type of tokenizer.""" + + spm = "spm" + tekken = "tekken" + + +@overload +def _maybe_remove_lang(text: str, skip_special_tokens: bool) -> str: ... +@overload +def _maybe_remove_lang(text: list[str], skip_special_tokens: bool) -> list[str]: ... +def _maybe_remove_lang(text: str | list[str], skip_special_tokens: bool) -> str | list[str]: + # in the specific case of Voxtral, the added f"lang:xx" (always a two char language code since it follows ISO 639-1 alpha-2 format) + # is not considered as a special token by mistral-common and is encoded/ decoded as normal text. + # Nevertheless we should remove it to ease users life. + if not skip_special_tokens: + return text + + if isinstance(text, str): + return re.sub(r"^lang:[a-z]{2}", "", text) + + return [re.sub(r"^lang:[a-z]{2}", "", string) for string in text] + + +_MAP_SPECIAL_TOKENS = { + "bos_token": SpecialTokens.bos.value, + "eos_token": SpecialTokens.eos.value, + "pad_token": SpecialTokens.pad.value, + "unk_token": SpecialTokens.unk.value, +} + +_VALID_INIT_KWARGS = {"_from_auto", "backend", "files_loaded"} + + +@requires(backends=("mistral-common",)) +class MistralCommonBackend(PreTrainedTokenizerBase): + """ + Class to wrap `mistral-common` tokenizers. + + `mistral-common` is the official tokenizer library for Mistral AI models. To use it, you need to install it with: + + ```bash + pip install transformers[mistral-common] + ``` + + Otherwise the tokenizer falls back to the Transformers implementation of the tokenizer. + + For more info on `mistral-common`, see [mistral-common](https://github.com/mistralai/mistral-common). + + This class is a wrapper around a `mistral_common.tokens.tokenizers.mistral.MistralTokenizer`. + It provides a Hugging Face compatible interface to tokenize using the official mistral-common tokenizer and inherits from the `PreTrainedTokenizerBase` class. + + Here are the key behavior differences with the `PythonBackend` class: + + - Pair of sequences are not supported. The signature has been kept for compatibility but all arguments related to pair of sequences are ignored. The return values for pairs are returned as `None`. + - The `is_split_into_words` argument is not supported. + - It is not possible to add new tokens to the tokenizer. Special tokens are handled differently from Transformers. In `mistral-common`, special tokens are never encoded directly. This means that: `tokenizer.encode("")` will not return the ID of the `` token. Instead, it will return a list of IDs corresponding to the tokenization of the string `""`. For more information, see the [mistral-common documentation](https://mistralai.github.io/mistral-common/usage/tokenizers/#special-tokens). + + If you have suggestions to improve this class, please open an issue on the [mistral-common GitHub repository](https://github.com/mistralai/mistral-common/issues) if it is related to the tokenizer or on the [Transformers GitHub repository](https://github.com/huggingface/transformers/issues) if it is related to the Hugging Face interface. + """ + + model_input_names: list[str] = ["input_ids", "attention_mask"] + padding_side: str = "left" + truncation_side: str = "right" + SPECIAL_TOKENS_ATTRIBUTES = [ + "bos_token", + "eos_token", + "unk_token", + "pad_token", + ] + + def __init__( + self, + tokenizer_path: str | os.PathLike | Path, + mode: ValidationMode = ValidationMode.test, + model_max_length: int = VERY_LARGE_INTEGER, + padding_side: str = "left", + truncation_side: str = "right", + model_input_names: list[str] | None = None, + clean_up_tokenization_spaces: bool = False, + **kwargs, + ): + """ + Constructs a `MistralCommonBackend`. + + - **model_input_names** (`list[str]`) -- A list of inputs expected in the forward pass of the model. + - **padding_side** (`str`) -- The default value for the side on which the model should have padding applied. + Should be `'right'` or `'left'`. + - **truncation_side** (`str`) -- The default value for the side on which the model should have truncation + applied. Should be `'right'` or `'left'`. + + Args: + tokenizer_path (`str` or `os.PathLike` or `Path`): + Path to the tokenizer file to load the `MistralTokenizer`. + mode (`Union[str, ValidationMode]`, *optional*, defaults to `ValidationMode.test`): + The mode to use for the tokenizer. This will be passed to the `MistralTokenizer` constructor. Possible values are: + - `"finetuning"` or `ValidationMode.finetuning`: The fine-tuning mode. + - `"test"` or `ValidationMode.test`: The test mode. + It changes how the tokenizer validates the input and prepares the request to the model. + model_max_length (`int`, *optional*): + The maximum length (in number of tokens) for the inputs to the transformer model. When the tokenizer is + loaded with [`~tokenization_utils_base.PreTrainedTokenizerBase.from_pretrained`], this will be set to the + value stored for the associated model in `max_model_input_sizes` (see above). If no value is provided, will + default to VERY_LARGE_INTEGER (`int(1e30)`). + 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. + truncation_side (`str`, *optional*): + The side on which the model should have truncation applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + model_input_names (`List[str]`, *optional*): + The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or + `"attention_mask"`). Default value is picked from the class attribute of the same name. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not the model should clean up the spaces that were added when splitting the input text during the + tokenization process. + """ + if kwargs and not set(kwargs.keys()).issubset(_VALID_INIT_KWARGS): + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported to init `MistralCommonBackend`.") + + self.init_kwargs = { + "tokenizer_path": tokenizer_path, + "mode": mode, + "model_max_length": model_max_length, + "padding_side": padding_side, + "truncation_side": truncation_side, + "model_input_names": model_input_names, + "clean_up_tokenization_spaces": clean_up_tokenization_spaces, + } + self._tokenizer_path = Path(tokenizer_path) + self._mode = self._get_validation_mode(mode) + + self.tokenizer: MistralTokenizer = MistralTokenizer.from_file(str(self._tokenizer_path), mode=self._mode) + self._tokenizer_type = ( + MistralTokenizerType.tekken + if isinstance(self.tokenizer.instruct_tokenizer.tokenizer, Tekkenizer) + else MistralTokenizerType.spm + ) + self._cache_get_vocab: dict[str, int] | None = None + + self._all_special_ids = self._get_all_special_ids() + self._all_special_tokens = self.convert_ids_to_tokens(self.all_special_ids) + + super().__init__( + truncation_side=truncation_side, + padding_side=padding_side, + model_max_length=model_max_length, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + extra_special_tokens=None, # Not used by this backend. + model_specific_special_tokens=None, # Not used by this backend. + model_input_names=model_input_names or self.model_input_names, + **_MAP_SPECIAL_TOKENS, + **kwargs, + ) + + @property + def mode(self) -> ValidationMode: + """ + `ValidationMode`: The mode used by the tokenizer. Possible values are: + - `"finetuning"` or `ValidationMode.finetuning`: The finetuning mode. + - `"test"` or `ValidationMode.test`: The test mode. + It changes how the tokenizer validates the input and prepares the request to the model. + """ + return self._mode + + @property + def all_special_ids(self) -> list[int]: + """ + `list[int]`: List the ids of the special tokens(`''`, `''`, etc.). + """ + return sorted(self._all_special_ids) + + @property + def all_special_tokens(self) -> list[str]: + """ + `list[str]`: A list of all unique special tokens. + """ + return self._all_special_tokens + + @property + def vocab_size(self) -> int: + """ + Returns the size of the vocabulary. + + `int`: Size of the vocabulary. + """ + return self.tokenizer.instruct_tokenizer.tokenizer.n_words + + def get_vocab(self) -> dict[str, int]: + """ + Returns the vocabulary as a dictionary of token to index. + + This is a lossy conversion. There may be multiple token ids that decode to the same + string due to partial UTF-8 byte sequences being converted to �. + + Returns: + `Dict[str, int]`: The vocabulary. + """ + if self._cache_get_vocab is None: + # We reverse the order to make sure that the first token is the one to be returned when there are multiple tokens with the same string representation. + vocab = self.tokenizer.instruct_tokenizer.tokenizer.vocab() + self._cache_get_vocab = {token: self._piece_to_id(token, False) for token in vocab} + # Order the dict. + self._cache_get_vocab = dict( + sorted(((k, v) for k, v in self._cache_get_vocab.items()), key=lambda x: x[1]) + ) + return self._cache_get_vocab + + def __len__(self): + """ + Size of the full vocabulary with the added tokens. + """ + return self.vocab_size + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, + """ + **kwargs: Not supported by `MistralCommonBackend.encode`. + Will raise an error if used. + """, + """ + Returns: + `list[int]`, `torch.Tensor`: The tokenized ids of the text. + """, + ) + def encode( + self, + text: TextInput | EncodedInput, + text_pair: None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + verbose: bool = True, + return_offsets_mapping: Literal[False] = False, + split_special_tokens: Literal[False] = False, + **kwargs, + ) -> list[int]: + """ + Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. + + Args: + text (`str` or `list[int]`): + The first sequence to be encoded. This can be a string or a list of integers (tokenized string ids). + text_pair (`None`, *optional*): + Not supported by `MistralCommonBackend.encode`. Kept to match `PreTrainedTokenizerBase.encode` signature. + """ + if return_offsets_mapping or split_special_tokens: + raise ValueError( + "`MistralCommonBackend` does not support `return_offsets_mapping` and `split_special_tokens`." + ) + + if truncation in [TruncationStrategy.ONLY_FIRST, TruncationStrategy.ONLY_SECOND, "only_first", "only_second"]: + raise ValueError( + "Truncation strategy `only_first` and `only_second` are not supported by `MistralCommonBackend`." + ) + + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.encode`.") + + if text_pair: + raise ValueError("`MistralCommonBackend.encode` does not support `text_pair`.") + + return super().encode( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + return_tensors=return_tensors, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + verbose=verbose, + ) + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> str: + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.decode`.") + + token_ids = to_py_obj(token_ids) + + if isinstance(token_ids, int): + token_ids = [token_ids] + + special_token_policy = SpecialTokenPolicy.IGNORE if skip_special_tokens else SpecialTokenPolicy.KEEP + + text = self.tokenizer.decode(token_ids, special_token_policy=special_token_policy) + + # Apply tokenizer-specific cleanup if available and requested + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + text = self.clean_up_tokenization(text) + + return _maybe_remove_lang(text=text, skip_special_tokens=skip_special_tokens) + + def decode( + self, + token_ids: Union[int, list[int], list[list[int]], np.ndarray, "torch.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> str | list[str]: + """ + Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special + tokens and clean up tokenization spaces. + + Args: + token_ids (`Union[int, list[int], list[list[int]], np.ndarray, torch.Tensor]`): + A single sequence or a batch (list of sequences) of tokenized input ids. Can be obtained using the + `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Not supported by `MistralCommonBackend.decode`. + Will raise an error if used. + + Returns: + `Union[str, list[str]]`: The decoded string for a single sequence, or a list of decoded strings for a + batch of sequences. + """ + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.decode`.") + + return super().decode( + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + ) + + def batch_decode( + self, + sequences: Union[list[int], list[list[int]], np.ndarray, "torch.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> list[str]: + """ + Convert a list of lists of token ids into a list of strings by calling decode. + + This method is provided for backwards compatibility. The `decode` method now handles batched input natively, + so you can use `decode` directly instead of `batch_decode`. + + Args: + sequences (`Union[list[int], list[list[int]], np.ndarray, torch.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Not supported by `MistralCommonBackend.batch_decode`. + Will raise an error if used. + + Returns: + `list[str]`: The list of decoded sentences. + """ + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.batch_decode`.") + + return super().batch_decode( + sequences=sequences, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + ) + + @overload + def convert_ids_to_tokens(self, ids: int, skip_special_tokens: bool = False) -> str: ... + @overload + def convert_ids_to_tokens(self, ids: list[int], skip_special_tokens: bool = False) -> list[str]: ... + def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `list[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `list[str]`: The decoded token(s). + """ + + if isinstance(ids, int): + return_int = True + ids = [ids] + else: + return_int = False + + tokens: list[str] = [] + for token_id in ids: + if self.tokenizer.instruct_tokenizer.tokenizer.is_special(token_id) and skip_special_tokens: + continue + tokens.append(self.tokenizer.instruct_tokenizer.tokenizer.id_to_piece(token_id)) + + if return_int and tokens == []: + raise ValueError(f"Invalid token id {ids[0]}.") + elif return_int: + return tokens[0] + + return tokens + + def _tekken_piece_to_id(self, piece: str, warn: bool) -> int: + tekken_tokenizer = self.tokenizer.instruct_tokenizer.tokenizer + assert isinstance(tekken_tokenizer, Tekkenizer), type(tekken_tokenizer) + + piece_bytes = piece.encode("utf-8") + shift = tekken_tokenizer.num_special_tokens + try: + return shift + tekken_tokenizer._tekken_token2id_nospecial[piece_bytes] + except KeyError: + piece_str = piece_bytes.decode("utf-8") + if piece_str in tekken_tokenizer._special_tokens_reverse_vocab: + return tekken_tokenizer._special_tokens_reverse_vocab[piece_str] + if warn: + logger.warning("Failed to convert token %s to id, replacing with ", piece_bytes) + return tekken_tokenizer.unk_id + + def _piece_to_id(self, piece: str, warn: bool) -> int: + if self._tokenizer_type == MistralTokenizerType.spm: + return self.tokenizer.instruct_tokenizer.tokenizer._model.piece_to_id(piece) + elif self._tokenizer_type == MistralTokenizerType.tekken: + return self._tekken_piece_to_id(piece, warn) + else: + raise ValueError(f"Unknown tokenizer type: {self._tokenizer_type}") + + def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: + """ + Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids), using the + vocabulary. + + Args: + tokens (`str` or `list[str]`): One or several token(s) to convert to token id(s). + + Returns: + `int` or `list[int]`: The token id or list of token ids. + """ + + if isinstance(tokens, str): + one_token = True + tokens = [tokens] + else: + one_token = False + + ids: list[int] = [] + for token in tokens: + ids.append(self._piece_to_id(token, True)) + + if one_token: + return ids[0] + return ids + + def _text_to_ids(self, text: TextInput, add_special_tokens: bool) -> list[int]: + """ + Converts a string into a sequence of tokens ids, using the tokenizer. + """ + add_eos = add_special_tokens and self._mode == ValidationMode.finetuning + tokens_ids = self.tokenizer.instruct_tokenizer.tokenizer.encode(text, bos=add_special_tokens, eos=add_eos) + return tokens_ids + + def tokenize( + self, + text: TextInput, + return_offsets_mapping: Literal[False] = False, + split_special_tokens: Literal[False] = False, + **kwargs, + ) -> list[str]: + """ + Converts a string into a sequence of tokens, using the tokenizer. + + Split in words for word-based vocabulary or sub-words for sub-word-based vocabularies. + + Args: + text (`str`): + The sequence to be encoded. + return_offsets_mapping (`Literal[False]`, *optional*): False, kept to match Transformers' signature. + split_special_tokens (`Literal[False]`, *optional*): False, kept to match Transformers' signature. + **kwargs (additional keyword arguments): + Not supported by `MistralCommonBackend.tokenize`. + Will raise an error if used. + + Returns: + `list[str]`: The list of tokens. + """ + if return_offsets_mapping or split_special_tokens: + raise ValueError( + "`MistralCommonBackend` does not support `return_offsets_mapping` and `split_special_tokens`." + ) + + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.tokenize`.") + + return self.convert_ids_to_tokens(self._text_to_ids(text, add_special_tokens=False), skip_special_tokens=False) + + def _get_all_special_ids(self) -> set[int]: + if self._tokenizer_type == MistralTokenizerType.tekken: + return self.tokenizer.instruct_tokenizer.tokenizer._special_token_ids + elif self._tokenizer_type == MistralTokenizerType.spm: + return { + token_id + for token_id in range(self.tokenizer.instruct_tokenizer.tokenizer.n_words) + if self.tokenizer.instruct_tokenizer.tokenizer.is_special(token_id) + } + else: + raise ValueError(f"Unknown tokenizer type: {self._tokenizer_type}") + + def get_special_tokens_mask( + self, token_ids_0: list[int], token_ids_1: None = None, already_has_special_tokens: bool = False + ) -> list[int]: + """ + Retrieves 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` or `encode_plus` methods. + + Args: + token_ids_0 (`list[int]`): List of ids of the sequence. + token_ids_1 (`None`, *optional*): None, kept to match Transformers' implementation. + 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: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if token_ids_1 is not None: + raise ValueError( + "`token_ids_1` is not supported by `MistralCommonBackend` and should be `None`, kept for compatibility." + ) + + if already_has_special_tokens: + return [1 if int(token_id) in self._all_special_ids else 0 for token_id in token_ids_0] + + if self.mode == ValidationMode.test: + # [BOS] seq0 + return [1] + ([0] * len(token_ids_0)) + else: + # [BOS] seq0 [EOS] + return [1] + ([0] * len(token_ids_0)) + [1] + + def _encode_plus( # type: ignore[override] + self, + text: TextInput | PreTokenizedInput | EncodedInput, + text_pair: 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, + 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_length: bool = False, + verbose: bool = True, + return_offsets_mapping: Literal[False] = False, + split_special_tokens: Literal[False] = False, + **kwargs, + ) -> BatchEncoding: + # Detect batched inputs (list of sequences) + if text_pair is not None: + raise ValueError("`MistralCommonBackend` does not support `text_pair != None` for `_encode_plus`.") + + if return_offsets_mapping or split_special_tokens: + raise ValueError( + "`MistralCommonBackend` does not support `return_offsets_mapping` and `split_special_tokens`." + ) + + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend._encode_plus`.") + + is_batched = isinstance(text, (list, tuple)) and ( + (not text and not is_split_into_words) + or (text and is_split_into_words and isinstance(text[0], (list, tuple))) + or (text and not is_split_into_words and isinstance(text[0], (str, list, tuple))) + ) + + if is_batched: + batch_outputs = {} + one_overflowed = False + for current_text in text: + current_output = self._encode_plus( + text=current_text, + text_pair=None, + add_special_tokens=add_special_tokens, + padding_strategy=PaddingStrategy.DO_NOT_PAD, # we pad in batch afterward + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=None, # we pad in batch afterward + padding_side=None, # we pad in batch afterward + return_tensors=None, # We convert the whole batch to tensors at the end + return_token_type_ids=return_token_type_ids, + return_attention_mask=False, # we pad in batch afterward + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + for key, value in current_output.items(): + batch_outputs.setdefault(key, []).append(value) + + # To ensure the list is built for each sample, we need to add this. + if return_overflowing_tokens and not return_tensors: + if "overflowing_tokens" not in current_output: + batch_outputs.setdefault("overflowing_tokens", []).append([0]) + batch_outputs.setdefault("num_truncated_tokens", []).append([0]) + else: + one_overflowed = True + + # Remove overflow-related keys before tensor conversion if return_tensors is set + # Slow tokenizers don't support returning these as tensors + if return_overflowing_tokens and (return_tensors or not one_overflowed): + batch_outputs.pop("overflowing_tokens", None) + batch_outputs.pop("num_truncated_tokens", None) + + batch_outputs = self.pad( + batch_outputs, + padding=padding_strategy.value, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def get_input_ids(text): + if isinstance(text, str): + return self._text_to_ids(text, False) + elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int): + return text + else: + raise ValueError(f"Input {text} is not valid. Should be a string, or a list/tuple of integers.") + + first_ids = get_input_ids(text) + + return self.prepare_for_model( + first_ids, + pair_ids=None, + add_special_tokens=add_special_tokens, + padding=padding_strategy.value, + truncation=truncation_strategy.value, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + prepend_batch_axis=True, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def prepare_for_model( + self, + ids: list[int], + pair_ids: None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + 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_length: bool = False, + verbose: bool = True, + prepend_batch_axis: bool = False, + return_offsets_mapping: Literal[False] = False, + split_special_tokens: Literal[False] = False, + **kwargs, + ) -> BatchEncoding: + """ + Prepares a sequence of input id so that it can be used by the model. It + adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens. + + Args: + ids (`list[int]`): + Tokenized input ids of the first sequence. + pair_ids (`None`, *optional*): + Not supported by `MistralCommonBackend`. Kept to match the interface of `PreTrainedTokenizerBase`. + """ + if return_offsets_mapping or split_special_tokens: + raise ValueError( + "`MistralCommonBackend` does not support `return_offsets_mapping` and `split_special_tokens`." + ) + + if pair_ids is not None: + raise ValueError( + "`pair_ids` is not supported by `MistralCommonBackend` and should be `None`, kept for compatibility." + ) + + if kwargs: + raise ValueError( + f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.prepare_for_model`." + ) + + padding_strategy, truncation_strategy, max_length, _ = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + # Validation + if ( + return_overflowing_tokens + and truncation_strategy == TruncationStrategy.LONGEST_FIRST + and pair_ids is not None + ): + raise ValueError( + "Not possible to return overflowing tokens for pair of sequences with the " + "`longest_first`. Please select another truncation strategy than `longest_first`, " + "for instance `only_second` or `only_first`." + ) + + # Defaults + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + # Truncation + num_special = self.num_special_tokens_to_add(pair=False) if add_special_tokens else 0 + total_len = len(ids) + len(pair_ids or []) + num_special + + overflowing_tokens = [] + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length: + ids, _, overflowing_tokens = self.truncate_sequences( + ids, + pair_ids=None, + num_tokens_to_remove=total_len - max_length, + truncation_strategy=truncation_strategy, + stride=stride, + ) + + # Add special tokens + if add_special_tokens: + sequence = self.build_inputs_with_special_tokens(ids, None) + token_type_ids = self.create_token_type_ids_from_sequences(ids, None) + else: + sequence = ids + token_type_ids = [0] * len(sequence) + + # Build output + encoded_inputs = {"input_ids": sequence} + if return_token_type_ids: + encoded_inputs["token_type_ids"] = token_type_ids + if return_special_tokens_mask: + encoded_inputs["special_tokens_mask"] = ( + self.get_special_tokens_mask(ids, None) if add_special_tokens else [0] * len(sequence) + ) + if return_overflowing_tokens and not return_tensors and overflowing_tokens: + encoded_inputs["overflowing_tokens"] = overflowing_tokens + encoded_inputs["num_truncated_tokens"] = total_len - max_length if max_length else 0 + + # Check sequence length and warn if needed + self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose) + + # Pad + if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask: + encoded_inputs = self.pad( + encoded_inputs, + max_length=max_length, + padding=padding_strategy.value, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + if return_length: + encoded_inputs["length"] = len(encoded_inputs["input_ids"]) + + return BatchEncoding(encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis) + + def truncate_sequences( # type: ignore[override] + self, + ids: list[int], + pair_ids: None = None, + num_tokens_to_remove: int = 0, + truncation_strategy: str | TruncationStrategy = "longest_first", + stride: int = 0, + **kwargs, + ) -> tuple[list[int], None, list[int]]: + """ + Truncates a sequence pair in-place following the strategy. + + Args: + ids (`list[int]`): + Tokenized input ids. Can be obtained from a string by chaining the `tokenize` and + `convert_tokens_to_ids` methods. + pair_ids (`None`, *optional*): + Not supported by `MistralCommonBackend`. Kept to match the signature of `PreTrainedTokenizerBase.truncate_sequences`. + num_tokens_to_remove (`int`, *optional*, defaults to 0): + Number of tokens to remove using the truncation strategy. + truncation_strategy (`str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `'longest_first'`): + The strategy to follow for truncation. Can be: + + - `'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. + - `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths greater + than the model maximum admissible input size). + stride (`int`, *optional*, defaults to 0): + If set to a positive number, the overflowing tokens returned will contain some tokens from the main + sequence returned. The value of this argument defines the number of additional tokens. + + Returns: + `Tuple[list[int], None, list[int]]`: The truncated `ids` and the list of + overflowing tokens. `None` is returned to match Transformers signature. + """ + + if pair_ids: + raise ValueError("`pair_ids` is not supported by `MistralCommonBackend.truncate_sequences`.") + + if not isinstance(truncation_strategy, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation_strategy) + + if truncation_strategy in [ + TruncationStrategy.ONLY_FIRST, + TruncationStrategy.ONLY_SECOND, + ]: + raise ValueError(f"{truncation_strategy=} is not supported by `MistralCommonBackend`.") + + if num_tokens_to_remove <= 0: + return ids, None, [] + + overflowing_tokens = [] + + if truncation_strategy == TruncationStrategy.LONGEST_FIRST: + window_len = min(len(ids), stride + num_tokens_to_remove) + if self.truncation_side == "left": + overflowing_tokens = ids[:window_len] + ids = ids[num_tokens_to_remove:] + else: + overflowing_tokens = ids[-window_len:] + ids = ids[:-num_tokens_to_remove] + + return ids, None, overflowing_tokens + + def apply_chat_template( # type: ignore[override] + self, + conversation: list[dict[str, str]] | list[list[dict[str, str]]], + tools: list[dict | Callable] | None = None, + add_generation_prompt: bool = False, + continue_final_message: bool = False, + tokenize: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool = False, + max_length: int | None = None, + return_tensors: str | TensorType | None = None, + return_dict: bool = True, + reasoning_effort: ReasoningEffort | None = None, + **kwargs, + ) -> str | list[int] | list[str] | list[list[int]] | BatchEncoding: + """ + Converts a list of dictionaries with `"role"` and `"content"` keys to a list of token + ids. + + Args: + conversation (Union[List[Dict[str, str]], List[List[Dict[str, str]]]]): A list of dicts + with "role" and "content" keys, representing the chat history so far. + tools (`List[Union[Dict, Callable]]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + add_generation_prompt (`bool`, *optional*): + This argument is a no-op for `MistralCommonBackend`. However, it cannot be used at the same time as `continue_final_message` to keep the API consistent. + If any conversation ends with an assistant message, it will raise an error. In such cases, use `continue_final_message` instead. + continue_final_message (bool, *optional*): + If this is set, the chat will be formatted so that the final + message in the chat is open-ended, without any EOS tokens. The model will continue this message + rather than starting a new one. This allows you to "prefill" part of + the model's response for it. Cannot be used at the same time as `add_generation_prompt`. + tokenize (`bool`, defaults to `True`): + Whether to tokenize the output. If `False`, the output will be a string. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, defaults to `False`): + Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`. + max_length (`int`, *optional*): + Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If + not specified, the tokenizer's `max_length` attribute will be used as a default. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable + values are: + - `'pt'`: Return PyTorch `torch.Tensor` objects. + return_dict (`bool`, defaults to `False`): + Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`. + If at least one conversation contains an image, its pixel values will be returned in the `pixel_values` key and image sizes in the `image_sizes` key. + reasoning_effort (`ReasoningEffort`, *optional*): + The reasoning effort to use for the chat completion for models that support it. Possible values are: + - `ReasoningEffort.none`: The model will not reason. + - `ReasoningEffort.high`: The model will use a reasoning approach. + If not specified, the default reasoning effort will be used. + + kwargs (additional keyword arguments, *optional*): + Not supported by `MistralCommonBackend.apply_chat_template`. + Will raise an error if used. + + Returns: + `Union[str, list[int], list[str], list[list[int]], BatchEncoding]`: The tokenized chat so far, including control tokens. This output is ready to pass to the model, either directly or via methods like `generate()`. + """ + if kwargs: + raise ValueError( + f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.apply_chat_template`." + ) + if not isinstance(truncation, bool): + raise TypeError("`truncation` must be a boolean for `apply_chat_template` method.") + + if add_generation_prompt and continue_final_message: + raise ValueError("Cannot use both `add_generation_prompt` and `continue_final_message`.") + + if isinstance(conversation, (list, tuple)) and ( + isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "messages") + ): + conversations = conversation + is_batched = True + else: + conversations = [conversation] + is_batched = False + + if add_generation_prompt: + for conversation in conversations: + last_message = conversation[-1] + if last_message.get("role") == "assistant": + raise ValueError( + "The last message in the conversation is already an assistant message. Consider using `continue_final_message` instead." + ) + + def _maybe_adapt_message(message: dict[str, Any]) -> None: + """Adapt message to `mistral-common` format and leave validation to `mistral-common`.""" + if not isinstance(message, dict): + return message + maybe_list_content: str | list[dict[str, str | dict[str, Any]]] | None = message.get("content") + if not maybe_list_content or isinstance(maybe_list_content, str): + return message + + normalized_content: list[dict[str, str | dict[str, Any]]] = [] + message = message.copy() + for content in maybe_list_content: + content_type = content.get("type", None) + if not content_type: + continue + elif content_type == "image": + maybe_url: str | None = content.get("url") + maybe_path: str | None = content.get("path") + maybe_base64: str | None = content.get("base64") + if maybe_url: + image_content = maybe_url + elif maybe_path: + if not maybe_path.startswith("file://"): + maybe_path = Path(maybe_path).resolve().as_uri() + image_content = maybe_path + elif maybe_base64: + if not maybe_base64.startswith("data:image"): + maybe_base64 = "data:image/unk;base64," + maybe_base64 + image_content = maybe_base64 + else: + raise ValueError("Image content must be specified.") + normalized_content.append({"type": "image_url", "image_url": {"url": image_content}}) + elif content_type == "audio": + maybe_url: str | None = content.get("url") + maybe_path: str | None = content.get("path") + maybe_base64: str | None = content.get("base64") + if maybe_url or maybe_path: + audio_data = load_audio_as(maybe_url or maybe_path, return_format="dict", force_mono=True) + normalized_content.append({"type": "input_audio", "input_audio": audio_data}) + continue + if not maybe_base64: + raise ValueError("Audio content must be specified.") + normalized_content.append({"type": "audio_url", "audio_url": {"url": maybe_base64}}) + else: + normalized_content.append(content) + message["content"] = normalized_content + return message + + outputs = [] + images: list[np.ndarray] = [] + audios: list[np.ndarray] = [] + + for conversation in conversations: + messages: list[dict[str, str | list[dict[str, str | dict[str, Any]]]]] = [] + for message in conversation: + message = _maybe_adapt_message(message) + messages.append(message) + + chat_request = ChatCompletionRequest.from_openai( + messages=messages, + tools=tools, + continue_final_message=continue_final_message, + reasoning_effort=reasoning_effort, + ) + + tokenized_request = self.tokenizer.encode_chat_completion(chat_request) + if tokenize: + outputs.append(tokenized_request.tokens) + else: + outputs.append(tokenized_request.text) + images.extend(tokenized_request.images) + audios.extend([el.audio_array for el in tokenized_request.audios]) + + if not is_batched: + outputs = outputs[0] + + if tokenize: + out = self( + outputs, + padding=padding, + truncation=truncation, + max_length=max_length, + add_special_tokens=False, + return_tensors=return_tensors, + ) + if return_dict: + if images: + pixel_values: list[np.ndarray] | np.ndarray | torch.Tensor + if return_tensors == "pt": + if not is_torch_available(): + raise ImportError( + "Unable to convert output to PyTorch tensors format, PyTorch is not installed." + ) + + pixel_values = torch.from_numpy(np.stack(images)) + elif return_tensors == "np": + pixel_values = np.array(images) + elif return_tensors is None: + pixel_values = images + else: + raise ValueError(f"Unsupported return_tensors type: {return_tensors}") + out.data["pixel_values"] = pixel_values + if images: + out.data["image_sizes"] = self._get_image_sizes_for_tensor(images, return_tensors) + if audios: + if return_tensors is not None: + raise NotImplementedError( + "When passing audio content in apply_chat_template, `return_tensors` must be None since we cannot batch the audio inputs. The returned audio will be a list of numpy arrays." + ) + # Transformers convention is audio for plural audio (audio does not take a "s") + out.data["audio"] = audios + return out + else: + return out["input_ids"] + + else: + logger.warning( + "`MistralCommonBackend.apply_chat_template(..., tokenize=False)` is unsafe and may lead to unexpected behavior." + " Please consider using `tokenize=True` instead and don't encode the output manually." + ) + return outputs + + def _get_image_sizes_for_tensor( + self, images: list[np.ndarray], return_tensors: str | TensorType | None + ) -> "list[list[int]] | np.ndarray | torch.Tensor": + """ + Convert image sizes to the appropriate format based on return_tensors. + + Args: + images: List of image arrays + return_tensors: The tensor type to return + + Returns: + Image sizes in the appropriate format + """ + image_sizes = [] + for image in images: + height, width = get_image_size(image) + image_sizes.append([height, width]) + + if return_tensors == "pt": + return torch.tensor(image_sizes, dtype=torch.long) + elif return_tensors == "np": + return np.array(image_sizes, dtype=np.int64) + else: + return image_sizes + + def build_inputs_with_special_tokens(self, token_ids_0: list[int], token_ids_1: None = None) -> list[int]: + """ + Build model inputs from a sequence by adding special tokens. + + This method dynamically builds inputs based on the tokenizer's `mode`: + - `"test"`: seq0 [EOS] + - `"finetuning"`: [BOS] seq0 + + Args: + token_ids_0 (`list[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`None`, *optional*): None, kept to match Transformers' signature. + + Returns: + `list[int]`: List of input IDs with the appropriate special tokens. + """ + if token_ids_1 is not None: + raise ValueError( + "`MistralCommonBackend` does not implement `token_ids_1 != None` for `build_inputs_with_special_tokens`." + ) + + if self.mode == ValidationMode.test: + # [BOS] seq0 + return [self.bos_token_id] + token_ids_0 + + else: + # [BOS] seq0 [EOS] + return [self.bos_token_id] + token_ids_0 + [self.eos_token_id] + + def create_token_type_ids_from_sequences(self, token_ids_0: list[int], token_ids_1: None = None) -> list[int]: + """ + Create a mask of zeroes from the token ids with special tokens added. + + Kept to match Transformers' implementation. + + Args: + token_ids_0 (`list[int]`): + List of IDs. + token_ids_1 (`None`, *optional*): None, kept to match Transformers' signature. + + + Returns: + `list[int]`: Token type IDs according to the configured pattern. + """ + if token_ids_1 is not None: + raise ValueError( + "`MistralCommonBackend` does not implement `token_ids_1 != None` for `create_token_type_ids_from_sequences`." + ) + + sequence = self.build_inputs_with_special_tokens(token_ids_0) + + return [0] * len(sequence) + + def num_special_tokens_to_add(self, pair: Literal[False] = False) -> int: + """ + Returns the number of added tokens when encoding a sequence with special tokens. + + + + This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put + this inside your training loop. + + + + Args: + pair (`Literal[False]`, *optional*): False, kept to match Transformer's signature. + + Returns: + `int`: Number of special tokens added to sequences. + """ + if pair: + raise ValueError( + "`MistralCommonBackend` does not implement `pair = True` for `num_special_tokens_to_add`." + ) + + return len(self.build_inputs_with_special_tokens([], None)) + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | EncodedInput | list[TextInput] | list[EncodedInput] | None = None, + text_pair: None = None, + text_target: None = None, + text_pair_target: None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + return_attention_mask: bool | None = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_length: bool = False, + verbose: bool = True, + return_offsets_mapping: Literal[False] = False, + split_special_tokens: Literal[False] = False, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences. + + Args: + text (`str`, `list[str]`, `list[list[str]]`, *optional*): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of int + (encoded strings). + text_pair (`None`, *optional*): + Not supported by `MistralCommonBackend`. Kept to match the signature of `PreTrainedTokenizerBase.__call__`. + text_target (`None`, *optional*): + Not supported by `MistralCommonBackend`. Kept to match the signature of `PreTrainedTokenizerBase.__call__`. + text_pair_target (`None`, *optional*): + Not supported by `MistralCommonBackend`. Kept to match the signature of `PreTrainedTokenizerBase.__call__`. + """ + if return_offsets_mapping or split_special_tokens: + raise ValueError( + "`MistralCommonBackend` does not support `return_offsets_mapping` and `split_special_tokens`." + ) + + if truncation in [TruncationStrategy.ONLY_FIRST, TruncationStrategy.ONLY_SECOND, "only_first", "only_second"]: + raise ValueError( + "Truncation strategy `only_first` and `only_second` are not supported by `MistralCommonBackend`." + ) + + if kwargs: + raise ValueError(f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.__call__`.") + + if text_pair or text_target or text_pair_target: + raise ValueError( + "`text_pair`, `text_target` and `text_pair_target` are not supported by `MistralCommonBackend`." + ) + + return super().__call__( + text=text, + text_pair=text_pair, + text_target=text_target, + 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_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | os.PathLike, + *init_inputs, + mode: str | ValidationMode = ValidationMode.test, + cache_dir: str | os.PathLike | None = None, + force_download: bool = False, + local_files_only: bool = False, + token: str | bool | None = None, + revision: str = "main", + model_max_length: int = VERY_LARGE_INTEGER, + padding_side: str = "left", + truncation_side: str = "right", + model_input_names: list[str] | None = None, + clean_up_tokenization_spaces: bool = False, + **kwargs, + ): + r""" + Instantiate a `MistralCommonBackend` from a predefined + tokenizer. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + + - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co. + - A path to a *directory* containing the tokenizer config, for instance saved + using the [`MistralCommonBackend.tokenization_mistral_common.save_pretrained`] method, e.g., + `./my_model_directory/`. + mode (`Union[str, ValidationMode]`, *optional*, defaults to `ValidationMode.test`): + Validation mode for the `MistralTokenizer` tokenizer. Possible values are: + - `"finetuning"` or `ValidationMode.finetuning`: The fine-tuning mode. + - `"test"` or `ValidationMode.test`: The test mode. + It changes how the tokenizer validates the input and prepares the request to the model. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded predefined tokenizer vocabulary files should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download the vocabulary files and override the cached versions if they + exist. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `hf auth login` (stored in `~/.huggingface`). + local_files_only (`bool`, *optional*, defaults to `False`): + Whether or not to only rely on local files and not to attempt to download any files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + 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. + padding_side (`str`, *optional*, defaults to `"left"`): + 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. + truncation_side (`str`, *optional*, defaults to `"right"`): + The side on which the model should have truncation applied. Should be selected between ['right', 'left']. + model_input_names (`List[str]`, *optional*): + The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or + `"attention_mask"`). Default value is picked from the class attribute of the same name. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not the model should clean up the spaces that were added when splitting the input text during the + tokenization process. + kwargs (additional keyword arguments, *optional*): + Not supported by `MistralCommonBackend.from_pretrained`. + Will raise an error if used. + """ + if init_inputs: + raise ValueError("`init_inputs` are not supported by `MistralCommonBackend.from_pretrained`.") + + # Handle kwargs and AutoTokenizer/AutoProcessor case + valid_kwargs = _VALID_INIT_KWARGS.union( + {"trust_remote_code", "_from_pipeline", "_commit_hash", "dtype", "subfolder"} + ) + if kwargs and not set(kwargs.keys()).issubset(valid_kwargs): + raise ValueError( + f"Some kwargs in {list(kwargs.keys())} are not supported by `MistralCommonBackend.from_pretrained`." + ) + + mode = cls._get_validation_mode(mode) + + if not os.path.isdir(pretrained_model_name_or_path): + tokenizer_path = download_tokenizer_from_hf_hub( + repo_id=pretrained_model_name_or_path, + cache_dir=cache_dir, + token=token, + revision=revision, + force_download=force_download, + local_files_only=local_files_only, + ) + else: + candidate_files = os.listdir(pretrained_model_name_or_path) + tokenizer_path = os.path.join(pretrained_model_name_or_path, get_one_valid_tokenizer_file(candidate_files)) + + return cls( + tokenizer_path=tokenizer_path, + mode=mode, + model_max_length=model_max_length, + padding_side=padding_side, + truncation_side=truncation_side, + model_input_names=model_input_names, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + ) + + def save_pretrained( # type: ignore[override] + self, + save_directory: str | os.PathLike | Path, + push_to_hub: bool = False, + token: str | bool | None = None, + commit_message: str | None = None, + repo_id: str | None = None, + private: bool | None = None, + **kwargs, + ) -> tuple[str, ...]: + """ + Save the full tokenizer state. + + + This method make sure the full tokenizer can then be re-loaded using the + [`~MistralCommonBackend.tokenization_mistral_common.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): The path to a directory where the tokenizer will be saved. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + token (`str` or *bool*, *optional*, defaults to `None`): + The token to use to push to the model hub. If `True`, will use the token in the `HF_TOKEN` environment + variable. + commit_message (`str`, *optional*): The commit message to use when pushing to the hub. + repo_id (`str`, *optional*): The name of the repository to which push to the Hub. + private (`bool`, *optional*): Whether the model repository is private or not. + kwargs (`Dict[str, Any]`, *optional*): + Not supported by `MistralCommonBackend.save_pretrained`. + Will raise an error if used. + + Returns: + A tuple of `str`: The files saved. + """ + if kwargs: + raise ValueError( + f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.save_pretrained`." + ) + + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + + shutil.copy(self._tokenizer_path, save_directory) + + if push_to_hub: + repo_id = repo_id or str(save_directory).split(os.path.sep)[-1] + repo_id = create_repo(repo_id, token=token, private=private, exist_ok=True).repo_id + files_timestamps = self._get_files_timestamps(save_directory) + + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + ) + + return (str(save_directory / self._tokenizer_path.name),) + + @staticmethod + def _get_validation_mode(mode: str | ValidationMode) -> ValidationMode: + """Get the validation mode from a string or a ValidationMode.""" + _invalid_mode_msg = ( + f"Invalid `mistral-common` tokenizer mode: {mode}. Possible values are 'finetuning' or 'test'." + ) + if isinstance(mode, str): + try: + mode = ValidationMode[mode] + except KeyError: + raise ValueError(_invalid_mode_msg) + elif not isinstance(mode, (str, ValidationMode)): + raise ValueError(_invalid_mode_msg) + + if mode not in [ValidationMode.finetuning, ValidationMode.test]: + raise ValueError(_invalid_mode_msg) + return mode + + def __repr__(self) -> str: + # MistralCommonBackend does not implement added_tokens_decoder, so we need a custom repr + return ( + f"{self.__class__.__name__}(name_or_path='{self.name_or_path}'," + f" vocab_size={self.vocab_size}, model_max_length={self.model_max_length}," + f" padding_side='{self.padding_side}', truncation_side='{self.truncation_side}'," + f" special_tokens={self.special_tokens_map})" + ) + + def added_tokens_decoder(self): + raise NotImplementedError("`MistralCommonBackend` does not implement `added_tokens_decoder`.") + + def add_special_tokens( + self, + special_tokens_dict: dict[str, str | AddedToken | Sequence[str | AddedToken]], + replace_extra_special_tokens: bool = True, + ): + r"""`MistralCommonBackend` does not implement `add_special_tokens` by design. + + If you would like this behaviour to be implemented, please open an issue in the `Transformers` or `mistral-common` repositories to request it. + """ + + raise NotImplementedError("`MistralCommonBackend` does not implement `add_special_tokens`.") + + def add_tokens( # type: ignore[override] + self, + special_tokens_dict: dict[str, str | AddedToken | Sequence[str | AddedToken]], + replace_extra_special_tokens: bool = True, + ): + """ + `MistralCommonBackend` does not implement `add_special_tokens` by design. + + If you would like this behaviour to be implemented, please open an issue in the `Transformers` or `mistral-common` repositories to request it. + """ + + raise NotImplementedError("`MistralCommonBackend` does not implement `add_tokens`.") + + def convert_added_tokens(cls, obj: AddedToken | Any, save: bool = False, add_type_field: bool = True): # type: ignore[override] + """ + `MistralCommonBackend` does not implement `convert_added_tokens` by design. + + If you would like this behaviour to be implemented, please open an issue in the `Transformers` or `mistral-common` repositories to request it. + """ + + raise NotImplementedError("`MistralCommonBackend` does not implement `convert_added_tokens`.") + + def get_chat_template(self, chat_template: str | None = None, tools: list[dict] | None = None) -> str: + """`MistralCommonBackend` does not implement `get_chat_template` by design as `mistral-common` does not use chat templates.""" + + raise NotImplementedError("`MistralCommonBackend` does not implement `get_chat_template`.") + + def save_chat_templates( + self, + save_directory: str | os.PathLike, + tokenizer_config: dict, + filename_prefix: str | None, + save_jinja_files: bool, + ): + """`MistralCommonBackend` does not implement `save_chat_templates` by design as `mistral-common` does not use chat templates.""" + + raise NotImplementedError("`MistralCommonBackend` does not implement `save_chat_templates`.") + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str, ...]: + """ + `MistralCommonBackend` does not implement `save_vocabulary` by design. + + This is because `mistral-common` is configured by one tokenizer file. If you'd like to save the vocabulary, please consider using the `save_pretrained` method instead. + """ + + raise NotImplementedError("`MistralCommonBackend` does not implement `save_vocabulary`.") + + +# Backward compatibility alias for codebases still importing the legacy name. +MistralCommonTokenizer = MistralCommonBackend diff --git a/third_party/transformers/src/transformers/tokenization_utils_base.py b/third_party/transformers/src/transformers/tokenization_utils_base.py new file mode 100644 index 0000000000000000000000000000000000000000..e61b20c3a4cf7250ccbd6e96370e1f77d331cec2 --- /dev/null +++ b/third_party/transformers/src/transformers/tokenization_utils_base.py @@ -0,0 +1,3557 @@ +# base +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Base classes common to both the slow and the fast tokenization classes: PreTrainedTokenizerBase (host all the user +fronting encoding methods) Special token mixing (host the special tokens logic) and BatchEncoding (wrap the dictionary +of output with special method for the Fast tokenizers) +""" + +from __future__ import annotations + +import copy +import json +import os +import re +import warnings +from collections import OrderedDict, UserDict +from collections.abc import Callable, Collection, Mapping, Sequence, Sized +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, NamedTuple, Union + +import numpy as np +from huggingface_hub import create_repo, is_offline_mode, list_repo_files +from packaging import version + +from . import __version__ +from .dynamic_module_utils import custom_object_save +from .utils import ( + CHAT_TEMPLATE_DIR, + CHAT_TEMPLATE_FILE, + ExplicitEnum, + PaddingStrategy, + PushToHubMixin, + TensorType, + add_end_docstrings, + cached_file, + copy_func, + extract_commit_hash, + is_mlx_available, + is_numpy_array, + is_protobuf_available, + is_tokenizers_available, + is_torch_available, + is_torch_device, + is_torch_tensor, + list_repo_templates, + logging, + requires_backends, + to_py_obj, +) +from .utils.chat_parsing_utils import recursive_parse +from .utils.chat_template_utils import render_jinja_template +from .utils.import_utils import PROTOBUF_IMPORT_ERROR + + +if TYPE_CHECKING: + if is_torch_available(): + import torch + + +def import_protobuf_decode_error(error_message=""): + if is_protobuf_available(): + from google.protobuf.message import DecodeError + + return DecodeError + else: + raise ImportError(PROTOBUF_IMPORT_ERROR.format(error_message)) + + +def flatten(arr: list): + res = [] + if len(arr) > 0: + for sub_arr in arr: + if isinstance(arr[0], (list, tuple)): + res.extend(flatten(sub_arr)) + else: + res.append(sub_arr) + return res + + +if is_tokenizers_available() or TYPE_CHECKING: + from tokenizers import Encoding as EncodingFast + +if is_tokenizers_available(): + from tokenizers import AddedToken +else: + + @dataclass(frozen=False, eq=True) + class AddedToken: + """ + AddedToken represents a token to be added to a Tokenizer An AddedToken can have special options defining the + way it should behave. + + The `normalized` will default to `not special` if it is not specified, similarly to the definition in + `tokenizers`. + """ + + def __init__( + self, content: str, single_word=False, lstrip=False, rstrip=False, special=False, normalized=None + ): + self.content = content + self.single_word = single_word + self.lstrip = lstrip + self.rstrip = rstrip + self.special = special + self.normalized = normalized if normalized is not None else not special + + def __getstate__(self): + return self.__dict__ + + def __str__(self): + return self.content + + +logger = logging.get_logger(__name__) + +VERY_LARGE_INTEGER = int(1e30) # This is used to set the max input length for a model with infinite size input +LARGE_INTEGER = int(1e20) # This is used when we need something big but slightly smaller than VERY_LARGE_INTEGER + +# Define type aliases and NamedTuples +TextInput = str +PreTokenizedInput = list[str] +EncodedInput = list[int] +TextInputPair = tuple[str, str] +PreTokenizedInputPair = tuple[list[str], list[str]] +EncodedInputPair = tuple[list[int], list[int]] + +# Define type aliases for text-related non-text modalities +AudioInput = Union[np.ndarray, "torch.Tensor", list[np.ndarray], list["torch.Tensor"]] + +# Slow tokenizers used to be saved in three separated files +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +ADDED_TOKENS_FILE = "added_tokens.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" + +# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file +FULL_TOKENIZER_FILE = "tokenizer.json" +_re_tokenizer_file = re.compile(r"tokenizer\.(.*)\.json") + + +class TruncationStrategy(ExplicitEnum): + """ + Possible values for the `truncation` argument in [`PreTrainedTokenizerBase.__call__`]. Useful for tab-completion in + an IDE. + """ + + ONLY_FIRST = "only_first" + ONLY_SECOND = "only_second" + LONGEST_FIRST = "longest_first" + DO_NOT_TRUNCATE = "do_not_truncate" + + +class CharSpan(NamedTuple): + """ + Character span in the original string. + + Args: + start (`int`): Index of the first character in the original string. + end (`int`): Index of the character following the last character in the original string. + """ + + start: int + end: int + + +class TokenSpan(NamedTuple): + """ + Token span in an encoded string (list of tokens). + + Args: + start (`int`): Index of the first token in the span. + end (`int`): Index of the token following the last token in the span. + """ + + start: int + end: int + + +class BatchEncoding(UserDict): + """ + Holds the output of the [`~tokenization_utils_base.PreTrainedTokenizerBase.__call__`], + [`~tokenization_utils_base.PreTrainedTokenizerBase.encode_plus`] and + [`~tokenization_utils_base.PreTrainedTokenizerBase.batch_encode_plus`] methods (tokens, attention_masks, etc). + + This class is derived from a python dictionary and can be used as a dictionary. In addition, this class exposes + utility methods to map from word/character space to token space. + + Args: + data (`dict`, *optional*): + Dictionary of lists/arrays/tensors returned by the `__call__`/`encode_plus`/`batch_encode_plus` methods + ('input_ids', 'attention_mask', etc.). + encoding (`tokenizers.Encoding` or `Sequence[tokenizers.Encoding]`, *optional*): + If the tokenizer is a fast tokenizer which outputs additional information like mapping from word/character + space to token space the `tokenizers.Encoding` instance or list of instance (for batches) hold this + information. + tensor_type (`Union[None, str, TensorType]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/Numpy Tensors at + initialization. + prepend_batch_axis (`bool`, *optional*, defaults to `False`): + Whether or not to add a batch axis when converting to tensors (see `tensor_type` above). Note that this + parameter has an effect if the parameter `tensor_type` is set, *otherwise has no effect*. + n_sequences (`Optional[int]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/Numpy Tensors at + initialization. + """ + + def __init__( + self, + data: dict[str, Any] | None = None, + encoding: EncodingFast | Sequence[EncodingFast] | None = None, + tensor_type: None | str | TensorType = None, + prepend_batch_axis: bool = False, + n_sequences: int | None = None, + ): + super().__init__(data) + + # If encoding is not None, the fast tokenization is used + if encoding is not None and isinstance(encoding, EncodingFast): + encoding = [encoding] + + self._encodings = encoding + + if n_sequences is None and encoding is not None and encoding: + n_sequences = encoding[0].n_sequences + + self._n_sequences = n_sequences + + self.convert_to_tensors(tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis) + + @property + def n_sequences(self) -> int | None: + """ + `Optional[int]`: The number of sequences used to generate each sample from the batch encoded in this + [`BatchEncoding`]. Currently can be one of `None` (unknown), `1` (a single sentence) or `2` (a pair of + sentences) + """ + return self._n_sequences + + def __getitem__(self, item: int | str) -> Any | EncodingFast: + """ + If the key is a string, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', + etc.). + + If the key is an integer, get the `tokenizers.Encoding` for batch item with index `key`. + + If the key is a slice, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', etc.) + with the constraint of slice. + """ + if isinstance(item, str): + return self.data[item] + elif self._encodings is not None: + return self._encodings[item] + elif isinstance(item, slice): + return {key: self.data[key][item] for key in self.data} + else: + raise KeyError( + "Invalid key. Only three types of key are available: " + "(1) string, (2) integers for backend Encoding, and (3) slices for data subsetting." + ) + + def __getattr__(self, item: str): + try: + return self.data[item] + except KeyError: + raise AttributeError + + def __getstate__(self): + return {"data": self.data, "encodings": self._encodings} + + def __setstate__(self, state): + if "data" in state: + self.data = state["data"] + + if "encodings" in state: + self._encodings = state["encodings"] + + # After this point: + # Extended properties and methods only available for fast (Rust-based) tokenizers + # provided by HuggingFace tokenizers library. + + @property + def is_fast(self) -> bool: + """ + TOOD: ita i will rm this `bool`: Whether or not this BatchEncoding was created by a fast tokenizer. + """ + return self._encodings is not None + + @property + def encodings(self) -> list[EncodingFast] | None: + """ + `Optional[list[tokenizers.Encoding]]`: The list all encodings from the tokenization process. Returns `None` if + the input was tokenized through Python (i.e., not a fast) tokenizer. + """ + return self._encodings + + def tokens(self, batch_index: int = 0) -> list[str]: + """ + Return the list of tokens (sub-parts of the input strings after word/subword splitting and before conversion to + integer indices) at a given batch index (only works for the output of a fast tokenizer). + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `list[str]`: The list of tokens at that index. + """ + if not self._encodings: + raise ValueError( + "tokens() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].tokens + + def sequence_ids(self, batch_index: int = 0) -> list[int | None]: + """ + Return a list mapping the tokens to the id of their original sentences: + + - `None` for special tokens added around or between sequences, + - `0` for tokens corresponding to words in the first sequence, + - `1` for tokens corresponding to words in the second sequence when a pair of sequences was jointly + encoded. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `list[Optional[int]]`: A list indicating the sequence id corresponding to each token. Special tokens added + by the tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding + sequence. + """ + if not self._encodings: + raise ValueError( + "sequence_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].sequence_ids + + def word_ids(self, batch_index: int = 0) -> list[int | None]: + """ + Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `list[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the + tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word + (several tokens will be mapped to the same word index if they are parts of that word). + """ + if not self._encodings: + raise ValueError( + "word_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].word_ids + + def token_to_sequence(self, batch_or_token_index: int, token_index: int | None = None) -> int: + """ + Get the index of the sequence represented by the given token. In the general use case, this method returns `0` + for a single sequence or the first sequence of a pair, and `1` for the second sequence of a pair + + Can be called as: + + - `self.token_to_sequence(token_index)` if batch size is 1 + - `self.token_to_sequence(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + + if not self._encodings: + raise ValueError("token_to_sequence() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_sequence(token_index) + + def token_to_word(self, batch_or_token_index: int, token_index: int | None = None) -> int: + """ + Get the index of the word corresponding (i.e. comprising) to an encoded token in a sequence of the batch. + + Can be called as: + + - `self.token_to_word(token_index)` if batch size is 1 + - `self.token_to_word(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + + if not self._encodings: + raise ValueError("token_to_word() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_word(token_index) + + def word_to_tokens( + self, batch_or_word_index: int, word_index: int | None = None, sequence_index: int = 0 + ) -> TokenSpan | None: + """ + Get the encoded token span corresponding to a word in a sequence of the batch. + + Token spans are returned as a [`~tokenization_utils_base.TokenSpan`] with: + + - **start** -- Index of the first token. + - **end** -- Index of the token following the last token. + + Can be called as: + + - `self.word_to_tokens(word_index, sequence_index: int = 0)` if batch size is 1 + - `self.word_to_tokens(batch_index, word_index, sequence_index: int = 0)` if batch size is greater or equal to + 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the word in the sequence. + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + ([`~tokenization_utils_base.TokenSpan`], *optional*): Span of tokens in the encoded sequence. Returns + `None` if no tokens correspond to the word. This can happen especially when the token is a special token + that has been used to format the tokenization. For example when we add a class token at the very beginning + of the tokenization. + """ + + if not self._encodings: + raise ValueError("word_to_tokens() is not available when using Python based tokenizers") + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if word_index < 0: + word_index = self._seq_len + word_index + span = self._encodings[batch_index].word_to_tokens(word_index, sequence_index) + return TokenSpan(*span) if span is not None else None + + def token_to_chars(self, batch_or_token_index: int, token_index: int | None = None) -> CharSpan | None: + """ + Get the character span corresponding to an encoded token in a sequence of the batch. + + Character spans are returned as a [`~tokenization_utils_base.CharSpan`] with: + + - **start** -- Index of the first character in the original string associated to the token. + - **end** -- Index of the character following the last character in the original string associated to the + token. + + Can be called as: + + - `self.token_to_chars(token_index)` if batch size is 1 + - `self.token_to_chars(batch_index, token_index)` if batch size is greater or equal to 1 + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token or tokens in + the sequence. + + Returns: + [`~tokenization_utils_base.CharSpan`]: Span of characters in the original string, or None, if the token + (e.g. , ) doesn't correspond to any chars in the origin string. + """ + + if not self._encodings: + raise ValueError("token_to_chars() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + span_indices = self._encodings[batch_index].token_to_chars(token_index) + + return CharSpan(*span_indices) if span_indices is not None else None + + def char_to_token(self, batch_or_char_index: int, char_index: int | None = None, sequence_index: int = 0) -> int: + """ + Get the index of the token in the encoded output comprising a character in the original string for a sequence + of the batch. + + Can be called as: + + - `self.char_to_token(char_index)` if batch size is 1 + - `self.char_to_token(batch_index, char_index)` if batch size is greater or equal to 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int`: Index of the token, or None if the char index refers to a whitespace only token and whitespace is + trimmed with `trim_offsets=True`. + """ + + if not self._encodings: + raise ValueError("char_to_token() is not available when using Python based tokenizers") + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_token(char_index, sequence_index) + + def word_to_chars( + self, batch_or_word_index: int, word_index: int | None = None, sequence_index: int = 0 + ) -> CharSpan: + """ + Get the character span in the original string corresponding to given word in a sequence of the batch. + + Character spans are returned as a CharSpan NamedTuple with: + + - start: index of the first character in the original string + - end: index of the character following the last character in the original string + + Can be called as: + + - `self.word_to_chars(word_index)` if batch size is 1 + - `self.word_to_chars(batch_index, word_index)` if batch size is greater or equal to 1 + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + `CharSpan` or `list[CharSpan]`: Span(s) of the associated character or characters in the string. CharSpan + are NamedTuple with: + + - start: index of the first character associated to the token in the original string + - end: index of the character following the last character associated to the token in the original + string + """ + + if not self._encodings: + raise ValueError("word_to_chars() is not available when using Python based tokenizers") + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + return CharSpan(*(self._encodings[batch_index].word_to_chars(word_index, sequence_index))) + + def char_to_word(self, batch_or_char_index: int, char_index: int | None = None, sequence_index: int = 0) -> int: + """ + Get the word in the original string corresponding to a character in the original string of a sequence of the + batch. + + Can be called as: + + - `self.char_to_word(char_index)` if batch size is 1 + - `self.char_to_word(batch_index, char_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the character in the original string. + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the character in the + original string. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int` or `list[int]`: Index or indices of the associated encoded token(s). + """ + + if not self._encodings: + raise ValueError("char_to_word() is not available when using Python based tokenizers") + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_word(char_index, sequence_index) + + def convert_to_tensors(self, tensor_type: str | TensorType | None = None, prepend_batch_axis: bool = False): + """ + Convert the inner content to tensors. + + Args: + tensor_type (`str` or [`~utils.TensorType`], *optional*): + The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If + `None`, no modification is done. + prepend_batch_axis (`int`, *optional*, defaults to `False`): + Whether or not to add the batch dimension during the conversion. + """ + if tensor_type is None: + return self + + # Convert to TensorType + if not isinstance(tensor_type, TensorType): + tensor_type = TensorType(tensor_type) + + if tensor_type == TensorType.PYTORCH: + if not is_torch_available(): + raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.") + import torch + + def as_tensor(value, dtype=None): + if isinstance(value, list) and len(value) > 0 and isinstance(value[0], np.ndarray): + return torch.from_numpy(np.array(value)) + if len(flatten(value)) == 0 and dtype is None: + dtype = torch.int64 + return torch.tensor(value, dtype=dtype) + + is_tensor = torch.is_tensor + + elif tensor_type == TensorType.MLX: + if not is_mlx_available(): + raise ImportError("Unable to convert output to MLX tensors format, MLX is not installed.") + import mlx.core as mx + + def as_tensor(value, dtype=None): + if len(flatten(value)) == 0 and dtype is None: + dtype = mx.int32 + return mx.array(value, dtype=dtype) + + def is_tensor(obj): + return isinstance(obj, mx.array) + else: + + def as_tensor(value, dtype=None): + if ( + isinstance(value, (list, tuple)) + and len(value) > 0 + and isinstance(value[0], (list, tuple, np.ndarray)) + ): + value_lens = [len(val) for val in value] + if len(set(value_lens)) > 1 and dtype is None: + # we have a ragged list so handle explicitly + value = as_tensor([np.asarray(val) for val in value], dtype=object) + if len(flatten(value)) == 0 and dtype is None: + dtype = np.int64 + return np.asarray(value, dtype=dtype) + + is_tensor = is_numpy_array + + # Do the tensor conversion in batch + for key, value in self.items(): + try: + if prepend_batch_axis: + value = [value] + + if not is_tensor(value): + tensor = as_tensor(value) + + # Removing this for now in favor of controlling the shape with `prepend_batch_axis` + # # at-least2d + # if tensor.ndim > 2: + # tensor = tensor.squeeze(0) + # elif tensor.ndim < 2: + # tensor = tensor[None, :] + + self[key] = tensor + except Exception as e: + if key == "overflowing_tokens": + raise ValueError( + "Unable to create tensor returning overflowing tokens of different lengths. " + "Please see if a fast version of this tokenizer is available to have this feature available." + ) from e + raise ValueError( + "Unable to create tensor, you should probably activate truncation and/or padding with" + " 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your" + f" features (`{key}` in this case) have excessive nesting (inputs type `list` where type `int` is" + " expected)." + ) from e + + return self + + def to(self, device: str | torch.device, *, non_blocking: bool = False) -> BatchEncoding: + """ + Send all values to device by calling `v.to(device, non_blocking=non_blocking)` (PyTorch only). + + Args: + device (`str` or `torch.device`): The device to put the tensors on. + non_blocking (`bool`): Whether to perform the copy asynchronously. + + Returns: + [`BatchEncoding`]: The same instance after modification. + """ + requires_backends(self, ["torch"]) + + # This check catches things like APEX blindly calling "to" on all inputs to a module + # Otherwise it passes the casts down and casts the LongTensor containing the token idxs + # into a HalfTensor + if isinstance(device, str) or is_torch_device(device) or isinstance(device, int): + self.data = { + k: v.to(device=device, non_blocking=non_blocking) if hasattr(v, "to") and callable(v.to) else v + for k, v in self.data.items() + } + else: + logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.") + return self + + +ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to add special tokens when encoding the sequences. This will use the underlying + `PretrainedTokenizerBase.build_inputs_with_special_tokens` function, which defines which tokens are + automatically added to the input ids. This is useful if you want to add `bos` or `eos` tokens + automatically. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence is 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. + is_split_into_words (`bool`, *optional*, defaults to `False`): + Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the + tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace) + which it will tokenize. This is useful for NER or token classification. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. Requires `padding` to be activated. + This is especially useful to enable the use of Tensor Cores 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_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. +""" + +ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) + + - **token_type_ids** -- List of token type ids to be fed to a model (when `return_token_type_ids=True` or + if *"token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) + + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) + + - **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`) +""" + + +INIT_TOKENIZER_DOCSTRING = r""" + Class attributes (overridden by derived classes) + + - **vocab_files_names** (`dict[str, str]`) -- A dictionary with, as keys, the `__init__` keyword name of each + vocabulary file required by the model, and as associated values, the filename for saving the associated file + (string). + - **pretrained_vocab_files_map** (`dict[str, dict[str, str]]`) -- A dictionary of dictionaries, with the + high-level keys being the `__init__` keyword name of each vocabulary file required by the model, the + low-level being the `short-cut-names` of the pretrained models with, as associated values, the `url` to the + associated pretrained vocabulary file. + - **model_input_names** (`list[str]`) -- A list of inputs expected in the forward pass of the model. + - **padding_side** (`str`) -- The default value for the side on which the model should have padding applied. + Should be `'right'` or `'left'`. + - **truncation_side** (`str`) -- The default value for the side on which the model should have truncation + applied. Should be `'right'` or `'left'`. + + Args: + model_max_length (`int`, *optional*): + The maximum length (in number of tokens) for the inputs to the transformer model. When the tokenizer is + loaded with [`~tokenization_utils_base.PreTrainedTokenizerBase.from_pretrained`], this will be set to the + value stored for the associated model in `max_model_input_sizes` (see above). If no value is provided, will + default to VERY_LARGE_INTEGER (`int(1e30)`). + 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. + truncation_side (`str`, *optional*): + The side on which the model should have truncation applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + chat_template (`str`, *optional*): + A Jinja template string that will be used to format lists of chat messages. See + https://huggingface.co/docs/transformers/chat_templating for a full description. + model_input_names (`list[string]`, *optional*): + The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or + `"attention_mask"`). Default value is picked from the class attribute of the same name. + bos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the beginning of a sentence. + eos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the end of a sentence. + unk_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing an out-of-vocabulary token. + sep_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token separating two different sentences in the same input (used by BERT for instance). + pad_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. + cls_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the class of the input (used by BERT for instance). + mask_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing a masked token (used by masked-language modeling pretraining objectives, like + BERT). Will be associated to `self.mask_token` and `self.mask_token_id`. + extra_special_tokens (list of `str` or `tokenizers.AddedToken`, *optional*): + A list of extra model-specific special tokens. Add them here to ensure they are skipped when decoding with + `skip_special_tokens` is set to True. If they are not part of the vocabulary, they will be added at the end + of the vocabulary. + split_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the special tokens should be split during the tokenization process. Passing will affect the + internal state of the tokenizer. The default behavior is to not split special tokens. This means that if + `` is the `bos_token`, then `tokenizer.tokenize("") = ['`]. Otherwise, if + `split_special_tokens=True`, then `tokenizer.tokenize("")` will be give `['<','s', '>']`. +""" + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizerBase(PushToHubMixin): + """ + Base class for all tokenizer backends. + """ + + vocab_files_names: dict[str, str] = {} + pretrained_vocab_files_map: dict[str, dict[str, str]] = {} + _auto_class: str | None = None + + # first name has to correspond to main model input name + # to make sure `tokenizer.pad(...)` works correctly + model_input_names: list[str] = ["input_ids", "attention_mask"] + padding_side: str = "right" + truncation_side: str = "right" + slow_tokenizer_class = None + + # Special tokens support (moved from SpecialTokensMixin) + # V5: Clean separation of named special tokens from extra special tokens + SPECIAL_TOKENS_ATTRIBUTES = [ + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", + ] + + def __init__(self, **kwargs): + self.init_inputs = () + for key in kwargs: + if hasattr(self, key) and callable(getattr(self, key)): + raise AttributeError(f"{key} conflicts with the method {key} in {self.__class__.__name__}") + + # V5: Convert deprecated additional_special_tokens to extra_special_tokens before storing init_kwargs + if "additional_special_tokens" in kwargs and "extra_special_tokens" not in kwargs: + kwargs["extra_special_tokens"] = kwargs.pop("additional_special_tokens") + + self.init_kwargs = copy.deepcopy(kwargs) + self.name_or_path = kwargs.pop("name_or_path", "") + self._processor_class = kwargs.pop("processor_class", None) + + self._pad_token_type_id = 0 + self.verbose = kwargs.pop("verbose", False) + + # V5: Separate storage for named special tokens and extra special tokens + self._special_tokens_map = dict.fromkeys(self.SPECIAL_TOKENS_ATTRIBUTES) + self._extra_special_tokens = [] # List of extra model-specific special tokens + + # V5: track both explicit and auto-detected model-specific tokens + explicit_model_specific_tokens = kwargs.pop("model_specific_special_tokens", None) + if explicit_model_specific_tokens is None: + explicit_model_specific_tokens = {} + elif not isinstance(explicit_model_specific_tokens, dict): + raise TypeError("model_specific_special_tokens must be a dictionary of token name to token value") + auto_model_specific_tokens = {} + + # Directly set hidden values to allow init with tokens not yet in vocab + for key in list(kwargs.keys()): + if key in self.SPECIAL_TOKENS_ATTRIBUTES: + value = kwargs.pop(key) + if value is None: + continue + if isinstance(value, (str, AddedToken)): + self._special_tokens_map[key] = value + else: + raise TypeError(f"Special token {key} has to be either str or AddedToken but got: {type(value)}") + elif key == "extra_special_tokens": + value = kwargs.pop(key) + if value is None: + continue + if isinstance(value, dict): + self._set_model_specific_special_tokens(special_tokens=value) + elif isinstance(value, (list, tuple)): + self._extra_special_tokens = list(value) + else: + raise TypeError("extra_special_tokens must be a list/tuple of tokens or a dict of named tokens") + elif ( + key.endswith("_token") + and key not in self.SPECIAL_TOKENS_ATTRIBUTES + and isinstance(kwargs[key], (str, AddedToken)) + ): + value = kwargs.pop(key) + if value is None: + continue + auto_model_specific_tokens[key] = value + + # For backward compatibility we fallback to set model_max_length from max_len if provided + model_max_length = kwargs.pop("model_max_length", kwargs.pop("max_len", None)) + self.model_max_length = model_max_length if model_max_length is not None else VERY_LARGE_INTEGER + + self.padding_side = kwargs.pop("padding_side", self.padding_side) + if self.padding_side not in ["right", "left"]: + raise ValueError( + f"Padding side should be selected between 'right' and 'left', current value: {self.padding_side}" + ) + + self.truncation_side = kwargs.pop("truncation_side", self.truncation_side) + if self.truncation_side not in ["right", "left"]: + raise ValueError( + f"Truncation side should be selected between 'right' and 'left', current value: {self.truncation_side}" + ) + + self.model_input_names = kwargs.pop("model_input_names", self.model_input_names) + + # By default, clean up tokenization spaces for both fast and slow tokenizers + self.clean_up_tokenization_spaces = kwargs.pop("clean_up_tokenization_spaces", False) + + # By default, do not split special tokens for both fast and slow tokenizers + self.split_special_tokens = kwargs.pop("split_special_tokens", False) + + self._in_target_context_manager = False + + self.chat_template = kwargs.pop("chat_template", None) + if isinstance(self.chat_template, (list, tuple)): + # Chat templates are stored as lists of dicts with fixed key names, + # we reconstruct that into a single dict while loading them. + self.chat_template = {template["name"]: template["template"] for template in self.chat_template} + + self.response_schema = kwargs.pop("response_schema", None) + + model_specific_tokens = {**auto_model_specific_tokens, **explicit_model_specific_tokens} + if model_specific_tokens: + self._set_model_specific_special_tokens(special_tokens=model_specific_tokens) + + self.deprecation_warnings = {} + + # Backend information (V5: tracking which backend and files were used) + self.backend = kwargs.pop("backend", None) + self.files_loaded = kwargs.pop("files_loaded", []) + + def _set_processor_class(self, processor_class: str): + """Sets processor class so it can be serialized in `tokenizer_config.json`.""" + self._processor_class = processor_class + + # ---- Special tokens API (moved from SpecialTokensMixin) ---- + def add_special_tokens( + self, + special_tokens_dict: dict[str, str | AddedToken | Sequence[str | AddedToken]], + replace_extra_special_tokens=True, + ) -> int: + """ + Add a dictionary of special tokens (eos, pad, cls, etc.) to the encoder and link them to class attributes. If + special tokens are NOT in the vocabulary, they are added to it (indexed starting from the last index of the + current vocabulary). + + When adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix of the + model so that its embedding matrix matches the tokenizer. + + In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method. + + Using `add_special_tokens` will ensure your special tokens can be used in several ways: + + - Special tokens can be skipped when decoding using `skip_special_tokens = True`. + - Special tokens are carefully handled by the tokenizer (they are never split), similar to `AddedTokens`. + - You can easily refer to special tokens using tokenizer class attributes like `tokenizer.cls_token`. This + makes it easy to develop model-agnostic training and fine-tuning scripts. + + When possible, special tokens are already registered for provided pretrained models (for instance + [`BertTokenizer`] `cls_token` is already registered to be `'[CLS]'` and XLM's one is also registered to be + `''`). + + Args: + special_tokens_dict (dictionary *str* to *str*, `tokenizers.AddedToken`, or `Sequence[Union[str, AddedToken]]`): + Keys should be in the list of predefined special attributes: [`bos_token`, `eos_token`, `unk_token`, + `sep_token`, `pad_token`, `cls_token`, `mask_token`, `extra_special_tokens`]. + + Tokens are only added if they are not already in the vocabulary (tested by checking if the tokenizer + assign the index of the `unk_token` to them). + replace_extra_special_tokens (`bool`, *optional*, defaults to `True`): + If `True`, the existing list of extra special tokens will be replaced by the list provided in + `special_tokens_dict`. Otherwise, `extra_special_tokens` will be extended. In the former + case, the tokens will NOT be removed from the tokenizer's full vocabulary - they are only being flagged + as non-special tokens. Remember, this only affects which tokens are skipped during decoding, not the + `added_tokens_encoder` and `added_tokens_decoder`. This means that the previous + `extra_special_tokens` are still added tokens, and will not be split by the model. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + + ```python + # Let's see how to add a new classification token to GPT-2 + tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2") + model = GPT2Model.from_pretrained("openai-community/gpt2") + + special_tokens_dict = {"cls_token": ""} + + num_added_toks = tokenizer.add_special_tokens(special_tokens_dict) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + + assert tokenizer.cls_token == "" + ```""" + if not special_tokens_dict: + return 0 + + # V5: Allowed keys are SPECIAL_TOKENS_ATTRIBUTES + "extra_special_tokens" + # Backward compatibility: convert "additional_special_tokens" to "extra_special_tokens" + special_tokens_dict = dict(special_tokens_dict) + if "additional_special_tokens" in special_tokens_dict: + special_tokens_dict.setdefault( + "extra_special_tokens", special_tokens_dict.pop("additional_special_tokens") + ) + + allowed_keys = set(self.SPECIAL_TOKENS_ATTRIBUTES) | {"extra_special_tokens"} + tokens_to_add = [] + for key, value in special_tokens_dict.items(): + if key not in allowed_keys: + raise ValueError(f"Key {key} is not a valid special token. Valid keys are: {allowed_keys}") + + if self.verbose: + logger.info(f"Assigning {value} to the {key} key of the tokenizer") + + if key == "extra_special_tokens": + if not isinstance(value, (list, tuple)) or not all(isinstance(t, (str, AddedToken)) for t in value): + raise ValueError(f"Tokens {value} for key {key} should all be str or AddedToken instances") + new_tokens = [ + ( + AddedToken(t, rstrip=False, lstrip=False, normalized=False, special=True) + if isinstance(t, str) + else t + ) + for t in value + if replace_extra_special_tokens or str(t) not in self.extra_special_tokens + ] + if replace_extra_special_tokens and new_tokens: + self._extra_special_tokens = list(new_tokens) + else: + self._extra_special_tokens.extend(new_tokens) + tokens_to_add.extend(new_tokens) + else: + if not isinstance(value, (str, AddedToken)): + raise ValueError(f"Token {value} for key {key} should be a str or an AddedToken instance") + if isinstance(value, str): + value = AddedToken(value, rstrip=False, lstrip=False, normalized=False, special=True) + setattr(self, key, value) + tokens_to_add.append(value) + + return self.add_tokens(tokens_to_add, special_tokens=True) + + def add_tokens( + self, new_tokens: str | AddedToken | Sequence[str | AddedToken], special_tokens: bool = False + ) -> int: + """ + #TODO remove this from here! PreTrainedTOkeniuzerBase should be agnostic of AddedToken. + + Add a list of new tokens. If the new tokens are not in the vocabulary, they are added to the end. Added tokens and + tokens from the vocabulary of the tokenization algorithm are therefore not treated in the same way. + + Args: + new_tokens (`str`, `tokenizers.AddedToken` or a sequence of *str* or `tokenizers.AddedToken`): + Tokens are only added if they are not already in the vocabulary. `tokenizers.AddedToken` wraps a string + token to let you personalize its behavior: whether this token should only match against a single word, + whether this token should strip all potential whitespaces on the left side, whether this token should + strip all potential whitespaces on the right side, etc. + special_tokens (`bool`, *optional*, defaults to `False`): + Specifies if the token is special. This mostly changes the normalization behavior + See details for `tokenizers.AddedToken` in HuggingFace tokenizers library. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + + ```python + # Let's see how to increase the vocabulary of Bert model and tokenizer + tokenizer = BertTokenizerFast.from_pretrained("google-bert/bert-base-uncased") + model = BertModel.from_pretrained("google-bert/bert-base-uncased") + + num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + ```""" + if not new_tokens: + return 0 + + if not isinstance(new_tokens, (list, tuple)): + new_tokens = [new_tokens] + return self._add_tokens(new_tokens, special_tokens=special_tokens) + + def _add_tokens(self, new_tokens: list[str] | list[AddedToken], special_tokens: bool = False) -> int: + raise NotImplementedError + + @property + def pad_token_type_id(self) -> int: + return self._pad_token_type_id + + def __setattr__(self, key, value): + # Handle _id/_ids suffix (eg. bos_token_id -> bos_token) + key_without_id = key.removesuffix("_ids").removesuffix("_id") if key.endswith(("_id", "_ids")) else key + + # Named special tokens (bos_token, eos_token, etc.) + if key_without_id in self.SPECIAL_TOKENS_ATTRIBUTES: + if key != key_without_id and value is not None: + value = self.convert_ids_to_tokens(value) + if value is not None and not isinstance(value, (str, AddedToken)): + raise ValueError(f"Cannot set a non-string value as the {key_without_id}") + self._special_tokens_map[key_without_id] = value + return + + # Extra special tokens: model-specific special tokens without standard names (eg. ) + if key_without_id == "extra_special_tokens": + if key != key_without_id and value is not None and isinstance(value, (list, tuple)): + value = [self.convert_ids_to_tokens(v) for v in value] + if not isinstance(value, (list, tuple)) and value is not None: + raise ValueError(f"extra_special_tokens must be a list or tuple, got {type(value)}") + self._extra_special_tokens = [] if value is None else list(value) + return + + super().__setattr__(key, value) + + def __getattr__(self, key): + # Handle _id/_ids suffix (eg. bos_token_id -> bos_token) + key_without_id = key.removesuffix("_ids").removesuffix("_id") if key.endswith(("_id", "_ids")) else key + + # Named special tokens (bos_token, eos_token, etc.) + if key_without_id in self.SPECIAL_TOKENS_ATTRIBUTES: + token_value = self._special_tokens_map.get(key_without_id) + if token_value is None: + if self.verbose: + logger.error(f"Using {key}, but it is not set yet.") + return None + return self.convert_tokens_to_ids(str(token_value)) if key != key_without_id else str(token_value) + + # Extra special tokens + if key_without_id == "extra_special_tokens": + tokens = [str(tok) for tok in self._extra_special_tokens] + return self.convert_tokens_to_ids(tokens) if key != key_without_id else tokens + + if key not in self.__dict__: + raise AttributeError(f"{self.__class__.__name__} has no attribute {key}") + return super().__getattr__(key) + + 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. + + For fast tokenizers, data collators call this with `already_has_special_tokens=True` to build a mask over an + already-formatted sequence. In that case, we compute the mask by checking membership in `all_special_ids`. + + Args: + token_ids_0: List of IDs for the (possibly already formatted) sequence. + token_ids_1: Unused when `already_has_special_tokens=True`. Must be None in that case. + already_has_special_tokens: Whether the sequence is already formatted with special tokens. + + Returns: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + if token_ids_1 is not None: + raise ValueError( + "You should not supply a second sequence if the provided sequence of ids is already formatted " + "with special tokens for the model." + ) + special_ids = set(self.all_special_ids) + return [1 if int(tid) in special_ids else 0 for tid in token_ids_0] + + # Default base implementation for non-formatted sequences is not provided here. + # Concrete tokenizer classes should override this for their specific formatting rules. + raise NotImplementedError( + f"{self.__class__.__name__} does not implement get_special_tokens_mask for non-formatted sequences" + ) + + @property + def special_tokens_map(self) -> dict[str, str]: + """ + `dict[str, str]`: A flat dictionary mapping named special token attributes to their string values. + + Only includes the standard named special tokens (bos_token, eos_token, etc.), not extra_special_tokens. + This provides a clean, flat structure without mixed types. + + Returns: + A dictionary with keys like 'bos_token', 'eos_token', etc., and string values. + + **V5 Change**: This now returns only named tokens. Use `extra_special_tokens` for the additional tokens. + """ + return { + attr: str(self._special_tokens_map[attr]) + for attr in self.SPECIAL_TOKENS_ATTRIBUTES + if self._special_tokens_map.get(attr) is not None + } + + # Note: extra_special_tokens and extra_special_tokens_ids are handled by __getattr__ and __setattr__ + # We don't define them as @property to keep the implementation simpler + + @property + def all_special_tokens(self) -> list[str]: + """ + `list[str]`: A list of all unique special tokens (named + extra) as strings. + + Includes both named special tokens (bos_token, eos_token, etc.) and extra special tokens. + Converts tokens of `tokenizers.AddedToken` type to string. + """ + seen = set() + all_toks = [] + + # Add named special tokens + for attr in self.SPECIAL_TOKENS_ATTRIBUTES: + value = self._special_tokens_map.get(attr) + if value is not None: + token_str = str(value) + if token_str not in seen: + all_toks.append(token_str) + seen.add(token_str) + + # Add extra special tokens + for token in self._extra_special_tokens: + token_str = str(token) + if token_str not in seen: + all_toks.append(token_str) + seen.add(token_str) + + return all_toks + + @property + def all_special_ids(self) -> list[int]: + """ + `list[int]`: List the ids of the special tokens(`''`, `''`, etc.) mapped to class attributes. + """ + return self.convert_tokens_to_ids(self.all_special_tokens) + + def _set_model_specific_special_tokens(self, special_tokens: dict[str, str | AddedToken]): + """ + Adds new model-specific special tokens (e.g., for multimodal models). + + These tokens are added to the named special tokens map and will be saved in tokenizer config. + For example: if the model tokenizer is multimodal, we can support special image or audio tokens. + + Args: + special_tokens: Dictionary of {token_name: token_value} + """ + self.SPECIAL_TOKENS_ATTRIBUTES = self.SPECIAL_TOKENS_ATTRIBUTES + list(special_tokens.keys()) + for key, value in special_tokens.items(): + if isinstance(value, (str, AddedToken)): + self._special_tokens_map[key] = value + else: + raise TypeError(f"Special token {key} has to be either str or AddedToken but got: {type(value)}") + + @property + def added_tokens_decoder(self) -> dict[int, AddedToken]: + raise NotImplementedError() + + def __repr__(self) -> str: + added_tokens_decoder_rep = "\n\t".join([f"{k}: {v.__repr__()}," for k, v in self.added_tokens_decoder.items()]) + if added_tokens_decoder_rep: + added_tokens_decoder_rep = f"\n\t{added_tokens_decoder_rep}\n" + return ( + f"{self.__class__.__name__}(name_or_path='{self.name_or_path}'," + f" vocab_size={self.vocab_size}, model_max_length={self.model_max_length}," + f" padding_side='{self.padding_side}', truncation_side='{self.truncation_side}'," + f" special_tokens={self.special_tokens_map}," + f" added_tokens_decoder={{{added_tokens_decoder_rep}}})" + ) + + def __len__(self) -> int: + raise NotImplementedError() + + @property + def vocab_size(self) -> int: + """ + `int`: Size of the base vocabulary (without the added tokens). + """ + raise NotImplementedError() + + def get_vocab(self) -> dict[str, int]: + """ + Returns the vocabulary as a dictionary of token to index. + + `tokenizer.get_vocab()[token]` is equivalent to `tokenizer.convert_tokens_to_ids(token)` when `token` is in the + vocab. + + Returns: + `dict[str, int]`: The vocabulary. + """ + raise NotImplementedError() + + def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: + """ + Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids), using the + vocabulary. + + Args: + tokens (`str` or `list[str]`): One or several token(s) to convert to token id(s). + + Returns: + `int` or `list[int]`: The token id or list of token ids. + """ + if isinstance(tokens, str): + return self._convert_token_to_id_with_added_voc(tokens) + + return [self._convert_token_to_id_with_added_voc(token) for token in tokens] + + def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `list[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `list[str]`: The decoded token(s). + """ + raise NotImplementedError() + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | os.PathLike, + *init_inputs, + cache_dir: str | os.PathLike | None = None, + force_download: bool = False, + local_files_only: bool = False, + token: str | bool | None = None, + revision: str = "main", + trust_remote_code=False, + **kwargs, + ): + r""" + Instantiate a [`~tokenization_utils_base.PreTrainedTokenizerBase`] (or a derived class) from a predefined + tokenizer. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + + - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co. + - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved + using the [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`] method, e.g., + `./my_model_directory/`. + - (**Deprecated**, not applicable to all derived classes) a path to a single saved vocabulary + file (if and only if the tokenizer only requires a single vocabulary file like Bert or XLNet), e.g., + `./my_model_directory/vocab.txt`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded predefined tokenizer vocabulary files should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download the vocabulary files and override the cached versions if they + exist. + proxies (`dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `hf auth login` (stored in `~/.huggingface`). + local_files_only (`bool`, *optional*, defaults to `False`): + Whether or not to only rely on local files and not to attempt to download any files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + subfolder (`str`, *optional*): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for + facebook/rag-token-base), specify it here. + inputs (additional positional arguments, *optional*): + Will be passed along to the Tokenizer `__init__` method. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether or not to allow for custom models defined on the Hub in their own modeling files. This option + should only be set to `True` for repositories you trust and in which you have read the code, as it will + execute code present on the Hub on your local machine. + kwargs (additional keyword arguments, *optional*): + Will be passed to the Tokenizer `__init__` method. Can be used to set special tokens like `bos_token`, + `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`, + `extra_special_tokens`. See parameters in the `__init__` for more details. + + + + Passing `token=True` is required when you want to use a private model. + + + + Examples: + + ```python + # We can't instantiate directly the base class *PreTrainedTokenizerBase* so let's show our examples on a derived class: BertTokenizer + # Download vocabulary from huggingface.co and cache. + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + + # Download vocabulary from huggingface.co (user-uploaded) and cache. + tokenizer = BertTokenizer.from_pretrained("dbmdz/bert-base-german-cased") + + # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*) + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/") + + # If the tokenizer uses a single vocabulary file, you can point directly to this file + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/my_vocab.txt") + + # You can link tokens to special vocabulary when instantiating + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased", unk_token="") + # You should be sure '' is in the vocabulary when doing that. + # Otherwise use tokenizer.add_special_tokens({'unk_token': ''}) instead) + assert tokenizer.unk_token == "" + ```""" + proxies = kwargs.pop("proxies", None) + subfolder = kwargs.pop("subfolder", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + commit_hash = kwargs.pop("_commit_hash", None) + gguf_file = kwargs.get("gguf_file") + + user_agent = {"file_type": "tokenizer", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + vocab_files = {} + additional_files_names = {} + init_configuration = {} + + is_local = os.path.isdir(pretrained_model_name_or_path) + single_file_id = None + if os.path.isfile(pretrained_model_name_or_path): + # For legacy support: allow single-file loading if: + # 1. Only one vocab file is required, OR + # 2. It's a fast tokenizer with tokenizer_file (which is optional), OR + # 3. It's a GGUF file + vocab_files_count = len(cls.vocab_files_names) + has_optional_tokenizer_file = vocab_files_count > 1 and "tokenizer_file" in cls.vocab_files_names + + if vocab_files_count > 1 and not gguf_file and not has_optional_tokenizer_file: + raise ValueError( + f"Calling {cls.__name__}.from_pretrained() with the path to a single file or url is not " + "supported for this tokenizer. Use a model identifier or the path to a directory instead." + ) + file_id = "vocab_file" + if pretrained_model_name_or_path.endswith("tokenizer.json"): + file_id = "tokenizer_file" + vocab_files[file_id] = pretrained_model_name_or_path + single_file_id = file_id + else: + if gguf_file: + vocab_files["vocab_file"] = gguf_file + else: + # At this point pretrained_model_name_or_path is either a directory or a model identifier name + additional_files_names = { + "added_tokens_file": ADDED_TOKENS_FILE, # kept only for legacy + "special_tokens_map_file": SPECIAL_TOKENS_MAP_FILE, # kept only for legacy + "tokenizer_config_file": TOKENIZER_CONFIG_FILE, + # tokenizer_file used to initialize a slow from a fast. Properly copy the `addedTokens` instead of adding in random orders + "tokenizer_file": FULL_TOKENIZER_FILE, + "chat_template_file": CHAT_TEMPLATE_FILE, + } + + vocab_files = {**cls.vocab_files_names, **additional_files_names} + + # Check for versioned tokenizer files + if "tokenizer_file" in vocab_files: + fast_tokenizer_file = FULL_TOKENIZER_FILE + resolved_config_file = cached_file( + pretrained_model_name_or_path, + TOKENIZER_CONFIG_FILE, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + local_files_only=local_files_only, + subfolder=subfolder, + user_agent=user_agent, + _raise_exceptions_for_missing_entries=False, + _commit_hash=commit_hash, + ) + if resolved_config_file is not None: + with open(resolved_config_file, encoding="utf-8") as reader: + tokenizer_config = json.load(reader) + if "fast_tokenizer_files" in tokenizer_config: + fast_tokenizer_file = get_fast_tokenizer_file(tokenizer_config["fast_tokenizer_files"]) + commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + vocab_files["tokenizer_file"] = fast_tokenizer_file + + # This block looks for any extra chat template files + if is_local: + template_dir = Path(pretrained_model_name_or_path, CHAT_TEMPLATE_DIR) + if template_dir.is_dir(): + for template_file in template_dir.glob("*.jinja"): + template_name = template_file.name.removesuffix(".jinja") + vocab_files[f"chat_template_{template_name}"] = f"{CHAT_TEMPLATE_DIR}/{template_file.name}" + else: + for template in list_repo_templates( + pretrained_model_name_or_path, + local_files_only=local_files_only, + revision=revision, + cache_dir=cache_dir, + token=token, + ): + template = template.removesuffix(".jinja") + vocab_files[f"chat_template_{template}"] = f"{CHAT_TEMPLATE_DIR}/{template}.jinja" + + remote_files = [] + if not is_local and not local_files_only: + try: + remote_files = list_repo_files(pretrained_model_name_or_path) + except Exception: + remote_files = [] + elif pretrained_model_name_or_path and os.path.isdir(pretrained_model_name_or_path): + remote_files = os.listdir(pretrained_model_name_or_path) + + if "tokenizer_file" in vocab_files and not re.search(vocab_files["tokenizer_file"], "".join(remote_files)): + # mistral tokenizer names are different, but we can still convert them if + # mistral common is not there + other_pattern = r"tekken\.json|tokenizer\.model\.*|tiktoken\.model" + "|".join( + getattr(cls, "VOCAB_FILES_NAMES", {}).keys() + ) + if match := re.search(other_pattern, "\n".join(remote_files)): + if "spm_file" in vocab_files: + vocab_files["spm_file"] = match.group() + else: + vocab_files["vocab_file"] = match.group() + + resolved_vocab_files = {} + for file_id, file_path in vocab_files.items(): + if file_path is None: + resolved_vocab_files[file_id] = None + elif single_file_id == file_id: + if os.path.isfile(file_path): + resolved_vocab_files[file_id] = file_path + else: + try: + resolved_vocab_files[file_id] = cached_file( + pretrained_model_name_or_path, + file_path, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + _commit_hash=commit_hash, + ) + except OSError: + # Re-raise any error raised by cached_file in order to get a helpful error message + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load tokenizer for '{pretrained_model_name_or_path}'. If you were trying to load it from " + "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " + f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " + f"containing all relevant files for a {cls.__name__} tokenizer." + ) + commit_hash = extract_commit_hash(resolved_vocab_files[file_id], commit_hash) + + for file_id, file_path in vocab_files.items(): + if file_id not in resolved_vocab_files: + continue + + return cls._from_pretrained( + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + _commit_hash=commit_hash, + _is_local=is_local, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def _from_pretrained( + cls, + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=None, + cache_dir=None, + local_files_only=False, + _commit_hash=None, + _is_local=False, + trust_remote_code=False, + **kwargs, + ): + # Prepare tokenizer initialization kwargs + # Did we saved some inputs and kwargs to reload ? + tokenizer_config_file = resolved_vocab_files.pop("tokenizer_config_file", None) + if tokenizer_config_file is not None: + with open(tokenizer_config_file, encoding="utf-8") as tokenizer_config_handle: + init_kwargs = json.load(tokenizer_config_handle) + # used in the past to check if the tokenizer class matches the class in the repo + init_kwargs.pop("tokenizer_class", None) + saved_init_inputs = init_kwargs.pop("init_inputs", ()) + if not init_inputs: + init_inputs = saved_init_inputs + else: + init_kwargs = init_configuration + + if resolved_vocab_files.get("tokenizer_file", None) is not None: + init_kwargs.pop("add_bos_token", None) + init_kwargs.pop("add_eos_token", None) + + # If independent chat template file(s) exist, they take priority over template entries in the tokenizer config + chat_templates = {} + chat_template_file = resolved_vocab_files.pop("chat_template_file", None) + extra_chat_templates = [key for key in resolved_vocab_files if key.startswith("chat_template_")] + if chat_template_file is not None: + with open(chat_template_file, encoding="utf-8") as chat_template_handle: + chat_templates["default"] = chat_template_handle.read() + for extra_chat_template in extra_chat_templates: + template_file = resolved_vocab_files.pop(extra_chat_template, None) + if template_file is None: + continue # I think this should never happen, but just in case + template_name = extra_chat_template.removeprefix("chat_template_") + with open(template_file) as chat_template_handle: + chat_templates[template_name] = chat_template_handle.read() + if len(chat_templates) == 1 and "default" in chat_templates: + init_kwargs["chat_template"] = chat_templates["default"] + elif chat_templates: + init_kwargs["chat_template"] = chat_templates + + if not _is_local: + if "auto_map" in init_kwargs: + # For backward compatibility with odl format. + if isinstance(init_kwargs["auto_map"], (tuple, list)): + init_kwargs["auto_map"] = {"AutoTokenizer": init_kwargs["auto_map"]} + + # Update with newly provided kwargs + init_kwargs.update(kwargs) + + # V5: Convert deprecated additional_special_tokens to extra_special_tokens + if "additional_special_tokens" in init_kwargs: + init_kwargs.setdefault("extra_special_tokens", init_kwargs.pop("additional_special_tokens")) + + # V5: Collect model-specific tokens (custom *_token keys not in standard attributes) + default_attrs = set(cls.SPECIAL_TOKENS_ATTRIBUTES) + model_specific_tokens = { + key: init_kwargs.pop(key) + for key in list(init_kwargs.keys()) + if key not in default_attrs and key.endswith("_token") and isinstance(init_kwargs[key], (str, AddedToken)) + } + # If extra_special_tokens is a dict, merge it into model_specific_tokens + if isinstance(init_kwargs.get("extra_special_tokens"), dict): + model_specific_tokens.update(init_kwargs.pop("extra_special_tokens")) + if model_specific_tokens: + init_kwargs["model_specific_special_tokens"] = model_specific_tokens + + # Merge resolved_vocab_files arguments in init_kwargs. + added_tokens_file = resolved_vocab_files.pop("added_tokens_file", None) + special_tokens_map_file = resolved_vocab_files.pop("special_tokens_map_file", None) + for args_name, file_path in resolved_vocab_files.items(): + if args_name not in init_kwargs or init_kwargs[args_name] is None: + init_kwargs[args_name] = file_path + tokenizer_file = resolved_vocab_files.get("tokenizer_file", None) + + init_kwargs["name_or_path"] = pretrained_model_name_or_path + init_kwargs["is_local"] = _is_local + + #### Handle tokenizer serialization of added and special tokens + added_tokens_decoder: dict[int, AddedToken] = {} + added_tokens_map: dict[str, AddedToken] = {} + # if we have info on the slow added tokens + if "added_tokens_decoder" in init_kwargs: + for idx, token in init_kwargs["added_tokens_decoder"].items(): + if isinstance(token, dict): + token = AddedToken(**token) + if isinstance(token, AddedToken): + added_tokens_decoder[int(idx)] = token + added_tokens_map[str(token)] = token + else: + raise TypeError( + f"Found a {token.__class__} in the saved `added_tokens_decoder`, should be a dictionary or an AddedToken instance" + ) + else: + # Legacy: read special_tokens_map.json and merge into init_kwargs + if special_tokens_map_file is not None: + with open(special_tokens_map_file, encoding="utf-8") as f: + special_tokens_map = json.load(f) + for key, value in special_tokens_map.items(): + if key in kwargs and kwargs[key]: + continue # User-provided kwargs take precedence + if isinstance(value, dict) and key != "extra_special_tokens": + value.pop("special", None) + value = AddedToken(**value, special=True) + elif key == "extra_special_tokens" and isinstance(value, list): + # Merge list tokens, converting dicts to AddedToken + existing = list(init_kwargs.get("extra_special_tokens") or []) + for tok in value: + tok = AddedToken(**tok, special=True) if isinstance(tok, dict) else tok + if tok not in existing: + existing.append(tok) + value = existing + init_kwargs[key] = value + # Convert dict extra_special_tokens to model_specific_special_tokens + if isinstance(init_kwargs.get("extra_special_tokens"), dict): + init_kwargs.setdefault("model_specific_special_tokens", {}).update( + init_kwargs.pop("extra_special_tokens") + ) + + # slow -> slow|fast, legacy: convert the `"added_tokens.json"` file to `added_tokens_decoder`. + # this is for legacy purpose. We don't add the tokens after init for efficiency. + if added_tokens_file is not None: + # V5: Check both named and extra special tokens + special_tokens = {str(init_kwargs[k]) for k in cls.SPECIAL_TOKENS_ATTRIBUTES if init_kwargs.get(k)} + special_tokens.update(str(t) for t in (init_kwargs.get("extra_special_tokens") or [])) + + with open(added_tokens_file, encoding="utf-8") as f: + added_tok_encoder = json.load(f) + for str_token, index in added_tok_encoder.items(): + is_special = str_token in special_tokens + added_tokens_decoder[index] = AddedToken( + str_token, rstrip=False, lstrip=False, normalized=not is_special, special=is_special + ) + added_tokens_map[str_token] = added_tokens_decoder[index] + + # allows converting a fast -> slow: add the `tokenizer.json`'s `"added_tokens"` to the slow tokenizer + # if `tokenizer_config.json` is `None` + if tokenizer_file is not None: + # This is for slow so can be done before + with open(tokenizer_file, encoding="utf-8") as tokenizer_file_handle: + tokenizer_file_handle = json.load(tokenizer_file_handle) + added_tokens = tokenizer_file_handle.pop("added_tokens") + for serialized_tokens in added_tokens: + idx = serialized_tokens.pop("id") + added_tokens_decoder[idx] = AddedToken(**serialized_tokens) + added_tokens_map[str(added_tokens_decoder[idx])] = added_tokens_decoder[idx] + # end legacy + + # Passing AddedTokens and not strings to the class to prevent it from casting the string to a different AddedToken + # convert {'__type': 'AddedToken', 'content': '', 'lstrip': False, 'normalized': True, ...} to AddedTokens + init_kwargs["added_tokens_decoder"] = added_tokens_decoder + init_kwargs = cls.convert_added_tokens(init_kwargs, save=False) + # V5: Map special tokens from added_tokens_map (named tokens only) + for key in cls.SPECIAL_TOKENS_ATTRIBUTES: + if key in init_kwargs and added_tokens_map != {} and init_kwargs[key] is not None: + init_kwargs[key] = added_tokens_map.get(str(init_kwargs[key]), init_kwargs[key]) + + # From pretrained with the legacy fixes + # for `tokenizers` based tokenizer, we actually want to have vocab and merges pre-extracted from whatever inputs + # for `none` (PythonBackend) based tokenizer, we also want the vocab file / merge files not extracted. + # for `sentencepiece` based tokenizer, we pass the sentencepiece model file directly. + init_kwargs = cls.convert_to_native_format(**init_kwargs) + + try: + tokenizer = cls(*init_inputs, **init_kwargs) + except import_protobuf_decode_error(): + raise RuntimeError( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(Google protobuf error: Tried to load SPM model with non-SPM vocab file).", + ) + except RuntimeError as e: + if "sentencepiece_processor.cc" in str(e): + raise RuntimeError( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(SentencePiece RuntimeError: Tried to load SPM model with non-SPM vocab file).", + ) from e + else: + raise e + except OSError: + raise OSError( + "Unable to load vocabulary from file. " + "Please check that the provided vocabulary is accessible and not corrupted." + ) + return tokenizer + + @classmethod + def convert_to_native_format(cls, **kwargs): + return kwargs + + @classmethod + def convert_added_tokens(cls, obj: AddedToken | Any, save=False, add_type_field=True): + if isinstance(obj, dict) and "__type" in obj and obj["__type"] == "AddedToken": + obj.pop("__type") + return AddedToken(**obj) + if isinstance(obj, AddedToken) and save: + obj = obj.__getstate__() + if add_type_field: + obj["__type"] = "AddedToken" + else: + # Don't save "special" for previous tokenizers + obj.pop("special") + return obj + elif isinstance(obj, (list, tuple)): + return [cls.convert_added_tokens(o, save=save, add_type_field=add_type_field) for o in obj] + elif isinstance(obj, dict): + return {k: cls.convert_added_tokens(v, save=save, add_type_field=add_type_field) for k, v in obj.items()} + return obj + + def save_pretrained( + self, + save_directory: str | os.PathLike, + legacy_format: bool | None = None, + filename_prefix: str | None = None, + push_to_hub: bool = False, + **kwargs, + ) -> tuple[str, ...]: + """ + Save the full tokenizer state. + + + This method make sure the full tokenizer can then be re-loaded using the + [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] class method.. + + Warning,None This won't save modifications you may have applied to the tokenizer after the instantiation (for + instance, modifying `tokenizer.do_lower_case` after creation). + + Args: + save_directory (`str` or `os.PathLike`): The path to a directory where the tokenizer will be saved. + legacy_format (`bool`, *optional*): + Only applicable for a fast tokenizer. If unset (default), will save the tokenizer in the unified JSON + format as well as in legacy format if it exists, i.e. with tokenizer specific vocabulary and a separate + added_tokens files. + + If `False`, will only save the tokenizer in the unified JSON format. This format is incompatible with + "slow" tokenizers (not powered by the *tokenizers* library), so the tokenizer will not be able to be + loaded in the corresponding "slow" tokenizer. + + If `True`, will save the tokenizer in legacy format. If the "slow" tokenizer doesn't exits, a value + error is raised. + filename_prefix (`str`, *optional*): + A prefix to add to the names of the files saved by the tokenizer. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + + Returns: + A tuple of `str`: The files saved. + """ + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, **kwargs).repo_id + files_timestamps = self._get_files_timestamps(save_directory) + + tokenizer_config_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + TOKENIZER_CONFIG_FILE + ) + + tokenizer_config = copy.deepcopy(self.init_kwargs) + tokenizer_config.pop("add_bos_token", None) + tokenizer_config.pop("add_eos_token", None) + + # Let's save the init kwargs + target_keys = set(self.init_kwargs.keys()) + target_keys.discard("add_bos_token") + target_keys.discard("add_eos_token") + # Let's save the special tokens map (only the strings) + target_keys.update(["model_max_length"]) + + for k in target_keys: + if hasattr(self, k): + tokenizer_config[k] = getattr(self, k) + + # Let's make sure we properly save the special tokens + # V5: Save both named tokens and extra tokens + tokenizer_config.update(self.special_tokens_map) + if self._extra_special_tokens: + tokenizer_config["extra_special_tokens"] = self.extra_special_tokens + + save_jinja_files = kwargs.get("save_jinja_files", True) + tokenizer_config, saved_raw_chat_template_files = self.save_chat_templates( + save_directory, tokenizer_config, filename_prefix, save_jinja_files + ) + + if getattr(self, "response_schema", None) is not None: + tokenizer_config["response_schema"] = self.response_schema + + if len(self.init_inputs) > 0: + tokenizer_config["init_inputs"] = copy.deepcopy(self.init_inputs) + for file_id in self.vocab_files_names: + tokenizer_config.pop(file_id, None) + + # no typefields, this way old fast and slow can load it + tokenizer_config = self.convert_added_tokens(tokenizer_config, add_type_field=True, save=True) + # Process added tokens separately: allows previous versions to ignore it! + added_tokens = {} + for key, value in self.added_tokens_decoder.items(): + added_tokens[key] = value.__getstate__() + tokenizer_config["added_tokens_decoder"] = added_tokens + + # Add tokenizer class to the tokenizer config to be able to reload it with from_pretrained + tokenizer_class = self.__class__.__name__ + + # tokenizers backend don't need to save added_tokens_decoder and additional_special_tokens + if any(base.__name__ == "TokenizersBackend" for base in self.__class__.__mro__): + tokenizer_config.pop("added_tokens_decoder", None) + tokenizer_config.pop("additional_special_tokens", None) + + # Remove the Fast at the end if we can save the slow tokenizer + if tokenizer_class.endswith("Fast") and getattr(self, "can_save_slow_tokenizer", False): + tokenizer_class = tokenizer_class[:-4] + tokenizer_config["tokenizer_class"] = tokenizer_class + if getattr(self, "_auto_map", None) is not None: + tokenizer_config["auto_map"] = self._auto_map + if getattr(self, "_processor_class", None) is not None: + tokenizer_config["processor_class"] = self._processor_class + tokenizer_config.pop("files_loaded", None) + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=tokenizer_config) + + # remove private information + if "name_or_path" in tokenizer_config: + tokenizer_config.pop("name_or_path") + tokenizer_config.pop("special_tokens_map_file", None) + tokenizer_config.pop("tokenizer_file", None) + if "device_map" in tokenizer_config: + tokenizer_config.pop("device_map") + if "slow_tokenizer_class" in tokenizer_config: + tokenizer_config.pop("slow_tokenizer_class") + + with open(tokenizer_config_file, "w", encoding="utf-8") as f: + out_str = json.dumps(tokenizer_config, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + logger.info(f"tokenizer config file saved in {tokenizer_config_file}") + + # Sanitize AddedTokens in special_tokens_map + + file_names = (tokenizer_config_file, *saved_raw_chat_template_files) + + save_files = self._save_pretrained( + save_directory=save_directory, + file_names=file_names, + legacy_format=legacy_format, + filename_prefix=filename_prefix, + ) + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + return save_files + + def _save_pretrained( + self, + save_directory: str | os.PathLike, + file_names: tuple[str, ...], + legacy_format: bool | None = None, + filename_prefix: str | None = None, + ) -> tuple[str, ...]: + """ + Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens. + + Fast tokenizers can also be saved in a unique JSON file containing {config + vocab + added-tokens} using the + specific [`~tokenization_utils_tokenizers.PreTrainedTokenizerFast._save_pretrained`] + """ + if legacy_format is False: + raise ValueError( + "Only fast tokenizers (instances of PreTrainedTokenizerFast) can be saved in non legacy format." + ) + + save_directory = str(save_directory) + + added_tokens_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + ADDED_TOKENS_FILE + ) + # the new get_added_vocab() also returns special tokens and tokens that have an index < vocab_size + added_vocab = {tok: index for tok, index in self.added_tokens_encoder.items() if index >= self.vocab_size} + if added_vocab: + with open(added_tokens_file, "w", encoding="utf-8") as f: + out_str = json.dumps(added_vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + logger.info(f"added tokens file saved in {added_tokens_file}") + + vocab_files = self.save_vocabulary(save_directory, filename_prefix=filename_prefix) + + return file_names + vocab_files + (added_tokens_file,) + + def clean_up_tokenization(self, text: str) -> str: + """ + Clean up tokenization spaces in a given text. + This method is mostly for remote code support. + + """ + + text = ( + text.replace(" .", ".") + .replace(" ?", "?") + .replace(" !", "!") + .replace(" ,", ",") + .replace(" ' ", "'") + .replace(" n't", "n't") + .replace(" 'm", "'m") + .replace(" 's", "'s") + .replace(" 've", "'ve") + .replace(" 're", "'re") + ) + return text + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str, ...]: + """ + Save only the vocabulary of the tokenizer (vocabulary + added tokens). + + This method won't save the configuration and special token mappings of the tokenizer. Use + [`~PreTrainedTokenizerFast._save_pretrained`] to save the whole state of the tokenizer. + + Args: + save_directory (`str`): + The directory in which to save the vocabulary. + filename_prefix (`str`, *optional*): + An optional prefix to add to the named of the saved files. + + Returns: + `tuple(str)`: Paths to the files saved. + """ + raise NotImplementedError + + def tokenize(self, text: str, pair: str | None = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + """ + Converts a string into a sequence of tokens, replacing unknown tokens with the `unk_token`. + + Args: + text (`str`): + The sequence to be encoded. + pair (`str`, *optional*): + A second sequence to be encoded with the first. + add_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to add the special tokens associated with the corresponding model. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific encode method. See details in + [`~PreTrainedTokenizerBase.__call__`] + + Returns: + `list[str]`: The list of tokens. + """ + raise NotImplementedError + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, + """ + **kwargs: Passed along to the `.tokenize()` method. + """, + """ + Returns: + `list[int]`, `torch.Tensor`, or `np.ndarray`: The tokenized ids of the text. + """, + ) + def encode( + self, + text: TextInput | PreTokenizedInput | EncodedInput, + text_pair: TextInput | PreTokenizedInput | EncodedInput | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + padding_side: str | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> list[int]: + """ + 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))`. + + Args: + 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). + """ + padding_strategy, truncation_strategy, max_length, kwargs_updated = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + **kwargs, + ) + + kwargs.update(kwargs_updated) + + encoded_inputs = self._encode_plus( + text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + padding_side=padding_side, + return_tensors=return_tensors, + **kwargs, + ) + + return encoded_inputs["input_ids"] + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + raise NotImplementedError + + @property + def max_len_single_sentence(self) -> int: + """ + `int`: The maximum length of a sentence that can be fed to the model. + """ + return self.model_max_length - self.num_special_tokens_to_add(pair=False) + + @max_len_single_sentence.setter + def max_len_single_sentence(self, value) -> None: + # For backward compatibility, allow to try to setup 'max_len_single_sentence'. + if value == self.model_max_length - self.num_special_tokens_to_add(pair=False) and self.verbose: + if not self.deprecation_warnings.get("max_len_single_sentence", False): + logger.warning( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_single_sentence"] = True + else: + raise ValueError( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + + @property + def max_len_sentences_pair(self) -> int: + """ + `int`: The maximum combined length of a pair of sentences that can be fed to the model. + """ + return self.model_max_length - self.num_special_tokens_to_add(pair=True) + + @max_len_sentences_pair.setter + def max_len_sentences_pair(self, value) -> None: + # For backward compatibility, allow to try to setup 'max_len_sentences_pair'. + if value == self.model_max_length - self.num_special_tokens_to_add(pair=True) and self.verbose: + if not self.deprecation_warnings.get("max_len_sentences_pair", False): + logger.warning( + "Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_sentences_pair"] = True + else: + raise ValueError("Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up.") + + def _get_padding_truncation_strategies( + self, padding=False, truncation=None, max_length=None, pad_to_multiple_of=None, verbose=True, **kwargs + ): + """ + Find the correct padding/truncation strategy + """ + + # Backward compatibility for previous behavior: + # If you only set max_length, it activates truncation for max_length + if max_length is not None and padding is False and truncation is None: + truncation = "longest_first" + + # Get padding strategy + if padding is not False: + if padding is True: + if verbose: + if max_length is not None and ( + truncation is None or truncation is False or truncation == "do_not_truncate" + ): + warnings.warn( + "`max_length` is ignored when `padding`=`True` and there is no truncation strategy. " + "To pad to max length, use `padding='max_length'`." + ) + padding_strategy = PaddingStrategy.LONGEST # Default to pad to the longest sequence in the batch + elif not isinstance(padding, PaddingStrategy): + padding_strategy = PaddingStrategy(padding) + elif isinstance(padding, PaddingStrategy): + padding_strategy = padding + else: + padding_strategy = PaddingStrategy.DO_NOT_PAD + + # Get truncation strategy + if truncation is not False and truncation is not None: + if truncation is True: + truncation_strategy = ( + TruncationStrategy.LONGEST_FIRST + ) # Default to truncate the longest sequences in pairs of inputs + elif not isinstance(truncation, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation) + elif isinstance(truncation, TruncationStrategy): + truncation_strategy = truncation + else: + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + + # Set max length if needed + if max_length is None: + if padding_strategy == PaddingStrategy.MAX_LENGTH: + if self.model_max_length > LARGE_INTEGER: + padding_strategy = PaddingStrategy.DO_NOT_PAD + else: + max_length = self.model_max_length + + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE: + if self.model_max_length > LARGE_INTEGER: + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + else: + max_length = self.model_max_length + + # Test if we have a padding token + if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.pad_token is None or self.pad_token_id < 0): + raise ValueError( + "Asking to pad but the tokenizer does not have a padding token. " + "Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` " + "or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`." + ) + + # Check that we will truncate to a multiple of pad_to_multiple_of if both are provided + if ( + truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE + and padding_strategy != PaddingStrategy.DO_NOT_PAD + and pad_to_multiple_of is not None + and max_length is not None + and (max_length % pad_to_multiple_of != 0) + ): + raise ValueError( + "Truncation and padding are both activated but " + f"truncation length ({max_length}) is not a multiple of pad_to_multiple_of ({pad_to_multiple_of})." + ) + + return padding_strategy, truncation_strategy, max_length, kwargs + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + text_pair: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + text_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + text_pair_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + add_special_tokens: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy | None = None, + max_length: int | None = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: int | None = None, + 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, + tokenizer_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences. + + Args: + text (`str`, `list[str]`, `list[list[str]]`, *optional*): + 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). + text_pair (`str`, `list[str]`, `list[list[str]]`, *optional*): + 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). + text_target (`str`, `list[str]`, `list[list[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. 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). + text_pair_target (`str`, `list[str]`, `list[list[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. 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). + tokenizer_kwargs (`dict[str, Any]`, *optional*): + Additional kwargs to pass to the tokenizer. These will be merged with the explicit parameters and + other kwargs, with explicit parameters taking precedence. + """ + # To avoid duplicating + all_kwargs = { + "add_special_tokens": add_special_tokens, + "padding": padding, + "truncation": truncation, + "max_length": max_length, + "stride": stride, + "is_split_into_words": is_split_into_words, + "pad_to_multiple_of": pad_to_multiple_of, + "padding_side": padding_side, + "return_tensors": return_tensors, + "return_token_type_ids": return_token_type_ids, + "return_attention_mask": return_attention_mask, + "return_overflowing_tokens": return_overflowing_tokens, + "return_special_tokens_mask": return_special_tokens_mask, + "return_offsets_mapping": return_offsets_mapping, + "return_length": return_length, + "split_special_tokens": kwargs.pop("split_special_tokens", self.split_special_tokens), + "verbose": verbose, + } + + max_target_length = kwargs.pop("max_target_length", None) + + # First merge tokenizer_kwargs, then other kwargs (explicit params take precedence) + if tokenizer_kwargs is not None: + all_kwargs.update(tokenizer_kwargs) + all_kwargs.update(kwargs) + if text is None and text_target is None: + raise ValueError("You need to specify either `text` or `text_target`.") + + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=all_kwargs.pop("padding", False), + truncation=all_kwargs.pop("truncation", None), + max_length=all_kwargs.pop("max_length", None), + pad_to_multiple_of=all_kwargs.get("pad_to_multiple_of"), + verbose=all_kwargs.get("verbose", True), + **kwargs, + ) + + 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._encode_plus( + text=text, + text_pair=text_pair, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + **all_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, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_target_length if max_target_length is not None else max_length, + **all_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 + + def _encode_plus( + self, + text: TextInput | PreTokenizedInput | EncodedInput, + text_pair: TextInput | PreTokenizedInput | EncodedInput | 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, + 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, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + raise NotImplementedError + + def pad( + self, + encoded_inputs: BatchEncoding + | list[BatchEncoding] + | dict[str, EncodedInput] + | dict[str, list[EncodedInput]] + | list[dict[str, EncodedInput]], + padding: bool | str | PaddingStrategy = True, + max_length: int | None = None, + pad_to_multiple_of: int | None = None, + padding_side: str | None = None, + return_attention_mask: bool | None = None, + return_tensors: str | TensorType | None = None, + verbose: bool = True, + ) -> BatchEncoding: + """ + Pad a single encoded input or a batch of encoded inputs up to predefined length or to the max sequence length + in the batch. + + Padding side (left/right) padding token ids are defined at the tokenizer level (with `self.padding_side`, + `self.pad_token_id` and `self.pad_token_type_id`). + + Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the + text followed by a call to the `pad` method to get a padded encoding. + + + + If the `encoded_inputs` passed are dictionary of numpy arrays, or PyTorch tensors, the + result will use the same type unless you provide a different tensor type with `return_tensors`. In the case of + PyTorch tensors, you will lose the specific device of your tensors however. + + + + Args: + encoded_inputs ([`BatchEncoding`], list of [`BatchEncoding`], `dict[str, list[int]]`, `dict[str, list[list[int]]` or `list[dict[str, list[int]]]`): + Tokenized inputs. Can represent one input ([`BatchEncoding`] or `dict[str, list[int]]`) or a batch of + tokenized inputs (list of [`BatchEncoding`], *dict[str, list[list[int]]]* or *list[dict[str, + list[int]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader + collate function. + + Instead of `list[int]` you can have tensors (numpy arrays, or PyTorch tensors), see + the note above for the return type. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'` (default): 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'`: No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`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 (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + """ + + # If we have a list of dicts, let's convert it in a dict of lists + # We do this to allow using this method as a collate_fn function in PyTorch Dataloader + if ( + isinstance(encoded_inputs, (list, tuple)) + and len(encoded_inputs) > 0 + and isinstance(encoded_inputs[0], Mapping) + ): + # Call .keys() explicitly for compatibility with TensorDict and other Mapping subclasses + encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()} + + # The model's main input name, usually `input_ids`, has been passed for padding + if self.model_input_names[0] not in encoded_inputs: + raise ValueError( + "You should supply an encoding or a list of encodings to this method " + f"that includes {self.model_input_names[0]}, but you provided {list(encoded_inputs.keys())}" + ) + + required_input = encoded_inputs[self.model_input_names[0]] + + if required_input is None or (isinstance(required_input, Sized) and len(required_input) == 0): + if return_attention_mask: + encoded_inputs["attention_mask"] = [] + return encoded_inputs + + # If we have PyTorch/NumPy tensors/arrays as inputs, we cast them as python objects + # and rebuild them afterwards if no return_tensors is specified + # Note that we lose the specific device the tensor may be on for PyTorch + + first_element = required_input[0] + if isinstance(first_element, (list, tuple)): + # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element. + for item in required_input: + if len(item) != 0: + first_element = item[0] + break + # At this state, if `first_element` is still a list/tuple, it's an empty one so there is nothing to do. + if not isinstance(first_element, (int, list, tuple)): + if is_torch_tensor(first_element): + return_tensors = "pt" if return_tensors is None else return_tensors + elif isinstance(first_element, np.ndarray): + return_tensors = "np" if return_tensors is None else return_tensors + else: + raise ValueError( + f"type of {first_element} unknown: {type(first_element)}. " + "Should be one of a python, numpy, or pytorch object." + ) + + for key, value in encoded_inputs.items(): + encoded_inputs[key] = to_py_obj(value) + + # Convert padding_strategy in PaddingStrategy + padding_strategy, _, max_length, _ = self._get_padding_truncation_strategies( + padding=padding, max_length=max_length, verbose=verbose + ) + + required_input = encoded_inputs[self.model_input_names[0]] + if required_input and not isinstance(required_input[0], (list, tuple)): + encoded_inputs = self._pad( + encoded_inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + return BatchEncoding(encoded_inputs, tensor_type=return_tensors) + + batch_size = len(required_input) + assert all(len(v) == batch_size for v in encoded_inputs.values()), ( + "Some items in the output dictionary have a different batch size than others." + ) + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = max(len(inputs) for inputs in required_input) + padding_strategy = PaddingStrategy.MAX_LENGTH + + batch_outputs = {} + for i in range(batch_size): + inputs = {k: v[i] for k, v in encoded_inputs.items()} + outputs = self._pad( + inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def _pad( + self, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, + 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 `padding_side` argument: + + - '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 "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 "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(f"Invalid padding strategy:{padding_side}") + + return encoded_inputs + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + """ + Converts a sequence of tokens in a single string. The most simple way to do it is `" ".join(tokens)` but we + often want to remove sub-word tokenization artifacts at the same time. + + Args: + tokens (`list[str]`): The token to join in a string. + + Returns: + `str`: The joined tokens. + """ + raise NotImplementedError + + def decode( + self, + token_ids: int | list[int] | list[list[int]] | np.ndarray | torch.Tensor, + skip_special_tokens: bool = False, + **kwargs, + ) -> str | list[str]: + """ + Converts a sequence of ids into a string, or a list of sequences into a list of strings, + using the tokenizer and vocabulary with options to remove special tokens and clean up + tokenization spaces. + + Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`. + + Args: + token_ids (`Union[int, list[int], list[list[int]], np.ndarray, torch.Tensor]`): + A single sequence or a batch (list of sequences) of tokenized input ids. Can be obtained using the + `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `Union[str, list[str]]`: The decoded string for a single sequence, or a list of decoded strings for a + batch of sequences. + """ + # Convert inputs to python lists + token_ids = to_py_obj(token_ids) + + # If we received batched input, decode each sequence + if isinstance(token_ids, (list, tuple)) and len(token_ids) > 0 and isinstance(token_ids[0], (list, tuple)): + clean_up_tokenization_spaces = kwargs.pop("clean_up_tokenization_spaces", False) + return [ + self._decode( + token_ids=seq, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + for seq in token_ids + ] + + return self._decode( + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + **kwargs, + ) + + def batch_decode( + self, + sequences: list[int] | list[list[int]] | np.ndarray | torch.Tensor, + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> list[str]: + """ + Convert a list of lists of token ids into a list of strings by calling decode. + + This method is provided for backwards compatibility. The `decode` method now handles batched input natively, + so you can use `decode` directly instead of `batch_decode`. + + Args: + sequences (`Union[list[int], list[list[int]], np.ndarray, torch.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `list[str]`: The list of decoded sentences. + """ + # Forward to decode() which now handles batched input natively + result = self.decode( + token_ids=sequences, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + # Ensure we always return a list for backwards compatibility + if isinstance(result, str): + return [result] + return result + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> str: + raise NotImplementedError + + def _eventual_warn_about_too_long_sequence(self, ids: list[int], max_length: int | None, verbose: bool): + """ + Depending on the input and internal state we might trigger a warning about a sequence that is too long for its + corresponding model + + Args: + ids (`list[str]`): The ids produced by the tokenization + max_length (`int`, *optional*): The max_length desired (does not trigger a warning if it is set) + verbose (`bool`): Whether or not to print more information and warnings. + + """ + if max_length is None and len(ids) > self.model_max_length and verbose and self.model_max_length != 0: + if not self.deprecation_warnings.get("sequence-length-is-longer-than-the-specified-maximum", False): + logger.warning( + "Token indices sequence length is longer than the specified maximum sequence length " + f"for this model ({len(ids)} > {self.model_max_length}). Running this sequence through the model " + "will result in indexing errors" + ) + self.deprecation_warnings["sequence-length-is-longer-than-the-specified-maximum"] = True + + @classmethod + def register_for_auto_class(cls, auto_class="AutoTokenizer"): + """ + Register this class with a given auto class. This should only be used for custom tokenizers as the ones in the + library are already mapped with `AutoTokenizer`. + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoTokenizer"`): + The auto class to register this new tokenizer with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def apply_chat_template( + self, + conversation: list[dict[str, str]] | list[list[dict[str, str]]], + tools: list[dict | Callable] | None = None, + documents: list[dict[str, str]] | None = None, + chat_template: str | None = None, + add_generation_prompt: bool = False, + continue_final_message: bool = False, + tokenize: bool = True, + padding: bool | str | PaddingStrategy = False, + truncation: bool = False, + max_length: int | None = None, + return_tensors: str | TensorType | None = None, + return_dict: bool = True, + return_assistant_tokens_mask: bool = False, + tokenizer_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> str | list[int] | list[str] | list[list[int]] | BatchEncoding: + """ + Converts a list of dictionaries with `"role"` and `"content"` keys to a list of token + ids. This method is intended for use with chat models, and will read the tokenizer's chat_template attribute to + determine the format and control tokens to use when converting. + + Args: + conversation (Union[list[dict[str, str]], list[list[dict[str, str]]]]): A list of dicts + with "role" and "content" keys, representing the chat history so far. + tools (`list[Union[Dict, Callable]]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [tool use guide](https://huggingface.co/docs/transformers/en/chat_extras#passing-tools) + for more information. + documents (`list[dict[str, str]]`, *optional*): + A list of dicts representing documents that will be accessible to the model if it is performing RAG + (retrieval-augmented generation). If the template does not support RAG, this argument will have no + effect. We recommend that each document should be a dict containing "title" and "text" keys. + chat_template (`str`, *optional*): + A Jinja template to use for this conversion. It is usually not necessary to pass anything to this + argument, as the model's template will be used by default. + add_generation_prompt (bool, *optional*): + If this is set, a prompt with the token(s) that indicate + the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model. + Note that this argument will be passed to the chat template, and so it must be supported in the + template for this argument to have any effect. + continue_final_message (bool, *optional*): + If this is set, the chat will be formatted so that the final + message in the chat is open-ended, without any EOS tokens. The model will continue this message + rather than starting a new one. This allows you to "prefill" part of + the model's response for it. Cannot be used at the same time as `add_generation_prompt`. + tokenize (`bool`, defaults to `True`): + Whether to tokenize the output. If `False`, the output will be a string. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, defaults to `False`): + Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`. + max_length (`int`, *optional*): + Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If + not specified, the tokenizer's `max_length` attribute will be used as a default. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable + values are: + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + return_dict (`bool`, defaults to `True`): + Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`. + tokenizer_kwargs (`dict[str: Any]`, *optional*): Additional kwargs to pass to the tokenizer. + return_assistant_tokens_mask (`bool`, defaults to `False`): + Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant, + the mask will contain 1. For user and system tokens, the mask will contain 0. + This functionality is only available for chat templates that support it via the `{% generation %}` keyword. + **kwargs: Additional kwargs to pass to the template renderer. Will be accessible by the chat template. + + Returns: + `Union[list[int], Dict]`: A list of token ids representing the tokenized chat so far, including control tokens. This + output is ready to pass to the model, either directly or via methods like `generate()`. If `return_dict` is + set, will return a dict of tokenizer outputs instead. + """ + + if not tokenize: + return_dict = False # dicts are only returned by the tokenizer anyway + + if return_assistant_tokens_mask and not (return_dict and tokenize): + raise ValueError("`return_assistant_tokens_mask=True` requires `return_dict=True` and `tokenize=True`") + + if tokenizer_kwargs is None: + tokenizer_kwargs = {} + + chat_template = self.get_chat_template(chat_template, tools) + + if isinstance(conversation, (list, tuple)) and ( + isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "messages") + ): + conversations = conversation + is_batched = True + else: + conversations = [conversation] + is_batched = False + + if continue_final_message: + if add_generation_prompt: + raise ValueError( + "continue_final_message and add_generation_prompt are not compatible. Use continue_final_message when you want the model to continue the final message, and add_generation_prompt when you want to add a header that will prompt it to start a new assistant message instead." + ) + if return_assistant_tokens_mask: + raise ValueError("continue_final_message is not compatible with return_assistant_tokens_mask.") + + template_kwargs = {**self.special_tokens_map, **kwargs} # kwargs overwrite special tokens if both are present + rendered_chat, generation_indices = render_jinja_template( + conversations=conversations, + tools=tools, + documents=documents, + chat_template=chat_template, + return_assistant_tokens_mask=return_assistant_tokens_mask, + continue_final_message=continue_final_message, + add_generation_prompt=add_generation_prompt, + **template_kwargs, + ) + + if not is_batched: + rendered_chat = rendered_chat[0] + + if tokenize: + out = self( + rendered_chat, + padding=padding, + truncation=truncation, + max_length=max_length, + add_special_tokens=False, + return_tensors=return_tensors, + **tokenizer_kwargs, + ) + if return_dict: + if return_assistant_tokens_mask: + assistant_masks = [] + if is_batched or return_tensors: + input_ids = out["input_ids"] + else: + input_ids = [out["input_ids"]] + for i in range(len(input_ids)): + current_mask = [0] * len(input_ids[i]) + for assistant_start_char, assistant_end_char in generation_indices[i]: + start_token = out.char_to_token(i, assistant_start_char) + end_token = out.char_to_token(i, assistant_end_char - 1) + if start_token is None: + # start_token is out of bounds maybe due to truncation. + break + for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])): + current_mask[token_id] = 1 + assistant_masks.append(current_mask) + + if not is_batched and not return_tensors: + assistant_masks = assistant_masks[0] + + out["assistant_masks"] = assistant_masks + + if return_tensors: + out.convert_to_tensors(tensor_type=return_tensors) + + return out + else: + return out["input_ids"] + else: + return rendered_chat + + def encode_message_with_chat_template( + self, + message: dict[str, str], + conversation_history: list[dict[str, str]] | None = None, + **kwargs, + ) -> list[int]: + """ + Tokenize a single message. This method is a convenience wrapper around `apply_chat_template` that allows you + to tokenize messages one by one. This is useful for things like token-by-token streaming. + This method is not guaranteed to be perfect. For some models, it may be impossible to robustly tokenize + single messages. For example, if the chat template adds tokens after each message, but also has a prefix that + is added to the entire chat, it will be impossible to distinguish a chat-start-token from a message-start-token. + In these cases, this method will do its best to find the correct tokenization, but it may not be perfect. + **Note:** This method does not support `add_generation_prompt`. If you want to add a generation prompt, + you should do it separately after tokenizing the conversation. + Args: + message (`dict`): + A dictionary with "role" and "content" keys, representing the message to tokenize. + conversation_history (`list[dict]`, *optional*): + A list of dicts with "role" and "content" keys, representing the chat history so far. If you are + tokenizing messages one by one, you should pass the previous messages in the conversation here. + **kwargs: + Additional kwargs to pass to the `apply_chat_template` method. + Returns: + `list[int]`: A list of token ids representing the tokenized message. + """ + if "add_generation_prompt" in kwargs: + raise ValueError( + "`encode_message_with_chat_template` does not support `add_generation_prompt`. Please add the generation prompt " + "separately." + ) + + if conversation_history is None or len(conversation_history) == 0: + return self.apply_chat_template( + [message], add_generation_prompt=False, tokenize=True, return_dict=False, **kwargs + ) + + conversation = conversation_history + [message] + tokens = self.apply_chat_template( + conversation, add_generation_prompt=False, tokenize=True, return_dict=False, **kwargs + ) + + prefix_tokens = self.apply_chat_template( + conversation_history, add_generation_prompt=False, tokenize=True, return_dict=False, **kwargs + ) + # It's possible that the prefix tokens are not a prefix of the full list of tokens. + # For example, if the prefix is `User: Hi` and the full conversation is `User: HiAssistant: Hello`. + # In this case, we can't simply find the prefix, so we have to do something a bit more subtle. + # We look for the first place where the tokens differ, and that's our split point. + # This is not perfect, but it's the best we can do without a token-level API. + # To make this more robust, we could do a diff and find the longest common subsequence, but this is + # a good first approximation. + # This is particularly important for models like Llama3 that have changed their chat template to include + # EOS tokens after user messages. + min_len = min(len(prefix_tokens), len(tokens)) + for i in range(min_len): + if prefix_tokens[i] != tokens[i]: + return tokens[i:] + return tokens[min_len:] + + def get_chat_template(self, chat_template: str | None = None, tools: list[dict] | None = None) -> str: + """ + Retrieve the chat template string used for tokenizing chat messages. This template is used + internally by the `apply_chat_template` method and can also be used externally to retrieve the model's chat + template for better generation tracking. + + Args: + chat_template (`str`, *optional*): + A Jinja template or the name of a template to use for this conversion. + It is usually not necessary to pass anything to this argument, + as the model's template will be used by default. + tools (`list[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + + Returns: + `str`: The chat template string. + """ + # First, handle the cases when the model has a dict of multiple templates + if isinstance(self.chat_template, dict): + template_dict = self.chat_template + if chat_template is not None and chat_template in template_dict: + # The user can pass the name of a template to the chat template argument instead of an entire template + chat_template = template_dict[chat_template] + elif chat_template is None: + if tools is not None and "tool_use" in template_dict: + chat_template = template_dict["tool_use"] + elif "default" in template_dict: + chat_template = template_dict["default"] + else: + raise ValueError( + "This model has multiple chat templates with no default specified! Please either pass a chat " + "template or the name of the template you wish to use to the `chat_template` argument. Available " + f"template names are {sorted(template_dict.keys())}." + ) + + elif chat_template is None: + # These are the cases when the model has a single template + # priority: `chat_template` argument > `tokenizer.chat_template` + if self.chat_template is not None: + chat_template = self.chat_template + else: + raise ValueError( + "Cannot use chat template functions because tokenizer.chat_template is not set and no template " + "argument was passed! For information about writing templates and setting the " + "tokenizer.chat_template attribute, please see the documentation at " + "https://huggingface.co/docs/transformers/main/en/chat_templating" + ) + + return chat_template + + def save_chat_templates( + self, + save_directory: str | os.PathLike, + tokenizer_config: dict, + filename_prefix: str | None, + save_jinja_files: bool, + ): + """ + Writes chat templates out to the save directory if we're using the new format, and removes them from + the tokenizer config if present. If we're using the legacy format, it doesn't write any files, and instead + writes the templates to the tokenizer config in the correct format. + """ + chat_template_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + CHAT_TEMPLATE_FILE + ) + chat_template_dir = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + CHAT_TEMPLATE_DIR + ) + + saved_raw_chat_template_files = [] + if save_jinja_files and isinstance(self.chat_template, str): + # New format for single templates is to save them as chat_template.jinja + with open(chat_template_file, "w", encoding="utf-8") as f: + f.write(self.chat_template) + logger.info(f"chat template saved in {chat_template_file}") + saved_raw_chat_template_files.append(chat_template_file) + if "chat_template" in tokenizer_config: + tokenizer_config.pop("chat_template") # To ensure it doesn't somehow end up in the config too + elif save_jinja_files and isinstance(self.chat_template, dict): + # New format for multiple templates is to save the default as chat_template.jinja + # and the other templates in the chat_templates/ directory + for template_name, template in self.chat_template.items(): + if template_name == "default": + with open(chat_template_file, "w", encoding="utf-8") as f: + f.write(self.chat_template["default"]) + logger.info(f"chat template saved in {chat_template_file}") + saved_raw_chat_template_files.append(chat_template_file) + else: + Path(chat_template_dir).mkdir(exist_ok=True) + template_filepath = os.path.join(chat_template_dir, f"{template_name}.jinja") + with open(template_filepath, "w", encoding="utf-8") as f: + f.write(template) + logger.info(f"chat template saved in {template_filepath}") + saved_raw_chat_template_files.append(template_filepath) + if "chat_template" in tokenizer_config: + tokenizer_config.pop("chat_template") # To ensure it doesn't somehow end up in the config too + elif isinstance(self.chat_template, dict): + # Legacy format for multiple templates: + # chat template dicts are saved to the config as lists of dicts with fixed key names. + tokenizer_config["chat_template"] = [{"name": k, "template": v} for k, v in self.chat_template.items()] + elif self.chat_template is not None: + # Legacy format for single templates: Just make them a key in tokenizer_config.json + tokenizer_config["chat_template"] = self.chat_template + return tokenizer_config, saved_raw_chat_template_files + + def parse_response( + self, + response: str | list[str | int | list[int]] | np.ndarray | torch.Tensor, + schema: list | dict | None = None, + ): + """ + Converts an output string created by generating text from a model into a parsed message dictionary. + This method is intended for use with chat models, and will read the tokenizer's `response_schema` attribute to + control parsing, although this can be overridden by passing a `response_schema` argument directly. + + Args: + response (`str`): + The output string generated by the model. This can be either a decoded string or list of strings, + or token IDs as a list/array. + schema (`Union[list, dict]`, *optional*): + A response schema that indicates the expected output format and how parsing should be performed. + If not provided, the tokenizer's `response_schema` attribute will be used. + """ + batched = ( + (isinstance(response, list) and not isinstance(response[0], int)) + or getattr(response, "ndim", 0) > 1 # For torch/numpy tensors + ) + + if schema is None: + if getattr(self, "response_schema", None) is None: + raise AttributeError("This tokenizer does not have a `response_schema` for parsing chat responses!") + schema = self.response_schema + if batched: + if not (isinstance(response, list) and isinstance(response[0], str)): + response = self.batch_decode(response) + return [recursive_parse(single_response, schema) for single_response in response] + else: + if not isinstance(response, str): + response = self.decode(response) + return recursive_parse(response, schema) + + +def get_fast_tokenizer_file(tokenization_files: list[str]) -> str: + """ + Get the tokenization file to use for this version of transformers. + + Args: + tokenization_files (`list[str]`): The list of available configuration files. + + Returns: + `str`: The tokenization file to use. + """ + tokenizer_files_map = {} + for file_name in tokenization_files: + search = _re_tokenizer_file.search(file_name) + if search is not None: + v = search.groups()[0] + tokenizer_files_map[v] = file_name + available_versions = sorted(tokenizer_files_map.keys()) + + # Defaults to FULL_TOKENIZER_FILE and then try to look at some newer versions. + tokenizer_file = FULL_TOKENIZER_FILE + transformers_version = version.parse(__version__) + for v in available_versions: + if version.parse(v) <= transformers_version: + tokenizer_file = tokenizer_files_map[v] + else: + # No point going further since the versions are sorted. + break + + return tokenizer_file + + +# Shared helper to locate a SentencePiece model file for a repo/path +def find_sentencepiece_model_file(pretrained_model_name_or_path, **kwargs): + """ + Find any .model file (SentencePiece model) in the model directory or Hub repo. + + Tries known filenames first ("tokenizer.model", "spm.model"), then scans local dir, + and as a last resort lists files on the Hub to find any .model. + + Returns the filename (str) relative to the repo root or directory if found, else None. + """ + from .utils.hub import has_file + + # Try common names first + for candidate in ("tokenizer.model", "spm.model"): + try: + if has_file( + pretrained_model_name_or_path, + candidate, + revision=kwargs.get("revision"), + token=kwargs.get("token"), + cache_dir=kwargs.get("cache_dir"), + local_files_only=kwargs.get("local_files_only", False), + ): + return candidate + except Exception: + # TODO: tighten to OSError / ProxyError + continue + + subfolder = kwargs.get("subfolder", "") + local_files_only = kwargs.get("local_files_only", False) + + # Local directory scan + if os.path.isdir(pretrained_model_name_or_path): + dir_path = ( + os.path.join(pretrained_model_name_or_path, subfolder) if subfolder else pretrained_model_name_or_path + ) + if os.path.isdir(dir_path): + for filename in os.listdir(dir_path): + if filename.endswith(".model"): + return filename if not subfolder else os.path.join(subfolder, filename) + + # Hub listing if allowed + if not local_files_only: + try: + from huggingface_hub import list_repo_tree + + entries = list_repo_tree( + repo_id=pretrained_model_name_or_path, + revision=kwargs.get("revision"), + path_in_repo=subfolder if subfolder else None, + recursive=False, + token=kwargs.get("token"), + ) + for entry in entries: + if entry.path.endswith(".model"): + return entry.path if not subfolder else entry.path.removeprefix(f"{subfolder}/") + except Exception as e: + # TODO: tighten exception class + logger.debug(f"Could not list Hub repository files: {e}") + + return None + + +def load_vocab_and_merges(pretrained_model_name_or_path, **kwargs): + """ + Resolve and load tokenizer vocabulary files from a repo/path. + + Priority order: + 1. Load ``vocab.json`` (WordLevel/WordPiece/BPE fast tokenizers) + 2. Load ``vocab.txt`` when only a WordPiece vocab is available + 3. Optionally load ``merges.txt`` (BPE tokenizers) + + Returns: + tuple (vocab: dict|None, merges: list[tuple[str,str]]|None, files_loaded: list[str]) + """ + files_loaded = [] + vocab = None + merges = None + try: + resolved_vocab_file = cached_file( + pretrained_model_name_or_path, + "vocab.json", + cache_dir=kwargs.get("cache_dir"), + force_download=kwargs.get("force_download", False), + proxies=kwargs.get("proxies"), + token=kwargs.get("token"), + revision=kwargs.get("revision"), + local_files_only=kwargs.get("local_files_only", False), + subfolder=kwargs.get("subfolder", ""), + ) + except Exception: + resolved_vocab_file = None + + if resolved_vocab_file is not None: + try: + with open(resolved_vocab_file, "r", encoding="utf-8") as vf: + vocab = json.load(vf) + files_loaded.append("vocab.json") + except Exception: + vocab = None + + # Fallback to vocab.txt (WordPiece-style vocabularies) + if vocab is None: + try: + resolved_vocab_txt = cached_file( + pretrained_model_name_or_path, + "vocab.txt", + cache_dir=kwargs.get("cache_dir"), + force_download=kwargs.get("force_download", False), + proxies=kwargs.get("proxies"), + token=kwargs.get("token"), + revision=kwargs.get("revision"), + local_files_only=kwargs.get("local_files_only", False), + subfolder=kwargs.get("subfolder", ""), + ) + except Exception: + resolved_vocab_txt = None + + if resolved_vocab_txt is not None: + try: + vocab = OrderedDict() + with open(resolved_vocab_txt, "r", encoding="utf-8") as vf: + for index, token in enumerate(vf): + token = token.rstrip("\n") + vocab[token] = index + files_loaded.append("vocab.txt") + except Exception: + vocab = None + + try: + resolved_merges_file = cached_file( + pretrained_model_name_or_path, + "merges.txt", + cache_dir=kwargs.get("cache_dir"), + force_download=kwargs.get("force_download", False), + proxies=kwargs.get("proxies"), + token=kwargs.get("token"), + revision=kwargs.get("revision"), + local_files_only=kwargs.get("local_files_only", False), + subfolder=kwargs.get("subfolder", ""), + ) + except Exception: + resolved_merges_file = None + + if resolved_merges_file is not None: + try: + merges = [] + with open(resolved_merges_file, "r", encoding="utf-8") as mf: + for line in mf: + line = line.strip() + if line and not line.startswith("#"): + parts = line.split() + if len(parts) == 2: + merges.append((parts[0], parts[1])) + files_loaded.append("merges.txt") + except Exception: + merges = None + + return vocab, merges, files_loaded + + +# To update the docstring, we need to copy the method, otherwise we change the original docstring. +PreTrainedTokenizerBase.push_to_hub = copy_func(PreTrainedTokenizerBase.push_to_hub) +if PreTrainedTokenizerBase.push_to_hub.__doc__ is not None: + PreTrainedTokenizerBase.push_to_hub.__doc__ = PreTrainedTokenizerBase.push_to_hub.__doc__.format( + object="tokenizer", object_class="AutoTokenizer", object_files="tokenizer files" + ) + + +def _get_prepend_scheme(add_prefix_space: bool, original_tokenizer) -> str: + if add_prefix_space: + prepend_scheme = "always" + if not getattr(original_tokenizer, "legacy", True): + prepend_scheme = "first" + else: + prepend_scheme = "never" + return prepend_scheme + + +def generate_merges(vocab, vocab_scores: dict[str, float] | None = None, skip_tokens: Collection[str] | None = None): + skip_tokens = set(skip_tokens) if skip_tokens is not None else set() + reverse = vocab_scores is not None + vocab_scores = dict(vocab_scores) if reverse else vocab + + merges = [] + for merge, piece_score in vocab_scores.items(): + if merge in skip_tokens: + continue + local = [] + for index in range(1, len(merge)): + piece_l, piece_r = merge[:index], merge[index:] + if piece_l in skip_tokens or piece_r in skip_tokens: + continue + if piece_l in vocab and piece_r in vocab: + local.append((piece_l, piece_r, piece_score)) + local = sorted(local, key=lambda x: (vocab[x[0]], vocab[x[1]])) + merges.extend(local) + + merges = sorted(merges, key=lambda val: (val[2], len(val[0]), len(val[1])), reverse=reverse) + merges = [(val[0], val[1]) for val in merges] + return merges diff --git a/third_party/transformers/src/transformers/tokenization_utils_tokenizers.py b/third_party/transformers/src/transformers/tokenization_utils_tokenizers.py new file mode 100644 index 0000000000000000000000000000000000000000..afca202127bed13358ddfcce99a1a1903fe72721 --- /dev/null +++ b/third_party/transformers/src/transformers/tokenization_utils_tokenizers.py @@ -0,0 +1,1388 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tokenization classes for fast tokenizers (provided by HuggingFace's tokenizers library). For slow (python) tokenizers +see tokenization_utils.py +""" + +import copy +import json +import os +from collections import defaultdict +from collections.abc import Iterable +from shutil import copyfile +from typing import Any + +import tokenizers.pre_tokenizers as pre_tokenizers_fast +from huggingface_hub import is_offline_mode +from tokenizers import AddedToken, processors +from tokenizers import Encoding as EncodingFast +from tokenizers import Tokenizer as TokenizerFast +from tokenizers.decoders import Decoder as DecoderFast +from tokenizers.models import BPE, Unigram +from tokenizers.trainers import BpeTrainer, UnigramTrainer, WordLevelTrainer, WordPieceTrainer + +from transformers.utils.hub import cached_file + +from .convert_slow_tokenizer import SpmConverter +from .integrations.ggml import convert_gguf_tokenizer +from .modeling_gguf_pytorch_utils import load_gguf_checkpoint +from .tokenization_utils_base import ( + INIT_TOKENIZER_DOCSTRING, + BatchEncoding, + PreTokenizedInput, + PreTrainedTokenizerBase, + TextInput, + TruncationStrategy, + generate_merges, +) +from .utils import PaddingStrategy, add_end_docstrings, logging + + +logger = logging.get_logger(__name__) + +# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file +TOKENIZER_FILE = "tokenizer.json" +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" +TIKTOKEN_VOCAB_FILE = "tokenizer.model" + +# Slow tokenizers have an additional added tokens files +ADDED_TOKENS_FILE = "added_tokens.json" + +INIT_TOKENIZER_DOCSTRING += """ + tokenizer_object ([`tokenizers.Tokenizer`]): + A [`tokenizers.Tokenizer`] object from 🤗 tokenizers to instantiate from. See [Using tokenizers from 🤗 + tokenizers](../fast_tokenizers) for more information. + tokenizer_file ([`str`]): + A path to a local JSON file representing a previously serialized [`tokenizers.Tokenizer`] object from 🤗 + tokenizers. +""" + +MODEL_TO_TRAINER_MAPPING = { + "BPE": BpeTrainer, + "Unigram": UnigramTrainer, + "WordLevel": WordLevelTrainer, + "WordPiece": WordPieceTrainer, +} + +VOCAB_FILES_NAMES = {"tokenizer_file": TOKENIZER_FILE, "vocab_file": TIKTOKEN_VOCAB_FILE} + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class TokenizersBackend(PreTrainedTokenizerBase): + """ + Base class for all fast tokenizers (wrapping HuggingFace tokenizers library). + + Inherits from [`~tokenization_utils_base.PreTrainedTokenizerBase`]. + + Handles all the shared methods for tokenization and special tokens, as well as methods for + downloading/caching/loading pretrained tokenizers, as well as adding tokens to the vocabulary. + + This class also contains the added tokens in a unified way on top of all tokenizers so we don't have to handle the + specific vocabulary augmentation methods of the various underlying dictionary structures (BPE, sentencepiece...). + """ + + vocab_files_names = VOCAB_FILES_NAMES + model = None + _tokenizer = None + + @classmethod + def convert_to_native_format(cls, trust_remote_code=False, **kwargs): + """s + Build a `tokenizers.Tokenizer` backend from the available serialization files (tokenizer.json, sentencepiece + models, tekken.json, vocab/merges). + """ + # Preserve kwargs for possible downstream use + local_kwargs = dict(kwargs) + fast_tokenizer_file = local_kwargs.pop("tokenizer_file", None) + + if ( + fast_tokenizer_file is not None + and os.path.isfile(fast_tokenizer_file) + and (cls is TokenizersBackend or "__init__" not in cls.__dict__ or trust_remote_code) + ): + local_kwargs["tokenizer_object"] = TokenizerFast.from_file(fast_tokenizer_file) + return local_kwargs + elif fast_tokenizer_file is not None and os.path.isfile(fast_tokenizer_file): + # we extract vocab/merges and pass decoder/pre_tokenizer/post_processor + # from the file so the reconstructed tokenizer matches the tokenizer.json + with open(fast_tokenizer_file, encoding="utf-8") as tokenizer_handle: + tokenizer_json = json.load(tokenizer_handle) + + # Build a minimal tokenizer (empty vocab/merges) to cheaply extract post_processor, + # padding and truncation as Rust objects — avoids parsing the full vocab via from_file. + # This optimization applies to BPE, WordPiece, and WordLevel only: + # - Unigram (SentencePiece) requires a non-empty vocab to initialize correctly in Rust + # (e.g. AlbertTokenizer, CamembertTokenizer, LlamaTokenizer, T5Tokenizer); passing an + # empty vocab causes "Unable to load vocab EmptyVocabulary". TODO: investigate if keeping + # just the UNK token is sufficient to make Unigram work with a minimal vocab. + # - Older tokenizer.json formats (e.g. XLNetTokenizer, DistilBertTokenizer) omit the + # "type" field in the "model" section, so we cannot determine the model type from JSON. + # In both cases we fall back to the original from_file path (no performance improvement). + model_type = tokenizer_json.get("model", {}).get("type") + if model_type not in (None, "Unigram"): + minimal_tokenizer_json = dict(tokenizer_json) + minimal_model = dict(tokenizer_json["model"]) + minimal_model["vocab"] = {} + if model_type == "BPE": + minimal_model["merges"] = [] + minimal_tokenizer_json["model"] = minimal_model + minimal_tokenizer_json["added_tokens"] = [] + tok_from_file = TokenizerFast.from_str(json.dumps(minimal_tokenizer_json)) + else: + tok_from_file = TokenizerFast.from_file(fast_tokenizer_file) + + local_kwargs["post_processor"] = tok_from_file.post_processor + local_kwargs["tokenizer_padding"] = tok_from_file.padding + local_kwargs["tokenizer_truncation"] = tok_from_file.truncation + # Preserve truncation and padding baked into tokenizer.json so that classes + # with a custom __init__ that rebuild the backend tokenizer from scratch + # can still access these settings. + if tok_from_file.truncation is not None: + local_kwargs["_json_truncation"] = tok_from_file.truncation + if tok_from_file.padding is not None: + local_kwargs["_json_padding"] = tok_from_file.padding + + # Extract precompiled SentencePiece charsmap from tokenizer.json normalizer + # when present (e.g. T5 tokenizers converted with SentencePiece >= 2.x). + normalizer_config = tokenizer_json.get("normalizer") + if normalizer_config: + if normalizer_config.get("type", None) == "Sequence": + normalizer_config = normalizer_config["normalizers"] + elif not isinstance(normalizer_config, list): + normalizer_config = [normalizer_config] + for normalizer in normalizer_config: + if normalizer.get("type") == "Precompiled" and "precompiled_charsmap" in normalizer: + import base64 + + local_kwargs["_spm_precompiled_charsmap"] = base64.b64decode( + normalizer["precompiled_charsmap"] + ) + break + + vocab = tokenizer_json.get("model", {}).get("vocab", None) + if cls.model is None: + if isinstance(vocab, list): + vocab = list(map(tuple, vocab)) # TODO just for now + elif cls.model.__name__ == "Unigram": + if isinstance(vocab, list) and vocab and isinstance(vocab[0], (list, tuple)): + vocab = [tuple(item) for item in vocab] + elif cls.model.__name__ == "WordLevel": + vocab = {token: i for i, token in enumerate(vocab)} + elif cls.model.__name__ == "BPE" or cls.model.__name__ == "WordPiece": + if isinstance(vocab, list): + vocab = {token[0] if isinstance(token, list) else token: i for i, token in enumerate(vocab)} + local_kwargs["vocab"] = vocab + + model_type = getattr(cls, "model", None) + if "merges" in tokenizer_json.get("model", {}) and (model_type and model_type.__name__ == "BPE"): + merges = tokenizer_json["model"]["merges"] + merges = [tuple(merge.split(" ")) if isinstance(merge, str) else tuple(merge) for merge in merges] + local_kwargs["merges"] = merges + + return local_kwargs + + vocab_file = local_kwargs.get("vocab_file") + merges_file = local_kwargs.get("merges_file") + vocab = local_kwargs.get("vocab") + merges = local_kwargs.get("merges") + + # Tekken converter (Mistral) + if isinstance(vocab_file, str) and vocab_file.endswith("tekken.json") and os.path.isfile(vocab_file): + from .convert_slow_tokenizer import MistralConverter + + local_kwargs["vocab"], local_kwargs["merges"] = MistralConverter( + vocab_file=vocab_file + ).extract_vocab_merges_from_model(vocab_file) + return local_kwargs + + # SentencePiece model (with TikToken fallback) + if isinstance(vocab_file, str) and os.path.isfile(vocab_file) and vocab_file.endswith(".model"): + try: + from .convert_slow_tokenizer import SentencePieceExtractor + + # 1. Extract vocab, merges, and spm_precompiled from the .model proto + extractor = SentencePieceExtractor(vocab_file) + local_kwargs = extractor.extract(cls.model, **local_kwargs) + + # 2. If a model-specific converter exists, use it. + try: + from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS + + converter_class = SLOW_TO_FAST_CONVERTERS.get(cls.__name__) + if converter_class is not None and hasattr(converter_class, "convert_from_spm"): + local_kwargs = converter_class.convert_from_spm(**local_kwargs) + except Exception as e: + logger.warning( + f"Could not reorder vocab using converter for {cls.__name__} due to {e}. Falling back to raw SentencePiece extraction." + ) + if hasattr(cls, "convert_from_spm_model"): + local_kwargs = cls.convert_from_spm_model(**local_kwargs) + + # 3. For non-model specific tokenizers (e.g. TokenizersBackend used + # for MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS), build a _tokenizer + # from the proto so normalizer/decoder are configured correctly. + if "tokenizer_object" not in local_kwargs and ( + cls is TokenizersBackend or "__init__" not in cls.__dict__ + ): + vocab = local_kwargs.pop("vocab", None) + merges = local_kwargs.pop("merges", None) + + # Replace placeholder tokens as specified in added_tokens_decoder + added_tokens_decoder = local_kwargs.get("added_tokens_decoder") or {} + if vocab is not None and added_tokens_decoder: + id_to_token = {token_id: token for token, token_id in vocab.items()} + for token_id, new_token in added_tokens_decoder.items(): + token_id = int(token_id) + new_token = str(new_token) + current_token = id_to_token.get(token_id) + if current_token and current_token != new_token and new_token not in vocab: + vocab[new_token] = vocab.pop(current_token) + id_to_token[token_id] = new_token + + tokenizer_object = SpmConverter.build_tokenizer_from_spm_proto( + proto=extractor.proto, + vocab=vocab, + merges=merges, + ) + if tokenizer_object is not None: + local_kwargs["tokenizer_object"] = tokenizer_object + # Set bos/eos tokens from proto spec if available. This is needed when + # building a tokenizer_object directly from a .model file because the + # tokenizer_object does not have bos/eos set. + proto_spec = extractor.proto.trainer_spec + if proto_spec.bos_id >= 0: + local_kwargs.setdefault("bos_token", proto_spec.bos_piece or "") + if proto_spec.eos_id >= 0: + local_kwargs.setdefault("eos_token", proto_spec.eos_piece or "") + if proto_spec.unk_id >= 0: + local_kwargs.setdefault("unk_token", proto_spec.unk_piece or "") + + except Exception as e: # TODO only catch deserialization error here! + logger.warning( + f"Could not extract SentencePiece model from {vocab_file} using sentencepiece library due to {e}. " + "Falling back to TikToken extractor." + ) + from .convert_slow_tokenizer import TikTokenConverter + + converter = TikTokenConverter( + vocab_file=vocab_file, extra_special_tokens=local_kwargs.get("extra_special_tokens") + ) + local_kwargs["tokenizer_object"] = converter.converted() + return local_kwargs + + # Fallback to standard vocab/merges files if they existed! + if vocab is None and isinstance(vocab_file, str) and os.path.isfile(vocab_file): + local_kwargs["vocab"] = vocab_file + vocab = local_kwargs["vocab"] + if merges is None and isinstance(merges_file, str) and os.path.isfile(merges_file): + local_kwargs["merges"] = merges_file + merges = local_kwargs["merges"] + + # Generate merges automatically when not provided for BPE tokenizers + if merges is None and cls.model is not None and cls.model.__name__ == "BPE" and isinstance(vocab, dict): + # Gather special tokens from kwargs to skip in merge generation + def _iter_special_tokens(values: Iterable[Any]) -> list[str]: + collected: list[str] = [] + for val in values: + if val is None: + continue + if isinstance(val, (list, tuple)): + collected.extend(_iter_special_tokens(val)) + else: + collected.append(str(val)) + return collected + + special_tokens_keys = [ + "pad_token", + "unk_token", + "bos_token", + "eos_token", + "sep_token", + "cls_token", + "mask_token", + "additional_special_tokens", + "extra_special_tokens", + ] + skip_tokens: set[str] = set() + for key in special_tokens_keys: + if key in local_kwargs: + skip_tokens.update(_iter_special_tokens([local_kwargs[key]])) + + merges = generate_merges(vocab, skip_tokens=skip_tokens) + local_kwargs["merges"] = merges + return local_kwargs + + def __init__(self, *args, **kwargs): + # Truncation/padding dicts extracted from tokenizer.json by convert_to_native_format + # when a class with a custom __init__ rebuilds the backend tokenizer from scratch. + _json_truncation = kwargs.pop("_json_truncation", None) + _json_padding = kwargs.pop("_json_padding", None) + # Precompiled SentencePiece charsmap is already used by model-specific tokenizers + # (before calling super().__init__) and should not be stored in `init_kwargs` to keep the tokenizer serializable. + kwargs.pop("_spm_precompiled_charsmap", None) + + tokenizer_object = kwargs.pop("tokenizer_object", None) + gguf_file = kwargs.pop("gguf_file", None) + fast_tokenizer_file = kwargs.pop("tokenizer_file", None) + # Note: added_tokens_decoder is NOT popped - it's passed to super().__init__() for processing + added_tokens_decoder = kwargs.get("added_tokens_decoder", {}) + # Store add_prefix_space before super().__init__() to ensure it's not overridden + add_prefix_space = kwargs.get("add_prefix_space", False) + vocab_file = kwargs.get("vocab_file") + + vocab = kwargs.get("vocab") + merges = kwargs.get("merges") + + fast_tokenizer = None + if tokenizer_object is not None: + fast_tokenizer = copy.deepcopy(tokenizer_object) + elif fast_tokenizer_file is not None and os.path.isfile(fast_tokenizer_file): + # We have a serialization from tokenizers which let us directly build the backend + fast_tokenizer = TokenizerFast.from_file(fast_tokenizer_file) + elif gguf_file is not None: + # We need to convert a slow tokenizer to build the backend + gguf_path = cached_file(kwargs.get("name_or_path", ""), gguf_file, **kwargs) + gguf_param = load_gguf_checkpoint(gguf_path) + architecture = gguf_param["config"]["model_type"] + tokenizer_dict = gguf_param["tokenizer"] + tokenizer_config = gguf_param["tokenizer_config"] + fast_tokenizer, additional_kwargs = convert_gguf_tokenizer(architecture, tokenizer_dict) + kwargs.update(tokenizer_config) + if len(additional_kwargs) > 0: + kwargs.update(additional_kwargs) + elif self._tokenizer is None and vocab is not None: + # Build from vocab/merges extracted by convert_to_native_format + if merges is not None: + vocab_dict = vocab if isinstance(vocab, dict) else {w: i for i, (w, _) in enumerate(vocab)} + fast_tokenizer = TokenizerFast(BPE(vocab=vocab_dict, merges=merges, fuse_unk=True, dropout=None)) + elif isinstance(vocab, dict): + fast_tokenizer = TokenizerFast(BPE(vocab=vocab, merges=[], fuse_unk=True, dropout=None)) + elif isinstance(vocab, list) and vocab and isinstance(vocab[0], (tuple, list)): + fast_tokenizer = TokenizerFast(Unigram(vocab=vocab, unk_id=kwargs.get("unk_id", 0))) + elif self._tokenizer is None: + raise ValueError( + "Couldn't instantiate the backend tokenizer from one of: \n" + "(1) a `tokenizers` library serialization file, \n" + "(2) a slow tokenizer instance to convert or \n" + "(3) an equivalent slow tokenizer class to instantiate and convert. \n" + "You need to have sentencepiece or tiktoken installed to convert a slow tokenizer to a fast one." + ) + # Only set defaults when creating TokenizersBackend from scratch + if fast_tokenizer_file is None and tokenizer_object is None and self._tokenizer is None: + kwargs.setdefault("bos_token", "") + kwargs.setdefault("eos_token", "") + + if fast_tokenizer is not None: + self._tokenizer = fast_tokenizer + + if self._tokenizer is None: + raise ValueError("The backend tokenizer is not correctly initialized.") + + _truncation = kwargs.pop("tokenizer_truncation", None) or self._tokenizer.truncation or _json_truncation + if _truncation is not None: + self._tokenizer.enable_truncation(**_truncation) + kwargs.setdefault("max_length", _truncation["max_length"]) + kwargs.setdefault("truncation_side", _truncation["direction"]) + kwargs.setdefault("stride", _truncation["stride"]) + kwargs.setdefault("truncation_strategy", _truncation["strategy"]) + else: + self._tokenizer.no_truncation() + + _padding = kwargs.pop("tokenizer_padding", None) or self._tokenizer.padding or _json_padding + if _padding is not None: + self._tokenizer.enable_padding(**_padding) + kwargs.setdefault("pad_token", _padding["pad_token"]) + kwargs.setdefault("pad_token_type_id", _padding["pad_type_id"]) + kwargs.setdefault("padding_side", _padding["direction"]) + kwargs.setdefault("max_length", _padding["length"]) + kwargs.setdefault("pad_to_multiple_of", _padding["pad_to_multiple_of"]) + + # Set backend to "tokenizers" if not already set + if "backend" not in kwargs: + kwargs["backend"] = "tokenizers" + + explicit_bos_eos_in_kwargs = "add_bos_token" in kwargs or "add_eos_token" in kwargs + self._add_bos_token = kwargs.get("add_bos_token", False) + self._add_eos_token = kwargs.get("add_eos_token", False) + if post_processor := kwargs.pop("post_processor", None): # most reliable way to get the post-processor + self._tokenizer.post_processor = post_processor + self._should_update_post_processor = explicit_bos_eos_in_kwargs or self._tokenizer.post_processor is None + # We call this after having initialized the backend tokenizer because we update it. + super().__init__(**kwargs) + + if vocab_file is not None: + self.vocab_file = vocab_file + # Ensure add_prefix_space is set correctly after parent init + self.add_prefix_space = add_prefix_space + self._tokenizer.encode_special_tokens = self.split_special_tokens + + added_tokens_decoder_hash = {hash(repr(token)) for token in self.added_tokens_decoder} + tokens_to_add = [ + token + for index, token in sorted(added_tokens_decoder.items(), key=lambda x: x[0]) + if hash(repr(token)) not in added_tokens_decoder_hash + ] + encoder = list(self.added_tokens_encoder.keys()) + [str(token) for token in tokens_to_add] + # if some of the special tokens are not already in the tokenizer, add them + # V5: Check both named special tokens and extra special tokens + # Iterate over _special_tokens_map to preserve AddedToken properties (lstrip, rstrip, etc.) + for special_token_value in self._special_tokens_map.values(): + if special_token_value is None: + continue + if str(special_token_value) not in encoder and special_token_value not in tokens_to_add: + tokens_to_add.append(special_token_value) + + # Also check extra special tokens + for token in self._extra_special_tokens: + if str(token) not in encoder and token not in tokens_to_add: + tokens_to_add.append(token) + + if len(tokens_to_add) > 0: + tokens = [] + all_named_tokens = [str(t) for t in self._special_tokens_map.values() if t] + for token in tokens_to_add: + if isinstance(token, str): + # Convert string to AddedToken, assuming it's special + token = AddedToken(token, special=True) + elif isinstance(token, AddedToken): + # Ensure the special flag is set correctly for special tokens + if not token.special and str(token) in all_named_tokens: + token.special = True + tokens.append(token) + if tokens: + # These tokens are from the special tokens map + self.add_tokens(tokens) + + try: + vocab_size = self._tokenizer.get_vocab_size() + except NotImplementedError: + vocab_size = 0 + + # Optionally patches mistral tokenizers with wrong regex + if vocab_size > 100000 and getattr(self._tokenizer, "pre_tokenizer", None) is not None: + kwargs.pop("tokenizer", None) + self._tokenizer = self._patch_mistral_regex( + self._tokenizer, + self.init_kwargs.get("name_or_path", None), + init_kwargs=self.init_kwargs, + fix_mistral_regex=kwargs.pop("fix_mistral_regex", None), + **kwargs, + ) + + self._should_update_post_processor = ( + self._should_update_post_processor or self._tokenizer.post_processor is None + ) + if self._should_update_post_processor: + self.update_post_processor() + + @property + def is_fast(self) -> bool: + return True + + @property + def can_save_slow_tokenizer(self) -> bool: + """ + `bool`: Whether or not the slow tokenizer can be saved. For a sentencepiece based slow tokenizer, this + can only be `True` if the original `"sentencepiece.model"` was not deleted. + """ + if "vocab_file" in self.vocab_files_names and self.vocab_files_names["vocab_file"].endswith(".model"): + if hasattr(self, "vocab_file") and self.vocab_file: + # If the vocab file is a sentencepiece model, we can save it + return os.path.isfile(self.vocab_file) + return False + else: + return True + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + if not os.path.isdir(save_directory): + logger.error(f"Vocabulary path ({save_directory}) should be a directory") + return + out_vocab_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file,) + + def update_post_processor(self): + """ + Updates the underlying post processor with the current `bos_token` and `eos_token`. + """ + bos = self.bos_token + bos_token_id = self.bos_token_id + if bos is None and self.add_bos_token: + self.add_bos_token = False + + eos = self.eos_token + eos_token_id = self.eos_token_id + if eos is None and self.add_eos_token: + self.add_eos_token = False + + single = f"{(bos + ':0 ') if self.add_bos_token else ''}$A:0{(' ' + eos + ':0') if self.add_eos_token else ''}" + pair = f"{single}{(' ' + bos + ':1') if self.add_bos_token else ''} $B:1{(' ' + eos + ':1') if self.add_eos_token else ''}" + + special_tokens = [] + if self.add_bos_token: + special_tokens.append((bos, bos_token_id)) + if self.add_eos_token: + special_tokens.append((eos, eos_token_id)) + self._tokenizer.post_processor = processors.TemplateProcessing( + single=single, pair=pair, special_tokens=special_tokens + ) + + @property + def add_eos_token(self): + return getattr(self, "_add_eos_token", False) + + @property + def add_bos_token(self): + return getattr(self, "_add_bos_token", False) + + @add_eos_token.setter + def add_eos_token(self, value): + object.__setattr__(self, "_add_eos_token", value) + self.update_post_processor() + + @add_bos_token.setter + def add_bos_token(self, value): + object.__setattr__(self, "_add_bos_token", value) + self.update_post_processor() + + def _post_init(self): + """ + Post-initialization hook that runs after the tokenizer is fully set up. + This is called by from_pretrained() after loading the tokenizer, which allows + us to add any special tokens that may have been passed as AddedToken objects. + + Child classes should call super()._post_init() if they override this method. + """ + tokens_to_add = [] + # V5: Check named special tokens + for token_value in self._special_tokens_map.values(): + if token_value is None: + continue + if isinstance(token_value, AddedToken): + tokens_to_add.append(token_value) + elif isinstance(token_value, str): + tokens_to_add.append(AddedToken(token_value, special=True, normalized=False)) + + # V5: Check extra special tokens + for token in self._extra_special_tokens: + if isinstance(token, AddedToken): + tokens_to_add.append(token) + elif isinstance(token, str): + tokens_to_add.append(AddedToken(token, special=True, normalized=False)) + + if tokens_to_add: + # Ensure special tokens are added as such to the backend + self.add_tokens(tokens_to_add, special_tokens=True) + + if getattr(self, "_should_update_post_processor", True) or self._tokenizer.post_processor is None: + self.update_post_processor() + + @property + def vocab_size(self) -> int: + """ + `int`: Size of the base vocabulary (without the added tokens). + """ + return self._tokenizer.get_vocab_size(with_added_tokens=False) + + def get_vocab(self) -> dict[str, int]: + return self._tokenizer.get_vocab(with_added_tokens=True) + + @property + def vocab(self) -> dict[str, int]: + return self.get_vocab() + + @property + def added_tokens_encoder(self) -> dict[str, int]: + """ + Returns the sorted mapping from string to index. The added tokens encoder is cached for performance + optimisation in `self._added_tokens_encoder` for the slow tokenizers. + """ + return {k.content: v for v, k in sorted(self.added_tokens_decoder.items(), key=lambda item: item[0])} + + @property + def added_tokens_decoder(self) -> dict[int, AddedToken]: + """ + Returns the added tokens in the vocabulary as a dictionary of index to AddedToken. + + Returns: + `dict[str, int]`: The added tokens. + """ + return self._tokenizer.get_added_tokens_decoder() + + # BC v5: expose ``_added_tokens_encoder`` / ``_added_tokens_decoder`` attrs for custom tokenizers that expect + # them from slow tokenizers. Only supports read, not write (won't sync to Rust backend, use add_tokens() instead + _added_tokens_encoder = added_tokens_encoder + _added_tokens_decoder = added_tokens_decoder + + def get_added_vocab(self) -> dict[str, int]: + """ + Returns the added tokens in the vocabulary as a dictionary of token to index. + + Returns: + `dict[str, int]`: The added tokens. + """ + return {k.content: v for v, k in sorted(self.added_tokens_decoder.items(), key=lambda item: item[0])} + + def __bool__(self) -> bool: + """ + Returns True, to avoid expensive `assert tokenizer` gotchas. + """ + return True + + def __len__(self) -> int: + """ + Size of the full vocabulary with the added tokens. + """ + return self._tokenizer.get_vocab_size(with_added_tokens=True) + + @property + def backend_tokenizer(self) -> TokenizerFast: + """ + `tokenizers.implementations.BaseTokenizer`: The Rust tokenizer used as a backend. + """ + return self._tokenizer + + @property + def decoder(self) -> DecoderFast: + """ + `tokenizers.decoders.Decoder`: The Rust decoder for this tokenizer. + """ + return self._tokenizer.decoder + + def _convert_encoding( + self, + encoding: EncodingFast, + 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, + ) -> tuple[dict[str, Any], list[EncodingFast]]: + """ + Convert the encoding representation (from low-level HuggingFace tokenizer output) to a python Dict and a list + of encodings, take care of building a batch from overflowing tokens. + + Overflowing tokens are converted to additional examples (like batches) so the output values of the dict are + lists (overflows) of lists (tokens). + + Output shape: (overflows, sequence length) + """ + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + if return_overflowing_tokens and encoding.overflowing is not None: + encodings = [encoding] + encoding.overflowing + else: + encodings = [encoding] + + encoding_dict = defaultdict(list) + for e in encodings: + encoding_dict["input_ids"].append(e.ids) + + if return_token_type_ids: + encoding_dict["token_type_ids"].append(e.type_ids) + if return_attention_mask: + encoding_dict["attention_mask"].append(e.attention_mask) + if return_special_tokens_mask: + encoding_dict["special_tokens_mask"].append(e.special_tokens_mask) + if return_offsets_mapping: + encoding_dict["offset_mapping"].append(e.offsets) + if return_length: + encoding_dict["length"].append(len(e.ids)) + + return encoding_dict, encodings + + def _convert_token_to_id_with_added_voc(self, token: str) -> int: + index = self._tokenizer.token_to_id(token) + if index is None: + return self.unk_token_id + return index + + def _convert_id_to_token(self, index: int) -> str | None: + return self._tokenizer.id_to_token(int(index)) + + def _add_tokens(self, new_tokens: list[str | AddedToken], special_tokens=False) -> int: + if special_tokens: + return self._tokenizer.add_special_tokens(new_tokens) + + return self._tokenizer.add_tokens(new_tokens) + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + """ + Returns the number of added tokens when encoding a sequence with special tokens. + + + + This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put + this inside your training loop. + + + + Args: + pair (`bool`, *optional*, defaults to `False`): + Whether the number of added tokens should be computed in the case of a sequence pair or a single + sequence. + + Returns: + `int`: Number of special tokens added to sequences. + """ + return self._tokenizer.num_special_tokens_to_add(pair) + + def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `list[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `list[str]`: The decoded token(s). + """ + if isinstance(ids, int): + return self._tokenizer.id_to_token(ids) + tokens = [] + # self.all_special_ids is an @property which may be slow, so only compute it once before the loop + ids_to_skip = set(self.all_special_ids) if skip_special_tokens else set() + for index in ids: + index = int(index) + if index in ids_to_skip: + continue + tokens.append(self._tokenizer.id_to_token(index)) + return tokens + + def tokenize(self, text: str, pair: str | None = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + return self._encode_plus(text=text, text_pair=pair, add_special_tokens=add_special_tokens, **kwargs).tokens() + + def set_truncation_and_padding( + self, + padding_strategy: PaddingStrategy, + truncation_strategy: TruncationStrategy, + max_length: int, + stride: int, + pad_to_multiple_of: int | None, + padding_side: str | None, + ): + """ + Define the truncation and the padding strategies for fast tokenizers (provided by HuggingFace tokenizers + library) and restore the tokenizer settings afterwards. + + The provided tokenizer has no padding / truncation strategy before the managed section. If your tokenizer set a + padding / truncation strategy before, then it will be reset to no padding / truncation when exiting the managed + section. + + Args: + padding_strategy ([`~utils.PaddingStrategy`]): + The kind of padding that will be applied to the input + truncation_strategy ([`~tokenization_utils_base.TruncationStrategy`]): + The kind of truncation that will be applied to the input + max_length (`int`): + The maximum size of a sequence. + stride (`int`): + The stride to use when handling overflow. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). + padding_side (`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. + """ + _truncation = self._tokenizer.truncation + _padding = self._tokenizer.padding + # Set truncation and padding on the backend tokenizer + if truncation_strategy == TruncationStrategy.DO_NOT_TRUNCATE: + if _truncation is not None: + self._tokenizer.no_truncation() + else: + target = { + "max_length": max_length, + "stride": stride, + "strategy": truncation_strategy.value, + "direction": self.truncation_side, + } + + # _truncation might contain more keys that the target `transformers` + # supports. Use only the target keys to trigger `enable_truncation`. + # This should enable this code to works on various `tokenizers` + # targets. + if _truncation is None: + current = None + else: + current = {k: _truncation.get(k, None) for k in target} + + if current != target: + self._tokenizer.enable_truncation(**target) + + if padding_strategy == PaddingStrategy.DO_NOT_PAD: + if _padding is not None: + self._tokenizer.no_padding() + else: + length = max_length if padding_strategy == PaddingStrategy.MAX_LENGTH else None + target = { + "length": length, + "direction": padding_side if padding_side is not None else self.padding_side, + "pad_id": self.pad_token_id, + "pad_token": self.pad_token, + "pad_type_id": self.pad_token_type_id, + "pad_to_multiple_of": pad_to_multiple_of, + } + if _padding != target: + self._tokenizer.enable_padding(**target) + + def _encode_plus( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + text_pair: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | 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, + is_split_into_words: bool = False, + 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, + split_special_tokens: bool | None = None, + **kwargs, + ) -> BatchEncoding: + # Input validation (from _call_one) + def _is_valid_text_input(t): + if isinstance(t, str): + return True + elif isinstance(t, (list, tuple)): + if len(t) == 0: + return True + elif isinstance(t[0], str): + return True + elif isinstance(t[0], (list, tuple)): + if len(t[0]) == 0 or isinstance(t[0][0], str): + return True + elif isinstance(t[0][0], (list, tuple)): + return len(t[0][0]) == 0 or isinstance(t[0][0][0], str) + else: + return False + else: + return False + else: + return False + + if not _is_valid_text_input(text): + raise ValueError( + "text input must be of type `str` (single example), `list[str]` (batch or single pretokenized example) " + "or `list[list[str]]` (batch of pretokenized examples) or `list[tuple[list[str], list[str]]]` (batch of pretokenized sequence pairs)." + ) + + if text_pair is not None and not _is_valid_text_input(text_pair): + raise ValueError( + "text input must be of type `str` (single example), `list[str]` (batch or single pretokenized example) " + "or `list[list[str]]` (batch of pretokenized examples) or `list[tuple[list[str], list[str]]]` (batch of pretokenized sequence pairs)." + ) + + # Batch detection (from _call_one) + if is_split_into_words: + is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)) + else: + is_batched = isinstance(text, (list, tuple)) + + if is_batched: + # Batch validation + if isinstance(text_pair, str): + raise TypeError( + "when tokenizing batches of text, `text_pair` must be a list or tuple with the same length as" + " `text`." + ) + 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 + else: + # Single input - convert to batch format + batch_text_or_text_pairs = [(text, text_pair)] if text_pair else [text] + + # Set tokenizer configuration (from _batch_encode_plus) + if not isinstance(batch_text_or_text_pairs, (tuple, list)): + raise TypeError( + f"batch_text_or_text_pairs has to be a list or a tuple (got {type(batch_text_or_text_pairs)})" + ) + + 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, + ) + + # Use self.split_special_tokens as default if not explicitly provided + if split_special_tokens is None: + split_special_tokens = self.split_special_tokens + + if self._tokenizer.encode_special_tokens != split_special_tokens: + self._tokenizer.encode_special_tokens = split_special_tokens + + # Direct rust backend call + encodings = self._tokenizer.encode_batch( + batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + is_pretokenized=is_split_into_words, + ) + + # Convert encodings to BatchEncoding format + 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=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + ) + for encoding in encodings + ] + + # Convert the output to have dict[list] from list[dict] + 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 + 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) + + batched_output = BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors) + + # If single input, remove the batch dimension (unless returning overflowing tokens) + if not is_batched and 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, + ) + + return batched_output + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return ( + self.backend_tokenizer.decoder.decode(tokens) + if self.backend_tokenizer.decoder is not None + else " ".join(tokens) + ) + + def _decode( + self, + token_ids: int | list[int], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool | None = None, + **kwargs, + ) -> str: + # Removed: use_source_tokenizer parameter (unused) + kwargs.pop("use_source_tokenizer", None) # Pop if present to avoid errors + + if isinstance(token_ids, int): + token_ids = [token_ids] + if isinstance(token_ids, dict): + token_ids = token_ids["input_ids"] + text = self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens) + + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + text = self.clean_up_tokenization(text) + + return text + + def _save_pretrained( + self, + save_directory: str | os.PathLike, + file_names: tuple[str, ...], + legacy_format: bool | None = None, + filename_prefix: str | None = None, + ) -> tuple[str, ...]: + save_directory = str(save_directory) + + tokenizer_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + TOKENIZER_FILE + ) + self.backend_tokenizer.save(tokenizer_file) + file_names = file_names + (tokenizer_file,) + + return file_names + + def train_new_from_iterator( + self, + text_iterator, + vocab_size, + length=None, + new_special_tokens=None, + special_tokens_map=None, + **kwargs, + ): + """ + Trains a tokenizer on a new corpus with the same defaults (in terms of special tokens or tokenization pipeline) + as the current one. + + Args: + text_iterator (generator of `list[str]`): + The training corpus. Should be a generator of batches of texts, for instance a list of lists of texts + if you have everything in memory. + vocab_size (`int`): + The size of the vocabulary you want for your tokenizer. + length (`int`, *optional*): + The total number of sequences in the iterator. This is used to provide meaningful progress tracking + new_special_tokens (list of `str` or `AddedToken`, *optional*): + A list of new special tokens to add to the tokenizer you are training. + special_tokens_map (`dict[str, str]`, *optional*): + If you want to rename some of the special tokens this tokenizer uses, pass along a mapping old special + token name to new special token name in this argument. + kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the trainer from the 🤗 Tokenizers library. + + Returns: + [`PreTrainedTokenizerFast`]: A new tokenizer of the same type as the original one, trained on + `text_iterator`. + + """ + tokenizer_json = json.loads(self._tokenizer.to_str()) + # Remove added tokens for now (uses IDs of tokens) + added_tokens = tokenizer_json.pop("added_tokens") + # Remove post processor for now (uses IDs of tokens) + post_processor = tokenizer_json.pop("post_processor") + + unk_token = None + # Remove vocab + if tokenizer_json["model"]["type"] == "BPE": + tokenizer_json["model"]["vocab"] = {} + tokenizer_json["model"]["merges"] = [] + elif tokenizer_json["model"]["type"] == "Unigram": + if tokenizer_json["model"]["unk_id"] is not None: + unk_id = tokenizer_json["model"]["unk_id"] + unk_token = tokenizer_json["model"]["vocab"][unk_id][0] + if special_tokens_map is not None and unk_token in special_tokens_map: + unk_token = special_tokens_map[unk_token] + tokenizer_json["model"]["unk_id"] = 0 + tokenizer_json["model"]["vocab"] = [[unk_token, 0.0]] + elif tokenizer_json["model"]["type"] in ["WordLevel", "WordPiece"]: + tokenizer_json["model"]["vocab"] = {} + else: + raise ValueError( + f"This method does not support this type of tokenizer (found {tokenizer_json['model']['type']}) " + "only BPE, Unigram, WordLevel and WordPiece." + ) + + if ( + special_tokens_map is not None + and "unk_token" in tokenizer_json["model"] + and tokenizer_json["model"]["unk_token"] in special_tokens_map + ): + tokenizer_json["model"]["unk_token"] = special_tokens_map[tokenizer_json["model"]["unk_token"]] + + tokenizer = TokenizerFast.from_str(json.dumps(tokenizer_json)) + + # Get the special tokens from the current tokenizer if none are specified. + special_tokens = [] + for added_token in added_tokens: + special = added_token.pop("special", None) + _ = added_token.pop("id", None) + if tokenizer_json["model"]["type"] != "Unigram" and not special: + continue + if special_tokens_map is not None and added_token["content"] in special_tokens_map: + added_token["content"] = special_tokens_map[added_token["content"]] + special_tokens.append(AddedToken(**added_token)) + + if new_special_tokens is not None: + special_tokens.extend(new_special_tokens) + + # Trainer needs to know the end of word / continuing subword thingies in BPE + if ( + tokenizer_json["model"]["type"] == "BPE" + and "continuing_subword_prefix" not in kwargs + and tokenizer_json["model"]["continuing_subword_prefix"] is not None + ): + kwargs["continuing_subword_prefix"] = tokenizer_json["model"]["continuing_subword_prefix"] + if ( + tokenizer_json["model"]["type"] == "BPE" + and "end_of_word_suffix" not in kwargs + and tokenizer_json["model"]["end_of_word_suffix"] is not None + ): + kwargs["end_of_word_suffix"] = tokenizer_json["model"]["end_of_word_suffix"] + if tokenizer_json["model"]["type"] == "Unigram" and unk_token is not None: + kwargs["unk_token"] = unk_token + if tokenizer_json["pre_tokenizer"] is not None: + if ( + tokenizer_json["pre_tokenizer"]["type"] == "ByteLevel" + or tokenizer_json["pre_tokenizer"]["type"] == "Sequence" + and "pretokenizers" in tokenizer_json["pre_tokenizer"] + and any( + pretokenizer["type"] == "ByteLevel" + for pretokenizer in tokenizer_json["pre_tokenizer"]["pretokenizers"] + ) + ): + kwargs["initial_alphabet"] = pre_tokenizers_fast.ByteLevel.alphabet() + + trainer_class = MODEL_TO_TRAINER_MAPPING[tokenizer_json["model"]["type"]] + trainer = trainer_class(vocab_size=vocab_size, special_tokens=special_tokens, **kwargs) + tokenizer.train_from_iterator(text_iterator, length=length, trainer=trainer) + + if post_processor is not None: + trained_tokenizer_json = json.loads(tokenizer.to_str()) + # Almost done, we just have to adjust the token IDs in the post processor + if "special_tokens" in post_processor: + for key in post_processor["special_tokens"]: + tokens = post_processor["special_tokens"][key]["tokens"] + if special_tokens_map is not None: + tokens = [special_tokens_map.get(token, token) for token in tokens] + post_processor["special_tokens"][key]["tokens"] = tokens + for token in tokens: + token_id = tokenizer.token_to_id(token) + if token_id is None: + raise ValueError( + "Attempted to set a token in the post processor that does not exist in the mapping" + ) + + post_processor["special_tokens"][key]["ids"] = [tokenizer.token_to_id(token) for token in tokens] + + for special_token in ["cls", "sep"]: + if special_token in post_processor: + token, _ = post_processor[special_token] + if special_tokens_map is not None and token in special_tokens_map: + token = special_tokens_map[token] + token_id = tokenizer.token_to_id(token) + if token_id is None: + raise ValueError( + "Attempted to set a token in the post processor that does not exist in the mapping" + ) + post_processor[special_token] = [token, token_id] + + trained_tokenizer_json["post_processor"] = post_processor + tokenizer = TokenizerFast.from_str(json.dumps(trained_tokenizer_json)) + + kwargs = self.init_kwargs.copy() + # V5: Map pad/cls/mask token at the Transformers level (named tokens only) + for token in PreTrainedTokenizerBase.SPECIAL_TOKENS_ATTRIBUTES: + if getattr(self, token) is not None: + special_token = getattr(self, token) + if special_tokens_map is not None and special_token in special_tokens_map: + special_token = special_tokens_map[special_token] + + special_token_full = self._special_tokens_map.get(token, None) + if isinstance(special_token_full, AddedToken): + # Create an added token with the same parameters except the content + kwargs[token] = AddedToken( + special_token, + single_word=special_token_full.single_word, + lstrip=special_token_full.lstrip, + rstrip=special_token_full.rstrip, + normalized=special_token_full.normalized, + special=True, + ) + else: + kwargs[token] = special_token + + # V5: Handle extra special tokens + extra_special_tokens = self.extra_special_tokens.copy() if self.extra_special_tokens else [] + if new_special_tokens is not None: + extra_special_tokens.extend(new_special_tokens) + if len(extra_special_tokens) > 0: + kwargs["extra_special_tokens"] = extra_special_tokens + + # Always try to pass tokenizer_object in kwargs first (standard TokenizersBackend usage) + # If the class creates its own tokenizer and passes it explicitly to super().__init__(), + # this will cause a TypeError, which we catch and handle by removing tokenizer_object + # from kwargs and setting _tokenizer directly after initialization. + kwargs["tokenizer_object"] = tokenizer + try: + return self.__class__(**kwargs) + except TypeError as e: + # Check if the error is due to multiple values for tokenizer_object + if "multiple values for keyword argument 'tokenizer_object'" in str(e): + # Class creates its own tokenizer and passes it explicitly (like LayoutLMv3Tokenizer) + # Remove tokenizer_object from kwargs and set _tokenizer directly + kwargs.pop("tokenizer_object", None) + new_tokenizer = self.__class__(**kwargs) + new_tokenizer._tokenizer = tokenizer + return new_tokenizer + else: + # Some other TypeError, re-raise it + raise + + @classmethod + def _patch_mistral_regex( + cls, + tokenizer, + pretrained_model_name_or_path, + token=None, + cache_dir=None, + local_files_only=False, + _commit_hash=None, + is_local=False, + init_kwargs=None, + fix_mistral_regex=None, + **kwargs, + ): + """ + Patches mistral related tokenizers with incorrect regex if detected + 1) Local file with an associated config saved next to it + >> Model type one of the mistral models (on older versions) + 2) Remote models on the hub from official mistral models + >> Tags including `base_model:.*mistralai` + """ + import re + + from huggingface_hub import model_info + from packaging import version + + from transformers.utils.hub import cached_file + + def is_base_mistral(model_id: str) -> bool: + model = model_info(model_id) + if model.tags is not None: + if re.search("base_model:.*mistralai", "".join(model.tags)): + return True + return False + + if is_offline_mode(): + is_local = True + + if pretrained_model_name_or_path is not None and ( + is_local or (not is_local and is_base_mistral(pretrained_model_name_or_path)) + ): + _config_file = cached_file( + pretrained_model_name_or_path, + "config.json", + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + _commit_hash=_commit_hash, + ) + + # Detected using a (local) mistral tokenizer + mistral_config_detected = False + if _config_file is not None: + with open(_config_file, encoding="utf-8") as f: + _config = json.load(f) + transformers_version = _config.get("transformers_version") + transformers_model_type = _config.get("model_type") + + # Detect if we can skip the mistral fix by + # a) having a non-mistral tokenizer + # b) fixed version of transformers + if transformers_version and version.parse(transformers_version) <= version.parse("4.57.2"): + if ( + is_local + and transformers_model_type is not None + and transformers_model_type + not in [ + "mistral", + "mistral3", + "voxtral", + "ministral", + "pixtral", + ] + ): + return tokenizer + elif transformers_version and version.parse(transformers_version) > version.parse("4.57.3"): + return tokenizer + + mistral_config_detected = True + + if mistral_config_detected or (not is_local and is_base_mistral(pretrained_model_name_or_path)): + # Expose the `fix_mistral_regex` flag on the tokenizer when provided, even if no correction is applied. + if init_kwargs and "fix_mistral_regex" in init_kwargs: + setattr(tokenizer, "fix_mistral_regex", init_kwargs["fix_mistral_regex"]) + + # only warn if its not explicitly passed + if fix_mistral_regex is None and not getattr(tokenizer, "fix_mistral_regex", False): + setattr(tokenizer, "fix_mistral_regex", False) + logger.warning( + f"The tokenizer you are loading from '{pretrained_model_name_or_path}'" + f" with an incorrect regex pattern: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/discussions/84#69121093e8b480e709447d5e." + " This will lead to incorrect tokenization. You should set the `fix_mistral_regex=True` flag when loading this tokenizer to fix this issue." + ) + elif fix_mistral_regex is True or getattr(tokenizer, "fix_mistral_regex", False): + setattr(tokenizer, "fix_mistral_regex", True) + import tokenizers + + split_pretokenizer = tokenizers.pre_tokenizers.Split( + pattern=tokenizers.Regex( + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+" + ), + behavior="isolated", + ) + current_pretokenizer = tokenizer.backend_tokenizer.pre_tokenizer + # Check if it's already a Sequence + if isinstance(current_pretokenizer, tokenizers.pre_tokenizers.Sequence): + # Replace the first element (the Split pattern) + tokenizer.backend_tokenizer.pre_tokenizer[0] = split_pretokenizer + else: + # Replace Metaspace with ByteLevel when adding Split, as Metaspace(split=False) doesn't + # work correctly with the Split pre-tokenizer and causes spaces to be lost during encoding + if isinstance(current_pretokenizer, tokenizers.pre_tokenizers.Metaspace): + current_pretokenizer = tokenizers.pre_tokenizers.ByteLevel( + add_prefix_space=False, use_regex=False + ) + + # Not a Sequence, so create one with Split + current pretokenizer + tokenizer.backend_tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Sequence( + [ + split_pretokenizer, + current_pretokenizer, + ] + ) + + return tokenizer + + +# Backward-compatible alias: allow referring to TokenizersBackend as PreTrainedTokenizerFast +PreTrainedTokenizerFast = TokenizersBackend diff --git a/third_party/transformers/src/transformers/trainer_seq2seq.py b/third_party/transformers/src/transformers/trainer_seq2seq.py new file mode 100644 index 0000000000000000000000000000000000000000..ada588adbd215cb6b2c1c1f970a48ffbb0275a5e --- /dev/null +++ b/third_party/transformers/src/transformers/trainer_seq2seq.py @@ -0,0 +1,392 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +from collections.abc import Callable +from copy import deepcopy +from pathlib import Path +from typing import TYPE_CHECKING, Any, Optional, Union + +import torch +from torch import nn +from torch.utils.data import Dataset + +from .generation.configuration_utils import GenerationConfig +from .integrations.deepspeed import is_deepspeed_zero3_enabled +from .integrations.fsdp import is_fsdp_managed_module +from .trainer import Trainer +from .utils import is_datasets_available, logging + + +if torch.distributed.is_available(): + from torch.distributed.fsdp import FullyShardedDataParallel + +if is_datasets_available(): + import datasets + +if TYPE_CHECKING: + from torch.utils.data import IterableDataset + + from .data.data_collator import DataCollator + from .feature_extraction_utils import FeatureExtractionMixin + from .image_processing_utils import BaseImageProcessor + from .modeling_utils import PreTrainedModel + from .processing_utils import ProcessorMixin + from .tokenization_utils_base import PreTrainedTokenizerBase + from .trainer_callback import TrainerCallback + from .trainer_utils import EvalPrediction, PredictionOutput + from .training_args import TrainingArguments + + +logger = logging.get_logger(__name__) + + +class Seq2SeqTrainer(Trainer): + def __init__( + self, + model: Union["PreTrainedModel", nn.Module] | None = None, + args: Optional["TrainingArguments"] = None, + data_collator: Optional["DataCollator"] = None, + train_dataset: Union[Dataset, "IterableDataset", "datasets.Dataset"] | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: Union[ + "PreTrainedTokenizerBase", "BaseImageProcessor", "FeatureExtractionMixin", "ProcessorMixin" + ] + | None = None, + model_init: Callable[[], "PreTrainedModel"] | None = None, + compute_loss_func: Callable | None = None, + compute_metrics: Callable[["EvalPrediction"], dict] | None = None, + callbacks: list["TrainerCallback"] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + ): + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Override self.model.generation_config if a GenerationConfig is specified in args. + # Priority: args.generation_config > model.generation_config > default GenerationConfig. + if self.args.generation_config is not None: + gen_config = self.load_generation_config(self.args.generation_config) + self.model.generation_config = gen_config + + @staticmethod + def load_generation_config(gen_config_arg: str | GenerationConfig) -> GenerationConfig: + """ + Loads a `~generation.GenerationConfig` from the `Seq2SeqTrainingArguments.generation_config` arguments. + + Args: + gen_config_arg (`str` or [`~generation.GenerationConfig]`): + `Seq2SeqTrainingArguments.generation_config` argument. + + Returns: + A `~generation.GenerationConfig`. + """ + + # GenerationConfig provided, nothing to do + if isinstance(gen_config_arg, GenerationConfig): + gen_config = deepcopy(gen_config_arg) + else: + # str or Path + pretrained_model_name = Path(gen_config_arg) if isinstance(gen_config_arg, str) else gen_config_arg + config_file_name = None + + # Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL + # This step is required in order to determine config_file_name + if pretrained_model_name.is_file(): + config_file_name = pretrained_model_name.name + pretrained_model_name = pretrained_model_name.parent + # dir path + elif pretrained_model_name.is_dir(): + pass + # model id or URL + else: + pretrained_model_name = gen_config_arg + + gen_config = GenerationConfig.from_pretrained(pretrained_model_name, config_file_name) + + # Strict validation to fail early. `GenerationConfig.save_pretrained()`, run at the end of training, throws + # an exception if there are warnings at validation time. + try: + gen_config.validate(strict=True) + except ValueError as exc: + raise ValueError(str(exc) + "\n\nFix these issues to train your model.") + + return gen_config + + def evaluate( + self, + eval_dataset: Dataset | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + **gen_kwargs, + ) -> dict[str, float]: + """ + Run evaluation and returns metrics. + + The calling script will be responsible for providing a method to compute metrics, as they are task-dependent + (pass it to the init `compute_metrics` argument). + + You can also subclass and override this method to inject custom behavior. + + Args: + eval_dataset (`Dataset`, *optional*): + Pass a dataset if you wish to override `self.eval_dataset`. If it is an [`~datasets.Dataset`], columns + not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__` + method. + ignore_keys (`list[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + Returns: + A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The + dictionary also contains the epoch number which comes from the training state. + """ + + gen_kwargs = gen_kwargs.copy() + + # Use legacy argument setting if a) the option is not explicitly passed; and b) the argument is set in the + # training args + if ( + gen_kwargs.get("max_length") is None + and gen_kwargs.get("max_new_tokens") is None + and self.args.generation_max_length is not None + ): + gen_kwargs["max_length"] = self.args.generation_max_length + if gen_kwargs.get("num_beams") is None and self.args.generation_num_beams is not None: + gen_kwargs["num_beams"] = self.args.generation_num_beams + # We don't want to drop samples in general + self.gather_function = self.accelerator.gather + self._gen_kwargs = gen_kwargs + return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def predict( + self, + test_dataset: Dataset, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "test", + **gen_kwargs, + ) -> "PredictionOutput": + """ + Run prediction and returns predictions and potential metrics. + + Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method + will also return metrics, like in `evaluate()`. + + Args: + test_dataset (`Dataset`): + Dataset to run the predictions on. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. Has to implement the method `__len__` + ignore_keys (`list[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + + + If your predictions or labels have different sequence lengths (for instance because you're doing dynamic + padding in a token classification task) the predictions will be padded (on the right) to allow for + concatenation into one array. The padding index is -100. + + + + Returns: *NamedTuple* A namedtuple with the following keys: + + - predictions (`np.ndarray`): The predictions on `test_dataset`. + - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some). + - metrics (`dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained + labels). + """ + + gen_kwargs = gen_kwargs.copy() + + # Use legacy argument setting if a) the option is not explicitly passed; and b) the argument is set in the + # training args + if ( + gen_kwargs.get("max_length") is None + and gen_kwargs.get("max_new_tokens") is None + and self.args.generation_max_length is not None + ): + gen_kwargs["max_length"] = self.args.generation_max_length + if gen_kwargs.get("num_beams") is None and self.args.generation_num_beams is not None: + gen_kwargs["num_beams"] = self.args.generation_num_beams + self.gather_function = self.accelerator.gather + self._gen_kwargs = gen_kwargs + + return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def prediction_step( + self, + model: nn.Module, + inputs: dict[str, torch.Tensor | Any], + prediction_loss_only: bool, + ignore_keys: list[str] | None = None, + **gen_kwargs, + ) -> tuple[float | None, torch.Tensor | None, torch.Tensor | None]: + """ + Perform an evaluation step on `model` using `inputs`. + + Subclass and override to inject custom behavior. + + Args: + model (`nn.Module`): + The model to evaluate. + inputs (`dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + prediction_loss_only (`bool`): + Whether or not to return the loss only. + gen_kwargs: + Additional `generate` specific kwargs. + + Return: + tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and + labels (each being optional). + """ + + if not self.args.predict_with_generate or prediction_loss_only: + return super().prediction_step( + model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys + ) + + has_labels = "labels" in inputs + inputs = self._prepare_inputs(inputs) + + # Priority (handled in generate): + # non-`None` gen_kwargs > model.generation_config > default GenerationConfig() + if len(gen_kwargs) == 0 and hasattr(self, "_gen_kwargs"): + gen_kwargs = self._gen_kwargs.copy() + if "num_beams" in gen_kwargs and gen_kwargs["num_beams"] is None: + gen_kwargs.pop("num_beams") + if "max_length" in gen_kwargs and gen_kwargs["max_length"] is None: + gen_kwargs.pop("max_length") + + default_synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self.model) + gen_kwargs["synced_gpus"] = gen_kwargs.get("synced_gpus", default_synced_gpus) + + generation_inputs = inputs.copy() + # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate + # (otherwise, it would continue generating from the padded `decoder_input_ids`) + if ( + "labels" in generation_inputs + and "decoder_input_ids" in generation_inputs + and generation_inputs["labels"].shape == generation_inputs["decoder_input_ids"].shape + ): + generation_inputs = { + k: v for k, v in inputs.items() if k not in ("decoder_input_ids", "decoder_attention_mask") + } + + summon_full_params_context = ( + FullyShardedDataParallel.summon_full_params(self.model) + if torch.distributed.is_available() and isinstance(self.model, FullyShardedDataParallel) + else contextlib.nullcontext() + ) + + with summon_full_params_context: + generated_tokens = self.model.generate(**generation_inputs, **gen_kwargs) + + # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop + # TODO: remove this hack when the legacy code that initializes generation_config from a model config is + # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183 + if self.model.generation_config._from_model_config: + self.model.generation_config._from_model_config = False + + # Retrieves GenerationConfig from model.generation_config + # Update with defaults because earlier the generation config used to be init + # with default values. Now we init it with `None` and keep defaults for BC + gen_config = self.model.generation_config + default_gen_config = gen_config._get_default_generation_params() + gen_config.update(**default_gen_config, defaults_only=True) + # in case the batch is shorter than max length, the output should be padded + if generated_tokens.shape[-1] < gen_config.max_length: + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_length) + elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1: + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_new_tokens + 1) + + with torch.no_grad(): + if has_labels: + with self.compute_loss_context_manager(): + outputs = model(**inputs) + if self.label_smoother is not None: + loss = self.label_smoother(outputs, inputs["labels"]).detach().mean() + else: + loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).detach().mean() + else: + loss = None + + if self.args.prediction_loss_only: + return loss, None, None + + if has_labels: + labels = inputs["labels"] + if labels.shape[-1] < gen_config.max_length: + labels = self._pad_tensors_to_max_len(labels, gen_config.max_length) + elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1: + labels = self._pad_tensors_to_max_len(labels, gen_config.max_new_tokens + 1) + else: + labels = None + + return loss, generated_tokens, labels + + def _pad_tensors_to_max_len(self, tensor, max_length): + if self.processing_class is not None and hasattr(self.processing_class, "pad_token_id"): + # If PAD token is not defined at least EOS token has to be defined + pad_token_id = ( + self.processing_class.pad_token_id + if self.processing_class.pad_token_id is not None + else self.processing_class.eos_token_id + ) + else: + if getattr(self.model.config, "pad_token_id", None) is not None: + pad_token_id = self.model.config.pad_token_id + else: + raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors") + + padded_tensor = pad_token_id * torch.ones( + (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device + ) + padded_tensor[:, : tensor.shape[-1]] = tensor + return padded_tensor diff --git a/third_party/transformers/src/transformers/training_args.py b/third_party/transformers/src/transformers/training_args.py new file mode 100644 index 0000000000000000000000000000000000000000..d4d4d49931cbe1af3117223a8c0d7be65442286c --- /dev/null +++ b/third_party/transformers/src/transformers/training_args.py @@ -0,0 +1,2829 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import json +import math +import os +import warnings +from dataclasses import asdict, dataclass, field, fields +from datetime import timedelta +from enum import Enum +from functools import cached_property +from typing import Any + +from .debug_utils import DebugOption +from .trainer_utils import ( + FSDPOption, + HubStrategy, + IntervalStrategy, + SaveStrategy, + SchedulerType, +) +from .utils import ( + ACCELERATE_MIN_VERSION, + ExplicitEnum, + is_accelerate_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_torch_available, + is_torch_bf16_gpu_available, + is_torch_cuda_available, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_mps_available, + is_torch_musa_available, + is_torch_neuron_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torch_xpu_available, + logging, + requires_backends, +) +from .utils.generic import strtobool +from .utils.import_utils import enable_tf32, is_optimum_neuron_available + + +logger = logging.get_logger(__name__) +log_levels = logging.get_log_levels_dict().copy() +trainer_log_levels = dict(**log_levels, passive=-1) + +if is_torch_available(): + import torch + import torch.distributed as dist + +if is_accelerate_available(): + from accelerate.state import AcceleratorState, PartialState + from accelerate.utils import DistributedType + + from .trainer_pt_utils import AcceleratorConfig + +if is_accelerate_available("1.10.1"): + from accelerate.parallelism_config import ParallelismConfig +else: + ParallelismConfig = Any + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + +if is_torch_neuroncore_available(check_device=False): + # torchrun support + # https://github.com/pytorch/xla/pull/3609 + if os.environ.get("TORCHELASTIC_RUN_ID"): + if is_optimum_neuron_available(): + logger.info( + "Make sure that you are performing the training with the NeuronTrainer from optimum[neuron], this " + "will fail otherwise." + ) + else: + logger.warning( + "Please use the NeuronTrainer from optimum[neuron] instead of the Transformers library to perform " + "training on AWS Trainium instances. More information here: " + "https://github.com/huggingface/optimum-neuron" + ) + import torch_xla.distributed.xla_backend as xbn + + if not isinstance(dist.group.WORLD, xbn.ProcessGroupXla): + dist.init_process_group(backend="xla") + if not isinstance(dist.group.WORLD, xbn.ProcessGroupXla): + raise AssertionError("Failed to initialize torch.distributed process group using XLA backend.") + + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + + smp.init() + + +class OptimizerNames(ExplicitEnum): + """ + Stores the acceptable string identifiers for optimizers. + """ + + ADAMW_TORCH = "adamw_torch" + ADAMW_TORCH_FUSED = "adamw_torch_fused" + ADAMW_TORCH_XLA = "adamw_torch_xla" + ADAMW_TORCH_NPU_FUSED = "adamw_torch_npu_fused" + ADAMW_APEX_FUSED = "adamw_apex_fused" + ADAFACTOR = "adafactor" + ADAMW_ANYPRECISION = "adamw_anyprecision" + ADAMW_TORCH_4BIT = "adamw_torch_4bit" + ADAMW_TORCH_8BIT = "adamw_torch_8bit" + ADEMAMIX = "ademamix" + SGD = "sgd" + ADAGRAD = "adagrad" + ADAMW_BNB = "adamw_bnb_8bit" + ADAMW_8BIT = "adamw_8bit" # just an alias for adamw_bnb_8bit + ADEMAMIX_8BIT = "ademamix_8bit" + LION_8BIT = "lion_8bit" + LION = "lion_32bit" + PAGED_ADAMW = "paged_adamw_32bit" + PAGED_ADAMW_8BIT = "paged_adamw_8bit" + PAGED_ADEMAMIX = "paged_ademamix_32bit" + PAGED_ADEMAMIX_8BIT = "paged_ademamix_8bit" + PAGED_LION = "paged_lion_32bit" + PAGED_LION_8BIT = "paged_lion_8bit" + RMSPROP = "rmsprop" + RMSPROP_BNB = "rmsprop_bnb" + RMSPROP_8BIT = "rmsprop_bnb_8bit" + RMSPROP_32BIT = "rmsprop_bnb_32bit" + GALORE_ADAMW = "galore_adamw" + GALORE_ADAMW_8BIT = "galore_adamw_8bit" + GALORE_ADAFACTOR = "galore_adafactor" + GALORE_ADAMW_LAYERWISE = "galore_adamw_layerwise" + GALORE_ADAMW_8BIT_LAYERWISE = "galore_adamw_8bit_layerwise" + GALORE_ADAFACTOR_LAYERWISE = "galore_adafactor_layerwise" + LOMO = "lomo" + ADALOMO = "adalomo" + GROKADAMW = "grokadamw" + SCHEDULE_FREE_RADAM = "schedule_free_radam" + SCHEDULE_FREE_ADAMW = "schedule_free_adamw" + SCHEDULE_FREE_SGD = "schedule_free_sgd" + APOLLO_ADAMW = "apollo_adamw" + APOLLO_ADAMW_LAYERWISE = "apollo_adamw_layerwise" + STABLE_ADAMW = "stable_adamw" + + +def _convert_str_dict(passed_value: dict): + "Safely checks that a passed value is a dictionary and converts any string values to their appropriate types." + for key, value in passed_value.items(): + if isinstance(value, dict): + passed_value[key] = _convert_str_dict(value) + elif isinstance(value, str): + # First check for bool and convert + if value.lower() in ("true", "false"): + passed_value[key] = value.lower() == "true" + # Check for digit + elif value.isdigit(): + passed_value[key] = int(value) + elif value.replace(".", "", 1).isdigit(): + passed_value[key] = float(value) + + return passed_value + + +@dataclass +class TrainingArguments: + """ + Configuration class for controlling all aspects of model training with the Trainer. + TrainingArguments centralizes all hyperparameters, optimization settings, logging preferences, and infrastructure choices needed for training. + + [`HfArgumentParser`] can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + output_dir (`str`, *optional*, defaults to `"trainer_output"`): + The output directory where the model predictions and checkpoints will be written. + + > Training Duration and Batch Size + + per_device_train_batch_size (`int`, *optional*, defaults to 8): + The batch size *per device*. The **global batch size** is computed as: + `per_device_train_batch_size * number_of_devices` in multi-GPU or distributed setups. + num_train_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents of + the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + Overrides `num_train_epochs`. If set to a positive number, the total number of training steps to perform. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + + > Learning Rate & Scheduler + + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate for the optimizer. This is typically the peak learning rate when using a scheduler with warmup. + lr_scheduler_type (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`): + The learning rate scheduler type to use. See [`SchedulerType`] for all possible values. Common choices: + - "linear" = [`get_linear_schedule_with_warmup`] + - "cosine" = [`get_cosine_schedule_with_warmup`] + - "constant" = [`get_constant_schedule`] + - "constant_with_warmup" = [`get_constant_schedule_with_warmup`] + lr_scheduler_kwargs (`dict` or `str`, *optional*, defaults to `None`): + The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values. + warmup_steps (`int` or `float`, *optional*, defaults to 0): + Number of steps for a linear warmup from 0 to `learning_rate`. Warmup helps stabilize training in the initial phase. Can be: + - An integer: exact number of warmup steps + - A float in range [0, 1): interpreted as ratio of total training steps + + > Optimizer + + optim (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"` (for torch>=2.8 `"adamw_torch_fused"`)): + The optimizer to use. Common options: + - `"adamw_torch"`: PyTorch's AdamW (recommended default) + - `"adamw_torch_fused"`: Fused AdamW kernel + - `"adamw_hf"`: HuggingFace's AdamW implementation + - `"sgd"`: Stochastic Gradient Descent with momentum + - `"adafactor"`: Memory-efficient optimizer for large models + - `"adamw_8bit"`: 8-bit AdamW (requires bitsandbytes) + See [`OptimizerNames`] for the complete list. + optim_args (`str`, *optional*): + Optional arguments that are supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore. + weight_decay (`float`, *optional*, defaults to 0): + Weight decay coefficient applied by the optimizer (not the loss function). Adds L2 + regularization to prevent overfitting by penalizing large weights. Automatically + excluded from bias and LayerNorm parameters. Typical values: 0.01 (standard), 0.1 + (stronger regularization), 0.0 (no regularization). + adam_beta1 (`float`, *optional*, defaults to 0.9): + The exponential decay rate for the first moment estimates (momentum) in Adam-based + optimizers. Controls how much history of gradients to retain. + adam_beta2 (`float`, *optional*, defaults to 0.999): + The exponential decay rate for the second moment estimates (variance) in Adam-based + optimizers. Controls adaptive learning rate scaling. + adam_epsilon (`float`, *optional*, defaults to 1e-8): + Epsilon value for numerical stability in Adam-based optimizers. Prevents division by + zero in the denominator of the update rule. + optim_target_modules (`Union[str, list[str]]`, *optional*): + The target modules to optimize, i.e. the module names that you would like to train. + Currently used for the [GaLore algorithm](https://huggingface.co/papers/2403.03507) and [APOLLO algorithm](https://huggingface.co/papers/2412.05270). + See [GaLore implementation](https://github.com/jiaweizzhao/GaLore) and [APOLLO implementation](https://github.com/zhuhanqing/APOLLO) for more details. + You need to make sure to pass a valid GaLore or APOLLO optimizer, e.g., one of: "apollo_adamw", "galore_adamw", "galore_adamw_8bit", "galore_adafactor" and make sure that the target modules are `nn.Linear` modules only. + + > Regularization & Training Stability + + gradient_accumulation_steps (`int`, *optional*, defaults to 1): + Number of update steps to accumulate gradients before performing a backward/update pass. + Simulates larger batch sizes without additional memory. Effective batch size = + `per_device_train_batch_size × num_devices × gradient_accumulation_steps`. + > [!TIP] + > When using gradient accumulation, one "step" is counted as one step with a backward pass. Therefore, logging, evaluation, and saving will occur every `gradient_accumulation_steps × xxx_step` training examples. + average_tokens_across_devices (`bool`, *optional*, defaults to `True`): + Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize + num_tokens_in_batch for precise loss calculation. Reference: + https://github.com/huggingface/transformers/issues/34242 + max_grad_norm (`float`, *optional*, defaults to 1.0): + Maximum gradient norm for gradient clipping. Applied after backward pass, before + optimizer step. Prevents gradient explosion by scaling down gradients when their global + norm exceeds this threshold. Set to 0 to disable clipping. Typical values: + 1.0 (standard), 0.5 (more conservative), 5.0 (less aggressive). + label_smoothing_factor (`float`, *optional*, defaults to 0.0): + Label smoothing factor to prevent overconfidence. Replaces hard 0/1 targets with soft + targets: 0 becomes `ε/num_labels` and 1 becomes `1 - ε + ε/num_labels`, where + ε = `label_smoothing_factor`. Zero means no smoothing. Typical range: 0.0 to 0.1. + + > Mixed Precision Training + + bf16 (`bool`, *optional*, defaults to `False`): + Enable bfloat16 (BF16) mixed precision training + Generally preferred over FP16 due to better numerical stability and no loss scaling required. + fp16 (`bool`, *optional*, defaults to `False`): + Enable float16 (FP16) mixed precision training. + Consider using BF16 instead if your hardware supports it. + bf16_full_eval (`bool`, *optional*, defaults to `False`): + Use full BF16 precision for evaluation (not just mixed precision). Faster and saves + memory but may affect metric values slightly. Only applies during evaluation. + fp16_full_eval (`bool`, *optional*, defaults to `False`): + Use full FP16 precision for evaluation (not just mixed precision). Faster and saves + memory but may affect metric values slightly. Only applies during evaluation. + tf32 (`bool`, *optional*): + Enable TensorFloat-32 (TF32) mode on Ampere and newer GPUs. TF32 uses 19-bit precision + for matrix multiplications (instead of FP32's 23-bit), providing up to 8x speedup with + negligible accuracy loss. Default depends on PyTorch version. See + [TF32 docs](https://huggingface.co/docs/transformers/perf_train_gpu_one#tf32). + + > Gradient Checkpointing + + gradient_checkpointing (`bool`, *optional*, defaults to `False`): + Enable gradient checkpointing to trade compute for memory. Reduces memory usage by + clearing activations during forward pass and recomputing them during backward pass. + Enables training larger models or batch sizes at the cost of ~20% slower training. + gradient_checkpointing_kwargs (`dict`, *optional*, defaults to `None`): + Keyword arguments passed to `gradient_checkpointing_enable()`. + + > Compilation + + torch_compile (`bool`, *optional*, defaults to `False`): + Compile the model using PyTorch 2.0's `torch.compile()` for faster training. Can provide + 20-50% speedup with no code changes. Uses default compilation settings unless + `torch_compile_backend` or `torch_compile_mode` are specified. + torch_compile_backend (`str`, *optional*): + Backend for `torch.compile()`. If set, automatically enables `torch_compile`. Options + include `"inductor"` (default), `"aot_eager"`, `"cudagraphs"`. Backends vary by PyTorch + version - see PyTorch docs for available options. + torch_compile_mode (`str`, *optional*): + Compilation mode for `torch.compile()`. If set, automatically enables `torch_compile`. + Options: `"default"`, `"reduce-overhead"` (minimize Python overhead), `"max-autotune"` + (aggressive optimization, slower compile time). + + > Kernels + + use_liger_kernel (`bool`, *optional*, defaults to `False`): + Enable [Liger Kernel](https://github.com/linkedin/Liger-Kernel) optimizations. Increases + multi-GPU throughput by ~20% and reduces memory usage by ~60%. Works with Flash Attention, + FSDP, and DeepSpeed. Currently supports Llama, Mistral, Mixtral, and Gemma models. + liger_kernel_config (`Optional[dict]`, *optional*): + Configuration for Liger Kernel. Passed as kwargs to `_apply_liger_kernel_to_instance()`. + Options typically include: `"rope"`, `"swiglu"`, `"cross_entropy"`, + `"fused_linear_cross_entropy"`, `"rms_norm"`. If `None`, uses default configuration. + + > Additional Optimizations + + use_cache (`bool`, *optional*, defaults to `False`): + Whether or not to enable cache for the model. For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`. + neftune_noise_alpha (`Optional[float]`): + If not `None`, this will activate NEFTune noise embeddings. This can drastically improve model performance + for instruction fine-tuning. Check out the [original paper](https://huggingface.co/papers/2310.05914) and the + [original code](https://github.com/neelsjain/NEFTune). Support transformers `PreTrainedModel` and also + `PeftModel` from peft. The original paper used values in the range [5.0, 15.0]. + torch_empty_cache_steps (`int`, *optional*): + Number of steps to wait before calling `torch..empty_cache()`. If left unset or set to None, cache will not be emptied. + This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about [10% slower performance](https://github.com/huggingface/transformers/issues/31372). + auto_find_batch_size (`bool`, *optional*, defaults to `False`) + Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding + CUDA Out-of-Memory errors. + + > Logging & Monitoring Training + + logging_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The logging strategy to adopt during training. Possible values are: + - `"no"`: No logging is done during training. + - `"epoch"`: Logging is done at the end of each epoch. + - `"steps"`: Logging is done every `logging_steps`. + logging_steps (`int` or `float`, *optional*, defaults to 500): + Number of update steps between two logs if `logging_strategy="steps"`. Should be an integer or a float in + range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps. + logging_first_step (`bool`, *optional*, defaults to `False`): + Whether to log the first `global_step` or not. + log_on_each_node (`bool`, *optional*, defaults to `True`): + In multinode distributed training, whether to log using `log_level` once per node, or only on the main + node. + logging_nan_inf_filter (`bool`, *optional*, defaults to `True`): + Filter out NaN and Inf losses when logging. If `True`, replaces NaN/Inf losses with the + average of recent valid losses. Does not affect gradient computation, only logging. + include_num_input_tokens_seen (`Optional[Union[str, bool]]`, *optional*, defaults to "no"): + Whether to track the number of input tokens seen. Must be one of ["all", "non_padding", "no"] or a boolean value which map to "all" or "no". + May be slower in distributed training as gather operations must be called. + + > Logging + + log_level (`str`, *optional*, defaults to `passive`): + Logging level for the main process. Options: `"debug"`, `"info"`, `"warning"`, `"error"`, + `"critical"`, or `"passive"` (doesn't change the current Transformers logging level, + which defaults to `"warning"`) + log_level_replica (`str`, *optional*, defaults to `"warning"`): + Logging level for replica processes in distributed training. Same options as `log_level`. + disable_tqdm (`bool`, *optional*): + Disable tqdm progress bars. Defaults to `True` if `log_level` is warning or lower, `False` otherwise. + + > Experiment Tracking Integration + + report_to (`str` or `list[str]`, *optional*, defaults to `"none"`): + The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, + `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"swanlab"`, + `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"` + for no integrations. + run_name (`str`, *optional*): + A descriptor for the run. Typically used for [trackio](https://github.com/gradio-app/trackio), + [wandb](https://www.wandb.com/), [mlflow](https://www.mlflow.org/), [comet](https://www.comet.com/site) and + [swanlab](https://swanlab.cn) logging. + project (`str`, *optional*, defaults to `"huggingface"`): + The name of the project to use for logging. Currently, only used by Trackio. + trackio_space_id (`str` or `None`, *optional*, defaults to `"trackio"`): + The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like + `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` in which case the Space will be + created in the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory. + Note that this Space will be public unless you set `hub_private_repo=True` or your organization's default + is to create private Spaces." + + > Evaluation + + eval_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`): + When to run evaluation. Options: + - `"no"`: No evaluation during training + - `"steps"`: Evaluate every `eval_steps` + - `"epoch"`: Evaluate at the end of each epoch + eval_steps (`int` or `float`, *optional*): + Number of update steps between two evaluations if `eval_strategy="steps"`. Will default to the same + value as `logging_steps` if not set. Should be an integer or a float in range `[0,1)`. If smaller than 1, + will be interpreted as ratio of total training steps. + eval_delay (`float`, *optional*): + Number of epochs or steps to wait for before the first evaluation can be performed, depending on the + eval_strategy. + per_device_eval_batch_size (`int`, *optional*, defaults to 8): + The batch size per device accelerator core/CPU for evaluation. + prediction_loss_only (`bool`, *optional*, defaults to `False`): + When performing evaluation and generating predictions, only returns the loss. + eval_on_start (`bool`, *optional*, defaults to `False`): + Whether to perform a evaluation step (sanity check) before the training to ensure the validation steps works correctly. + eval_do_concat_batches (`bool`, *optional*, defaults to `True`): + Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, + will instead store them as lists, with each batch kept separate. + eval_use_gather_object (`bool`, *optional*, defaults to `False`): + Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices. This should only be enabled if users are not just returning tensors, and this is actively discouraged by PyTorch. + This is useful when the labels structure is non standard, like in computer vision tasks. + eval_accumulation_steps (`int`, *optional*): + Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If + left unset, the whole predictions are accumulated on the device accelerator before being moved to the CPU (faster but + requires more memory). + + > Metrics Computation + + include_for_metrics (`list[str]`, *optional*, defaults to `[]`): + Include additional data in the `compute_metrics` function if needed for metrics computation. + Possible options to add to `include_for_metrics` list: + - `"inputs"`: Input data passed to the model, intended for calculating input dependent metrics. + - `"loss"`: Loss values computed during evaluation, intended for calculating loss dependent metrics. + batch_eval_metrics (`bool`, *optional*, defaults to `False`): + If set to `True`, evaluation will call compute_metrics at the end of each batch to accumulate statistics + rather than saving all eval logits in memory. When set to `True`, you must pass a compute_metrics function + that takes a boolean argument `compute_result`, which when passed `True`, will trigger the final global + summary statistics from the batch-level summary statistics you've accumulated over the evaluation set. + + > Checkpointing & Saving + + save_only_model (`bool`, *optional*, defaults to `False`): + Save only model weights, not optimizer/scheduler/RNG state. Significantly reduces + checkpoint size but prevents resuming training from the checkpoint. Use when you only + need the trained model for inference, not continued training. + You can only load the model using `from_pretrained` with this option set to `True`. + save_strategy (`str` or [`~trainer_utils.SaveStrategy`], *optional*, defaults to `"steps"`): + The checkpoint save strategy to adopt during training. Possible values are: + - `"no"`: No save is done during training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps`. + - `"best"`: Save is done whenever a new `best_metric` is achieved. + save_steps (`int` or `float`, *optional*, defaults to 500): + Number of updates steps before two checkpoint saves if `save_strategy="steps"`. Should be an integer or a + float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps. + save_on_each_node (`bool`, *optional*, defaults to `False`): + When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on + the main one. + This should not be activated when the different nodes use the same storage as the files will be saved with + the same names for each node. + save_total_limit (`int`, *optional*): + Maximum number of checkpoints to keep. Deletes older checkpoints in `output_dir`. When + `load_best_model_at_end=True`, the best checkpoint is always retained plus the most + recent ones. For example, `save_total_limit=5` keeps the 4 most recent plus the best + enable_jit_checkpoint (`bool`, *optional*, defaults to `False`): + Enable Just-In-Time checkpointing on SIGTERM signal for graceful termination on + preemptible workloads. **Important**: Configure your orchestrator's graceful shutdown + period to allow sufficient time. For Kubernetes, set `terminationGracePeriodSeconds` + (default 30s is usually insufficient). For Slurm, use `--signal=USR1@`. + Required grace period ≥ longest iteration time + checkpoint save time. + + > Hugging Face Hub Integration + + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push the model to the Hub every time the model is saved. If this is activated, + `output_dir` will begin a git directory synced with the repo (determined by `hub_model_id`) and the content + will be pushed each time a save is triggered (depending on your `save_strategy`). Calling + [`~Trainer.save_model`] will also trigger a push. + hub_token (`str`, *optional*): + The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with + `hf auth login`. + hub_private_repo (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's + default is private. This value is ignored if the repo already exists. If reporting to Trackio with + deployment to Hugging Face Spaces enabled, the same logic determines whether the Space is private. + hub_model_id (`str`, *optional*): + The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in + which case the model will be pushed in your namespace. Otherwise it should be the whole repository name, + for instance `"user_name/model"`, which allows you to push to an organization you are a member of with + `"organization_name/model"`. Will default to `user_name/output_dir_name` with *output_dir_name* being the + name of `output_dir`. + hub_strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`): + Defines what and when to push to Hub. Options: + - `"end"`: Push only at the end of training + - `"every_save"`: Push on each save (async to not block training) + - `"checkpoint"`: Like `"every_save"` plus push latest checkpoint to `"last-checkpoint"` subfolder for easy resuming + - `"all_checkpoints"`: Push all checkpoints as they appear + hub_always_push (`bool`, *optional*, defaults to `False`): + Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not finished. + hub_revision (`str`, *optional*): + The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash. + + > Best Model Tracking + + load_best_model_at_end (`bool`, *optional*, defaults to `False`): + Load the best checkpoint at the end of training. Requires `eval_strategy` to be set. + When enabled, the best checkpoint is always saved (see `save_total_limit`). + + When `True`, `save_strategy` must match `eval_strategy` (unless `save_strategy` is `"best"`), and if using `"steps"`, + `save_steps` must be a multiple of `eval_steps`. + + metric_for_best_model (`str`, *optional*): + Metric to use for comparing models when `load_best_model_at_end=True`. Must be a metric + name returned by evaluation, with or without the `"eval_"` prefix. Defaults to `"loss"`. + If you set this, `greater_is_better` will default to `True` unless the name ends with + `"loss"`. Examples: `"accuracy"`, `"f1"`, `"eval_bleu"`. + greater_is_better (`bool`, *optional*): + Whether higher metric values are better. Defaults based on `metric_for_best_model`: + `True` if the metric name doesn't end in `"loss"`, `False` otherwise. + + > Resuming Training + + ignore_data_skip (`bool`, *optional*, defaults to `False`): + When resuming training, skip fast-forwarding through the dataset to reach the previous + state. If `True`, training starts from the beginning of the dataset (faster resume but + results won't match interrupted training). If `False`, skips seen data (slower resume + but exact continuation). + restore_callback_states_from_checkpoint (`bool`, *optional*, defaults to `False`): + Restore callback states from checkpoint when resuming. If `True`, will override callbacks + passed to Trainer if they exist in the checkpoint. + + > Reproducibility + + full_determinism (`bool`, *optional*, defaults to `False`) + If `True`, [`enable_full_determinism`] is called instead of [`set_seed`] to ensure reproducible results in + distributed training. Important: this will negatively impact the performance, so only use it for debugging. + seed (`int`, *optional*, defaults to 42): + Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the + [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized parameters. + data_seed (`int`, *optional*): + Random seed to be used with data samplers. If not set, random generators for data sampling will use the + same seed as `seed`. This can be used to ensure reproducibility of data sampling, independent of the model + seed. + + > Hardware Configuration + + use_cpu (`bool`, *optional*, defaults to `False`): + Whether or not to use cpu. If set to False, we will use the available torch device/backend. + + > Accelerate Configuration + + accelerator_config (`str`, `dict`, or `AcceleratorConfig`, *optional*): + Configuration for the internal Accelerate integration. Can be: + - Path to JSON config file: `"accelerator_config.json"` + - Dictionary with config options + - `AcceleratorConfig` instance + Key options: + - `split_batches` (`bool`, defaults to `False`): Whether to split batches across devices. + If `True`, actual batch size is the same on all devices (total must be divisible by + num_processes). If `False`, each device gets the specified batch size. + - `dispatch_batches` (`bool`): If `True`, only main process iterates through dataloader + and dispatches batches to devices. Defaults to `True` for `IterableDataset`, `False` + otherwise. + - `even_batches` (`bool`, defaults to `True`): Duplicate samples from dataset start to + ensure all workers get equal batch sizes. + - `use_seedable_sampler` (`bool`, defaults to `True`): Use fully seedable random sampler + for reproducibility. + - `use_configured_state` (`bool`, defaults to `False`): Use pre-initialized + `AcceleratorState`/`PartialState` instead of creating new one. May cause issues with + hyperparameter tuning. + + parallelism_config (`ParallelismConfig`, *optional*): + Parallelism configuration for the training run. Requires Accelerate `1.10.1` + + > Dataloader + + dataloader_drop_last (`bool`, *optional*, defaults to `False`): + Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) + or not. + dataloader_num_workers (`int`, *optional*, defaults to 0): + Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the + main process. + dataloader_pin_memory (`bool`, *optional*, defaults to `True`): + Whether you want to pin memory in data loaders or not. Will default to `True`. + dataloader_persistent_workers (`bool`, *optional*, defaults to `False`): + If True, the data loader will not shut down the worker processes after a dataset has been consumed once. + This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will + increase RAM usage. Will default to `False`. + dataloader_prefetch_factor (`int`, *optional*): + Number of batches loaded in advance by each worker. + 2 means there will be a total of 2 * num_workers batches prefetched across all workers. + remove_unused_columns (`bool`, *optional*, defaults to `True`): + Whether or not to automatically remove the columns unused by the model forward method. + label_names (`list[str]`, *optional*): + The list of keys in your dictionary of inputs that correspond to the labels. + Will eventually default to the list of argument names accepted by the model that contain the word "label", + except if the model used is one of the `XxxForQuestionAnswering` in which case it will also include the + `["start_positions", "end_positions"]` keys. + You should only specify `label_names` if you're using custom label names or if your model's `forward` consumes multiple label tensors (e.g., extractive QA). + train_sampling_strategy (`str`, *optional*, defaults to `"random"`): + The sampler to use for the training dataloader. Possible values are: + + - `"random"`: Uses `RandomSampler` (default). + - `"sequential"`: Uses `SequentialSampler`. + - `"group_by_length"`: Uses `LengthGroupedSampler` to group samples of roughly the same length + together (to minimize padding and be more efficient). + + Note: When using an `IterableDataset`, this argument is ignored. + length_column_name (`str`, *optional*, defaults to `"length"`): + Column name for precomputed lengths. If the column exists, grouping by length will use these values rather + than computing them on train startup. Ignored unless `train_sampling_strategy` is `"group_by_length"` and the dataset + is an instance of `Dataset`. + + > DDP (DistributedDataParallel) + + ddp_find_unused_parameters (`bool`, *optional*): + When using distributed training, the value of the flag `find_unused_parameters` passed to + `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise. + ddp_bucket_cap_mb (`int`, *optional*): + When using distributed training, the value of the flag `bucket_cap_mb` passed to `DistributedDataParallel`. + ddp_broadcast_buffers (`bool`, *optional*): + When using distributed training, the value of the flag `broadcast_buffers` passed to + `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise. + ddp_backend (`str`, *optional*): + The backend to use for distributed training. Must be one of `"nccl"`, `"mpi"`, `"xccl"`, `"gloo"`, `"hccl"`. + ddp_timeout (`int`, *optional*, defaults to 1800): + The timeout for `torch.distributed.init_process_group` calls, used to avoid GPU socket timeouts when + performing slow operations in distributed runnings. Please refer to the [PyTorch documentation](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more + information. + + > FSDP (Fully Sharded Data Parallel) + + fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `None`): + Enable PyTorch Fully Sharded Data Parallel (FSDP) for distributed training. Options: + - `"full_shard"`: Shard parameters, gradients, and optimizer states (most memory efficient) + - `"shard_grad_op"`: Shard only optimizer states and gradients (ZeRO-2) + - `"hybrid_shard"`: Full shard within nodes, replicate across nodes + - `"hybrid_shard_zero2"`: Shard gradients/optimizer within nodes, replicate across nodes + - `"offload"`: Offload parameters and gradients to CPU (only with `"full_shard"` or + `"shard_grad_op"`) + - `"auto_wrap"`: Automatically wrap layers using `default_auto_wrap_policy` + fsdp_config (`str` or `dict`, *optional*): + Config to be used with fsdp (Pytorch Distributed Parallel Training). The value is either a location of + fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`. + + A List of config and its options: + - fsdp_version (`int`, *optional*, defaults to `1`): + The version of FSDP to use. Defaults to 1. + - min_num_params (`int`, *optional*, defaults to `0`): + FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `fsdp` field is + passed). + - transformer_layer_cls_to_wrap (`list[str]`, *optional*): + List of transformer layer class names (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`, + `T5Block` .... (useful only when `fsdp` flag is passed). + - backward_prefetch (`str`, *optional*) + FSDP's backward prefetch mode. Controls when to prefetch next set of parameters (useful only when + `fsdp` field is passed). + + A list of options along the following: + + - `"backward_pre"` : Prefetches the next set of parameters before the current set of parameter's + gradient computation. + - `"backward_post"` : This prefetches the next set of parameters after the current set of + parameter's gradient computation. + - forward_prefetch (`bool`, *optional*, defaults to `False`) + FSDP's forward prefetch mode (useful only when `fsdp` field is passed). + If `"True"`, then FSDP explicitly prefetches the next upcoming all-gather while executing in the + forward pass. + - limit_all_gathers (`bool`, *optional*, defaults to `False`) + FSDP's limit_all_gathers (useful only when `fsdp` field is passed). + If `"True"`, FSDP explicitly synchronizes the CPU thread to prevent too many in-flight + all-gathers. + - use_orig_params (`bool`, *optional*, defaults to `True`) + If `"True"`, allows non-uniform `requires_grad` during init, which means support for interspersed + frozen and trainable parameters. Useful in cases such as parameter-efficient fine-tuning. Please + refer this + [blog](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019 + - sync_module_states (`bool`, *optional*, defaults to `True`) + If `"True"`, each individually wrapped FSDP unit will broadcast module parameters from rank 0 to + ensure they are the same across all ranks after initialization + - cpu_ram_efficient_loading (`bool`, *optional*, defaults to `False`) + If `"True"`, only the first process loads the pretrained model checkpoint while all other processes + have empty weights. When this setting as `"True"`, `sync_module_states` also must to be `"True"`, + otherwise all the processes except the main process would have random weights leading to unexpected + behaviour during training. + - activation_checkpointing (`bool`, *optional*, defaults to `False`): + If `"True"`, activation checkpointing is a technique to reduce memory usage by clearing activations of + certain layers and recomputing them during a backward pass. Effectively, this trades extra + computation time for reduced memory usage. + - xla (`bool`, *optional*, defaults to `False`): + Whether to use PyTorch/XLA Fully Sharded Data Parallel Training. This is an experimental feature + and its API may evolve in the future. + - xla_fsdp_settings (`dict`, *optional*) + The value is a dictionary which stores the XLA FSDP wrapping parameters. + + For a complete list of options, please see [here]( + https://github.com/pytorch/xla/blob/master/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py). + - xla_fsdp_grad_ckpt (`bool`, *optional*, defaults to `False`): + Will use gradient checkpointing over each nested XLA FSDP wrapped layer. This setting can only be + used when the xla flag is set to true, and an auto wrapping policy is specified through + fsdp_min_num_params or fsdp_transformer_layer_cls_to_wrap. + + > DeepSpeed + + deepspeed (`str` or `dict`, *optional*): + Enable [DeepSpeed](https://github.com/deepspeedai/DeepSpeed) integration. Value is either: + - Path to DeepSpeed JSON config file: `"ds_config.json"` + - Loaded config as dictionary + > [!TIP] + > If using ZeRO initialization, instantiate your model *after* initializing + `TrainingArguments`, otherwise ZeRO won't be applied. + + > Debugging & Profiling (Experimental) + + debug (`str` or list of [`~debug_utils.DebugOption`], *optional*, defaults to `""`): + Enable one or more debug features. This is an experimental feature. + Possible options are: + - "underflow_overflow": detects overflow in model's input/outputs and reports the last frames that led to + the event + - "tpu_metrics_debug": print debug metrics on TPU + skip_memory_metrics (`bool`, *optional*, defaults to `True`): + Whether to skip adding of memory profiler reports to metrics. This is skipped by default because it slows + down the training and evaluation speed. + + > External Script Flags (not used by Trainer) + + do_train (`bool`, *optional*, defaults to `False`): + Whether to run training or not. This argument is not directly used by [`Trainer`], it's intended to be used + by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_eval (`bool`, *optional*): + Whether to run evaluation on the validation set or not. Will be set to `True` if `eval_strategy` is + different from `"no"`. This argument is not directly used by [`Trainer`], it's intended to be used by your + training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_predict (`bool`, *optional*, defaults to `False`): + Whether to run predictions on the test set or not. This argument is not directly used by [`Trainer`], it's + intended to be used by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + resume_from_checkpoint (`str`, *optional*): + The path to a folder with a valid checkpoint for your model. This argument is not directly used by + [`Trainer`], it's intended to be used by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + """ + + # Fields that accept dict values via CLI as JSON strings (e.g., '{"key": "value"}'). + # Any new dict-typed arg must be added here and typed as `dict | str | None`. + _VALID_DICT_FIELDS = [ + "accelerator_config", + "fsdp_config", + "deepspeed", + "gradient_checkpointing_kwargs", + "lr_scheduler_kwargs", + ] + + # --- Output --- + output_dir: str | None = field( + default=None, + metadata={"help": "The output directory where the model predictions and checkpoints will be written."}, + ) + + # --- Training Duration and Batch Size --- + per_device_train_batch_size: int = field(default=8, metadata={"help": "The batch size per device for training."}) + num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."}) + max_steps: int = field( + default=-1, + metadata={ + "help": "Overrides `num_train_epochs`. If set to a positive number, the total number of training steps to perform." + }, + ) + + # --- Learning Rate & Scheduler --- + learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for the optimizer."}) + lr_scheduler_type: SchedulerType | str = field( + default="linear", + metadata={"help": "The learning rate scheduler type to use. See `SchedulerType` for all possible values."}, + ) + lr_scheduler_kwargs: dict | str | None = field( + default=None, + metadata={ + "help": "The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values." + }, + ) + warmup_steps: float = field( + default=0, + metadata={ + "help": "Number of steps for a linear warmup from 0 to `learning_rate`. Can be an integer (exact steps) or a float in [0, 1) (ratio of total steps)." + }, + ) + + # --- Optimizer --- + default_optim = "adamw_torch" + if is_torch_available(): + from .pytorch_utils import is_torch_greater_or_equal_than_2_8 + + if is_torch_greater_or_equal_than_2_8: + default_optim = "adamw_torch_fused" + optim: OptimizerNames | str = field( + default=default_optim, + metadata={"help": "The optimizer to use. See `OptimizerNames` for the complete list."}, + ) + optim_args: str | None = field( + default=None, + metadata={ + "help": "Optional arguments supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore." + }, + ) + weight_decay: float = field( + default=0.0, + metadata={ + "help": "Weight decay coefficient applied by the optimizer. Automatically excluded from bias and LayerNorm parameters." + }, + ) + adam_beta1: float = field( + default=0.9, + metadata={ + "help": "The exponential decay rate for the first moment estimates (momentum) in Adam-based optimizers." + }, + ) + adam_beta2: float = field( + default=0.999, + metadata={ + "help": "The exponential decay rate for the second moment estimates (variance) in Adam-based optimizers." + }, + ) + adam_epsilon: float = field( + default=1e-8, metadata={"help": "Epsilon value for numerical stability in Adam-based optimizers."} + ) + optim_target_modules: None | str | list[str] = field( + default=None, + metadata={"help": "The target modules to optimize. Currently used for the GaLore and APOLLO algorithms."}, + ) + + # --- Regularization & Training Stability --- + gradient_accumulation_steps: int = field( + default=1, + metadata={ + "help": ( + "Number of update steps to accumulate gradients before performing a backward/update pass." + " Effective batch size = per_device_train_batch_size * num_devices * gradient_accumulation_steps." + ) + }, + ) + average_tokens_across_devices: bool = field( + default=True, + metadata={ + "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to " + "synchronize num_tokens_in_batch for precise loss calculation. Reference: " + "https://github.com/huggingface/transformers/issues/34242" + }, + ) + max_grad_norm: float = field( + default=1.0, metadata={"help": "Maximum gradient norm for gradient clipping. Set to 0 to disable."} + ) + label_smoothing_factor: float = field( + default=0.0, metadata={"help": "Label smoothing factor to prevent overconfidence. Zero means no smoothing."} + ) + + # --- Mixed Precision --- + bf16: bool = field( + default=False, + metadata={ + "help": "Enable bfloat16 (BF16) mixed precision training. Generally preferred over FP16 due to better numerical stability." + }, + ) + fp16: bool = field( + default=False, + metadata={ + "help": "Enable float16 (FP16) mixed precision training. Consider using BF16 instead if your hardware supports it." + }, + ) + bf16_full_eval: bool = field( + default=False, + metadata={ + "help": "Use full BF16 precision for evaluation (not just mixed precision). Faster and saves memory." + }, + ) + fp16_full_eval: bool = field( + default=False, + metadata={ + "help": "Use full FP16 precision for evaluation (not just mixed precision). Faster and saves memory." + }, + ) + tf32: bool | None = field( + default=None, + metadata={ + "help": "Enable TF32 mode on Ampere and newer GPUs. Provides up to 8x speedup with negligible accuracy loss." + }, + ) + + # --- Gradient Checkpointing --- + gradient_checkpointing: bool = field( + default=False, + metadata={ + "help": "Enable gradient checkpointing to trade compute for memory. Reduces memory at the cost of ~20%% slower training." + }, + ) + gradient_checkpointing_kwargs: dict[str, Any] | str | None = field( + default=None, + metadata={"help": "Keyword arguments passed to `gradient_checkpointing_enable()`."}, + ) + + # --- Compilation --- + torch_compile: bool = field( + default=False, metadata={"help": "Compile the model using `torch.compile()` for faster training."} + ) + torch_compile_backend: str | None = field( + default=None, + metadata={ + "help": "Backend for `torch.compile()`. If set, automatically enables `torch_compile`.", + }, + ) + torch_compile_mode: str | None = field( + default=None, + metadata={ + "help": "Compilation mode for `torch.compile()`. If set, automatically enables `torch_compile`.", + }, + ) + + # --- Kernels --- + use_liger_kernel: bool = field( + default=False, + metadata={ + "help": "Enable Liger Kernel optimizations. Increases throughput by ~20%% and reduces memory by ~60%%." + }, + ) + liger_kernel_config: dict[str, bool] | None = field( + default=None, + metadata={ + "help": "Configuration for Liger Kernel. Passed as kwargs to `_apply_liger_kernel_to_instance()`. If None, uses default configuration." + }, + ) + + # --- Additional Optimizations --- + use_cache: bool = field( + default=False, + metadata={ + "help": "Whether or not to use cache for the model For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`." + }, + ) + neftune_noise_alpha: float | None = field( + default=None, + metadata={ + "help": "If not None, activates NEFTune noise embeddings. Can drastically improve performance for instruction fine-tuning. Typical range: [5.0, 15.0]." + }, + ) + torch_empty_cache_steps: int | None = field( + default=None, + metadata={ + "help": "Number of steps to wait before calling `torch..empty_cache()`. Helps avoid CUDA OOM at a cost of ~10%% slower performance. If None, cache will not be emptied." + }, + ) + auto_find_batch_size: bool = field( + default=False, + metadata={ + "help": "Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding CUDA Out-of-Memory errors." + }, + ) + + # --- Logging & Monitoring --- + logging_strategy: IntervalStrategy | str = field( + default="steps", + metadata={"help": "The logging strategy to adopt during training. Options: 'no', 'epoch', 'steps'."}, + ) + logging_steps: float = field( + default=500, + metadata={ + "help": ( + "Log every X updates steps. Should be an integer or a float in range `[0,1)`. " + "If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + logging_first_step: bool = field( + default=False, metadata={"help": "Whether to log the first `global_step` or not."} + ) + log_on_each_node: bool = field( + default=True, + metadata={ + "help": ( + "When doing a multinode distributed training, whether to log once per node or just once on the main" + " node." + ) + }, + ) + logging_nan_inf_filter: bool = field( + default=True, + metadata={ + "help": "Filter out NaN and Inf losses when logging. Does not affect gradient computation, only logging." + }, + ) + include_num_input_tokens_seen: str | bool = field( + default="no", + metadata={ + "help": ( + "Whether to track the number of input tokens seen. " + "Must be one of [`all`, `non_padding`, `no`] or a boolean value which map to `all` or `no`" + ) + }, + ) + + # --- Log Levels --- + log_level: str = field( + default="passive", + metadata={ + "help": "Logging level for the main process. Options: 'debug', 'info', 'warning', 'error', 'critical', 'passive'.", + "choices": trainer_log_levels.keys(), + }, + ) + log_level_replica: str = field( + default="warning", + metadata={ + "help": "Logging level for replica processes in distributed training. Same options as `log_level`.", + "choices": trainer_log_levels.keys(), + }, + ) + disable_tqdm: bool | None = field( + default=None, + metadata={"help": "Disable tqdm progress bars. Defaults to True if log_level is warning or lower."}, + ) + + # --- Experiment Tracking --- + report_to: None | str | list[str] = field( + default="none", + metadata={ + "help": "The list of integrations to report the results and logs to. Use 'all' for all installed integrations, 'none' for no integrations." + }, + ) + run_name: str | None = field( + default=None, + metadata={ + "help": ( + "An optional descriptor for the run. Notably used for trackio, wandb, mlflow comet and swanlab " + "logging." + ) + }, + ) + project: str = field( + default="huggingface", + metadata={"help": "The name of the project to use for logging. Currently, only used by Trackio."}, + ) + trackio_space_id: str | None = field( + default="trackio", + metadata={ + "help": "The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like " + "'username/reponame' or 'orgname/reponame', or just 'reponame' in which case the Space will be created in " + "the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory. Note " + "that this Space will be public unless you set `hub_private_repo=True` or your organization's " + "default is to create private Spaces." + }, + ) + + # --- Evaluation --- + eval_strategy: IntervalStrategy | str = field( + default="no", + metadata={"help": "When to run evaluation. Options: 'no', 'steps', 'epoch'."}, + ) + eval_steps: float | None = field( + default=None, + metadata={ + "help": ( + "Number of update steps between evaluations if `eval_strategy='steps'`. Defaults to `logging_steps` if not set." + " Should be an integer or a float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + eval_delay: float = field( + default=0, + metadata={ + "help": ( + "Number of epochs or steps to wait for before the first evaluation can be performed, depending on the" + " eval_strategy." + ) + }, + ) + per_device_eval_batch_size: int = field( + default=8, metadata={"help": "The batch size per device (GPU/TPU core/CPU) for evaluation."} + ) + prediction_loss_only: bool = field( + default=False, + metadata={"help": "When performing evaluation and generating predictions, only returns the loss."}, + ) + eval_on_start: bool = field( + default=False, + metadata={ + "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check." + }, + ) + eval_do_concat_batches: bool = field( + default=True, + metadata={ + "help": "Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, will instead store them as lists, with each batch kept separate." + }, + ) + eval_use_gather_object: bool = field( + default=False, + metadata={ + "help": "Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices." + }, + ) + eval_accumulation_steps: int | None = field( + default=None, + metadata={ + "help": "Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If unset, predictions are accumulated on the accelerator before being moved to the CPU." + }, + ) + + # --- Metrics --- + include_for_metrics: list[str] = field( + default_factory=list, + metadata={"help": "Include additional data in the `compute_metrics` function. Options: 'inputs', 'loss'."}, + ) + batch_eval_metrics: bool = field( + default=False, + metadata={"help": "Break eval metrics calculation into batches to save memory."}, + ) + + # --- Checkpointing & Saving --- + save_only_model: bool = field( + default=False, + metadata={ + "help": "Save only model weights, not optimizer/scheduler/RNG state. Prevents resuming training from checkpoint." + }, + ) + save_strategy: SaveStrategy | str = field( + default="steps", + metadata={ + "help": "The checkpoint save strategy to adopt during training. Options: 'no', 'epoch', 'steps', 'best'." + }, + ) + save_steps: float = field( + default=500, + metadata={ + "help": ( + "Save checkpoint every X updates steps. Should be an integer or a float in range `[0,1)`. " + "If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + save_on_each_node: bool = field( + default=False, + metadata={ + "help": ( + "When doing multi-node distributed training, whether to save models and checkpoints on each node, or" + " only on the main one" + ) + }, + ) + save_total_limit: int | None = field( + default=None, + metadata={ + "help": "Maximum number of checkpoints to keep. Deletes older checkpoints in `output_dir`. The best checkpoint is always retained when `load_best_model_at_end=True`." + }, + ) + enable_jit_checkpoint: bool = field( + default=False, + metadata={ + "help": "Enable JIT checkpointing on SIGTERM signal for graceful termination on preemptible workloads. Configure your orchestrator's graceful shutdown period accordingly." + }, + ) + + # --- Hub Integration --- + push_to_hub: bool = field( + default=False, metadata={"help": "Whether or not to push the model to the Hub every time the model is saved."} + ) + hub_token: str | None = field( + default=None, + metadata={ + "help": "The token to use to push the model to the Hub. Defaults to the token from `hf auth login`." + }, + ) + hub_private_repo: bool | None = field( + default=None, + metadata={ + "help": "Whether to make the repo private. If `None` (default), the repo will be public unless the " + "organization's default is private. This value is ignored if the repo already exists. If reporting to " + "Trackio with deployment to Hugging Face Spaces enabled, the same logic determines whether the Space is " + "private." + }, + ) + hub_model_id: str | None = field( + default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."} + ) + hub_strategy: HubStrategy | str = field( + default="every_save", + metadata={ + "help": "Defines what and when to push to Hub. Options: 'end', 'every_save', 'checkpoint', 'all_checkpoints'." + }, + ) + hub_always_push: bool = field( + default=False, + metadata={"help": "Unless `True`, the Trainer will skip pushes if the previous one wasn't finished yet."}, + ) + hub_revision: str | None = field( + default=None, + metadata={ + "help": "The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash." + }, + ) + + # --- Best Model Tracking --- + load_best_model_at_end: bool = field( + default=False, + metadata={"help": "Load the best checkpoint at the end of training. Requires `eval_strategy` to be set."}, + ) + metric_for_best_model: str | None = field( + default=None, + metadata={ + "help": "Metric to use for comparing models when `load_best_model_at_end=True`. Defaults to 'loss'." + }, + ) + greater_is_better: bool | None = field( + default=None, + metadata={"help": "Whether higher metric values are better. Defaults based on `metric_for_best_model`."}, + ) + + # --- Resuming Training --- + ignore_data_skip: bool = field( + default=False, + metadata={ + "help": "When resuming training, skip fast-forwarding through the dataset to reach the previous state. If True, training starts from the beginning of the dataset." + }, + ) + restore_callback_states_from_checkpoint: bool = field( + default=False, + metadata={ + "help": "Whether to restore the callback states from the checkpoint. If `True`, will override callbacks passed to the `Trainer` if they exist in the checkpoint." + }, + ) + + # --- Reproducibility --- + full_determinism: bool = field( + default=False, + metadata={ + "help": ( + "Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed" + " training. Important: this will negatively impact the performance, so only use it for debugging." + ) + }, + ) + seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."}) + data_seed: int | None = field( + default=None, + metadata={"help": "Random seed to be used with data samplers. If not set, uses the same seed as `seed`."}, + ) + + # --- Hardware --- + use_cpu: bool = field( + default=False, + metadata={ + "help": "Whether or not to use cpu. If set to False, we will use the available torch device/backend." + }, + ) + + # --- Accelerate --- + accelerator_config: dict | str | None = field( + default=None, + metadata={ + "help": "Configuration for the internal Accelerate integration. Can be a path to a JSON config file or a dict." + }, + ) + parallelism_config: ParallelismConfig | None = field( + default=None, + metadata={"help": "Parallelism configuration for the training run. Requires Accelerate `1.10.1`."}, + ) + + # --- Dataloader --- + dataloader_drop_last: bool = field( + default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."} + ) + dataloader_num_workers: int = field( + default=0, + metadata={ + "help": ( + "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded" + " in the main process." + ) + }, + ) + dataloader_pin_memory: bool = field( + default=True, metadata={"help": "Whether or not to pin memory for DataLoader."} + ) + dataloader_persistent_workers: bool = field( + default=False, + metadata={ + "help": "If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage." + }, + ) + dataloader_prefetch_factor: int | None = field( + default=None, + metadata={ + "help": ( + "Number of batches loaded in advance by each worker. " + "2 means there will be a total of 2 * num_workers batches prefetched across all workers. " + ) + }, + ) + remove_unused_columns: bool = field( + default=True, + metadata={"help": "Whether or not to automatically remove the columns unused by the model forward method."}, + ) + label_names: list[str] | None = field( + default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."} + ) + train_sampling_strategy: str = field( + default="random", + metadata={ + "help": "Sampler for training: 'random' (default), 'sequential', or 'group_by_length'.", + "choices": ["random", "sequential", "group_by_length"], + }, + ) + length_column_name: str = field( + default="length", + metadata={ + "help": "Column name for precomputed lengths. Ignored unless `train_sampling_strategy` is 'group_by_length'." + }, + ) + + # --- DDP --- + ddp_find_unused_parameters: bool | None = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `find_unused_parameters` passed to " + "`DistributedDataParallel`." + ) + }, + ) + ddp_bucket_cap_mb: int | None = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `bucket_cap_mb` passed to " + "`DistributedDataParallel`." + ) + }, + ) + ddp_broadcast_buffers: bool | None = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `broadcast_buffers` passed to " + "`DistributedDataParallel`." + ) + }, + ) + ddp_backend: str | None = field( + default=None, + metadata={ + "help": "The backend to use for distributed training. Must be one of 'nccl', 'mpi', 'xccl', 'gloo', 'hccl'.", + "choices": ["nccl", "gloo", "mpi", "xccl", "hccl", "cncl", "mccl"], + }, + ) + ddp_timeout: int = field( + default=1800, + metadata={"help": "The timeout for `torch.distributed.init_process_group` calls (in seconds)."}, + ) + + # --- FSDP --- + fsdp: list[FSDPOption] | str | None = field( + default=None, + metadata={ + "help": "Enable PyTorch FSDP for distributed training. Options: 'full_shard', 'shard_grad_op', 'hybrid_shard', 'hybrid_shard_zero2', 'offload', 'auto_wrap'.", + }, + ) + fsdp_config: dict[str, Any] | str | None = field( + default=None, + metadata={ + "help": ( + "Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a " + "fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`." + ) + }, + ) + + # --- DeepSpeed --- + deepspeed: dict | str | None = field( + default=None, + metadata={"help": "Enable DeepSpeed integration. Value is a path to a JSON config file or a dict."}, + ) + + # --- Debugging --- + debug: str | list[DebugOption] = field( + default="", + metadata={ + "help": "Enable one or more debug features. Options: 'underflow_overflow' (detect overflow in model I/O), 'tpu_metrics_debug' (print TPU metrics)." + }, + ) + skip_memory_metrics: bool = field( + default=True, + metadata={ + "help": "Whether to skip adding memory profiler reports to metrics. Skipped by default because it slows down training." + }, + ) + + # --- External Script Flags --- + do_train: bool = field( + default=False, + metadata={ + "help": "Whether to run training. Not directly used by Trainer; intended for training/evaluation scripts." + }, + ) + do_eval: bool = field( + default=False, + metadata={ + "help": "Whether to run evaluation. Not directly used by Trainer; intended for training/evaluation scripts." + }, + ) + do_predict: bool = field( + default=False, + metadata={ + "help": "Whether to run predictions on the test set. Not directly used by Trainer; intended for training/evaluation scripts." + }, + ) + resume_from_checkpoint: str | None = field( + default=None, + metadata={ + "help": "Path to a folder with a valid checkpoint for your model. Not directly used by Trainer; intended for training/evaluation scripts." + }, + ) + + # --- Deprecated / Internal --- + warmup_ratio: float | None = field( + default=None, + metadata={ + "help": "This argument is deprecated and will be removed in v5.2. Use `warmup_steps` instead as it also works with float values." + }, + ) + logging_dir: str | None = field( + default=None, + metadata={ + "help": "Deprecated and will be removed in v5.2. Set env var `TENSORBOARD_LOGGING_DIR` instead. TensorBoard log directory." + }, + ) + local_rank: int = field( + default=-1, + metadata={ + "help": "When using torch.distributed.launch (Deprecated), it will pass `local_rank` in the script, so we need this for the parser. To get the local rank, prefer using the property `local_process_index`" + }, + ) + + def __post_init__(self): + # ── 1. Defaults & Normalization ── + if self.output_dir is None: + self.output_dir = "trainer_output" + logger.info( + "No output directory specified, defaulting to 'trainer_output'. " + "To change this behavior, specify --output_dir when creating TrainingArguments." + ) + + # Parse JSON string dict args from CLI (e.g., '{"key": "value"}'). + # Only parses strings starting with '{'; other strings are treated as file paths. + for valid_field in self._VALID_DICT_FIELDS: + passed_value = getattr(self, valid_field) + if isinstance(passed_value, str) and passed_value.startswith("{"): + loaded_dict = json.loads(passed_value) + loaded_dict = _convert_str_dict(loaded_dict) + setattr(self, valid_field, loaded_dict) + + # Expand ~ in paths so os.makedirs works correctly (#10628) + if self.output_dir is not None: + self.output_dir = os.path.expanduser(self.output_dir) + + if self.disable_tqdm is None: + self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN + + if self.warmup_ratio is not None: + logger.warning("warmup_ratio is deprecated and will be removed in v5.2. Use `warmup_steps` instead.") + self.warmup_steps = self.warmup_ratio + + if self.logging_dir is not None: + logger.warning( + "`logging_dir` is deprecated and will be removed in v5.2. Please set `TENSORBOARD_LOGGING_DIR` instead." + ) + + if isinstance(self.include_num_input_tokens_seen, bool): + self.include_num_input_tokens_seen = "all" if self.include_num_input_tokens_seen else "no" + + # ── 2. Enum / Type Conversions ── + self.eval_strategy = IntervalStrategy(self.eval_strategy) + self.logging_strategy = IntervalStrategy(self.logging_strategy) + self.save_strategy = SaveStrategy(self.save_strategy) + self.hub_strategy = HubStrategy(self.hub_strategy) + self.lr_scheduler_type = SchedulerType(self.lr_scheduler_type) + self.optim = OptimizerNames(self.optim) + + if isinstance(self.debug, str): + self.debug = [DebugOption(s) for s in self.debug.split()] + elif self.debug is None: + self.debug = [] + + # ── 3. Auto-derived Values ── + if self.do_eval is False and self.eval_strategy != IntervalStrategy.NO: + self.do_eval = True + + # Fall back to logging_steps if eval_steps is unset + if self.eval_strategy == IntervalStrategy.STEPS and (self.eval_steps is None or self.eval_steps == 0): + if self.logging_steps > 0: + logger.info(f"using `logging_steps` to initialize `eval_steps` to {self.logging_steps}") + self.eval_steps = self.logging_steps + else: + raise ValueError( + f"evaluation strategy {self.eval_strategy} requires either non-zero --eval_steps or" + " --logging_steps" + ) + + if ( + self.load_best_model_at_end + or self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU + or self.lr_scheduler_type == SchedulerType.GREEDY + ) and self.metric_for_best_model is None: + self.metric_for_best_model = "loss" + if self.greater_is_better is None and self.metric_for_best_model is not None: + self.greater_is_better = not self.metric_for_best_model.endswith("loss") + + if self.report_to == "all" or self.report_to == ["all"]: + from .integrations import get_available_reporting_integrations + + self.report_to = get_available_reporting_integrations() + elif self.report_to == "none" or self.report_to == ["none"]: + self.report_to = [] + elif not isinstance(self.report_to, list): + self.report_to = [self.report_to] + + # Auto-enable Kubeflow integration when running inside a Kubeflow TrainJob + from .integrations import is_kubeflow_available + + if is_kubeflow_available() and "kubeflow" not in self.report_to: + self.report_to = list(self.report_to) + ["kubeflow"] + + # ── 4. Validation ── + self._validate_args() + + # ── 5. Mixed Precision ── + # Read from env first; DeepSpeed may override this later + self.mixed_precision = os.environ.get("ACCELERATE_MIXED_PRECISION", "no") + if self.fp16: + self.mixed_precision = "fp16" + elif self.bf16: + self.mixed_precision = "bf16" + + # ── 6. Torch Compile ── + if (self.torch_compile_mode is not None or self.torch_compile_backend is not None) and not self.torch_compile: + self.torch_compile = True + if self.torch_compile and self.torch_compile_backend is None: + if not self.use_cpu and is_torch_hpu_available(): + self.torch_compile_backend = "hpu_backend" + else: + self.torch_compile_backend = "inductor" + + if self.torch_compile: + # TODO: remove env var fallback once minimum accelerate >= 1.2.0 + if not is_accelerate_available("1.2.0"): + os.environ["ACCELERATE_DYNAMO_BACKEND"] = self.torch_compile_backend + if self.torch_compile_mode is not None: + os.environ["ACCELERATE_DYNAMO_MODE"] = self.torch_compile_mode + + # ── 7. Accelerator Config (must come before self.device) ── + if is_accelerate_available(): + if not isinstance(self.accelerator_config, AcceleratorConfig): + if self.accelerator_config is None: + self.accelerator_config = AcceleratorConfig() + elif isinstance(self.accelerator_config, dict): + self.accelerator_config = AcceleratorConfig(**self.accelerator_config) + # Reject uninstantiated class (e.g. AcceleratorConfig instead of AcceleratorConfig()) + elif isinstance(self.accelerator_config, type): + raise NotImplementedError( + "Tried passing in a callable to `accelerator_config`, but this is not supported. " + "Please pass in a fully constructed `AcceleratorConfig` object instead." + ) + else: + self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config) + if self.accelerator_config.split_batches: + logger.info( + "Using `split_batches=True` in `accelerator_config` will override the `per_device_train_batch_size` " + "Batches will be split across all processes equally when using `split_batches=True`." + ) + + # ── 8. Device Init ── + if is_torch_available(): + self.device + + # ── 9. TF32 ── + if is_torch_available() and self.torch_compile: + if is_torch_tf32_available(): + if self.tf32 is None and not self.fp16 or self.bf16: + device_str = "MUSA" if is_torch_musa_available() else "CUDA" + logger.info( + f"Setting TF32 in {device_str} backends to speedup torch compile, you won't see any improvement" + " otherwise." + ) + enable_tf32(True) + else: + logger.warning( + "The speedups for torchdynamo mostly come with GPU Ampere or higher and which is not detected here." + ) + if is_torch_available() and self.tf32 is not None: + if self.tf32: + if is_torch_tf32_available(): + enable_tf32(True) + else: + raise ValueError("--tf32 requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7") + else: + if is_torch_tf32_available(): + enable_tf32(False) + # TF32 not available, nothing to disable + + # ── 10. Hardware Overrides ── + if self.use_cpu: + self.dataloader_pin_memory = False + + # ── 11. FSDP ── + # Store args only (not the plugin itself) to avoid pickle issues + self.fsdp_plugin_args = self._process_fsdp_args() + + # ── 12. DeepSpeed (must be last) ── + self.deepspeed_plugin = None + if self.deepspeed: + from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig + + # Leave self.deepspeed unmodified; users may rely on the original value + self.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.deepspeed) + self.hf_deepspeed_config.trainer_config_process(self) + + from accelerate.utils import DeepSpeedPlugin + + self.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.hf_deepspeed_config) + elif strtobool(os.environ.get("ACCELERATE_USE_DEEPSPEED", "false")): + from accelerate.utils import DeepSpeedPlugin + + self.deepspeed_plugin = DeepSpeedPlugin() + self.deepspeed_plugin.set_mixed_precision(self.mixed_precision) + self.deepspeed_plugin.set_deepspeed_weakref() + + def _validate_args(self): + """Validate argument combinations and value constraints.""" + if self.torch_empty_cache_steps is not None: + if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0): + raise ValueError( + f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}." + ) + + # logging_steps must be non-zero when logging_strategy="steps" + if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps == 0: + raise ValueError(f"logging strategy {self.logging_strategy} requires non-zero --logging_steps") + + if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps > 1: + if self.logging_steps != int(self.logging_steps): + raise ValueError(f"--logging_steps must be an integer if bigger than 1: {self.logging_steps}") + self.logging_steps = int(self.logging_steps) + if self.eval_strategy == IntervalStrategy.STEPS and self.eval_steps > 1: + if self.eval_steps != int(self.eval_steps): + raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}") + self.eval_steps = int(self.eval_steps) + if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1: + if self.save_steps != int(self.save_steps): + raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}") + self.save_steps = int(self.save_steps) + + # load_best_model_at_end requires compatible save and eval strategies + if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST: + if self.eval_strategy != self.save_strategy: + raise ValueError( + '--load_best_model_at_end requires the save and eval strategy to match, except when --save_strategy="best", but found\n- Evaluation ' + f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}" + ) + if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0: + if self.eval_steps < 1 or self.save_steps < 1: + if not (self.eval_steps < 1 and self.save_steps < 1): + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation " + "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps " + f"{self.save_steps} and eval_steps {self.eval_steps}." + ) + # Use integer arithmetic to avoid floating point precision issues + LARGE_MULTIPLIER = 1_000_000 + if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0: + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation " + f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}." + ) + else: + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation " + f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}." + ) + + if is_torch_available(): + if self.bf16 or self.bf16_full_eval: + if not self.use_cpu and not is_torch_bf16_gpu_available() and not is_torch_xla_available(): + error_message = "Your setup doesn't support bf16/gpu. You need to assign use_cpu if you want to train the model on CPU." + if is_torch_cuda_available(): + error_message += " You need Ampere+ GPU with cuda>=11.0." + raise ValueError(error_message) + + if self.fp16 and self.bf16: + raise ValueError("At most one of fp16 and bf16 can be True, but not both") + + if self.fp16_full_eval and self.bf16_full_eval: + raise ValueError("At most one of fp16 and bf16 can be True for full eval, but not both") + + if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU: + if self.eval_strategy == IntervalStrategy.NO: + raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires an eval strategy") + if not is_torch_available(): + raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires torch>=0.2.0") + + if self.lr_scheduler_type == SchedulerType.GREEDY: + if self.eval_strategy == IntervalStrategy.NO: + raise ValueError("lr_scheduler_type greedy requires an eval strategy") + + if self.warmup_steps < 0: + raise ValueError("warmup_steps must be an integer or a float") + + if self.dataloader_num_workers == 0 and self.dataloader_prefetch_factor is not None: + raise ValueError( + "--dataloader_prefetch_factor can only be set when data is loaded in a different process, i.e." + " when --dataloader_num_workers > 0." + ) + + def __str__(self): + self_as_dict = asdict(self) + + self_as_dict = {k: f"<{k.upper()}>" if k.endswith("_token") else v for k, v in self_as_dict.items()} + + attrs_as_str = [f"{k}={v},\n" for k, v in sorted(self_as_dict.items())] + return f"{self.__class__.__name__}(\n{''.join(attrs_as_str)})" + + __repr__ = __str__ + + @property + def train_batch_size(self) -> int: + """ + The actual batch size for training. + """ + train_batch_size = self.per_device_train_batch_size * max(1, self.n_gpu) + return train_batch_size + + @property + def eval_batch_size(self) -> int: + """ + The actual batch size for evaluation. + """ + eval_batch_size = self.per_device_eval_batch_size * max(1, self.n_gpu) + return eval_batch_size + + @property + def ddp_timeout_delta(self) -> timedelta: + """ + The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable. + """ + return timedelta(seconds=self.ddp_timeout) + + @cached_property + def _setup_devices(self) -> "torch.device": + requires_backends(self, ["torch"]) + logger.info("PyTorch: setting up devices") + if not is_sagemaker_mp_enabled(): + if not is_accelerate_available(): + raise ImportError( + f"Using the `Trainer` with `PyTorch` requires `accelerate>={ACCELERATE_MIN_VERSION}`: " + f"Please run `pip install transformers[torch]` or `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`" + ) + # Build kwargs for PartialState; actual init happens below + accelerator_state_kwargs: dict[str, Any] = {"enabled": True, "use_configured_state": False} + if isinstance(self.accelerator_config, AcceleratorConfig): + accelerator_state_kwargs["use_configured_state"] = self.accelerator_config.pop( + "use_configured_state", False + ) + if accelerator_state_kwargs["use_configured_state"]: + if PartialState._shared_state == {}: + raise ValueError( + "Passing `'use_configured_state':True` to the AcceleratorConfig requires a pre-configured " + "`AcceleratorState` or `PartialState` to be defined before calling `TrainingArguments`. " + ) + self.distributed_state = PartialState(cpu=self.use_cpu) + if self.deepspeed and self.distributed_state.distributed_type != DistributedType.DEEPSPEED: + raise RuntimeError( + "Tried to use an already configured `Accelerator` or `PartialState` that was not initialized for DeepSpeed, " + "but also passed in a `deepspeed` configuration to the `TrainingArguments`. Please set " + "`use_configured_state:False` instead or setup your `Accelerator` or `PartialState` properly." + ) + else: + AcceleratorState._reset_state(reset_partial_state=True) + self.distributed_state = None + + self._n_gpu = 1 + if self.use_cpu or strtobool(os.environ.get("ACCELERATE_USE_CPU", "False")): + accelerator_state_kwargs["cpu"] = True + accelerator_state_kwargs["backend"] = self.ddp_backend + self._n_gpu = 0 + elif is_sagemaker_mp_enabled(): + accelerator_state_kwargs["enabled"] = False + device = torch.device("cuda", smp.local_rank()) + torch.cuda.set_device(device) + elif is_sagemaker_dp_enabled(): + accelerator_state_kwargs["_use_sagemaker_dp"] = True + elif self.deepspeed: + accelerator_state_kwargs["use_deepspeed"] = True + accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout) + else: + accelerator_state_kwargs["backend"] = self.ddp_backend + accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout) + + # Initialize PartialState with the accumulated kwargs + if accelerator_state_kwargs.pop("enabled", False) and not accelerator_state_kwargs.pop( + "use_configured_state", False + ): + # Temporarily set env var so Accelerate detects DeepSpeed + use_deepspeed = accelerator_state_kwargs.pop("use_deepspeed", False) + if use_deepspeed: + os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" + self.distributed_state = PartialState(**accelerator_state_kwargs) + if use_deepspeed: + del os.environ["ACCELERATE_USE_DEEPSPEED"] + if not is_sagemaker_mp_enabled(): + device = self.distributed_state.device + if dist.is_available() and dist.is_initialized() and self.parallel_mode != ParallelMode.DISTRIBUTED: + logger.warning( + "torch.distributed process group is initialized, but parallel_mode != ParallelMode.DISTRIBUTED. " + "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch" + ) + + if is_torch_xla_available(): + device = self.distributed_state.device + self._n_gpu = 0 + elif is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled(): + pass # _n_gpu already set above + elif self.distributed_state.distributed_type == DistributedType.NO: + if self.use_cpu: + device = torch.device("cpu") + elif is_torch_mps_available(): + device = torch.device("mps") + elif is_torch_xpu_available(): + device = torch.device("xpu:0") + torch.xpu.set_device(device) + elif is_torch_mlu_available(): + device = torch.device("mlu:0") + torch.mlu.set_device(device) + elif is_torch_musa_available(): + device = torch.device("musa:0") + torch.musa.set_device(device) + elif is_torch_npu_available(): + device = torch.device("npu:0") + torch.npu.set_device(device) + elif is_torch_hpu_available(): + device = torch.device("hpu:0") + torch.hpu.set_device(device) + elif is_torch_neuron_available(): + device = torch.device("neuron:0") + torch.neuron.set_device(device) + else: + # Default to cuda:0 (respects CUDA_VISIBLE_DEVICES); nn.DataParallel handles n_gpu > 1 + device = torch.device( + "cuda:0" if torch.cuda.is_available() else os.environ.get("ACCELERATE_TORCH_DEVICE", "cpu") + ) + # _n_gpu may not have been set yet if _setup_devices is called early + self._n_gpu = torch.cuda.device_count() + if device.type == "cuda": + torch.cuda.set_device(device) + return device + + @property + def device(self) -> "torch.device": + """ + The device used by this process. + """ + requires_backends(self, ["torch"]) + return self._setup_devices + + @property + def n_gpu(self): + """ + The number of GPUs used by this process. + + Note: + This will only be greater than one when you have multiple GPUs available but are not using distributed + training. For distributed training, it will always be 1. + """ + requires_backends(self, ["torch"]) + # Ensure _setup_devices has been called + if not hasattr(self, "_n_gpu"): + _ = self._setup_devices + return self._n_gpu + + @property + def parallel_mode(self): + """ + The current mode used for parallelism if multiple GPUs/TPU cores are available. One of: + + - `ParallelMode.NOT_PARALLEL`: no parallelism (CPU or one GPU). + - `ParallelMode.NOT_DISTRIBUTED`: several GPUs in one single process (uses `torch.nn.DataParallel`). + - `ParallelMode.DISTRIBUTED`: several GPUs, each having its own process (uses + `torch.nn.DistributedDataParallel`). + - `ParallelMode.TPU`: several TPU cores. + """ + requires_backends(self, ["torch"]) + if is_torch_xla_available(): + return ParallelMode.TPU + elif is_sagemaker_mp_enabled(): + return ParallelMode.SAGEMAKER_MODEL_PARALLEL + elif is_sagemaker_dp_enabled(): + return ParallelMode.SAGEMAKER_DATA_PARALLEL + elif self.distributed_state is not None and self.distributed_state.distributed_type != DistributedType.NO: + return ParallelMode.DISTRIBUTED + elif self.n_gpu > 1: + return ParallelMode.NOT_DISTRIBUTED + else: + return ParallelMode.NOT_PARALLEL + + @property + def world_size(self): + """ + The number of processes used in parallel. + """ + requires_backends(self, ["torch"]) + if self.distributed_state is not None: + return self.distributed_state.num_processes + elif is_sagemaker_mp_enabled(): + return smp.dp_size() if not smp.state.cfg.prescaled_batch else smp.rdp_size() + return 1 + + @property + def process_index(self): + """ + The index of the current process used. + """ + requires_backends(self, ["torch"]) + if self.distributed_state is not None: + return self.distributed_state.process_index + elif is_sagemaker_mp_enabled(): + return smp.dp_rank() if not smp.state.cfg.prescaled_batch else smp.rdp_rank() + return 0 + + @property + def local_process_index(self): + """ + The index of the local process used. + """ + requires_backends(self, ["torch"]) + + if self.distributed_state is not None: + return self.distributed_state.local_process_index + elif is_sagemaker_mp_enabled(): + return smp.local_rank() + return 0 + + @property + def should_log(self): + """ + Whether or not the current process should produce log. + """ + if self.log_on_each_node: + return self.local_process_index == 0 + else: + if is_sagemaker_mp_enabled(): + return smp.rank() == 0 + else: + return self.process_index == 0 + + @property + def should_save(self): + """ + Whether or not the current process should write to disk, e.g., to save models and checkpoints. + """ + if self.save_on_each_node: + return self.local_process_index == 0 + else: + if is_sagemaker_mp_enabled(): + return smp.rank() == 0 + else: + return self.process_index == 0 + + def get_process_log_level(self): + """ + Returns the log level to be used depending on whether this process is the main process of node 0, main process + of node non-0, or a non-main process. + + For the main process the log level defaults to the logging level set (`logging.WARNING` if you didn't do + anything) unless overridden by `log_level` argument. + + For the replica processes the log level defaults to `logging.WARNING` unless overridden by `log_level_replica` + argument. + + The choice between the main and replica process settings is made according to the return value of `should_log`. + """ + + # convert to int + log_level = trainer_log_levels[self.log_level] + log_level_replica = trainer_log_levels[self.log_level_replica] + + log_level_main_node = logging.get_verbosity() if log_level == -1 else log_level + log_level_replica_node = logging.get_verbosity() if log_level_replica == -1 else log_level_replica + return log_level_main_node if self.should_log else log_level_replica_node + + @property + def place_model_on_device(self) -> bool | None: + """ + Can be subclassed and overridden for some specific integrations. + """ + return None + + @property + def _no_sync_in_gradient_accumulation(self): + """ + Whether or not to use no_sync for the gradients when doing gradient accumulation. + """ + return not ( + self.deepspeed or is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled() or is_torch_neuroncore_available() + ) + + @contextlib.contextmanager + def main_process_first(self, local=True, desc="work"): + """ + A context manager for torch distributed environment where on needs to do something on the main process, while + blocking replicas, and when it's finished releasing the replicas. + + One such use is for `datasets`'s `map` feature which to be efficient should be run once on the main process, + which upon completion saves a cached version of results and which then automatically gets loaded by the + replicas. + + Args: + local (`bool`, *optional*, defaults to `True`): + if `True` first means process of rank 0 of each node if `False` first means process of rank 0 of node + rank 0 In multi-node environment with a shared filesystem you most likely will want to use + `local=False` so that only the main process of the first node will do the processing. If however, the + filesystem is not shared, then the main process of each node will need to do the processing, which is + the default behavior. + desc (`str`, *optional*, defaults to `"work"`): + a work description to be used in debug logs + + """ + if is_torch_available() and self.world_size > 1: + main_process_desc = "main local process" if local else "main process" + if self.distributed_state is not None: + is_main_process = ( + self.distributed_state.is_local_main_process if local else self.distributed_state.is_main_process + ) + elif is_sagemaker_mp_enabled(): + is_main_process = smp.rank() == 0 + + try: + if not is_main_process: + # tell all replicas to wait + logger.debug(f"{self.process_index}: waiting for the {main_process_desc} to perform {desc}") + + if is_torch_xla_available(): + xm.rendezvous(desc) + else: + dist.barrier() + yield + finally: + if is_main_process: + # the wait is over + logger.debug(f"{self.process_index}: {main_process_desc} completed {desc}, releasing all replicas") + if is_torch_xla_available(): + xm.rendezvous(desc) + else: + dist.barrier() + else: + yield + + def get_warmup_steps(self, num_training_steps: int): + """ + Get number of steps used for a linear warmup. + """ + warmup_steps = ( + int(self.warmup_steps) if self.warmup_steps >= 1 else math.ceil(num_training_steps * self.warmup_steps) + ) + return warmup_steps + + def _dict_dtype_to_str(self, d: dict[str, Any]) -> None: + """ + Checks whether the passed dictionary and its nested dicts have a *dtype* key and if it's not None, + converts torch.dtype to a string of just the type. For example, `torch.float32` get converted into *"float32"* + string, which can then be stored in the json format. + """ + if d.get("dtype") is not None and not isinstance(d["dtype"], str): + d["dtype"] = str(d["dtype"]).split(".")[1] + for value in d.values(): + if isinstance(value, dict): + self._dict_dtype_to_str(value) + + def to_dict(self): + """ + Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates + the token values by removing their value. + """ + # Exclude non-init fields (they aren't user-facing config) + d = {field.name: getattr(self, field.name) for field in fields(self) if field.init} + + for k, v in d.items(): + if isinstance(v, Enum): + d[k] = v.value + if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum): + d[k] = [x.value for x in v] + if k.endswith("_token"): + d[k] = f"<{k.upper()}>" + # Serialize AcceleratorConfig to dict + if is_accelerate_available() and isinstance(v, AcceleratorConfig): + d[k] = v.to_dict() + # Serialize quantization_config if nested inside model_init_kwargs + if k == "model_init_kwargs" and isinstance(v, dict) and "quantization_config" in v: + quantization_config = v.get("quantization_config") + if quantization_config and not isinstance(quantization_config, dict): + d[k]["quantization_config"] = quantization_config.to_dict() + if k == "parallelism_config" and v is not None: + d[k] = v.to_json() + + self._dict_dtype_to_str(d) + + return d + + def to_json_string(self): + """ + Serializes this instance to a JSON string. + """ + return json.dumps(self.to_dict(), indent=2) + + def to_sanitized_dict(self) -> dict[str, Any]: + """ + Sanitized serialization to use with TensorBoard's hparams + """ + d = self.to_dict() + d = {**d, "train_batch_size": self.train_batch_size, "eval_batch_size": self.eval_batch_size} + + valid_types = [bool, int, float, str] + if is_torch_available(): + valid_types.append(torch.Tensor) + + return {k: v if type(v) in valid_types else str(v) for k, v in d.items()} + + # Convenience setters for grouped configuration + def set_training( + self, + learning_rate: float = 5e-5, + batch_size: int = 8, + weight_decay: float = 0, + num_epochs: float = 3, + max_steps: int = -1, + gradient_accumulation_steps: int = 1, + seed: int = 42, + gradient_checkpointing: bool = False, + ): + """ + A method that regroups all basic arguments linked to the training. + + + + Calling this method will automatically set `self.do_train` to `True`. + + + + Args: + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate for the optimizer. + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for training. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in the + optimizer. + num_train_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents + of the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + gradient_accumulation_steps (`int`, *optional*, defaults to 1): + Number of updates steps to accumulate the gradients for, before performing a backward/update pass. + + + + When using gradient accumulation, one step is counted as one step with backward pass. Therefore, + logging, evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training + examples. + + + + seed (`int`, *optional*, defaults to 42): + Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use + the [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized + parameters. + gradient_checkpointing (`bool`, *optional*, defaults to `False`): + If True, use gradient checkpointing to save memory at the expense of slower backward pass. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_training(learning_rate=1e-4, batch_size=32) + >>> args.learning_rate + 1e-4 + ``` + """ + self.do_train = True + self.learning_rate = learning_rate + self.per_device_train_batch_size = batch_size + self.weight_decay = weight_decay + self.num_train_epochs = num_epochs + self.max_steps = max_steps + self.gradient_accumulation_steps = gradient_accumulation_steps + self.seed = seed + self.gradient_checkpointing = gradient_checkpointing + return self + + def set_evaluate( + self, + strategy: str | IntervalStrategy = "no", + steps: int = 500, + batch_size: int = 8, + accumulation_steps: int | None = None, + delay: float | None = None, + loss_only: bool = False, + ): + """ + A method that regroups all arguments linked to evaluation. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`): + The evaluation strategy to adopt during training. Possible values are: + + - `"no"`: No evaluation is done during training. + - `"steps"`: Evaluation is done (and logged) every `steps`. + - `"epoch"`: Evaluation is done at the end of each epoch. + + Setting a `strategy` different from `"no"` will set `self.do_eval` to `True`. + steps (`int`, *optional*, defaults to 500): + Number of update steps between two evaluations if `strategy="steps"`. + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for evaluation. + accumulation_steps (`int`, *optional*): + Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. + If left unset, the whole predictions are accumulated on GPU/TPU before being moved to the CPU (faster + but requires more memory). + delay (`float`, *optional*): + Number of epochs or steps to wait for before the first evaluation can be performed, depending on the + eval_strategy. + loss_only (`bool`, *optional*, defaults to `False`): + Ignores all outputs except the loss. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_evaluate(strategy="steps", steps=100) + >>> args.eval_steps + 100 + ``` + """ + self.eval_strategy = IntervalStrategy(strategy) + if self.eval_strategy == IntervalStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.do_eval = self.eval_strategy != IntervalStrategy.NO + self.eval_steps = steps + self.per_device_eval_batch_size = batch_size + self.eval_accumulation_steps = accumulation_steps + self.eval_delay = delay + self.prediction_loss_only = loss_only + return self + + def set_testing( + self, + batch_size: int = 8, + loss_only: bool = False, + ): + """ + A method that regroups all basic arguments linked to testing on a held-out dataset. + + + + Calling this method will automatically set `self.do_predict` to `True`. + + + + Args: + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for testing. + loss_only (`bool`, *optional*, defaults to `False`): + Ignores all outputs except the loss. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_testing(batch_size=32) + >>> args.per_device_eval_batch_size + 32 + ``` + """ + self.do_predict = True + self.per_device_eval_batch_size = batch_size + self.prediction_loss_only = loss_only + return self + + def set_save( + self, + strategy: str | IntervalStrategy = "steps", + steps: int = 500, + total_limit: int | None = None, + on_each_node: bool = False, + ): + """ + A method that regroups all arguments linked to checkpoint saving. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The checkpoint save strategy to adopt during training. Possible values are: + + - `"no"`: No save is done during training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps`. + + steps (`int`, *optional*, defaults to 500): + Number of updates steps before two checkpoint saves if `strategy="steps"`. + total_limit (`int`, *optional*): + If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in + `output_dir`. + on_each_node (`bool`, *optional*, defaults to `False`): + When doing multi-node distributed training, whether to save models and checkpoints on each node, or + only on the main one. + + This should not be activated when the different nodes use the same storage as the files will be saved + with the same names for each node. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_save(strategy="steps", steps=100) + >>> args.save_steps + 100 + ``` + """ + self.save_strategy = SaveStrategy(strategy) + if self.save_strategy == SaveStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.save_steps = steps + self.save_total_limit = total_limit + self.save_on_each_node = on_each_node + return self + + def set_logging( + self, + strategy: str | IntervalStrategy = "steps", + steps: int = 500, + report_to: str | list[str] = "none", + level: str = "passive", + first_step: bool = False, + nan_inf_filter: bool = False, + on_each_node: bool = False, + replica_level: str = "passive", + ): + """ + A method that regroups all arguments linked to logging. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The logging strategy to adopt during training. Possible values are: + + - `"no"`: No logging is done during training. + - `"epoch"`: Logging is done at the end of each epoch. + - `"steps"`: Logging is done every `logging_steps`. + + steps (`int`, *optional*, defaults to 500): + Number of update steps between two logs if `strategy="steps"`. + level (`str`, *optional*, defaults to `"passive"`): + Logger log level to use on the main process. Possible choices are the log levels as strings: `"debug"`, + `"info"`, `"warning"`, `"error"` and `"critical"`, plus a `"passive"` level which doesn't set anything + and lets the application set the level. + report_to (`str` or `list[str]`, *optional*, defaults to `"none"`): + The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, + `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, + `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations + installed, `"none"` for no integrations. + first_step (`bool`, *optional*, defaults to `False`): + Whether to log and evaluate the first `global_step` or not. + nan_inf_filter (`bool`, *optional*, defaults to `True`): + Whether to filter `nan` and `inf` losses for logging. If set to `True` the loss of every step that is + `nan` or `inf` is filtered and the average loss of the current logging window is taken instead. + + + + `nan_inf_filter` only influences the logging of loss values, it does not change the behavior the + gradient is computed or applied to the model. + + + + on_each_node (`bool`, *optional*, defaults to `True`): + In multinode distributed training, whether to log using `log_level` once per node, or only on the main + node. + replica_level (`str`, *optional*, defaults to `"passive"`): + Logger log level to use on replicas. Same choices as `log_level` + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_logging(strategy="steps", steps=100) + >>> args.logging_steps + 100 + ``` + """ + self.logging_strategy = IntervalStrategy(strategy) + if self.logging_strategy == IntervalStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.logging_steps = steps + self.report_to = report_to + self.log_level = level + self.logging_first_step = first_step + self.logging_nan_inf_filter = nan_inf_filter + self.log_on_each_node = on_each_node + self.log_level_replica = replica_level + return self + + def set_push_to_hub( + self, + model_id: str, + strategy: str | HubStrategy = "every_save", + token: str | None = None, + private_repo: bool | None = None, + always_push: bool = False, + revision: str | None = None, + ): + """ + A method that regroups all arguments linked to synchronizing checkpoints with the Hub. + + + + Calling this method will set `self.push_to_hub` to `True`, which means the `output_dir` will begin a git + directory synced with the repo (determined by `model_id`) and the content will be pushed each time a save is + triggered (depending on your `self.save_strategy`). Calling [`~Trainer.save_model`] will also trigger a push. + + + + Args: + model_id (`str`): + The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in + which case the model will be pushed in your namespace. Otherwise it should be the whole repository + name, for instance `"user_name/model"`, which allows you to push to an organization you are a member of + with `"organization_name/model"`. + strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`): + Defines the scope of what is pushed to the Hub and when. Possible values are: + + - `"end"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) and a + draft of a model card when the [`~Trainer.save_model`] method is called. + - `"every_save"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) + and + a draft of a model card each time there is a model save. The pushes are asynchronous to not block + training, and in case the save are very frequent, a new push is only attempted if the previous one is + finished. A last push is made with the final model at the end of training. + - `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named + last-checkpoint, allowing you to resume training easily with + `trainer.train(resume_from_checkpoint="last-checkpoint")`. + - `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the + output + folder (so you will get one checkpoint folder per folder in your final repository) + + token (`str`, *optional*): + The token to use to push the model to the Hub. Will default to the token in the cache folder obtained + with `hf auth login`. + private_repo (`bool`, *optional*, defaults to `False`): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + always_push (`bool`, *optional*, defaults to `False`): + Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not + finished. + revision (`str`, *optional*): + The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_push_to_hub("me/awesome-model") + >>> args.hub_model_id + 'me/awesome-model' + ``` + """ + self.push_to_hub = True + self.hub_model_id = model_id + self.hub_strategy = HubStrategy(strategy) + self.hub_token = token + self.hub_private_repo = private_repo + self.hub_always_push = always_push + self.hub_revision = revision + return self + + def set_optimizer( + self, + name: str | OptimizerNames = "adamw_torch", + learning_rate: float = 5e-5, + weight_decay: float = 0, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-8, + args: str | None = None, + ): + """ + A method that regroups all arguments linked to the optimizer and its hyperparameters. + + Args: + name (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"`): + The optimizer to use: `"adamw_torch"`, `"adamw_torch_fused"`, `"adamw_apex_fused"`, + `"adamw_anyprecision"` or `"adafactor"`. + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights. + beta1 (`float`, *optional*, defaults to 0.9): + The beta1 hyperparameter for the adam optimizer or its variants. + beta2 (`float`, *optional*, defaults to 0.999): + The beta2 hyperparameter for the adam optimizer or its variants. + epsilon (`float`, *optional*, defaults to 1e-8): + The epsilon hyperparameter for the adam optimizer or its variants. + args (`str`, *optional*): + Optional arguments that are supplied to AnyPrecisionAdamW (only useful when + `optim="adamw_anyprecision"`). + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_optimizer(name="adamw_torch", beta1=0.8) + >>> args.optim + 'adamw_torch' + ``` + """ + self.optim = OptimizerNames(name) + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.adam_beta1 = beta1 + self.adam_beta2 = beta2 + self.adam_epsilon = epsilon + self.optim_args = args + return self + + def set_lr_scheduler( + self, + name: str | SchedulerType = "linear", + num_epochs: float = 3.0, + max_steps: int = -1, + warmup_steps: float = 0, + warmup_ratio: float | None = None, + ): + """ + A method that regroups all arguments linked to the learning rate scheduler and its hyperparameters. + + Args: + name (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`): + The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values. + num_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents + of the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + warmup_steps (`float`, *optional*, defaults to 0): + Number of steps used for a linear warmup from 0 to `learning_rate`. Should be an integer or a float in range `[0,1)`. + If smaller than 1, will be interpreted as ratio of steps used for a linear warmup from 0 to `learning_rate`. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_lr_scheduler(name="cosine", warmup_steps=0.05) + >>> args.warmup_steps + 0.05 + ``` + """ + if warmup_ratio is not None: + logger.warning("warmup_ratio is deprecated and will be removed in v5.2 . Use `warmup_steps` instead.") + warmup_steps = warmup_ratio + + self.lr_scheduler_type = SchedulerType(name) + self.num_train_epochs = num_epochs + self.max_steps = max_steps + self.warmup_steps = warmup_steps + return self + + def set_dataloader( + self, + train_batch_size: int = 8, + eval_batch_size: int = 8, + drop_last: bool = False, + num_workers: int = 0, + pin_memory: bool = True, + persistent_workers: bool = False, + prefetch_factor: int | None = None, + auto_find_batch_size: bool = False, + ignore_data_skip: bool = False, + sampler_seed: int | None = None, + ): + """ + A method that regroups all arguments linked to the dataloaders creation. + + Args: + drop_last (`bool`, *optional*, defaults to `False`): + Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch + size) or not. + num_workers (`int`, *optional*, defaults to 0): + Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in + the main process. + pin_memory (`bool`, *optional*, defaults to `True`): + Whether you want to pin memory in data loaders or not. Will default to `True`. + persistent_workers (`bool`, *optional*, defaults to `False`): + If True, the data loader will not shut down the worker processes after a dataset has been consumed + once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, + but will increase RAM usage. Will default to `False`. + prefetch_factor (`int`, *optional*): + Number of batches loaded in advance by each worker. + 2 means there will be a total of 2 * num_workers batches prefetched across all workers. + auto_find_batch_size (`bool`, *optional*, defaults to `False`) + Whether to find a batch size that will fit into memory automatically through exponential decay, + avoiding CUDA Out-of-Memory errors. Requires accelerate to be installed (`pip install accelerate`) + ignore_data_skip (`bool`, *optional*, defaults to `False`): + When resuming training, whether or not to skip the epochs and batches to get the data loading at the + same stage as in the previous training. If set to `True`, the training will begin faster (as that + skipping step can take a long time) but will not yield the same results as the interrupted training + would have. + sampler_seed (`int`, *optional*): + Random seed to be used with data samplers. If not set, random generators for data sampling will use the + same seed as `self.seed`. This can be used to ensure reproducibility of data sampling, independent of + the model seed. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_dataloader(train_batch_size=16, eval_batch_size=64) + >>> args.per_device_train_batch_size + 16 + ``` + """ + self.per_device_train_batch_size = train_batch_size + self.per_device_eval_batch_size = eval_batch_size + self.dataloader_drop_last = drop_last + self.dataloader_num_workers = num_workers + self.dataloader_pin_memory = pin_memory + self.dataloader_persistent_workers = persistent_workers + self.dataloader_prefetch_factor = prefetch_factor + self.auto_find_batch_size = auto_find_batch_size + self.ignore_data_skip = ignore_data_skip + self.data_seed = sampler_seed + return self + + def _process_fsdp_args(self): + if not self.fsdp: + self.fsdp = [] + elif self.fsdp is True: + self.fsdp = [FSDPOption.FULL_SHARD] + elif isinstance(self.fsdp, str): + self.fsdp = [FSDPOption(s) for s in self.fsdp.split()] + + if self.fsdp == [FSDPOption.OFFLOAD]: + raise ValueError( + "`--fsdp offload` can't work on its own. It needs to be added to `--fsdp full_shard` or " + '`--fsdp shard_grad_op`. For example, `--fsdp "full_shard offload"`.' + ) + elif FSDPOption.FULL_SHARD in self.fsdp and FSDPOption.SHARD_GRAD_OP in self.fsdp: + raise ValueError("`--fsdp full_shard` is not compatible with `--fsdp shard_grad_op`.") + + if self.gradient_checkpointing and ( + FSDPOption.FULL_SHARD in self.fsdp or FSDPOption.HYBRID_SHARD in self.fsdp + ): + logger.warning( + "When using FSDP full shard, instead of using `gradient_checkpointing` in TrainingArguments, please" + " use `activation_checkpointing` in `fsdp_config`. The former introduces a redundant AllGather" + " operation in backward pass. Reference: https://github.com/huggingface/transformers/issues/30404" + ) + + if self.fsdp_config is None: + self.fsdp_config = {} + + if isinstance(self.fsdp_config, str): + if len(self.fsdp) == 0: + warnings.warn("`--fsdp_config` is useful only when `--fsdp` is specified.") + with open(self.fsdp_config, encoding="utf-8") as f: + self.fsdp_config = json.load(f) + + if self.fsdp_config is not None and isinstance(self.fsdp_config, dict): + for k in list(self.fsdp_config.keys()): + if k.startswith("fsdp_"): + v = self.fsdp_config.pop(k) + self.fsdp_config[k[5:]] = v + + self.fsdp_config["min_num_params"] = self.fsdp_config.get("min_num_params", 0) + + # Normalize transformer_layer_cls_to_wrap from string to list + if isinstance(self.fsdp_config.get("transformer_layer_cls_to_wrap", None), str): + self.fsdp_config["transformer_layer_cls_to_wrap"] = [self.fsdp_config["transformer_layer_cls_to_wrap"]] + + if len(self.fsdp) == 0 and self.fsdp_config["min_num_params"] > 0: + warnings.warn("`min_num_params` is useful only when `--fsdp` is specified.") + + if len(self.fsdp) == 0 and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None: + warnings.warn("`transformer_layer_cls_to_wrap` is useful only when `--fsdp` is specified.") + + if ( + len(self.fsdp) > 0 + and self.fsdp_config["min_num_params"] > 0 + and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None + ): + raise ValueError("`min_num_params` and `transformer_layer_cls_to_wrap` are mutually exclusive.") + self.fsdp_config["xla"] = self.fsdp_config.get("xla", False) + self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False) + self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) + if self.fsdp_config["xla"]: + if len(self.fsdp) > 0: + # Copy to avoid mutating the original (needed for JSON serialization) + self.xla_fsdp_config = self.fsdp_config.get("xla_fsdp_settings", {}).copy() + # Convert string dtype names to torch.dtype + if "compute_dtype" in self.xla_fsdp_config: + self.xla_fsdp_config["compute_dtype"] = getattr(torch, self.xla_fsdp_config["compute_dtype"]) + if "buffer_dtype" in self.xla_fsdp_config: + self.xla_fsdp_config["buffer_dtype"] = getattr(torch, self.xla_fsdp_config["buffer_dtype"]) + else: + warnings.warn("XLA FSDP can be used only when `--fsdp` is specified.") + else: + if self.fsdp_config["xla_fsdp_grad_ckpt"]: + warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.") + + # Build kwargs for Accelerate's FSDPPlugin + fsdp_plugin_args = None + if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: + from accelerate.utils.constants import ( + FSDP_AUTO_WRAP_POLICY, + FSDP_SHARDING_STRATEGY, + ) + + fsdp_plugin_args = {} + fsdp_sharding = None + for fsdp_option in self.fsdp: + if fsdp_option.upper() in FSDP_SHARDING_STRATEGY: + fsdp_sharding = fsdp_option + elif fsdp_option == FSDPOption.OFFLOAD: + fsdp_plugin_args["cpu_offload"] = True + elif fsdp_option == FSDPOption.AUTO_WRAP: + fsdp_plugin_args["auto_wrap_policy"] = FSDP_AUTO_WRAP_POLICY[0] + if self.fsdp_config["min_num_params"] > 0: + fsdp_plugin_args["min_num_params"] = self.fsdp_config["min_num_params"] + fsdp_plugin_args["auto_wrap_policy"] = FSDP_AUTO_WRAP_POLICY[1] + elif self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None: + fsdp_plugin_args["transformer_cls_names_to_wrap"] = ",".join( + self.fsdp_config["transformer_layer_cls_to_wrap"] + ) + fsdp_version = int(self.fsdp_config.get("version", 1)) + fsdp_plugin_args["fsdp_version"] = fsdp_version + prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH") + if fsdp_version == 2: + # full_shard → True (reshard after forward), shard_grad_op → False + default_reshard = fsdp_sharding != "shard_grad_op" if fsdp_sharding else True + fsdp_plugin_args["reshard_after_forward"] = str_to_bool( + str(self.fsdp_config.get("reshard_after_forward", default_reshard)).lower() + ) + else: + fsdp_plugin_args["forward_prefetch"] = str_to_bool( + str(self.fsdp_config.get("forward_prefetch", "false")).lower() + ) + fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper() + # Pass sharding strategy as reshard_after_forward (accelerate converts it to ShardingStrategy) + default_reshard = fsdp_sharding.upper() if fsdp_sharding else "FULL_SHARD" + fsdp_plugin_args["reshard_after_forward"] = str( + self.fsdp_config.get("reshard_after_forward", default_reshard) + ).lower() + fsdp_plugin_args["use_orig_params"] = str_to_bool( + str(self.fsdp_config.get("use_orig_params", "true")).lower() + ) + + sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower() + cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower() + if sync_module_states == "false" and cpu_ram_efficient_loading == "true": + # Without sync, non-main processes would have random weights + raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`') + + # Set env var to suppress Accelerate warning and for transformers to read + fsdp_plugin_args["cpu_ram_efficient_loading"] = str_to_bool(cpu_ram_efficient_loading) + os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading + + fsdp_plugin_args["sync_module_states"] = str_to_bool(sync_module_states) + + return fsdp_plugin_args + + +class ParallelMode(Enum): + NOT_PARALLEL = "not_parallel" + NOT_DISTRIBUTED = "not_distributed" + DISTRIBUTED = "distributed" + SAGEMAKER_MODEL_PARALLEL = "sagemaker_model_parallel" + SAGEMAKER_DATA_PARALLEL = "sagemaker_data_parallel" + TPU = "tpu" + + +def str_to_bool(value, to_bool: bool = True) -> int | bool: + """ + Converts a string representation of truth to `True` (1) or `False` (0). + + True values are `y`, `yes`, `t`, `true`, `on`, and `1`; False value are `n`, `no`, `f`, `false`, `off`, and `0`; + """ + value = value.lower() + if value in ("y", "yes", "t", "true", "on", "1"): + return 1 if not to_bool else True + elif value in ("n", "no", "f", "false", "off", "0"): + return 0 if not to_bool else False + else: + raise ValueError(f"invalid truth value {value}") diff --git a/third_party/transformers/src/transformers/training_args_seq2seq.py b/third_party/transformers/src/transformers/training_args_seq2seq.py new file mode 100644 index 0000000000000000000000000000000000000000..eb5ff6119a841b834555dc796c0c13b2ad12e255 --- /dev/null +++ b/third_party/transformers/src/transformers/training_args_seq2seq.py @@ -0,0 +1,94 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from dataclasses import dataclass, field +from pathlib import Path + +from .generation.configuration_utils import GenerationConfig +from .training_args import TrainingArguments +from .utils import add_start_docstrings + + +logger = logging.getLogger(__name__) + + +@dataclass +@add_start_docstrings(TrainingArguments.__doc__) +class Seq2SeqTrainingArguments(TrainingArguments): + """ + sortish_sampler (`bool`, *optional*, defaults to `False`): + Whether to use a *sortish sampler* or not. Only possible if the underlying datasets are *Seq2SeqDataset* + for now but will become generally available in the near future. + + It sorts the inputs according to lengths in order to minimize the padding size, with a bit of randomness + for the training set. + predict_with_generate (`bool`, *optional*, defaults to `False`): + Whether to use generate to calculate generative metrics (ROUGE, BLEU). + generation_max_length (`int`, *optional*): + The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default to the + `max_length` value of the model configuration. + generation_num_beams (`int`, *optional*): + The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default to the + `num_beams` value of the model configuration. + generation_config (`str` or `Path` or [`~generation.GenerationConfig`], *optional*): + Allows to load a [`~generation.GenerationConfig`] from the `from_pretrained` method. This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~GenerationConfig.save_pretrained`] method, e.g., `./my_model_directory/`. + - a [`~generation.GenerationConfig`] object. + """ # fmt: skip # Prevent Ruff from altering the indentation + + sortish_sampler: bool = field(default=False, metadata={"help": "Whether to use SortishSampler or not."}) + predict_with_generate: bool = field( + default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."} + ) + generation_max_length: int | None = field( + default=None, + metadata={ + "help": ( + "The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default " + "to the `max_length` value of the model configuration." + ) + }, + ) + generation_num_beams: int | None = field( + default=None, + metadata={ + "help": ( + "The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default " + "to the `num_beams` value of the model configuration." + ) + }, + ) + generation_config: str | Path | GenerationConfig | None = field( + default=None, + metadata={ + "help": "Model id, file path or url pointing to a GenerationConfig json file, to use during prediction." + }, + ) + + def to_dict(self): + """ + Serializes this instance while replace `Enum` by their values and `GenerationConfig` by dictionaries (for JSON + serialization support). It obfuscates the token values by removing their value. + """ + # filter out fields that are defined as field(init=False) + d = super().to_dict() + for k, v in d.items(): + if isinstance(v, GenerationConfig): + d[k] = v.to_dict() + return d diff --git a/third_party/transformers/src/transformers/video_processing_utils.py b/third_party/transformers/src/transformers/video_processing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e8cede50a99350cf3e99dfcfeddb9e48ce25dc86 --- /dev/null +++ b/third_party/transformers/src/transformers/video_processing_utils.py @@ -0,0 +1,845 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import warnings +from collections.abc import Callable +from functools import partial +from typing import Any + +import numpy as np +from huggingface_hub import create_repo, is_offline_mode +from huggingface_hub.dataclasses import validate_typed_dict + +from .dynamic_module_utils import custom_object_save +from .image_processing_backends import TorchvisionBackend +from .image_processing_utils import BatchFeature +from .image_utils import ( + ChannelDimension, + SizeDict, + is_vision_available, + validate_kwargs, +) +from .processing_utils import Unpack, VideosKwargs +from .utils import ( + IMAGE_PROCESSOR_NAME, + PROCESSOR_NAME, + VIDEO_PROCESSOR_NAME, + TensorType, + add_start_docstrings, + copy_func, + is_torch_available, + is_torchcodec_available, + is_torchvision_v2_available, + logging, + safe_load_json_file, +) +from .utils.hub import cached_file +from .utils.import_utils import requires +from .video_utils import ( + VideoInput, + VideoMetadata, + group_videos_by_shape, + infer_channel_dimension_format, + is_valid_video, + load_video, + make_batched_metadata, + make_batched_videos, + reorder_videos, +) + + +if is_torch_available(): + import torch + +if is_torchvision_v2_available(): + import torchvision.transforms.v2.functional as tvF + +if is_vision_available(): + from .image_utils import PILImageResampling + + +logger = logging.get_logger(__name__) + + +BASE_VIDEO_PROCESSOR_DOCSTRING = r""" + Args: + do_resize (`bool`, *optional*, defaults to `self.do_resize`): + Whether to resize the video's (height, width) dimensions to the specified `size`. Can be overridden by the + `do_resize` parameter in the `preprocess` method. + size (`dict`, *optional*, defaults to `self.size`): + Size of the output video after resizing. Can be overridden by the `size` parameter in the `preprocess` + method. + size_divisor (`int`, *optional*, defaults to `self.size_divisor`): + The size by which to make sure both the height and width can be divided. + default_to_square (`bool`, *optional*, defaults to `self.default_to_square`): + Whether to default to a square video when resizing, if size is an int. + resample (`PILImageResampling`, *optional*, defaults to `self.resample`): + Resampling filter to use if resizing the video. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `resample` parameter in the `preprocess` method. + do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`): + Whether to center crop the video to the specified `crop_size`. Can be overridden by `do_center_crop` in the + `preprocess` method. + crop_size (`dict[str, int]` *optional*, defaults to `self.crop_size`): + Size of the output video after applying `center_crop`. Can be overridden by `crop_size` in the `preprocess` + method. + do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): + Whether to rescale the video by the specified scale `rescale_factor`. Can be overridden by the + `do_rescale` parameter in the `preprocess` method. + rescale_factor (`int` or `float`, *optional*, defaults to `self.rescale_factor`): + Scale factor to use if rescaling the video. Only has an effect if `do_rescale` is set to `True`. Can be + overridden by the `rescale_factor` parameter in the `preprocess` method. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the video. Can be overridden by the `do_normalize` parameter in the `preprocess` + method. Can be overridden by the `do_normalize` parameter in the `preprocess` method. + image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`): + Mean to use if normalizing the video. This is a float or list of floats the length of the number of + channels in the video. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be + overridden by the `image_mean` parameter in the `preprocess` method. + image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`): + Standard deviation to use if normalizing the video. This is a float or list of floats the length of the + number of channels in the video. Can be overridden by the `image_std` parameter in the `preprocess` method. + Can be overridden by the `image_std` parameter in the `preprocess` method. + do_convert_rgb (`bool`, *optional*, defaults to `self.image_std`): + Whether to convert the video to RGB. + video_metadata (`VideoMetadata`, *optional*): + Metadata of the video containing information about total duration, fps and total number of frames. + do_sample_frames (`int`, *optional*, defaults to `self.do_sample_frames`): + Whether to sample frames from the video before processing or to process the whole video. + num_frames (`int`, *optional*, defaults to `self.num_frames`): + Maximum number of frames to sample when `do_sample_frames=True`. + fps (`int` or `float`, *optional*, defaults to `self.fps`): + Target frames to sample per second when `do_sample_frames=True`. + return_tensors (`str` or `TensorType`, *optional*): + Returns stacked tensors if set to `pt, otherwise returns a list of tensors. + data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`): + The channel dimension format for the output video. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: video in (height, width, num_channels) format. + - Unset: Use the channel dimension format of the input video. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input video. If unset, the channel dimension format is inferred + from the input video. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: video in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: video in (height, width) format. + device (`torch.device`, *optional*): + The device to process the videos on. If unset, the device is inferred from the input videos. + return_metadata (`bool`, *optional*): + Whether to return video metadata or not. + """ + + +@add_start_docstrings( + "Constructs a base VideoProcessor.", + BASE_VIDEO_PROCESSOR_DOCSTRING, +) +@requires(backends=("vision", "torchvision")) +class BaseVideoProcessor(TorchvisionBackend): + _auto_class = None + + resample = None + image_mean = None + image_std = None + size = None + size_divisor = None + default_to_square = True + crop_size = None + do_resize = None + do_center_crop = None + do_rescale = None + rescale_factor = 1 / 255 + do_normalize = None + do_convert_rgb = None + do_sample_frames = None + fps = None + num_frames = None + video_metadata = None + return_metadata = False + valid_kwargs = VideosKwargs + model_input_names = ["pixel_values_videos"] + + def __init__(self, **kwargs: Unpack[VideosKwargs]) -> None: + super().__init__(**kwargs) + + def __call__(self, videos, **kwargs) -> BatchFeature: + return self.preprocess(videos, **kwargs) + + def convert_to_rgb( + self, + video: "torch.Tensor", + ) -> VideoInput: + """ + Converts a video to RGB format. + + Args: + video (`"torch.Tensor"`): + The video to convert. + + Returns: + `torch.Tensor`: The converted video. + """ + + video = tvF.grayscale_to_rgb(video) + if video.shape[-3] == 3 or not (video[..., 3, :, :] < 255).any(): + return video + + # There is a transparency layer, blend it with a white background. + # Calculate the alpha proportion for blending. + alpha = video[..., 3, :, :] / 255.0 + video = (1 - alpha[..., None, :, :]) * 255 + alpha[..., None, :, :] * video[..., :3, :, :] + return video + + def sample_frames( + self, + metadata: VideoMetadata, + num_frames: int | None = None, + fps: int | float | None = None, + **kwargs, + ): + """ + Default sampling function which uniformly samples the desired number of frames between 0 and total number of frames. + If `fps` is passed along with metadata, `fps` frames per second are sampled uniformty. Arguments `num_frames` + and `fps` are mutually exclusive. + + Args: + metadata (`VideoMetadata`): + Metadata of the video containing information about total duration, fps and total number of frames. + num_frames (`int`, *optional*): + Maximum number of frames to sample. Defaults to `self.num_frames`. + fps (`int` or `float`, *optional*): + Target frames to sample per second. Defaults to `self.fps`. + + Returns: + np.ndarray: + Indices to sample video frames. + """ + if fps is not None and num_frames is not None: + raise ValueError( + "`num_frames`, `fps`, and `sample_indices_fn` are mutually exclusive arguments, please use only one!" + ) + + num_frames = num_frames if num_frames is not None else self.num_frames + fps = fps if fps is not None else self.fps + total_num_frames = metadata.total_num_frames + + # If num_frames is not given but fps is, calculate num_frames from fps + if num_frames is None and fps is not None: + if metadata is None or metadata.fps is None: + raise ValueError( + "Asked to sample `fps` frames per second but no video metadata was provided which is required when sampling with `fps`. " + "Please pass in `VideoMetadata` object or use a fixed `num_frames` per input video" + ) + num_frames = int(total_num_frames / metadata.fps * fps) + + if num_frames > total_num_frames: + raise ValueError( + f"Video can't be sampled. The `num_frames={num_frames}` exceeds `total_num_frames={total_num_frames}`. " + ) + + if num_frames is not None: + indices = torch.arange(0, total_num_frames, total_num_frames / num_frames).int() + else: + indices = torch.arange(0, total_num_frames).int() + return indices + + def _decode_and_sample_videos( + self, + videos: VideoInput, + video_metadata: VideoMetadata | dict, + do_sample_frames: bool | None = None, + sample_indices_fn: Callable | None = None, + ) -> list["torch.Tensor"]: + """ + Decode input videos and sample frames if needed. + """ + videos = make_batched_videos(videos) + video_metadata = make_batched_metadata(videos, video_metadata=video_metadata) + + # Only sample frames if an array video is passed, otherwise first decode -> then sample + if is_valid_video(videos[0]) and do_sample_frames: + sampled_videos = [] + sampled_metadata = [] + for video, metadata in zip(videos, video_metadata): + indices = sample_indices_fn(metadata=metadata) + metadata.frames_indices = indices + sampled_videos.append(video[indices]) + sampled_metadata.append(metadata) + videos = sampled_videos + video_metadata = sampled_metadata + elif not is_valid_video(videos[0]): + if isinstance(videos[0], list): + # Videos sometimes are passed as a list of image URLs, especially through templates + videos = [ + torch.stack([tvF.pil_to_tensor(image) for image in images], dim=0) + for images in self.fetch_images(videos) + ] + if do_sample_frames: + raise ValueError( + "Sampling frames from a list of images is not supported! Set `do_sample_frames=False`." + ) + else: + videos, video_metadata = self.fetch_videos(videos, sample_indices_fn=sample_indices_fn) + + return videos, video_metadata + + def _prepare_input_videos( + self, + videos: VideoInput, + input_data_format: str | ChannelDimension | None = None, + device: str | None = None, + ) -> list["torch.Tensor"]: + """ + Prepare the input videos for processing. + """ + processed_videos = [] + for video in videos: + # `make_batched_videos` always returns a 4D array per video + if isinstance(video, np.ndarray): + # not using tvF.to_tensor as it doesn't handle (C, H, W) numpy arrays + video = torch.from_numpy(video).contiguous() + + # Infer the channel dimension format if not provided + if input_data_format is None: + input_data_format = infer_channel_dimension_format(video) + + if input_data_format == ChannelDimension.LAST: + video = video.permute(0, 3, 1, 2).contiguous() + + if device is not None: + video = video.to(device) + + processed_videos.append(video) + return processed_videos + + @add_start_docstrings( + BASE_VIDEO_PROCESSOR_DOCSTRING, + ) + def preprocess( + self, + videos: VideoInput, + **kwargs: Unpack[VideosKwargs], + ) -> BatchFeature: + validate_kwargs( + captured_kwargs=kwargs.keys(), + valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"], + ) + + # Perform type validation on received kwargs + validate_typed_dict(self.valid_kwargs, kwargs) + + # Set default kwargs from self. This ensures that if a kwarg is not provided + # by the user, it gets its default value from the instance, or is set to None. + for kwarg_name in self.valid_kwargs.__annotations__: + kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None)) + + input_data_format = kwargs.pop("input_data_format") + do_sample_frames = kwargs.pop("do_sample_frames") + device = kwargs.pop("device") + video_metadata = kwargs.pop("video_metadata") + + sample_indices_fn = partial(self.sample_frames, **kwargs) if do_sample_frames else None + videos, video_metadata = self._decode_and_sample_videos( + videos, + video_metadata=video_metadata, + do_sample_frames=do_sample_frames, + sample_indices_fn=sample_indices_fn, + ) + videos = self._prepare_input_videos(videos=videos, input_data_format=input_data_format, device=device) + + kwargs = self._standardize_kwargs(**kwargs) + self._validate_preprocess_kwargs(**kwargs) + + # Pop kwargs that are not needed in _preprocess + kwargs.pop("data_format") + return_metadata = kwargs.pop("return_metadata") + + preprocessed_videos = self._preprocess(videos=videos, **kwargs) + if return_metadata: + preprocessed_videos["video_metadata"] = video_metadata + return preprocessed_videos + + def _preprocess( + self, + videos: list["torch.Tensor"], + do_convert_rgb: bool, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + # Group videos by size for batched resizing + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + for shape, stacked_videos in grouped_videos.items(): + if do_convert_rgb: + stacked_videos = self.convert_to_rgb(stacked_videos) + if do_resize: + stacked_videos = self.resize(stacked_videos, size=size, resample=resample) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + # Group videos by size for further processing + # Needed in case do_resize is False, or resize returns videos with different sizes + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + for shape, stacked_videos in grouped_videos.items(): + if do_center_crop: + stacked_videos = self.center_crop(stacked_videos, crop_size) + # Fused rescale and normalize + stacked_videos = self.rescale_and_normalize( + stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_videos_grouped[shape] = stacked_videos + + processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index) + + return BatchFeature(data={"pixel_values_videos": processed_videos}, tensor_type=return_tensors) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | os.PathLike, + cache_dir: str | os.PathLike | None = None, + force_download: bool = False, + local_files_only: bool = False, + token: str | bool | None = None, + revision: str = "main", + **kwargs, + ): + r""" + Instantiate a type of [`~video_processing_utils.VideoProcessorBase`] from an video processor. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained video hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a video processor file saved using the + [`~video_processing_utils.VideoProcessorBase.save_pretrained`] method, e.g., + `./my_model_directory/`. + - a path to a saved video processor JSON *file*, e.g., + `./my_model_directory/video_preprocessor_config.json`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded pretrained model video processor should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the video processor files and override the cached versions if + they exist. + proxies (`dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `hf auth login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + return_unused_kwargs (`bool`, *optional*, defaults to `False`): + If `False`, then this function returns just the final video processor object. If `True`, then this + functions returns a `Tuple(video_processor, unused_kwargs)` where *unused_kwargs* is a dictionary + consisting of the key/value pairs whose keys are not video processor attributes: i.e., the part of + `kwargs` which has not been used to update `video_processor` and is otherwise ignored. + subfolder (`str`, *optional*, defaults to `""`): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can + specify the folder name here. + kwargs (`dict[str, Any]`, *optional*): + The values in kwargs of any keys which are video processor attributes will be used to override the + loaded values. Behavior concerning key/value pairs whose keys are *not* video processor attributes is + controlled by the `return_unused_kwargs` keyword parameter. + + Returns: + A video processor of type [`~video_processing_utils.ImagVideoProcessorBase`]. + + Examples: + + ```python + # We can't instantiate directly the base class *VideoProcessorBase* so let's show the examples on a + # derived class: *LlavaOnevisionVideoProcessor* + video_processor = LlavaOnevisionVideoProcessor.from_pretrained( + "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" + ) # Download video_processing_config from huggingface.co and cache. + video_processor = LlavaOnevisionVideoProcessor.from_pretrained( + "./test/saved_model/" + ) # E.g. video processor (or model) was saved using *save_pretrained('./test/saved_model/')* + video_processor = LlavaOnevisionVideoProcessor.from_pretrained("./test/saved_model/video_preprocessor_config.json") + video_processor = LlavaOnevisionVideoProcessor.from_pretrained( + "llava-hf/llava-onevision-qwen2-0.5b-ov-hf", do_normalize=False, foo=False + ) + assert video_processor.do_normalize is False + video_processor, unused_kwargs = LlavaOnevisionVideoProcessor.from_pretrained( + "llava-hf/llava-onevision-qwen2-0.5b-ov-hf", do_normalize=False, foo=False, return_unused_kwargs=True + ) + assert video_processor.do_normalize is False + assert unused_kwargs == {"foo": False} + ```""" + kwargs["cache_dir"] = cache_dir + kwargs["force_download"] = force_download + kwargs["local_files_only"] = local_files_only + kwargs["revision"] = revision + + if token is not None: + kwargs["token"] = token + + video_processor_dict, kwargs = cls.get_video_processor_dict(pretrained_model_name_or_path, **kwargs) + + return cls.from_dict(video_processor_dict, **kwargs) + + def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): + """ + Save an video processor object to the directory `save_directory`, so that it can be re-loaded using the + [`~video_processing_utils.VideoProcessorBase.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the video processor JSON file will be saved (will be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, **kwargs).repo_id + files_timestamps = self._get_files_timestamps(save_directory) + + # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self) + + # If we save using the predefined names, we can load using `from_pretrained` + output_video_processor_file = os.path.join(save_directory, VIDEO_PROCESSOR_NAME) + + self.to_json_file(output_video_processor_file) + logger.info(f"Video processor saved in {output_video_processor_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + return [output_video_processor_file] + + @classmethod + def get_video_processor_dict( + cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + """ + From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a + video processor of type [`~video_processing_utils.VideoProcessorBase`] using `from_dict`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + The identifier of the pre-trained checkpoint from which we want the dictionary of parameters. + subfolder (`str`, *optional*, defaults to `""`): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can + specify the folder name here. + + Returns: + `tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the video processor object. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", "") + + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + + user_agent = {"file_type": "video processor", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isfile(pretrained_model_name_or_path): + resolved_video_processor_file = pretrained_model_name_or_path + resolved_processor_file = None + is_local = True + else: + video_processor_file = VIDEO_PROCESSOR_NAME + try: + # Try to load with a new config name first and if not successful try with the old file name + # NOTE: we save all processor configs as nested dict in PROCESSOR_NAME from v5, which is the standard + resolved_processor_file = cached_file( + pretrained_model_name_or_path, + filename=PROCESSOR_NAME, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + ) + resolved_video_processor_files = [ + resolved_file + for filename in [video_processor_file, IMAGE_PROCESSOR_NAME] + if ( + resolved_file := cached_file( + pretrained_model_name_or_path, + filename=filename, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + ) + ) + is not None + ] + resolved_video_processor_file = ( + resolved_video_processor_files[0] if resolved_video_processor_files else None + ) + except OSError: + # Raise any OS error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load video processor for '{pretrained_model_name_or_path}'. If you were trying to load" + " it from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a {video_processor_file} file" + ) + + # Load video_processor dict. Priority goes as (nested config if found -> video processor config -> image processor config) + # We are downloading both configs because almost all models have a `processor_config.json` but + # not all of these are nested. We need to check if it was saved recebtly as nested or if it is legacy style + video_processor_dict = None + if resolved_processor_file is not None: + processor_dict = safe_load_json_file(resolved_processor_file) + if "video_processor" in processor_dict: + video_processor_dict = processor_dict["video_processor"] + + if resolved_video_processor_file is not None and video_processor_dict is None: + video_processor_dict = safe_load_json_file(resolved_video_processor_file) + + if video_processor_dict is None: + raise OSError( + f"Can't load video processor for '{pretrained_model_name_or_path}'. If you were trying to load" + " it from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a {video_processor_file} file" + ) + + if is_local: + logger.info(f"loading configuration file {resolved_video_processor_file}") + else: + logger.info( + f"loading configuration file {video_processor_file} from cache at {resolved_video_processor_file}" + ) + + return video_processor_dict, kwargs + + @classmethod + def from_dict(cls, video_processor_dict: dict[str, Any], **kwargs): + """ + Instantiates a type of [`~video_processing_utils.VideoProcessorBase`] from a Python dictionary of parameters. + + Args: + video_processor_dict (`dict[str, Any]`): + Dictionary that will be used to instantiate the video processor object. Such a dictionary can be + retrieved from a pretrained checkpoint by leveraging the + [`~video_processing_utils.VideoProcessorBase.to_dict`] method. + kwargs (`dict[str, Any]`): + Additional parameters from which to initialize the video processor object. + + Returns: + [`~video_processing_utils.VideoProcessorBase`]: The video processor object instantiated from those + parameters. + """ + video_processor_dict = video_processor_dict.copy() + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + video_processor_dict.update({k: v for k, v in kwargs.items() if k in cls.valid_kwargs.__annotations__}) + video_processor = cls(**video_processor_dict) + + # Apply extra kwargs to instance (BC for remote code, e.g. phi4_multimodal) + extra_keys = [] + for key in reversed(list(kwargs.keys())): + if hasattr(video_processor, key) and key not in cls.valid_kwargs.__annotations__: + setattr(video_processor, key, kwargs.pop(key, None)) + extra_keys.append(key) + if extra_keys: + logger.warning_once( + f"Image processor {cls.__name__}: kwargs {extra_keys} were applied for backward compatibility. " + f"To avoid this warning, add them to valid_kwargs: create a custom TypedDict extending " + f"ImagesKwargs with these keys and set it as the `valid_kwargs` class attribute." + ) + + logger.info(f"Video processor {video_processor}") + if return_unused_kwargs: + return video_processor, kwargs + else: + return video_processor + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + `dict[str, Any]`: Dictionary of all the attributes that make up this video processor instance. + """ + filtered_dict = super().to_dict() + filtered_dict.pop("image_processor_type", None) + filtered_dict["video_processor_type"] = self.__class__.__name__ + + return filtered_dict + + def to_json_string(self) -> str: + """ + Serializes this instance to a JSON string. + + Returns: + `str`: String containing all the attributes that make up this feature_extractor instance in JSON format. + """ + dictionary = self.to_dict() + + for key, value in dictionary.items(): + if isinstance(value, np.ndarray): + dictionary[key] = value.tolist() + + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: str | os.PathLike): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this image_processor instance's parameters will be saved. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + @classmethod + def from_json_file(cls, json_file: str | os.PathLike): + """ + Instantiates a video processor of type [`~video_processing_utils.VideoProcessorBase`] from the path to a JSON + file of parameters. + + Args: + json_file (`str` or `os.PathLike`): + Path to the JSON file containing the parameters. + + Returns: + A video processor of type [`~video_processing_utils.VideoProcessorBase`]: The video_processor object + instantiated from that JSON file. + """ + with open(json_file, "r", encoding="utf-8") as reader: + text = reader.read() + video_processor_dict = json.loads(text) + return cls(**video_processor_dict) + + @classmethod + def register_for_auto_class(cls, auto_class="AutoVideoProcessor"): + """ + Register this class with a given auto class. This should only be used for custom video processors as the ones + in the library are already mapped with `AutoVideoProcessor `. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoVideoProcessor "`): + The auto class to register this new video processor with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def fetch_videos(self, video_url_or_urls: str | list[str] | list[list[str]], sample_indices_fn=None): + """ + Convert a single or a list of urls into the corresponding `np.array` objects. + + If a single url is passed, the return value will be a single object. If a list is passed a list of objects is + returned. + """ + backend = "torchcodec" + if not is_torchcodec_available(): + warnings.warn( + "`torchcodec` is not installed and cannot be used to decode the video by default. " + "Falling back to `torchvision`. Note that `torchvision` decoding is deprecated and will be removed in future versions. " + ) + backend = "torchvision" + + if isinstance(video_url_or_urls, list): + return list(zip(*[self.fetch_videos(x, sample_indices_fn=sample_indices_fn) for x in video_url_or_urls])) + else: + return load_video(video_url_or_urls, backend=backend, sample_indices_fn=sample_indices_fn) + + +BaseVideoProcessor.push_to_hub = copy_func(BaseVideoProcessor.push_to_hub) +if BaseVideoProcessor.push_to_hub.__doc__ is not None: + BaseVideoProcessor.push_to_hub.__doc__ = BaseVideoProcessor.push_to_hub.__doc__.format( + object="video processor", object_class="AutoVideoProcessor", object_files="video processor file" + ) diff --git a/third_party/transformers/src/transformers/video_utils.py b/third_party/transformers/src/transformers/video_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..971e4fc08905b67d97118a4030544c4ea3b9185e --- /dev/null +++ b/third_party/transformers/src/transformers/video_utils.py @@ -0,0 +1,891 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import warnings +from collections.abc import Callable, Iterable, Mapping +from contextlib import redirect_stdout +from dataclasses import dataclass, fields +from io import BytesIO +from typing import NewType, Union +from urllib.parse import urlparse + +import httpx +import numpy as np + +from .image_transforms import PaddingMode, to_channel_dimension_format +from .image_utils import ChannelDimension, infer_channel_dimension_format, is_valid_image +from .utils import ( + is_av_available, + is_cv2_available, + is_decord_available, + is_numpy_array, + is_torch_available, + is_torch_tensor, + is_torchcodec_available, + is_torchvision_available, + is_vision_available, + is_yt_dlp_available, + logging, + requires_backends, +) + + +if is_vision_available(): + import PIL.Image + + if is_torchvision_available(): + from torchvision import io as torchvision_io + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + +URL = NewType("URL", str) +Path = NewType("Path", str) + +VideoInput = Union[ + list["PIL.Image.Image"], + np.ndarray, + "torch.Tensor", + list[np.ndarray], + list["torch.Tensor"], + list[list["PIL.Image.Image"]], + list[list[np.ndarray]], + list[list["torch.Tensor"]], + URL, + list[URL], + list[list[URL]], + Path, + list[Path], + list[list[Path]], +] + + +@dataclass +class VideoMetadata(Mapping): + total_num_frames: int + fps: float | None = None + width: int | None = None + height: int | None = None + duration: float | None = None + video_backend: str | None = None + frames_indices: list[int] | None = None + + def __iter__(self): + return (f.name for f in fields(self)) + + def __len__(self): + return len(fields(self)) + + def __getitem__(self, item): + return getattr(self, item) + + def __setitem__(self, key, value): + return setattr(self, key, value) + + @property + def timestamps(self) -> list[float]: + "Timestamps of the sampled frames in seconds." + if self.fps is None or self.frames_indices is None: + raise ValueError("Cannot infer video `timestamps` when `fps` or `frames_indices` is None.") + return [frame_idx / self.fps for frame_idx in self.frames_indices] + + @property + def sampled_fps(self) -> float: + "FPS of the sampled video." + if self.frames_indices is None or self.total_num_frames is None or self.fps is None: + return self.fps or 24 + return len(self.frames_indices) / self.total_num_frames * self.fps + + def update(self, dictionary): + for key, value in dictionary.items(): + if hasattr(self, key): + setattr(self, key, value) + + +VideoMetadataType = VideoMetadata | dict | list[dict | VideoMetadata] | list[list[dict | VideoMetadata]] + + +def is_valid_video_frame(frame): + return isinstance(frame, PIL.Image.Image) or ( + (is_numpy_array(frame) or is_torch_tensor(frame)) and frame.ndim == 3 + ) + + +def is_valid_video(video): + if not isinstance(video, (list, tuple)): + return (is_numpy_array(video) or is_torch_tensor(video)) and video.ndim == 4 + return video and all(is_valid_video_frame(frame) for frame in video) + + +def valid_videos(videos): + # If we have a list of videos, it could be either one video as list of frames or a batch + if isinstance(videos, (list, tuple)): + for video_or_frame in videos: + if not (is_valid_video(video_or_frame) or is_valid_video_frame(video_or_frame)): + return False + # If not a list, then we have a single 4D video or 5D batched tensor + elif not is_valid_video(videos) or videos.ndim == 5: + return False + return True + + +def is_batched_video(videos): + if isinstance(videos, (list, tuple)): + return is_valid_video(videos[0]) + elif (is_numpy_array(videos) or is_torch_tensor(videos)) and videos.ndim == 5: + return True + return False + + +def is_scaled_video(video: np.ndarray) -> bool: + """ + Checks to see whether the pixel values have already been rescaled to [0, 1]. + """ + # It's possible the video has pixel values in [0, 255] but is of floating type + return np.min(video) >= 0 and np.max(video) <= 1 + + +def convert_pil_frames_to_video(videos: list[VideoInput]) -> list[Union[np.ndarray, "torch.Tensor"]]: + """ + Given a batch of videos, converts each video to a 4D array. If video is already in array type, + it is simply returned. We assume that all inputs in the list are in the same format, based on the type of the first element. + + Args: + videos (`VideoInput`): + Video inputs to turn into a list of videos. + """ + + if not (isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0])): + return videos + + video_converted = [] + for video in videos: + video = [np.array(frame) for frame in video] + video = np.stack(video) + video_converted.append(video) + return video_converted + + +def make_batched_videos(videos) -> list[Union[np.ndarray, "torch.Tensor", "URL", "Path"]]: + """ + Ensure that the input is a list of videos. If the input is a single video, it is converted to a list of length 1. + If the input is a batch of videos, it is converted to a list of 4D video arrays. Videos passed as list `PIL.Image` + frames are converted to 4D arrays. + + We assume that all inputs in the list are in the same format, based on the type of the first element. + + Args: + videos (`VideoInput`): + Video inputs to turn into a list of videos. + """ + # Early exit for deeply nested list of image frame paths. We shouldn't flatten them + try: + if isinstance(videos[0][0], list) and isinstance(videos[0][0][0], str): + return [image_paths for sublist in videos for image_paths in sublist] + except (IndexError, TypeError): + pass + + if is_batched_video(videos): + return convert_pil_frames_to_video(list(videos)) + elif isinstance(videos, str) or is_valid_video(videos): + return convert_pil_frames_to_video([videos]) + # only one frame passed, thus we unsqueeze time dim + elif is_valid_image(videos): + if isinstance(videos, PIL.Image.Image): + videos = np.array(videos) + return [videos[None, ...]] + elif not isinstance(videos, list): + raise ValueError( + f"Invalid video input. Expected either a list of video frames or an input of 4 or 5 dimensions, but got" + f" type {type(videos)}." + ) + + # Recursively flatten any nested structure + flat_videos_list = [] + for item in videos: + if isinstance(item, str) or is_valid_video(item): + flat_videos_list.append(item) + elif isinstance(item, list) and item: + flat_videos_list.extend(make_batched_videos(item)) + + flat_videos_list = convert_pil_frames_to_video(flat_videos_list) + return flat_videos_list + + +def make_batched_metadata(videos: VideoInput, video_metadata: VideoMetadataType) -> list[VideoMetadata]: + if video_metadata is None: + # Create default metadata and fill attributes we can infer from given video + video_metadata = [ + { + "total_num_frames": len(video), + "fps": None, + "duration": None, + "frames_indices": list(range(len(video))), + "height": get_video_size(video)[0] if is_valid_video(video) else None, + "width": get_video_size(video)[1] if is_valid_video(video) else None, + } + for video in videos + ] + + if isinstance(video_metadata, list): + # Flatten if nested list + if isinstance(video_metadata[0], list): + video_metadata = [ + VideoMetadata(**metadata) for metadata_list in video_metadata for metadata in metadata_list + ] + # Simply wrap in VideoMetadata if simple dict + elif isinstance(video_metadata[0], dict): + video_metadata = [VideoMetadata(**metadata) for metadata in video_metadata] + else: + # Create a batched list from single object + video_metadata = [VideoMetadata(**video_metadata)] + return video_metadata + + +def get_video_size(video: np.ndarray, channel_dim: ChannelDimension | None = None) -> tuple[int, int]: + """ + Returns the (height, width) dimensions of the video. + + Args: + video (`np.ndarray`): + The video to get the dimensions of. + channel_dim (`ChannelDimension`, *optional*): + Which dimension the channel dimension is in. If `None`, will infer the channel dimension from the video. + + Returns: + A tuple of the video's height and width. + """ + if channel_dim is None: + channel_dim = infer_channel_dimension_format(video, num_channels=(1, 3, 4)) + + if channel_dim == ChannelDimension.FIRST: + return video.shape[-2], video.shape[-1] + elif channel_dim == ChannelDimension.LAST: + return video.shape[-3], video.shape[-2] + else: + raise ValueError(f"Unsupported data format: {channel_dim}") + + +def get_uniform_frame_indices(total_num_frames: int, num_frames: int | None = None): + """ + Creates a numpy array for uniform sampling of `num_frame` frames from `total_num_frames` + when loading a video. + + Args: + total_num_frames (`int`): + Total number of frames that a video has. + num_frames (`int`, *optional*): + Number of frames to sample uniformly. If not specified, all frames are sampled. + + Returns: + np.ndarray: np array of frame indices that will be sampled. + """ + if num_frames is not None: + indices = np.arange(0, total_num_frames, total_num_frames / num_frames).astype(int) + else: + indices = np.arange(0, total_num_frames).astype(int) + return indices + + +def default_sample_indices_fn(metadata: VideoMetadata, num_frames=None, fps=None, **kwargs): + """ + A default sampling function that replicates the logic used in get_uniform_frame_indices, + while optionally handling `fps` if `num_frames` is not provided. + + Args: + metadata (`VideoMetadata`): + `VideoMetadata` object containing metadata about the video, such as "total_num_frames" or "fps". + num_frames (`int`, *optional*): + Number of frames to sample uniformly. + fps (`int` or `float`, *optional*): + Desired frames per second. Takes priority over num_frames if both are provided. + + Returns: + `np.ndarray`: Array of frame indices to sample. + """ + total_num_frames = metadata.total_num_frames + video_fps = metadata.fps + + # If num_frames is not given but fps is, calculate num_frames from fps + if num_frames is None and fps is not None: + num_frames = int(total_num_frames / video_fps * fps) + if num_frames > total_num_frames: + raise ValueError( + f"When loading the video with fps={fps}, we computed num_frames={num_frames} " + f"which exceeds total_num_frames={total_num_frames}. Check fps or video metadata." + ) + + if num_frames is not None: + indices = np.arange(0, total_num_frames, total_num_frames / num_frames, dtype=int) + else: + indices = np.arange(0, total_num_frames, dtype=int) + return indices + + +def read_video_opencv( + video_path: Union["URL", "Path"], + sample_indices_fn: Callable, + **kwargs, +) -> tuple[np.ndarray, VideoMetadata]: + """ + Decode a video using the OpenCV backend. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + tuple[`np.ndarray`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import cv2 + requires_backends(read_video_opencv, ["cv2"]) + import cv2 + + video = cv2.VideoCapture(video_path) + total_num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + video_fps = video.get(cv2.CAP_PROP_FPS) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), + fps=float(video_fps), + duration=float(duration), + video_backend="opencv", + height=int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)), + width=int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + index = 0 + frames = [] + while video.isOpened(): + success, frame = video.read() + if not success: + break + if index in indices: + height, width, channel = frame.shape + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame[0:height, 0:width, 0:channel]) + if success: + index += 1 + if index >= total_num_frames: + break + + video.release() + metadata.frames_indices = indices + return np.stack(frames), metadata + + +def read_video_decord( + video_path: Union["URL", "Path"], + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode a video using the Decord backend. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import from decord + requires_backends(read_video_decord, ["decord"]) + from decord import VideoReader, cpu + + vr = VideoReader(uri=video_path, ctx=cpu(0)) # decord has problems with gpu + video_fps = vr.get_avg_fps() + total_num_frames = len(vr) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), + fps=float(video_fps), + duration=float(duration), + video_backend="decord", + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + video = vr.get_batch(indices).asnumpy() + + metadata.update( + { + "frames_indices": indices, + "height": video.shape[1], + "width": video.shape[2], + } + ) + return video, metadata + + +def read_video_pyav( + video_path: Union["URL", "Path"], + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode the video with PyAV decoder. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import av + requires_backends(read_video_pyav, ["av"]) + import av + + container = av.open(video_path) + total_num_frames = container.streams.video[0].frames + video_fps = container.streams.video[0].average_rate # should we better use `av_guess_frame_rate`? + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), + fps=float(video_fps), + duration=float(duration), + video_backend="pyav", + height=container.streams.video[0].height, + width=container.streams.video[0].width, + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + frames = [] + container.seek(0) + end_index = indices[-1] + for i, frame in enumerate(container.decode(video=0)): + if i > end_index: + break + if i >= 0 and i in indices: + frames.append(frame) + + video = np.stack([x.to_ndarray(format="rgb24") for x in frames]) + metadata.frames_indices = indices + return video, metadata + + +def read_video_torchvision( + video_path: Union["URL", "Path"], + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode the video with torchvision decoder. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + tuple[`torch.Tensor`, `VideoMetadata`]: A tuple containing: + - Torch tensor of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + warnings.warn( + "Using `torchvision` for video decoding is deprecated and will be removed in future versions. " + "Please use `torchcodec` instead." + ) + video, _, info = torchvision_io.read_video( + video_path, + start_pts=0.0, + end_pts=None, + pts_unit="sec", + output_format="TCHW", + ) + video_fps = info["video_fps"] + total_num_frames = video.size(0) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), + fps=float(video_fps), + duration=float(duration), + video_backend="torchvision", + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + video = video[indices].contiguous() + metadata.update( + { + "frames_indices": indices, + "height": video.shape[2], + "width": video.shape[3], + } + ) + return video, metadata + + +def read_video_torchcodec( + video_path: Union["URL", "Path"], + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode the video with torchcodec decoder. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`torch.Tensor`, `VideoMetadata`]: A tuple containing: + - Torch tensor of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import torchcodec + requires_backends(read_video_torchcodec, ["torchcodec"]) + from torchcodec.decoders import VideoDecoder + + # VideoDecoder expects a string for device, default to "cpu" if None + + decoder = VideoDecoder( + video_path, + # Interestingly `exact` mode takes less than approximate when we load the whole video + seek_mode="exact", + # Allow FFmpeg decide on the number of threads for efficiency + num_ffmpeg_threads=0, + device=kwargs.get("device", "cpu"), + ) + total_num_frames = decoder.metadata.num_frames + video_fps = decoder.metadata.average_fps + metadata = VideoMetadata( + total_num_frames=total_num_frames, + fps=video_fps, + duration=decoder.metadata.duration_seconds, + video_backend="torchcodec", + height=decoder.metadata.height, + width=decoder.metadata.width, + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + video = decoder.get_frames_at(indices=indices).data.contiguous() + metadata.frames_indices = indices + return video, metadata + + +VIDEO_DECODERS = { + "decord": read_video_decord, + "opencv": read_video_opencv, + "pyav": read_video_pyav, + "torchvision": read_video_torchvision, + "torchcodec": read_video_torchcodec, +} + + +def load_video( + video: VideoInput, + num_frames: int | None = None, + fps: int | float | None = None, + backend: str = "pyav", + sample_indices_fn: Callable | None = None, + **kwargs, +) -> np.ndarray: + """ + Loads `video` to a numpy array. + + Args: + video (`VideoInput`): + The video to convert to the numpy array format. Can be a link to video or local path. + num_frames (`int`, *optional*): + Number of frames to sample uniformly. If not passed, the whole video is loaded. + fps (`int` or `float`, *optional*): + Number of frames to sample per second. Should be passed only when `num_frames=None`. + If not specified and `num_frames==None`, all frames are sampled. + backend (`str`, *optional*, defaults to `"pyav"`): + The backend to use when loading the video. Can be any of ["decord", "pyav", "opencv", "torchvision", "torchcodec"]. Defaults to "pyav". + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniformt sampling with fps is performed, otherwise `sample_indices_fn` has priority over other args. + The function expects at input the all args along with all kwargs passed to `load_video` and should output valid + indices at which the video should be sampled. For example: + + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + tuple[`np.ndarray`, Dict]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - Metadata dictionary. + """ + + # If `sample_indices_fn` is given, we can accept any args as those might be needed by custom `sample_indices_fn` + if fps is not None and num_frames is not None and sample_indices_fn is None: + raise ValueError( + "`num_frames`, `fps`, and `sample_indices_fn` are mutually exclusive arguments, please use only one!" + ) + + # If user didn't pass a sampling function, create one on the fly with default logic + if sample_indices_fn is None: + + def sample_indices_fn_func(metadata, **fn_kwargs): + return default_sample_indices_fn(metadata, num_frames=num_frames, fps=fps, **fn_kwargs) + + sample_indices_fn = sample_indices_fn_func + + # Early exit if provided an array or `PIL` frames + if not isinstance(video, str): + metadata = [None] * len(video) + return video, metadata + + if urlparse(video).netloc in ["www.youtube.com", "youtube.com"]: + if not is_yt_dlp_available(): + raise ImportError("To load a video from YouTube url you have to install `yt_dlp` first.") + # Lazy import from yt_dlp + requires_backends(load_video, ["yt_dlp"]) + from yt_dlp import YoutubeDL + + buffer = BytesIO() + with redirect_stdout(buffer), YoutubeDL() as f: + f.download([video]) + bytes_obj = buffer.getvalue() + file_obj = BytesIO(bytes_obj) + elif video.startswith("http://") or video.startswith("https://"): + file_obj = BytesIO(httpx.get(video, follow_redirects=True).content) + elif os.path.isfile(video): + file_obj = video + else: + raise TypeError("Incorrect format used for video. Should be an url linking to an video or a local path.") + + # can also load with decord, but not cv2/torchvision + # both will fail in case of url links + video_is_url = video.startswith("http://") or video.startswith("https://") + if video_is_url and backend == "opencv": + raise ValueError("If you are trying to load a video from URL, you cannot use 'opencv' as backend") + + if ( + (not is_decord_available() and backend == "decord") + or (not is_av_available() and backend == "pyav") + or (not is_cv2_available() and backend == "opencv") + or (not is_torchvision_available() and backend == "torchvision") + or (not is_torchcodec_available() and backend == "torchcodec") + ): + raise ImportError( + f"You chose backend={backend} for loading the video but the required library is not found in your environment " + f"Make sure to install {backend} before loading the video." + ) + + video_decoder = VIDEO_DECODERS[backend] + video, metadata = video_decoder(file_obj, sample_indices_fn, **kwargs) + return video, metadata + + +def convert_to_rgb( + video: np.ndarray, + input_data_format: str | ChannelDimension | None = None, +) -> np.ndarray: + """ + Convert video to RGB by blending the transparency layer if it's in RGBA format, otherwise simply returns it. + + Args: + video (`np.ndarray`): + The video to convert. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input video. If unset, will use the inferred format from the input. + """ + if not isinstance(video, np.ndarray): + raise TypeError(f"Video has to be a numpy array to convert to RGB format, but found {type(video)}") + + # np.array usually comes with ChannelDimension.LAST so let's convert it + if input_data_format is None: + input_data_format = infer_channel_dimension_format(video) + video = to_channel_dimension_format(video, ChannelDimension.FIRST, input_channel_dim=input_data_format) + + # 3 channels for RGB already + if video.shape[-3] == 3: + return video + + # Grayscale video so we repeat it 3 times for each channel + if video.shape[-3] == 1: + return video.repeat(3, -3) + + if not (video[..., 3, :, :] < 255).any(): + return video + + # There is a transparency layer, blend it with a white background. + # Calculate the alpha proportion for blending. + alpha = video[..., 3, :, :] / 255.0 + video = (1 - alpha[..., None, :, :]) * 255 + alpha[..., None, :, :] * video[..., 3, :, :] + return video + + +def pad( + video: np.ndarray, + padding: int | tuple[int, int] | Iterable[tuple[int, int]], + mode: PaddingMode = PaddingMode.CONSTANT, + constant_values: float | Iterable[float] = 0.0, + data_format: str | ChannelDimension | None = None, + input_data_format: str | ChannelDimension | None = None, +) -> np.ndarray: + """ + Pads the `video` with the specified (height, width) `padding` and `mode`. + + Args: + video (`np.ndarray`): + The video to pad. + padding (`int` or `tuple[int, int]` or `Iterable[tuple[int, int]]`): + Padding to apply to the edges of the height, width axes. Can be one of three formats: + - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis. + - `((before, after),)` yields same before and after pad for height and width. + - `(pad,)` or int is a shortcut for before = after = pad width for all axes. + mode (`PaddingMode`): + The padding mode to use. Can be one of: + - `"constant"`: pads with a constant value. + - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the + vector along each axis. + - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis. + - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array. + constant_values (`float` or `Iterable[float]`, *optional*): + The value to use for the padding if `mode` is `"constant"`. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output video. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_frames, num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: video in (num_frames, height, width, num_channels) format. + If unset, will use same as the input video. + input_data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the input video. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_frames, num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: video in (num_frames, height, width, num_channels) format. + If unset, will use the inferred format of the input video. + + Returns: + `np.ndarray`: The padded video. + + """ + if input_data_format is None: + input_data_format = infer_channel_dimension_format(video) + + def _expand_for_data_format(values): + """ + Convert values to be in the format expected by np.pad based on the data format. + """ + if isinstance(values, (int, float)): + values = ((values, values), (values, values)) + elif isinstance(values, tuple) and len(values) == 1: + values = ((values[0], values[0]), (values[0], values[0])) + elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], int): + values = (values, values) + elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], tuple): + pass + else: + raise ValueError(f"Unsupported format: {values}") + + # add 0 for channel dimension + values = ( + ((0, 0), (0, 0), *values) if input_data_format == ChannelDimension.FIRST else ((0, 0), *values, (0, 0)) + ) + + # Add additional padding if there's a batch dimension + values = (0, *values) if video.ndim == 5 else values + return values + + padding_map = { + PaddingMode.CONSTANT: "constant", + PaddingMode.REFLECT: "reflect", + PaddingMode.REPLICATE: "replicate", + PaddingMode.SYMMETRIC: "symmetric", + } + padding = _expand_for_data_format(padding) + + pad_kwargs = {} + if mode not in padding_map: + raise ValueError(f"Invalid padding mode: {mode}") + elif mode == PaddingMode.CONSTANT: + pad_kwargs["constant_values"] = _expand_for_data_format(constant_values) + + video = np.pad(video, padding, mode=padding_map[mode], **pad_kwargs) + video = to_channel_dimension_format(video, data_format, input_data_format) if data_format is not None else video + return video + + +def group_videos_by_shape( + videos: list["torch.Tensor"], +) -> tuple[dict[tuple[int, int], "torch.Tensor"], dict[int, tuple[tuple[int, int], int]]]: + """ + Groups videos by shape. + Returns a dictionary with the shape as key and a list of videos with that shape as value, + and a dictionary with the index of the video in the original list as key and the shape and index in the grouped list as value. + """ + grouped_videos = {} + grouped_videos_index = {} + for i, video in enumerate(videos): + shape = video.shape[-2::] + num_frames = video.shape[-4] # video format BTCHW + shape = (num_frames, *shape) + if shape not in grouped_videos: + grouped_videos[shape] = [] + grouped_videos[shape].append(video) + grouped_videos_index[i] = (shape, len(grouped_videos[shape]) - 1) + # stack videos with the same size and number of frames + grouped_videos = {shape: torch.stack(videos, dim=0) for shape, videos in grouped_videos.items()} + return grouped_videos, grouped_videos_index + + +def reorder_videos( + processed_videos: dict[tuple[int, int], "torch.Tensor"], + grouped_videos_index: dict[int, tuple[tuple[int, int], int]], +) -> list["torch.Tensor"]: + """ + Reconstructs a list of videos in the original order. + """ + return [ + processed_videos[grouped_videos_index[i][0]][grouped_videos_index[i][1]] + for i in range(len(grouped_videos_index)) + ]