Visual question answering
-
-
Transformers acts as the model-definition framework for state-of-the-art machine learning models in text, computer
vision, audio, video, and multimodal model, for both inference and training.
@@ -35,6 +34,10 @@ There are over 1M+ Transformers [model checkpoints](https://huggingface.co/model
Explore the [Hub](https://huggingface.com/) today to find a model and use Transformers to help you get started right away.
+Explore the [Models Timeline](./models_timeline) to discover the latest text, vision, audio and multimodal model architectures in Transformers.
+
+
+
## Features
Transformers provides everything you need for inference or training with state-of-the-art pretrained models. Some of the main features include:
@@ -61,4 +64,4 @@ Transformers is designed for developers and machine learning engineers and resea
## Learn
-If you're new to Transformers or want to learn more about transformer models, we recommend starting with the [LLM course](https://huggingface.co/learn/llm-course/chapter1/1?fw=pt). This comprehensive course covers everything from the fundamentals of how transformer models work to practical applications across various tasks. You'll learn the complete workflow, from curating high-quality datasets to fine-tuning large language models and implementing reasoning capabilities. The course contains both theoretical and hands-on exercises to build a solid foundational knowledge of transformer models as you learn.
\ No newline at end of file
+If you're new to Transformers or want to learn more about transformer models, we recommend starting with the [LLM course](https://huggingface.co/learn/llm-course/chapter1/1?fw=pt). This comprehensive course covers everything from the fundamentals of how transformer models work to practical applications across various tasks. You'll learn the complete workflow, from curating high-quality datasets to fine-tuning large language models and implementing reasoning capabilities. The course contains both theoretical and hands-on exercises to build a solid foundational knowledge of transformer models as you learn.
diff --git a/docs/source/en/internal/file_utils.md b/docs/source/en/internal/file_utils.md
index 31fbc5b88110..63db5756a622 100644
--- a/docs/source/en/internal/file_utils.md
+++ b/docs/source/en/internal/file_utils.md
@@ -20,7 +20,6 @@ This page lists all of Transformers general utility functions that are found in
Most of those are only useful if you are studying the general code in the library.
-
## Enums and namedtuples
[[autodoc]] utils.ExplicitEnum
diff --git a/docs/source/en/internal/generation_utils.md b/docs/source/en/internal/generation_utils.md
index d47eba82d8cc..87b0111ff053 100644
--- a/docs/source/en/internal/generation_utils.md
+++ b/docs/source/en/internal/generation_utils.md
@@ -65,7 +65,6 @@ values. Here, for instance, it has two keys that are `sequences` and `scores`.
We document here all output types.
-
[[autodoc]] generation.GenerateDecoderOnlyOutput
[[autodoc]] generation.GenerateEncoderDecoderOutput
@@ -74,13 +73,11 @@ We document here all output types.
[[autodoc]] generation.GenerateBeamEncoderDecoderOutput
-
## LogitsProcessor
A [`LogitsProcessor`] can be used to modify the prediction scores of a language model head for
generation.
-
[[autodoc]] AlternatingCodebooksLogitsProcessor
- __call__
@@ -174,8 +171,6 @@ generation.
[[autodoc]] WatermarkLogitsProcessor
- __call__
-
-
## StoppingCriteria
A [`StoppingCriteria`] can be used to change when to stop generation (other than EOS token). Please note that this is exclusively available to our PyTorch implementations.
@@ -300,7 +295,6 @@ A [`Constraint`] can be used to force the generation to include specific tokens
- to_legacy_cache
- from_legacy_cache
-
## Watermark Utils
[[autodoc]] WatermarkingConfig
diff --git a/docs/source/en/internal/import_utils.md b/docs/source/en/internal/import_utils.md
index 0d76c2bbe33a..4a9915378a1f 100644
--- a/docs/source/en/internal/import_utils.md
+++ b/docs/source/en/internal/import_utils.md
@@ -22,8 +22,8 @@ worked around. We don't want for all users of `transformers` to have to install
we therefore mark those as soft dependencies rather than hard dependencies.
The transformers toolkit is not made to error-out on import of a model that has a specific dependency; instead, an
-object for which you are lacking a dependency will error-out when calling any method on it. As an example, if
-`torchvision` isn't installed, the fast image processors will not be available.
+object for which you are lacking a dependency will error-out when calling any method on it. As an example, if
+`torchvision` isn't installed, the fast image processors will not be available.
This object is still importable:
@@ -60,7 +60,7 @@ PyTorch dependency
**Tokenizers**: All files starting with `tokenization_` and ending with `_fast` have an automatic `tokenizers` dependency
-**Vision**: All files starting with `image_processing_` have an automatic dependency to the `vision` dependency group;
+**Vision**: All files starting with `image_processing_` have an automatic dependency to the `vision` dependency group;
at the time of writing, this only contains the `pillow` dependency.
**Vision + Torch + Torchvision**: All files starting with `image_processing_` and ending with `_fast` have an automatic
@@ -71,7 +71,7 @@ All of these automatic dependencies are added on top of the explicit dependencie
### Explicit Object Dependencies
We add a method called `requires` that is used to explicitly specify the dependencies of a given object. As an
-example, the `Trainer` class has two hard dependencies: `torch` and `accelerate`. Here is how we specify these
+example, the `Trainer` class has two hard dependencies: `torch` and `accelerate`. Here is how we specify these
required dependencies:
```python
diff --git a/docs/source/en/internal/model_debugging_utils.md b/docs/source/en/internal/model_debugging_utils.md
index 262113575f42..553a5ce56845 100644
--- a/docs/source/en/internal/model_debugging_utils.md
+++ b/docs/source/en/internal/model_debugging_utils.md
@@ -21,10 +21,8 @@ provides for it.
Most of those are only useful if you are adding new models in the library.
-
## Model addition debuggers
-
### 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
@@ -72,7 +70,6 @@ with model_addition_debugger_context(
```
-
### Reading results
The debugger generates two files from the forward call, both with the same base name, but ending either with
@@ -221,9 +218,9 @@ path reference to the associated `.safetensors` file. Each tensor is written to
the state dictionary. File names are constructed using the `module_path` as a prefix with a few possible postfixes that
are built recursively.
-* Module inputs are denoted with the `_inputs` and outputs by `_outputs`.
-* `list` and `tuple` instances, such as `args` or function return values, will be postfixed with `_{index}`.
-* `dict` instances will be postfixed with `_{key}`.
+* Module inputs are denoted with the `_inputs` and outputs by `_outputs`.
+* `list` and `tuple` instances, such as `args` or function return values, will be postfixed with `_{index}`.
+* `dict` instances will be postfixed with `_{key}`.
### Comparing between implementations
@@ -231,10 +228,8 @@ Once the forward passes of two models have been traced by the debugger, one can
below: we can see slight differences between these two implementations' key projection layer. Inputs are mostly
identical, but not quite. Looking through the file differences makes it easier to pinpoint which layer is wrong.
-
-
### Limitations and scope
This feature will only work for torch-based models, and would require more work and case-by-case approach for say
@@ -254,13 +249,14 @@ layers.
This small util is a power user tool intended for model adders and maintainers. It lists all test methods
existing in `test_modeling_common.py`, inherited by all model tester classes, and scans the repository to measure
-how many tests are being skipped and for which models.
+how many tests are being skipped and for which models.
### Rationale
When porting models to transformers, tests fail as they should, and sometimes `test_modeling_common` feels irreconcilable with the peculiarities of our brand new model. But how can we be sure we're not breaking everything by adding a seemingly innocent skip?
This utility:
+
- scans all test_modeling_common methods
- looks for times where a method is skipped
- returns a summary json you can load as a DataFrame/inspect
@@ -269,8 +265,7 @@ This utility:
-
-### Usage
+### Usage
You can run the skipped test analyzer in two ways:
@@ -286,7 +281,7 @@ python utils/scan_skipped_tests.py --output_dir path/to/output
**Example output:**
-```
+```text
🔬 Parsing 331 model test files once each...
📝 Aggregating 224 tests...
(224/224) test_update_candidate_strategy_with_matches_1es_3d_is_nonecodet_schedule_fa_kwargs
diff --git a/docs/source/en/internal/pipelines_utils.md b/docs/source/en/internal/pipelines_utils.md
index 6ea6de9a61b8..23856e5639c3 100644
--- a/docs/source/en/internal/pipelines_utils.md
+++ b/docs/source/en/internal/pipelines_utils.md
@@ -20,7 +20,6 @@ This page lists all the utility functions the library provides for pipelines.
Most of those are only useful if you are studying the code of the models in the library.
-
## Argument handling
[[autodoc]] pipelines.ArgumentHandler
diff --git a/docs/source/en/jan.md b/docs/source/en/jan.md
index ff580496c81b..95309f46cd04 100644
--- a/docs/source/en/jan.md
+++ b/docs/source/en/jan.md
@@ -25,7 +25,7 @@ You are now ready to chat!
To conclude this example, let's look into a more advanced use-case. If you have a beefy machine to serve models with, but prefer using Jan on a different device, you need to add port forwarding. If you have `ssh` access from your Jan machine into your server, this can be accomplished by typing the following to your Jan machine's terminal
-```
+```bash
ssh -N -f -L 8000:localhost:8000 your_server_account@your_server_IP -p port_to_ssh_into_your_server
```
diff --git a/docs/source/en/kv_cache.md b/docs/source/en/kv_cache.md
index f0a781cba4fc..f318c73d28a9 100644
--- a/docs/source/en/kv_cache.md
+++ b/docs/source/en/kv_cache.md
@@ -67,7 +67,7 @@ out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_valu
## Fixed-size cache
-The default [`DynamicCache`] prevents you from taking advantage of most just-in-time (JIT) optimizations because the cache size isn't fixed. JIT optimizations enable you to maximize latency at the expense of memory usage. All of the following cache types are compatible with JIT optimizations like [torch.compile](./llm_optims#static-kv-cache-and-torchcompile) to accelerate generation.
+The default [`DynamicCache`] prevents you from taking advantage of most just-in-time (JIT) optimizations because the cache size isn't fixed. JIT optimizations enable you to maximize latency at the expense of memory usage. All of the following cache types are compatible with JIT optimizations like [torch.compile](./llm_optims#static-kv-cache-and-torchcompile) to accelerate generation.
A fixed-size cache ([`StaticCache`]) pre-allocates a specific maximum cache size for the kv pairs. You can generate up to the maximum cache size without needing to modify it. However, having a fixed (usually large) size for the key/value states means that while generating, a lot of tokens will actually be masked as they should not take part in the attention. So this trick allows to easily `compile` the decoding stage, but it incurs a waste of tokens in the attention computation. As all things, it's then a trade-off which should be very good if you generate with several sequence of more or less the same lengths, but may be sub-optimal if you have for example 1 very large sequence, and then only short sequences (as the fix cache size would be large, a lot would be wasted for the short sequences). Make sure you understand the impact if you use it!
@@ -213,7 +213,7 @@ A cache can also work in iterative generation settings where there is back-and-f
For iterative generation with a cache, start by initializing an empty cache class and then you can feed in your new prompts. Keep track of dialogue history with a [chat template](./chat_templating).
-The following example demonstrates [Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf). If you’re using a different chat-style model, [`~PreTrainedTokenizer.apply_chat_template`] may process messages differently. It might cut out important tokens depending on how the Jinja template is written.
+The following example demonstrates [Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf). If you're using a different chat-style model, [`~PreTrainedTokenizer.apply_chat_template`] may process messages differently. It might cut out important tokens depending on how the Jinja template is written.
For example, some models use special `
@@ -108,7 +106,6 @@ To put this into perspective, on an NVIDIA A100 and when generating 400 semantic
At batch size 8, on an NVIDIA A100, Flash Attention 2 is also 10% faster than Better Transformer, and at batch size 16, 25%.
-
#### Combining optimization techniques
You can combine optimization techniques, and use CPU offload, half-precision and Flash Attention 2 (or 🤗 Better Transformer) all at once.
@@ -147,7 +144,7 @@ These presets are also uploaded in the hub [here](https://huggingface.co/suno/ba
>>> audio_array = audio_array.cpu().numpy().squeeze()
```
-Bark can generate highly realistic, **multilingual** speech as well as other audio - including music, background noise and simple sound effects.
+Bark can generate highly realistic, **multilingual** speech as well as other audio - including music, background noise and simple sound effects.
```python
>>> # Multilingual speech - simplified Chinese
@@ -165,7 +162,6 @@ Bark can generate highly realistic, **multilingual** speech as well as other aud
The model can also produce **nonverbal communications** like laughing, sighing and crying.
-
```python
>>> # Adding non-speech cues to the input text
>>> inputs = processor("Hello uh ... [clears throat], my dog is cute [laughter]")
@@ -235,4 +231,3 @@ To save the audio, simply take the sample rate from the model config and some sc
[[autodoc]] BarkSemanticConfig
- all
-
diff --git a/docs/source/en/model_doc/bart.md b/docs/source/en/model_doc/bart.md
index b0252ea92311..daa65d6afc0c 100644
--- a/docs/source/en/model_doc/bart.md
+++ b/docs/source/en/model_doc/bart.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2019-10-29 and added to Hugging Face Transformers on 2020-11-16.*
-
@@ -24,7 +23,7 @@ rendered properly in your Markdown viewer.
# BART
-[BART](https://huggingface.co/papers/1910.13461) is a sequence-to-sequence model that combines the pretraining objectives from BERT and GPT. It’s pretrained by corrupting text in different ways like deleting words, shuffling sentences, or masking tokens and learning how to fix it. The encoder encodes the corrupted document and the corrupted text is fixed by the decoder. As it learns to recover the original text, BART gets really good at both understanding and generating language.
+[BART](https://huggingface.co/papers/1910.13461) is a sequence-to-sequence model that combines the pretraining objectives from BERT and GPT. It's pretrained by corrupting text in different ways like deleting words, shuffling sentences, or masking tokens and learning how to fix it. The encoder encodes the corrupted document and the corrupted text is fixed by the decoder. As it learns to recover the original text, BART gets really good at both understanding and generating language.
You can find all the original BART checkpoints under the [AI at Meta](https://huggingface.co/facebook?search_models=bart) organization.
@@ -46,6 +45,7 @@ pipeline = pipeline(
pipeline("Plants create through a process known as photosynthesis.")
```
+
@@ -89,7 +89,7 @@ echo -e "Plants create through a process known as photosynthesis." | tran
- Inputs should be padded on the right because BERT uses absolute position embeddings.
- The [facebook/bart-large-cnn](https://huggingface.co/facebook/bart-large-cnn) checkpoint doesn't include `mask_token_id` which means it can't perform mask-filling tasks.
-- BART doesn’t use `token_type_ids` for sequence classification. Use [`BartTokenizer`] or [`~PreTrainedTokenizerBase.encode`] to get the proper splitting.
+- BART doesn't use `token_type_ids` for sequence classification. Use [`BartTokenizer`] or [`~PreTrainedTokenizerBase.encode`] to get the proper splitting.
- The forward pass of [`BartModel`] creates the `decoder_input_ids` if they're not passed. This can be different from other model APIs, but it is a useful feature for mask-filling tasks.
- Model predictions are intended to be identical to the original implementation when `forced_bos_token_id=0`. This only works if the text passed to `fairseq.encode` begins with a space.
- [`~GenerationMixin.generate`] should be used for conditional generation tasks like summarization.
diff --git a/docs/source/en/model_doc/barthez.md b/docs/source/en/model_doc/barthez.md
index 43b6521f1013..f7a100a4208c 100644
--- a/docs/source/en/model_doc/barthez.md
+++ b/docs/source/en/model_doc/barthez.md
@@ -31,7 +31,6 @@ You can find all of the original BARThez checkpoints under the [BARThez](https:/
> This model was contributed by [moussakam](https://huggingface.co/moussakam).
> Refer to the [BART](./bart) docs for more usage examples.
-
The example below demonstrates how to predict the `` token with [`Pipeline`], [`AutoModel`], and from the command line.
diff --git a/docs/source/en/model_doc/bartpho.md b/docs/source/en/model_doc/bartpho.md
index 9e86a1b615d0..15e96c57669f 100644
--- a/docs/source/en/model_doc/bartpho.md
+++ b/docs/source/en/model_doc/bartpho.md
@@ -33,12 +33,9 @@ You can find all the original checkpoints under the [VinAI](https://huggingface.
The example below demonstrates how to summarize text with [`Pipeline`] or the [`AutoModel`] class.
-
-
-
```python
import torch
from transformers import pipeline
@@ -98,8 +95,6 @@ transformers run --task summarization --model vinai/bartpho-word --device 0
-
-
## Notes
- BARTpho uses the large architecture of BART with an additional layer-normalization layer on top of the encoder and decoder. The BART-specific classes should be replaced with the mBART-specific classes.
diff --git a/docs/source/en/model_doc/beit.md b/docs/source/en/model_doc/beit.md
index b66021ec8d98..ee516a935ed4 100644
--- a/docs/source/en/model_doc/beit.md
+++ b/docs/source/en/model_doc/beit.md
@@ -87,7 +87,7 @@ page for more information.
SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import BeitForImageClassification
model = BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224", attn_implementation="sdpa", dtype=torch.float16)
...
@@ -123,6 +123,7 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- See also: [Image classification task guide](../tasks/image_classification)
**Semantic segmentation**
+
- [Semantic segmentation task guide](../tasks/semantic_segmentation)
If you're interested in submitting a resource to be included here, please feel free to open a Pull Request and we'll review it! The resource should ideally demonstrate something new instead of duplicating an existing resource.
diff --git a/docs/source/en/model_doc/bert-generation.md b/docs/source/en/model_doc/bert-generation.md
index 38cbe2137eb7..d57734b069ba 100644
--- a/docs/source/en/model_doc/bert-generation.md
+++ b/docs/source/en/model_doc/bert-generation.md
@@ -13,6 +13,7 @@ specific language governing permissions and limitations under the License.
rendered properly in your Markdown viewer.
-->
+*This model was released on 2019-07-29 and added to Hugging Face Transformers on 2020-11-16.*
-
# Cohere
Cohere [Command-R](https://cohere.com/blog/command-r) is a 35B parameter multilingual large language model designed for long context tasks like retrieval-augmented generation (RAG) and calling external APIs and tools. The model is specifically trained for grounded generation and supports both single-step and multi-step tool use. It supports a context length of 128K tokens.
You can find all the original Command-R checkpoints under the [Command Models](https://huggingface.co/collections/CohereForAI/command-models-67652b401665205e17b192ad) collection.
-
> [!TIP]
> Click on the Cohere models in the right sidebar for more examples of how to apply Cohere to different language tasks.
@@ -123,9 +121,9 @@ visualizer("Plants create energy through a process known as")
-
## Notes
-- Don’t use the dtype parameter in [`~AutoModel.from_pretrained`] if you’re using FlashAttention-2 because it only supports fp16 or bf16. You should use [Automatic Mixed Precision](https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html), set fp16 or bf16 to True if using [`Trainer`], or use [torch.autocast](https://pytorch.org/docs/stable/amp.html#torch.autocast).
+
+- Don't use the dtype parameter in [`~AutoModel.from_pretrained`] if you're using FlashAttention-2 because it only supports fp16 or bf16. You should use [Automatic Mixed Precision](https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html), set fp16 or bf16 to True if using [`Trainer`], or use [torch.autocast](https://pytorch.org/docs/stable/amp.html#torch.autocast).
## CohereConfig
@@ -145,7 +143,6 @@ visualizer("Plants create energy through a process known as")
[[autodoc]] CohereModel
- forward
-
## CohereForCausalLM
[[autodoc]] CohereForCausalLM
diff --git a/docs/source/en/model_doc/cohere2.md b/docs/source/en/model_doc/cohere2.md
index bcfa05e98d19..52555d6ae558 100644
--- a/docs/source/en/model_doc/cohere2.md
+++ b/docs/source/en/model_doc/cohere2.md
@@ -22,7 +22,6 @@ rendered properly in your Markdown viewer.
-
# Cohere 2
[Cohere Command R7B](https://cohere.com/blog/command-r7b) is an open weights research release of a 7B billion parameter model. It is a multilingual model trained on 23 languages and has a context window of 128k. The model features three layers with sliding window attention and ROPE for efficient local context modeling and relative positional encoding. A fourth layer uses global attention without positional embeddings, enabling unrestricted token interactions across the entire sequence.
@@ -31,7 +30,6 @@ This model is optimized for speed, cost-performance, and compute resources.
You can find all the original Command-R checkpoints under the [Command Models](https://huggingface.co/collections/CohereForAI/command-models-67652b401665205e17b192ad) collection.
-
> [!TIP]
> Click on the Cohere models in the right sidebar for more examples of how to apply Cohere to different language tasks.
@@ -136,7 +134,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
[[autodoc]] Cohere2Model
- forward
-
## Cohere2ForCausalLM
[[autodoc]] Cohere2ForCausalLM
diff --git a/docs/source/en/model_doc/cohere2_vision.md b/docs/source/en/model_doc/cohere2_vision.md
index 2e12ff3e4767..e466ce6a5f09 100644
--- a/docs/source/en/model_doc/cohere2_vision.md
+++ b/docs/source/en/model_doc/cohere2_vision.md
@@ -113,6 +113,7 @@ outputs = pipe(text=messages, max_new_tokens=300, return_full_text=False)
print(outputs)
```
+
diff --git a/docs/source/en/model_doc/cpm.md b/docs/source/en/model_doc/cpm.md
index ccfa1596bad4..275f5629db13 100644
--- a/docs/source/en/model_doc/cpm.md
+++ b/docs/source/en/model_doc/cpm.md
@@ -42,7 +42,6 @@ NLP tasks in the settings of few-shot (even zero-shot) learning.*
This model was contributed by [canwenxu](https://huggingface.co/canwenxu). The original implementation can be found
here: https://github.com/TsinghuaAI/CPM-Generate
-
@@ -24,7 +23,7 @@ rendered properly in your Markdown viewer.
@@ -155,4 +156,4 @@ print(tokenizer.decode(outputs[0]))
## BertGenerationDecoder
[[autodoc]] BertGenerationDecoder
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/bert-japanese.md b/docs/source/en/model_doc/bert-japanese.md
index 812e5a455ad5..6599efa73e08 100644
--- a/docs/source/en/model_doc/bert-japanese.md
+++ b/docs/source/en/model_doc/bert-japanese.md
@@ -81,7 +81,6 @@ API reference information.
-
## BertJapaneseTokenizer
[[autodoc]] BertJapaneseTokenizer
diff --git a/docs/source/en/model_doc/bertweet.md b/docs/source/en/model_doc/bertweet.md
index 4dffe29168d3..20206da87e43 100644
--- a/docs/source/en/model_doc/bertweet.md
+++ b/docs/source/en/model_doc/bertweet.md
@@ -24,8 +24,7 @@ rendered properly in your Markdown viewer.
## BERTweet
-[BERTweet](https://huggingface.co/papers/2005.10200) shares the same architecture as [BERT-base](./bert), but it’s pretrained like [RoBERTa](./roberta) on English Tweets. It performs really well on Tweet-related tasks like part-of-speech tagging, named entity recognition, and text classification.
-
+[BERTweet](https://huggingface.co/papers/2005.10200) shares the same architecture as [BERT-base](./bert), but it's pretrained like [RoBERTa](./roberta) on English Tweets. It performs really well on Tweet-related tasks like part-of-speech tagging, named entity recognition, and text classification.
You can find all the original BERTweet checkpoints under the [VinAI Research](https://huggingface.co/vinai?search_models=BERTweet) organization.
@@ -49,6 +48,7 @@ pipeline = pipeline(
)
pipeline("Plants create through a process known as photosynthesis.")
```
+
@@ -88,7 +88,8 @@ echo -e "Plants create through a process known as photosynthesis." | tran
## Notes
-- Use the [`AutoTokenizer`] or [`BertweetTokenizer`] because it’s preloaded with a custom vocabulary adapted to tweet-specific tokens like hashtags (#), mentions (@), emojis, and common abbreviations. Make sure to also install the [emoji](https://pypi.org/project/emoji/) library.
+
+- Use the [`AutoTokenizer`] or [`BertweetTokenizer`] because it's preloaded with a custom vocabulary adapted to tweet-specific tokens like hashtags (#), mentions (@), emojis, and common abbreviations. Make sure to also install the [emoji](https://pypi.org/project/emoji/) library.
- Inputs should be padded on the right (`padding="max_length"`) because BERT uses absolute position embeddings.
## BertweetTokenizer
diff --git a/docs/source/en/model_doc/big_bird.md b/docs/source/en/model_doc/big_bird.md
index 2d3b6d545faf..b4bfeefa516a 100644
--- a/docs/source/en/model_doc/big_bird.md
+++ b/docs/source/en/model_doc/big_bird.md
@@ -47,6 +47,7 @@ pipeline = pipeline(
)
pipeline("Plants create [MASK] through a process known as photosynthesis.")
```
+
@@ -81,10 +82,12 @@ print(f"The predicted token is: {predicted_token}")
```bash
!echo -e "Plants create [MASK] through a process known as photosynthesis." | transformers-cli run --task fill-mask --model google/bigbird-roberta-base --device 0
```
+
## Notes
+
- Inputs should be padded on the right because BigBird uses absolute position embeddings.
- BigBird supports `original_full` and `block_sparse` attention. If the input sequence length is less than 1024, it is recommended to use `original_full` since sparse patterns don't offer much benefit for smaller inputs.
- The current implementation uses window size of 3 blocks and 2 global blocks, only supports the ITC-implementation, and doesn't support `num_random_blocks=0`.
diff --git a/docs/source/en/model_doc/bigbird_pegasus.md b/docs/source/en/model_doc/bigbird_pegasus.md
index cae1e8f779d4..c4a6d54b9442 100644
--- a/docs/source/en/model_doc/bigbird_pegasus.md
+++ b/docs/source/en/model_doc/bigbird_pegasus.md
@@ -52,6 +52,7 @@ Through photosynthesis, plants capture energy from sunlight using a green pigmen
These ingredients are then transformed into glucose, a type of sugar that serves as a source of chemical energy, and oxygen, which is released as a byproduct into the atmosphere. The glucose produced during photosynthesis is not just used immediately; plants also store it as starch or convert it into other organic compounds like cellulose, which is essential for building their cellular structure.
This energy reserve allows them to grow, develop leaves, produce flowers, bear fruit, and carry out various physiological processes throughout their lifecycle.""")
```
+
@@ -77,6 +78,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = model.generate(**input_ids, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
diff --git a/docs/source/en/model_doc/biogpt.md b/docs/source/en/model_doc/biogpt.md
index 60b84f015122..82c2cb0e8cd0 100644
--- a/docs/source/en/model_doc/biogpt.md
+++ b/docs/source/en/model_doc/biogpt.md
@@ -135,31 +135,26 @@ print(output)
[[autodoc]] BioGptConfig
-
## BioGptTokenizer
[[autodoc]] BioGptTokenizer
- save_vocabulary
-
## BioGptModel
[[autodoc]] BioGptModel
- forward
-
## BioGptForCausalLM
[[autodoc]] BioGptForCausalLM
- forward
-
## BioGptForTokenClassification
[[autodoc]] BioGptForTokenClassification
- forward
-
## BioGptForSequenceClassification
[[autodoc]] BioGptForSequenceClassification
diff --git a/docs/source/en/model_doc/bit.md b/docs/source/en/model_doc/bit.md
index 5a6630566fca..5ed3b8f816ab 100644
--- a/docs/source/en/model_doc/bit.md
+++ b/docs/source/en/model_doc/bit.md
@@ -36,6 +36,7 @@ The original code can be found [here](https://github.com/google-research/big_tra
## Usage tips
- BiT models are equivalent to ResNetv2 in terms of architecture, except that: 1) all batch normalization layers are replaced by [group normalization](https://huggingface.co/papers/1803.08494),
+
2) [weight standardization](https://huggingface.co/papers/1903.10520) is used for convolutional layers. The authors show that the combination of both is useful for training with large batch sizes, and has a significant
impact on transfer learning.
@@ -72,4 +73,4 @@ If you're interested in submitting a resource to be included here, please feel f
## BitForImageClassification
[[autodoc]] BitForImageClassification
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/bitnet.md b/docs/source/en/model_doc/bitnet.md
index 6946ec65d437..c674f51fc305 100644
--- a/docs/source/en/model_doc/bitnet.md
+++ b/docs/source/en/model_doc/bitnet.md
@@ -35,33 +35,29 @@ Several versions of the model weights are available on Hugging Face:
* [**`microsoft/bitnet-b1.58-2B-4T-gguf`**](https://huggingface.co/microsoft/bitnet-b1.58-2B-4T-gguf): Contains the model weights in GGUF format, compatible with the `bitnet.cpp` library for CPU inference.
-
### Model Details
-
* **Architecture:** Transformer-based, modified with `BitLinear` layers (BitNet framework).
- * Uses Rotary Position Embeddings (RoPE).
- * Uses squared ReLU (ReLU²) activation in FFN layers.
- * Employs [`subln`](https://proceedings.mlr.press/v202/wang23u.html) normalization.
- * No bias terms in linear or normalization layers.
+ * Uses Rotary Position Embeddings (RoPE).
+ * Uses squared ReLU (ReLU²) activation in FFN layers.
+ * Employs [`subln`](https://proceedings.mlr.press/v202/wang23u.html) normalization.
+ * No bias terms in linear or normalization layers.
* **Quantization:** Native 1.58-bit weights and 8-bit activations (W1.58A8).
- * Weights are quantized to ternary values {-1, 0, +1} using absmean quantization during the forward pass.
- * Activations are quantized to 8-bit integers using absmax quantization (per-token).
- * **Crucially, the model was *trained from scratch* with this quantization scheme, not post-training quantized.**
+ * Weights are quantized to ternary values {-1, 0, +1} using absmean quantization during the forward pass.
+ * Activations are quantized to 8-bit integers using absmax quantization (per-token).
+ * **Crucially, the model was *trained from scratch* with this quantization scheme, not post-training quantized.**
* **Parameters:** ~2 Billion
* **Training Tokens:** 4 Trillion
-* **Context Length:** Maximum sequence length of **4096 tokens**.
- * *Recommendation:* For optimal performance on tasks requiring very long contexts (beyond the pre-training length or for specialized long-reasoning tasks), we recommend performing intermediate long-sequence adaptation/training before the final fine-tuning stage.
+* **Context Length:** Maximum sequence length of **4096 tokens**.
+ * *Recommendation:* For optimal performance on tasks requiring very long contexts (beyond the pre-training length or for specialized long-reasoning tasks), we recommend performing intermediate long-sequence adaptation/training before the final fine-tuning stage.
* **Training Stages:**
- 1. **Pre-training:** Large-scale training on public text/code and synthetic math data using a two-stage learning rate and weight decay schedule.
- 2. **Supervised Fine-tuning (SFT):** Fine-tuned on instruction-following and conversational datasets using sum loss aggregation and specific hyperparameter tuning.
- 3. **Direct Preference Optimization (DPO):** Aligned with human preferences using preference pairs.
+ 1. **Pre-training:** Large-scale training on public text/code and synthetic math data using a two-stage learning rate and weight decay schedule.
+ 2. **Supervised Fine-tuning (SFT):** Fine-tuned on instruction-following and conversational datasets using sum loss aggregation and specific hyperparameter tuning.
+ 3. **Direct Preference Optimization (DPO):** Aligned with human preferences using preference pairs.
* **Tokenizer:** LLaMA 3 Tokenizer (vocab size: 128,256).
-
## Usage tips
-
**VERY IMPORTANT NOTE ON EFFICIENCY**
> Please do NOT expect performance efficiency gains (in terms of speed, latency, or energy consumption) when using this model with the standard transformers library.
@@ -106,7 +102,6 @@ response = tokenizer.decode(chat_outputs[0][chat_input.shape[-1]:], skip_special
print("\nAssistant Response:", response)
```
-
## BitNetConfig
[[autodoc]] BitNetConfig
diff --git a/docs/source/en/model_doc/blenderbot-small.md b/docs/source/en/model_doc/blenderbot-small.md
index 1967013208b0..830db710e039 100644
--- a/docs/source/en/model_doc/blenderbot-small.md
+++ b/docs/source/en/model_doc/blenderbot-small.md
@@ -55,7 +55,6 @@ found [here](https://github.com/facebookresearch/ParlAI).
Blenderbot Small is a model with absolute position embeddings so it's usually advised to pad the inputs on the right rather than
the left.
-
## Resources
- [Causal language modeling task guide](../tasks/language_modeling)
diff --git a/docs/source/en/model_doc/blenderbot.md b/docs/source/en/model_doc/blenderbot.md
index 99149c5d948f..168c744235d8 100644
--- a/docs/source/en/model_doc/blenderbot.md
+++ b/docs/source/en/model_doc/blenderbot.md
@@ -71,7 +71,6 @@ An example:
`facebook/blenderbot_small_90M`, have a different architecture and consequently should be used with
[BlenderbotSmall](blenderbot-small).
-
## Resources
- [Causal language modeling task guide](../tasks/language_modeling)
diff --git a/docs/source/en/model_doc/blip-2.md b/docs/source/en/model_doc/blip-2.md
index fe4e939c2dc8..faaaee7b0840 100644
--- a/docs/source/en/model_doc/blip-2.md
+++ b/docs/source/en/model_doc/blip-2.md
@@ -26,14 +26,14 @@ rendered properly in your Markdown viewer.
The BLIP-2 model was proposed in [BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models](https://huggingface.co/papers/2301.12597) by
Junnan Li, Dongxu Li, Silvio Savarese, Steven Hoi. BLIP-2 leverages frozen pre-trained image encoders and large language models (LLMs) by training a lightweight, 12-layer Transformer
encoder in between them, achieving state-of-the-art performance on various vision-language tasks. Most notably, BLIP-2 improves upon [Flamingo](https://huggingface.co/papers/2204.14198), an 80 billion parameter model, by 8.7%
-on zero-shot VQAv2 with 54x fewer trainable parameters.
+on zero-shot VQAv2 with 54x fewer trainable parameters.
The abstract from the paper is the following:
*The cost of vision-and-language pre-training has become increasingly prohibitive due to end-to-end training of large-scale models. This paper proposes BLIP-2, a generic and efficient pre-training strategy that bootstraps vision-language pre-training from off-the-shelf frozen pre-trained image encoders and frozen large language models. BLIP-2 bridges the modality gap with a lightweight Querying Transformer, which is pre-trained in two stages. The first stage bootstraps vision-language representation learning from a frozen image encoder. The second stage bootstraps vision-to-language generative learning from a frozen language model. BLIP-2 achieves state-of-the-art performance on various vision-language tasks, despite having significantly fewer trainable parameters than existing methods. For example, our model outperforms Flamingo80B by 8.7% on zero-shot VQAv2 with 54x fewer trainable parameters. We also demonstrate the model's emerging capabilities of zero-shot image-to-text generation that can follow natural language instructions.*
+alt="drawing" width="600"/>
BLIP-2 architecture. Taken from the original paper.
diff --git a/docs/source/en/model_doc/blip.md b/docs/source/en/model_doc/blip.md
index 13a2a5731a5f..5e727050f6ee 100644
--- a/docs/source/en/model_doc/blip.md
+++ b/docs/source/en/model_doc/blip.md
@@ -25,7 +25,6 @@ rendered properly in your Markdown viewer.
[BLIP](https://huggingface.co/papers/2201.12086) (Bootstrapped Language-Image Pretraining) is a vision-language pretraining (VLP) framework designed for *both* understanding and generation tasks. Most existing pretrained models are only good at one or the other. It uses a captioner to generate captions and a filter to remove the noisy captions. This increases training data quality and more effectively uses the messy web data.
-
You can find all the original BLIP checkpoints under the [BLIP](https://huggingface.co/collections/Salesforce/blip-models-65242f40f1491fbf6a9e9472) collection.
> [!TIP]
@@ -129,7 +128,7 @@ Refer to this [notebook](https://github.com/huggingface/notebooks/blob/main/exam
## BlipTextLMHeadModel
[[autodoc]] BlipTextLMHeadModel
-- forward
+ - forward
## BlipVisionModel
diff --git a/docs/source/en/model_doc/bloom.md b/docs/source/en/model_doc/bloom.md
index 805379338e32..51e2970c25f6 100644
--- a/docs/source/en/model_doc/bloom.md
+++ b/docs/source/en/model_doc/bloom.md
@@ -43,17 +43,19 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- [`BloomForCausalLM`] is supported by this [causal language modeling example script](https://github.com/huggingface/transformers/tree/main/examples/pytorch/language-modeling#gpt-2gpt-and-causal-language-modeling) and [notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/language_modeling.ipynb).
See also:
+
- [Causal language modeling task guide](../tasks/language_modeling)
- [Text classification task guide](../tasks/sequence_classification)
- [Token classification task guide](../tasks/token_classification)
- [Question answering task guide](../tasks/question_answering)
-
⚡️ Inference
+
- A blog on [Optimization story: Bloom inference](https://huggingface.co/blog/bloom-inference-optimization).
- A blog on [Incredibly Fast BLOOM Inference with DeepSpeed and Accelerate](https://huggingface.co/blog/bloom-inference-pytorch-scripts).
⚙️ Training
+
- A blog on [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed).
## BloomConfig
diff --git a/docs/source/en/model_doc/blt.md b/docs/source/en/model_doc/blt.md
new file mode 100644
index 000000000000..254cf6c0f44a
--- /dev/null
+++ b/docs/source/en/model_doc/blt.md
@@ -0,0 +1,97 @@
+
+*This model was released on 2024-12-13 and added to Hugging Face Transformers on 2025-09-19.*
+
+
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+model = AutoModelForCausalLM.from_pretrained(
+ "itazap/blt-1b-hf",
+ device_map="auto",
+)
+
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+prompt = "my name is"
+generated_ids = model.generate(
+ **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, use_cache=False
+)
+
+print(tokenizer.decode(generated_ids[0]))
+```
+
+
+
+This model was contributed by [itazap](https://huggingface.co/).
+The original code can be found [here]().
+
+## BltConfig
+
+[[autodoc]] BltConfig
+
+[[autodoc]] BltModel
+ - forward
+
+## BltForCausalLM
+
+[[autodoc]] BltForCausalLM
+ - forward
diff --git a/docs/source/en/model_doc/bridgetower.md b/docs/source/en/model_doc/bridgetower.md
index 6a2b09e263ab..861dd32c16fe 100644
--- a/docs/source/en/model_doc/bridgetower.md
+++ b/docs/source/en/model_doc/bridgetower.md
@@ -26,7 +26,7 @@ rendered properly in your Markdown viewer.
The BridgeTower model was proposed in [BridgeTower: Building Bridges Between Encoders in Vision-Language Representative Learning](https://huggingface.co/papers/2206.08657) by Xiao Xu, Chenfei Wu, Shachar Rosenman, Vasudev Lal, Wanxiang Che, Nan Duan. The goal of this model is to build a
bridge between each uni-modal encoder and the cross-modal encoder to enable comprehensive and detailed interaction at each layer of the cross-modal encoder thus achieving remarkable performance on various downstream tasks with almost negligible additional performance and computational costs.
-This paper has been accepted to the [AAAI'23](https://aaai.org/Conferences/AAAI-23/) conference.
+This paper has been accepted to the [AAAI'23](https://aaai.org/Conferences/AAAI-23/) conference.
The abstract from the paper is the following:
@@ -54,6 +54,7 @@ The [`BridgeTowerProcessor`] wraps [`RobertaTokenizer`] and [`BridgeTowerImagePr
encode the text and prepare the images respectively.
The following example shows how to run contrastive learning using [`BridgeTowerProcessor`] and [`BridgeTowerForContrastiveLearning`].
+
```python
>>> from transformers import BridgeTowerProcessor, BridgeTowerForContrastiveLearning
>>> import requests
@@ -76,6 +77,7 @@ The following example shows how to run contrastive learning using [`BridgeTowerP
```
The following example shows how to run image-text retrieval using [`BridgeTowerProcessor`] and [`BridgeTowerForImageAndTextRetrieval`].
+
```python
>>> from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
>>> import requests
@@ -130,7 +132,6 @@ Tips:
- Please refer to [Table 5](https://huggingface.co/papers/2206.08657) for BridgeTower's performance on Image Retrieval and other down stream tasks.
- The PyTorch version of this model is only available in torch 1.10 and higher.
-
## BridgeTowerConfig
[[autodoc]] BridgeTowerConfig
@@ -177,4 +178,3 @@ Tips:
[[autodoc]] BridgeTowerForImageAndTextRetrieval
- forward
-
diff --git a/docs/source/en/model_doc/bros.md b/docs/source/en/model_doc/bros.md
index aeb3dd76e52b..4ef3d3737ae2 100644
--- a/docs/source/en/model_doc/bros.md
+++ b/docs/source/en/model_doc/bros.md
@@ -57,7 +57,6 @@ def expand_and_normalize_bbox(bboxes, doc_width, doc_height):
- [`~transformers.BrosForTokenClassification.forward`, `~transformers.BrosSpadeEEForTokenClassification.forward`, `~transformers.BrosSpadeEEForTokenClassification.forward`] require not only `input_ids` and `bbox` but also `box_first_token_mask` for loss calculation. It is a mask to filter out non-first tokens of each box. You can obtain this mask by saving start token indices of bounding boxes when creating `input_ids` from words. You can make `box_first_token_mask` with following code,
-
```python
def make_box_first_token_mask(bboxes, words, tokenizer, max_seq_length=512):
@@ -102,7 +101,6 @@ def make_box_first_token_mask(bboxes, words, tokenizer, max_seq_length=512):
[[autodoc]] BrosModel
- forward
-
## BrosForTokenClassification
[[autodoc]] BrosForTokenClassification
diff --git a/docs/source/en/model_doc/camembert.md b/docs/source/en/model_doc/camembert.md
index ddce66f2dedb..8affbd73a570 100644
--- a/docs/source/en/model_doc/camembert.md
+++ b/docs/source/en/model_doc/camembert.md
@@ -16,10 +16,10 @@ rendered properly in your Markdown viewer.
*This model was released on 2019-11-10 and added to Hugging Face Transformers on 2020-11-16.*
.")
```
+
@@ -72,6 +73,7 @@ predicted_token = tokenizer.decode(predicted_token_id)
print(f"The predicted token is: {predicted_token}")
```
+
@@ -84,7 +86,6 @@ echo -e "Le camembert est un délicieux fromage ." | transformers run --ta
-
Quantization reduces the memory burden of large models by representing weights in lower precision. Refer to the [Quantization](../quantization/overview) overview for available options.
The example below uses [bitsandbytes](../quantization/bitsandbytes) quantization to quantize the weights to 8-bits.
diff --git a/docs/source/en/model_doc/canine.md b/docs/source/en/model_doc/canine.md
index e1d8bb7f7f68..29a926c305cd 100644
--- a/docs/source/en/model_doc/canine.md
+++ b/docs/source/en/model_doc/canine.md
@@ -23,7 +23,7 @@ rendered properly in your Markdown viewer.
# CANINE
-[CANINE](https://huggingface.co/papers/2103.06874) is a tokenization-free Transformer. It skips the usual step of splitting text into subwords or wordpieces and processes text character by character. That means it works directly with raw Unicode, making it especially useful for languages with complex or inconsistent tokenization rules and even noisy inputs like typos. Since working with characters means handling longer sequences, CANINE uses a smart trick. The model compresses the input early on (called downsampling) so the transformer doesn’t have to process every character individually. This keeps things fast and efficient.
+[CANINE](https://huggingface.co/papers/2103.06874) is a tokenization-free Transformer. It skips the usual step of splitting text into subwords or wordpieces and processes text character by character. That means it works directly with raw Unicode, making it especially useful for languages with complex or inconsistent tokenization rules and even noisy inputs like typos. Since working with characters means handling longer sequences, CANINE uses a smart trick. The model compresses the input early on (called downsampling) so the transformer doesn't have to process every character individually. This keeps things fast and efficient.
You can find all the original CANINE checkpoints under the [Google](https://huggingface.co/google?search_models=canine) organization.
@@ -86,6 +86,7 @@ echo -e "Plant create energy through a process known as photosynthesis." | trans
inputs = ["Life is like a box of chocolates.", "You never know what you gonna get."]
encoding = tokenizer(inputs, padding="longest", truncation=True, return_tensors="pt")
```
+
- CANINE is primarily designed to be fine-tuned on a downstream task. The pretrained model can be used for either masked language modeling or next sentence prediction.
## CanineConfig
diff --git a/docs/source/en/model_doc/chameleon.md b/docs/source/en/model_doc/chameleon.md
index eb71349115ed..dc573faa1112 100644
--- a/docs/source/en/model_doc/chameleon.md
+++ b/docs/source/en/model_doc/chameleon.md
@@ -28,7 +28,6 @@ rendered properly in your Markdown viewer.
The Chameleon model was proposed in [Chameleon: Mixed-Modal Early-Fusion Foundation Models
](https://huggingface.co/papers/2405.09818) by META AI Chameleon Team. Chameleon is a Vision-Language Model that use vector quantization to tokenize images which enables the model to generate multimodal output. The model takes images and texts as input, including an interleaved format, and generates textual response. Image generation module is not released yet.
-
The abstract from the paper is the following:
*We present Chameleon, a family of early-fusion token-based mixed-modal models capable of understanding and generating images and text in any arbitrary sequence. We outline a stable training
@@ -43,7 +42,6 @@ including Gemini Pro and GPT-4V, according to human judgments on a new long-form
generation evaluation, where either the prompt or outputs contain mixed sequences of both images and
text. Chameleon marks a significant step forward in unified modeling of full multimodal documents*
-
@@ -52,7 +50,6 @@ alt="drawing" width="600"/>
This model was contributed by [joaogante](https://huggingface.co/joaogante) and [RaushanTurganbay](https://huggingface.co/RaushanTurganbay).
The original code can be found [here](https://github.com/facebookresearch/chameleon).
-
## Usage tips
- We advise users to use `padding_side="left"` when computing batched generation as it leads to more accurate results. Simply make sure to set `processor.tokenizer.padding_side = "left"` before generating.
diff --git a/docs/source/en/model_doc/chinese_clip.md b/docs/source/en/model_doc/chinese_clip.md
index 7ed4d503c00f..96b094ccd91b 100644
--- a/docs/source/en/model_doc/chinese_clip.md
+++ b/docs/source/en/model_doc/chinese_clip.md
@@ -119,4 +119,4 @@ Currently, following scales of pretrained Chinese-CLIP models are available on
## ChineseCLIPVisionModel
[[autodoc]] ChineseCLIPVisionModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/clipseg.md b/docs/source/en/model_doc/clipseg.md
index e27d49ffe484..099fd4fb1bac 100644
--- a/docs/source/en/model_doc/clipseg.md
+++ b/docs/source/en/model_doc/clipseg.md
@@ -47,7 +47,7 @@ can be formulated. Finally, we find our system to adapt well
to generalized queries involving affordances or properties*
+alt="drawing" width="600"/>
CLIPSeg overview. Taken from the original paper.
@@ -106,4 +106,4 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
## CLIPSegForImageSegmentation
[[autodoc]] CLIPSegForImageSegmentation
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/clvp.md b/docs/source/en/model_doc/clvp.md
index 926438a3c1f5..eead4a546435 100644
--- a/docs/source/en/model_doc/clvp.md
+++ b/docs/source/en/model_doc/clvp.md
@@ -29,29 +29,25 @@ The abstract from the paper is the following:
*In recent years, the field of image generation has been revolutionized by the application of autoregressive transformers and DDPMs. These approaches model the process of image generation as a step-wise probabilistic processes and leverage large amounts of compute and data to learn the image distribution. This methodology of improving performance need not be confined to images. This paper describes a way to apply advances in the image generative domain to speech synthesis. The result is TorToise - an expressive, multi-voice text-to-speech system.*
-
This model was contributed by [Susnato Dhar](https://huggingface.co/susnato).
The original code can be found [here](https://github.com/neonbjb/tortoise-tts).
-
## Usage tips
1. CLVP is an integral part of the Tortoise TTS model.
2. CLVP can be used to compare different generated speech candidates with the provided text, and the best speech tokens are forwarded to the diffusion model.
3. The use of the [`ClvpModelForConditionalGeneration.generate()`] method is strongly recommended for tortoise usage.
-4. Note that the CLVP model expects the audio to be sampled at 22.05 kHz contrary to other audio models which expects 16 kHz.
-
+4. Note that the CLVP model expects the audio to be sampled at 22.05 kHz contrary to other audio models which expects 16 kHz.
## Brief Explanation:
- The [`ClvpTokenizer`] tokenizes the text input, and the [`ClvpFeatureExtractor`] extracts the log mel-spectrogram from the desired audio.
- [`ClvpConditioningEncoder`] takes those text tokens and audio representations and converts them into embeddings conditioned on the text and audio.
- The [`ClvpForCausalLM`] uses those embeddings to generate multiple speech candidates.
-- Each speech candidate is passed through the speech encoder ([`ClvpEncoder`]) which converts them into a vector representation, and the text encoder ([`ClvpEncoder`]) converts the text tokens into the same latent space.
-- At the end, we compare each speech vector with the text vector to see which speech vector is most similar to the text vector.
+- Each speech candidate is passed through the speech encoder ([`ClvpEncoder`]) which converts them into a vector representation, and the text encoder ([`ClvpEncoder`]) converts the text tokens into the same latent space.
+- At the end, we compare each speech vector with the text vector to see which speech vector is most similar to the text vector.
- [`ClvpModelForConditionalGeneration.generate()`] compresses all of the logic described above into a single method.
-
Example :
```python
@@ -74,7 +70,6 @@ Example :
>>> generated_output = model.generate(**processor_output)
```
-
## ClvpConfig
[[autodoc]] ClvpConfig
@@ -128,4 +123,3 @@ Example :
## ClvpDecoder
[[autodoc]] ClvpDecoder
-
diff --git a/docs/source/en/model_doc/code_llama.md b/docs/source/en/model_doc/code_llama.md
index 60e9cb4c3cf2..a46e1f05b32a 100644
--- a/docs/source/en/model_doc/code_llama.md
+++ b/docs/source/en/model_doc/code_llama.md
@@ -143,6 +143,7 @@ visualizer("""def func(a, b):
- Infilling is only available in the 7B and 13B base models, and not in the Python, Instruct, 34B, or 70B models.
- Use the `` token where you want your input to be filled. The tokenizer splits this token to create a formatted input string that follows the [original training pattern](https://github.com/facebookresearch/codellama/blob/cb51c14ec761370ba2e2bc351374a79265d0465e/llama/generation.py#L402). This is more robust than preparing the pattern yourself.
+
```py
from transformers import LlamaForCausalLM, CodeLlamaTokenizer
@@ -158,6 +159,7 @@ visualizer("""def func(a, b):
filling = tokenizer.batch_decode(generated_ids[:, input_ids.shape[1]:], skip_special_tokens = True)[0]
print(PROMPT.replace("", filling))
```
+
- Use `bfloat16` for further training or fine-tuning and `float16` for inference.
- The `BOS` character is not used for infilling when encoding the prefix or suffix, but only at the beginning of each prompt.
- The tokenizer is a byte-pair encoding model based on [SentencePiece](https://github.com/google/sentencepiece). During decoding, if the first token is the start of the word (for example, “Banana”), the tokenizer doesn’t prepend the prefix space to the string.
diff --git a/docs/source/en/model_doc/codegen.md b/docs/source/en/model_doc/codegen.md
index e5ad3863b67c..c341154921e3 100644
--- a/docs/source/en/model_doc/codegen.md
+++ b/docs/source/en/model_doc/codegen.md
@@ -29,7 +29,7 @@ CodeGen is an autoregressive language model for program synthesis trained sequen
The abstract from the paper is the following:
-*Program synthesis strives to generate a computer program as a solution to a given problem specification. We propose a conversational program synthesis approach via large language models, which addresses the challenges of searching over a vast program space and user intent specification faced in prior approaches. Our new approach casts the process of writing a specification and program as a multi-turn conversation between a user and a system. It treats program synthesis as a sequence prediction problem, in which the specification is expressed in natural language and the desired program is conditionally sampled. We train a family of large language models, called CodeGen, on natural language and programming language data. With weak supervision in the data and the scaling up of data size and model size, conversational capacities emerge from the simple autoregressive language modeling. To study the model behavior on conversational program synthesis, we develop a multi-turn programming benchmark (MTPB), where solving each problem requires multi-step synthesis via multi-turn conversation between the user and the model. Our findings show the emergence of conversational capabilities and the effectiveness of the proposed conversational program synthesis paradigm. In addition, our model CodeGen (with up to 16B parameters trained on TPU-v4) outperforms OpenAI's Codex on the HumanEval benchmark. We make the training library JaxFormer including checkpoints available as open source contribution: [this https URL](https://github.com/salesforce/codegen).*
+*Program synthesis strives to generate a computer program as a solution to a given problem specification. We propose a conversational program synthesis approach via large language models, which addresses the challenges of searching over a vast program space and user intent specification faced in prior approaches. Our new approach casts the process of writing a specification and program as a multi-turn conversation between a user and a system. It treats program synthesis as a sequence prediction problem, in which the specification is expressed in natural language and the desired program is conditionally sampled. We train a family of large language models, called CodeGen, on natural language and programming language data. With weak supervision in the data and the scaling up of data size and model size, conversational capacities emerge from the simple autoregressive language modeling. To study the model behavior on conversational program synthesis, we develop a multi-turn programming benchmark (MTPB), where solving each problem requires multi-step synthesis via multi-turn conversation between the user and the model. Our findings show the emergence of conversational capabilities and the effectiveness of the proposed conversational program synthesis paradigm. In addition, our model CodeGen (with up to 16B parameters trained on TPU-v4) outperforms OpenAI's Codex on the HumanEval benchmark. We make the training library JaxFormer including checkpoints available as open source contribution: [this https URL](https://github.com/salesforce/codegen).*
This model was contributed by [Hiroaki Hayashi](https://huggingface.co/rooa).
The original code can be found [here](https://github.com/salesforce/codegen).
@@ -39,7 +39,7 @@ The original code can be found [here](https://github.com/salesforce/codegen).
* CodeGen model [checkpoints](https://huggingface.co/models?other=codegen) are available on different pre-training data with variable sizes.
* The format is: `Salesforce/codegen-{size}-{data}`, where
* `size`: `350M`, `2B`, `6B`, `16B`
- * `data`:
+ * `data`:
* `nl`: Pre-trained on the Pile
* `multi`: Initialized with `nl`, then further pre-trained on multiple programming languages data
* `mono`: Initialized with `multi`, then further pre-trained on Python data
diff --git a/docs/source/en/model_doc/cohere.md b/docs/source/en/model_doc/cohere.md
index 9fc6d266d69a..022a178b5cfa 100644
--- a/docs/source/en/model_doc/cohere.md
+++ b/docs/source/en/model_doc/cohere.md
@@ -22,14 +22,12 @@ rendered properly in your Markdown viewer.
+alt="drawing" width="600"/>
BLIP-2 architecture. Taken from the original paper.
diff --git a/docs/source/en/model_doc/blip.md b/docs/source/en/model_doc/blip.md
index 13a2a5731a5f..5e727050f6ee 100644
--- a/docs/source/en/model_doc/blip.md
+++ b/docs/source/en/model_doc/blip.md
@@ -25,7 +25,6 @@ rendered properly in your Markdown viewer.
[BLIP](https://huggingface.co/papers/2201.12086) (Bootstrapped Language-Image Pretraining) is a vision-language pretraining (VLP) framework designed for *both* understanding and generation tasks. Most existing pretrained models are only good at one or the other. It uses a captioner to generate captions and a filter to remove the noisy captions. This increases training data quality and more effectively uses the messy web data.
-
You can find all the original BLIP checkpoints under the [BLIP](https://huggingface.co/collections/Salesforce/blip-models-65242f40f1491fbf6a9e9472) collection.
> [!TIP]
@@ -129,7 +128,7 @@ Refer to this [notebook](https://github.com/huggingface/notebooks/blob/main/exam
## BlipTextLMHeadModel
[[autodoc]] BlipTextLMHeadModel
-- forward
+ - forward
## BlipVisionModel
diff --git a/docs/source/en/model_doc/bloom.md b/docs/source/en/model_doc/bloom.md
index 805379338e32..51e2970c25f6 100644
--- a/docs/source/en/model_doc/bloom.md
+++ b/docs/source/en/model_doc/bloom.md
@@ -43,17 +43,19 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- [`BloomForCausalLM`] is supported by this [causal language modeling example script](https://github.com/huggingface/transformers/tree/main/examples/pytorch/language-modeling#gpt-2gpt-and-causal-language-modeling) and [notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/language_modeling.ipynb).
See also:
+
- [Causal language modeling task guide](../tasks/language_modeling)
- [Text classification task guide](../tasks/sequence_classification)
- [Token classification task guide](../tasks/token_classification)
- [Question answering task guide](../tasks/question_answering)
-
⚡️ Inference
+
- A blog on [Optimization story: Bloom inference](https://huggingface.co/blog/bloom-inference-optimization).
- A blog on [Incredibly Fast BLOOM Inference with DeepSpeed and Accelerate](https://huggingface.co/blog/bloom-inference-pytorch-scripts).
⚙️ Training
+
- A blog on [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed).
## BloomConfig
diff --git a/docs/source/en/model_doc/blt.md b/docs/source/en/model_doc/blt.md
new file mode 100644
index 000000000000..254cf6c0f44a
--- /dev/null
+++ b/docs/source/en/model_doc/blt.md
@@ -0,0 +1,97 @@
+
+*This model was released on 2024-12-13 and added to Hugging Face Transformers on 2025-09-19.*
+
+
+
+
+# Byte Lantet Transformer (BLT)
+
+## Overview
+
+The BLT model was proposed in [Byte Latent Transformer: Patches Scale Better Than Tokens](https://huggingface.co/papers/2412.09871) by Artidoro Pagnoni, Ram Pasunuru, Pedro Rodriguez, John Nguyen, Benjamin Muller, Margaret Li1, Chunting Zhou, Lili Yu, Jason Weston, Luke Zettlemoyer, Gargi Ghosh, Mike Lewis, Ari Holtzman†, Srinivasan Iyer.
+BLT is a byte-level LLM that achieves tokenization-level performance through entropy-based dynamic patching.
+
+The abstract from the paper is the following:
+
+*We introduce the Byte Latent Transformer (BLT), a new byte-level LLM architecture that, for the first time, matches tokenization-based LLM performance at scale with significant improvements in inference
+efficiency and robustness. BLT encodes bytes into dynamically sized patches, which serve as the primary units of computation. Patches are segmented based on the entropy of the next byte, allocating
+more compute and model capacity where increased data complexity demands it. We present the first flop controlled scaling study of byte-level models up to 8B parameters and 4T training bytes. Our results demonstrate the feasibility of scaling models trained on raw bytes without a fixed vocabulary. Both training and inference efficiency improve due to dynamically selecting long patches when data is predictable, along with qualitative improvements on reasoning and long tail generalization. Overall, for fixed inference costs, BLT shows significantly better scaling than tokenization-based models, by simultaneously growing both patch and model size.*
+
+## Usage Tips:
+
+- **Dual Model Architecture**: BLT consists of two separate trained models:
+ - **Patcher (Entropy Model)**: A smaller transformer model that predicts byte-level entropy to determine patch boundaries and segment input.
+ - **Main Transformer Model**: The primary model that processes the patches through a Local Encoder, Global Transformer, and Local Decoder.
+
+- **Dynamic Patching**: The model uses entropy-based dynamic patching where:
+ - High-entropy regions (complex data) get shorter patches with more computational attention
+ - Low-entropy regions (predictable data) get longer patches for efficiency
+ - This allows the model to allocate compute resources where they're most needed
+
+- **Local Encoder**: Processes byte sequences with cross-attention to patch embeddings
+- **Global Transformer**: Processes patch-level representations with full attention across patches
+- **Local Decoder**: Generates output with cross-attention back to the original byte sequence
+
+- **Byte-Level Tokenizer**: Unlike traditional tokenizers that use learned vocabularies, BLT's tokenizer simply converts text to UTF-8 bytes and maps each byte to a token ID. There is no need for a vocabulary.
+
+The model can be loaded via:
+
+
+
+ 
+ 
+ 
+
+
+
+
+
+
-
# CamemBERT
@@ -50,6 +50,7 @@ from transformers import pipeline
pipeline = pipeline("fill-mask", model="camembert-base", dtype=torch.float16, device=0)
pipeline("Le camembert est un délicieux fromage
-
+
+
+
-
+
-
@@ -52,7 +50,6 @@ alt="drawing" width="600"/>
This model was contributed by [joaogante](https://huggingface.co/joaogante) and [RaushanTurganbay](https://huggingface.co/RaushanTurganbay).
The original code can be found [here](https://github.com/facebookresearch/chameleon).
-
## Usage tips
- We advise users to use `padding_side="left"` when computing batched generation as it leads to more accurate results. Simply make sure to set `processor.tokenizer.padding_side = "left"` before generating.
diff --git a/docs/source/en/model_doc/chinese_clip.md b/docs/source/en/model_doc/chinese_clip.md
index 7ed4d503c00f..96b094ccd91b 100644
--- a/docs/source/en/model_doc/chinese_clip.md
+++ b/docs/source/en/model_doc/chinese_clip.md
@@ -119,4 +119,4 @@ Currently, following scales of pretrained Chinese-CLIP models are available on
## ChineseCLIPVisionModel
[[autodoc]] ChineseCLIPVisionModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/clipseg.md b/docs/source/en/model_doc/clipseg.md
index e27d49ffe484..099fd4fb1bac 100644
--- a/docs/source/en/model_doc/clipseg.md
+++ b/docs/source/en/model_doc/clipseg.md
@@ -47,7 +47,7 @@ can be formulated. Finally, we find our system to adapt well
to generalized queries involving affordances or properties*
+alt="drawing" width="600"/>
CLIPSeg overview. Taken from the original paper.
@@ -106,4 +106,4 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
## CLIPSegForImageSegmentation
[[autodoc]] CLIPSegForImageSegmentation
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/clvp.md b/docs/source/en/model_doc/clvp.md
index 926438a3c1f5..eead4a546435 100644
--- a/docs/source/en/model_doc/clvp.md
+++ b/docs/source/en/model_doc/clvp.md
@@ -29,29 +29,25 @@ The abstract from the paper is the following:
*In recent years, the field of image generation has been revolutionized by the application of autoregressive transformers and DDPMs. These approaches model the process of image generation as a step-wise probabilistic processes and leverage large amounts of compute and data to learn the image distribution. This methodology of improving performance need not be confined to images. This paper describes a way to apply advances in the image generative domain to speech synthesis. The result is TorToise - an expressive, multi-voice text-to-speech system.*
-
This model was contributed by [Susnato Dhar](https://huggingface.co/susnato).
The original code can be found [here](https://github.com/neonbjb/tortoise-tts).
-
## Usage tips
1. CLVP is an integral part of the Tortoise TTS model.
2. CLVP can be used to compare different generated speech candidates with the provided text, and the best speech tokens are forwarded to the diffusion model.
3. The use of the [`ClvpModelForConditionalGeneration.generate()`] method is strongly recommended for tortoise usage.
-4. Note that the CLVP model expects the audio to be sampled at 22.05 kHz contrary to other audio models which expects 16 kHz.
-
+4. Note that the CLVP model expects the audio to be sampled at 22.05 kHz contrary to other audio models which expects 16 kHz.
## Brief Explanation:
- The [`ClvpTokenizer`] tokenizes the text input, and the [`ClvpFeatureExtractor`] extracts the log mel-spectrogram from the desired audio.
- [`ClvpConditioningEncoder`] takes those text tokens and audio representations and converts them into embeddings conditioned on the text and audio.
- The [`ClvpForCausalLM`] uses those embeddings to generate multiple speech candidates.
-- Each speech candidate is passed through the speech encoder ([`ClvpEncoder`]) which converts them into a vector representation, and the text encoder ([`ClvpEncoder`]) converts the text tokens into the same latent space.
-- At the end, we compare each speech vector with the text vector to see which speech vector is most similar to the text vector.
+- Each speech candidate is passed through the speech encoder ([`ClvpEncoder`]) which converts them into a vector representation, and the text encoder ([`ClvpEncoder`]) converts the text tokens into the same latent space.
+- At the end, we compare each speech vector with the text vector to see which speech vector is most similar to the text vector.
- [`ClvpModelForConditionalGeneration.generate()`] compresses all of the logic described above into a single method.
-
Example :
```python
@@ -74,7 +70,6 @@ Example :
>>> generated_output = model.generate(**processor_output)
```
-
## ClvpConfig
[[autodoc]] ClvpConfig
@@ -128,4 +123,3 @@ Example :
## ClvpDecoder
[[autodoc]] ClvpDecoder
-
diff --git a/docs/source/en/model_doc/code_llama.md b/docs/source/en/model_doc/code_llama.md
index 60e9cb4c3cf2..a46e1f05b32a 100644
--- a/docs/source/en/model_doc/code_llama.md
+++ b/docs/source/en/model_doc/code_llama.md
@@ -143,6 +143,7 @@ visualizer("""def func(a, b):
- Infilling is only available in the 7B and 13B base models, and not in the Python, Instruct, 34B, or 70B models.
- Use the `
-
# Deformable DETR
diff --git a/docs/source/en/model_doc/deit.md b/docs/source/en/model_doc/deit.md
index b40db07365a1..185a741d5b44 100644
--- a/docs/source/en/model_doc/deit.md
+++ b/docs/source/en/model_doc/deit.md
@@ -86,7 +86,7 @@ page for more information.
SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import DeiTForImageClassification
model = DeiTForImageClassification.from_pretrained("facebook/deit-base-distilled-patch16-224", attn_implementation="sdpa", dtype=torch.float16)
...
diff --git a/docs/source/en/model_doc/deplot.md b/docs/source/en/model_doc/deplot.md
index 651ddcef7fe9..5a7d4d12dcd6 100644
--- a/docs/source/en/model_doc/deplot.md
+++ b/docs/source/en/model_doc/deplot.md
@@ -21,7 +21,7 @@ rendered properly in your Markdown viewer.
-
-
+
-
+
+
+
-## Overview
+## Overview
DePlot was proposed in the paper [DePlot: One-shot visual language reasoning by plot-to-table translation](https://huggingface.co/papers/2212.10505) from Fangyu Liu, Julian Martin Eisenschlos, Francesco Piccinno, Syrine Krichene, Chenxi Pang, Kenton Lee, Mandar Joshi, Wenhu Chen, Nigel Collier, Yasemin Altun.
@@ -36,8 +36,7 @@ DePlot is a Visual Question Answering subset of `Pix2Struct` architecture. It re
Currently one checkpoint is available for DePlot:
-- `google/deplot`: DePlot fine-tuned on ChartQA dataset
-
+- `google/deplot`: DePlot fine-tuned on ChartQA dataset
```python
from transformers import AutoProcessor, Pix2StructForConditionalGeneration
@@ -57,6 +56,7 @@ print(processor.decode(predictions[0], skip_special_tokens=True))
## Fine-tuning
To fine-tune DePlot, refer to the pix2struct [fine-tuning notebook](https://github.com/huggingface/notebooks/blob/main/examples/image_captioning_pix2struct.ipynb). For `Pix2Struct` models, we have found out that fine-tuning the model with Adafactor and cosine learning rate scheduler leads to faster convergence:
+
```python
from transformers.optimization import Adafactor, get_cosine_schedule_with_warmup
@@ -68,4 +68,4 @@ scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=1000, nu
DePlot is a model trained using `Pix2Struct` architecture. For API reference, see [`Pix2Struct` documentation](pix2struct).
-
\ No newline at end of file
+
diff --git a/docs/source/en/model_doc/depth_anything.md b/docs/source/en/model_doc/depth_anything.md
index 5ac7007595ff..44774c961eaa 100644
--- a/docs/source/en/model_doc/depth_anything.md
+++ b/docs/source/en/model_doc/depth_anything.md
@@ -86,4 +86,4 @@ Image.fromarray(depth.astype("uint8"))
## DepthAnythingForDepthEstimation
[[autodoc]] DepthAnythingForDepthEstimation
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/depth_anything_v2.md b/docs/source/en/model_doc/depth_anything_v2.md
index e8637ba6192c..fbcf2248f658 100644
--- a/docs/source/en/model_doc/depth_anything_v2.md
+++ b/docs/source/en/model_doc/depth_anything_v2.md
@@ -110,4 +110,4 @@ If you're interested in submitting a resource to be included here, please feel f
## DepthAnythingForDepthEstimation
[[autodoc]] DepthAnythingForDepthEstimation
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/depth_pro.md b/docs/source/en/model_doc/depth_pro.md
index 85423359ceb0..c19703cdccc3 100644
--- a/docs/source/en/model_doc/depth_pro.md
+++ b/docs/source/en/model_doc/depth_pro.md
@@ -84,12 +84,13 @@ alt="drawing" width="600"/>
The `DepthProForDepthEstimation` model uses a `DepthProEncoder`, for encoding the input image and a `FeatureFusionStage` for fusing the output features from encoder.
The `DepthProEncoder` further uses two encoders:
+
- `patch_encoder`
- - Input image is scaled with multiple ratios, as specified in the `scaled_images_ratios` configuration.
- - Each scaled image is split into smaller **patches** of size `patch_size` with overlapping areas determined by `scaled_images_overlap_ratios`.
- - These patches are processed by the **`patch_encoder`**
+ - Input image is scaled with multiple ratios, as specified in the `scaled_images_ratios` configuration.
+ - Each scaled image is split into smaller **patches** of size `patch_size` with overlapping areas determined by `scaled_images_overlap_ratios`.
+ - These patches are processed by the **`patch_encoder`**
- `image_encoder`
- - Input image is also rescaled to `patch_size` and processed by the **`image_encoder`**
+ - Input image is also rescaled to `patch_size` and processed by the **`image_encoder`**
Both these encoders can be configured via `patch_model_config` and `image_model_config` respectively, both of which are separate `Dinov2Model` by default.
@@ -102,12 +103,14 @@ The network is supplemented with a focal length estimation head. A small convolu
The `use_fov_model` parameter in `DepthProConfig` controls whether **FOV prediction** is enabled. By default, it is set to `False` to conserve memory and computation. When enabled, the **FOV encoder** is instantiated based on the `fov_model_config` parameter, which defaults to a `Dinov2Model`. The `use_fov_model` parameter can also be passed when initializing the `DepthProForDepthEstimation` model.
The pretrained model at checkpoint `apple/DepthPro-hf` uses the FOV encoder. To use the pretrained-model without FOV encoder, set `use_fov_model=False` when loading the model, which saves computation.
+
```py
>>> from transformers import DepthProForDepthEstimation
>>> model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf", use_fov_model=False)
```
To instantiate a new model with FOV encoder, set `use_fov_model=True` in the config.
+
```py
>>> from transformers import DepthProConfig, DepthProForDepthEstimation
>>> config = DepthProConfig(use_fov_model=True)
@@ -115,6 +118,7 @@ To instantiate a new model with FOV encoder, set `use_fov_model=True` in the con
```
Or set `use_fov_model=True` when initializing the model, which overrides the value in config.
+
```py
>>> from transformers import DepthProConfig, DepthProForDepthEstimation
>>> config = DepthProConfig()
@@ -123,13 +127,13 @@ Or set `use_fov_model=True` when initializing the model, which overrides the val
### Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
```py
@@ -156,8 +160,8 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- Official Implementation: [apple/ml-depth-pro](https://github.com/apple/ml-depth-pro)
- DepthPro Inference Notebook: [DepthPro Inference](https://github.com/qubvel/transformers-notebooks/blob/main/notebooks/DepthPro_inference.ipynb)
- DepthPro for Super Resolution and Image Segmentation
- - Read blog on Medium: [Depth Pro: Beyond Depth](https://medium.com/@raoarmaghanshakir040/depth-pro-beyond-depth-9d822fc557ba)
- - Code on Github: [geetu040/depthpro-beyond-depth](https://github.com/geetu040/depthpro-beyond-depth)
+ - Read blog on Medium: [Depth Pro: Beyond Depth](https://medium.com/@raoarmaghanshakir040/depth-pro-beyond-depth-9d822fc557ba)
+ - Code on Github: [geetu040/depthpro-beyond-depth](https://github.com/geetu040/depthpro-beyond-depth)
If you're interested in submitting a resource to be included here, please feel free to open a Pull Request and we'll review it! The resource should ideally demonstrate something new instead of duplicating an existing resource.
diff --git a/docs/source/en/model_doc/detr.md b/docs/source/en/model_doc/detr.md
index 425ab0f04c51..46c9d3dadce6 100644
--- a/docs/source/en/model_doc/detr.md
+++ b/docs/source/en/model_doc/detr.md
@@ -16,9 +16,9 @@ rendered properly in your Markdown viewer.
*This model was released on 2020-05-26 and added to Hugging Face Transformers on 2021-06-09.*
-
# DETR
@@ -113,6 +113,7 @@ DETR can be naturally extended to perform panoptic segmentation (which unifies s
There are three other ways to instantiate a DETR model (depending on what you prefer):
- Option 1: Instantiate DETR with pre-trained weights for entire model
+
```python
from transformers import DetrForObjectDetection
@@ -120,6 +121,7 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
```
- Option 2: Instantiate DETR with randomly initialized weights for Transformer, but pre-trained weights for backbone
+
```python
from transformers import DetrConfig, DetrForObjectDetection
@@ -128,6 +130,7 @@ model = DetrForObjectDetection(config)
```
- Option 3: Instantiate DETR with randomly initialized weights for backbone + Transformer
+
```python
config = DetrConfig(use_pretrained_backbone=False)
model = DetrForObjectDetection(config)
@@ -144,7 +147,7 @@ As a summary, consider the following table:
| **Postprocessing** (i.e. converting the output of the model to Pascal VOC format) | [`~transformers.DetrImageProcessor.post_process`] | [`~transformers.DetrImageProcessor.post_process_segmentation`] | [`~transformers.DetrImageProcessor.post_process_segmentation`], [`~transformers.DetrImageProcessor.post_process_panoptic`] |
| **evaluators** | `CocoEvaluator` with `iou_types="bbox"` | `CocoEvaluator` with `iou_types="bbox"` or `"segm"` | `CocoEvaluator` with `iou_tupes="bbox"` or `"segm"`, `PanopticEvaluator` |
-- In short, one should prepare the data either in COCO detection or COCO panoptic format, then use [`~transformers.DetrImageProcessor`] to create `pixel_values`, `pixel_mask` and optional `labels`, which can then be used to train (or fine-tune) a model.
+- In short, one should prepare the data either in COCO detection or COCO panoptic format, then use [`~transformers.DetrImageProcessor`] to create `pixel_values`, `pixel_mask` and optional `labels`, which can then be used to train (or fine-tune) a model.
- For evaluation, one should first convert the outputs of the model using one of the postprocessing methods of [`~transformers.DetrImageProcessor`]. These can be provided to either `CocoEvaluator` or `PanopticEvaluator`, which allow you to calculate metrics like mean Average Precision (mAP) and Panoptic Quality (PQ). The latter objects are implemented in the [original repository](https://github.com/facebookresearch/detr). See the [example notebooks](https://github.com/NielsRogge/Transformers-Tutorials/tree/master/DETR) for more info regarding evaluation.
## Resources
diff --git a/docs/source/en/model_doc/dia.md b/docs/source/en/model_doc/dia.md
index 1a07e8831ee7..bab0cb4a72d3 100644
--- a/docs/source/en/model_doc/dia.md
+++ b/docs/source/en/model_doc/dia.md
@@ -117,11 +117,9 @@ out = model(**inputs)
out.loss.backward()
```
-
This model was contributed by [Jaeyong Sung](https://huggingface.co/buttercrab), [Arthur Zucker](https://huggingface.co/ArthurZ),
and [Anton Vlasjuk](https://huggingface.co/AntonV). The original code can be found [here](https://github.com/nari-labs/dia/).
-
## DiaConfig
[[autodoc]] DiaConfig
diff --git a/docs/source/en/model_doc/diffllama.md b/docs/source/en/model_doc/diffllama.md
index 406bae43c5f2..79b8314d0ae2 100644
--- a/docs/source/en/model_doc/diffllama.md
+++ b/docs/source/en/model_doc/diffllama.md
@@ -35,7 +35,6 @@ The abstract from the paper is the following:
### Usage tips
The hyperparameters of this model is the same as Llama model.
-
## DiffLlamaConfig
[[autodoc]] DiffLlamaConfig
diff --git a/docs/source/en/model_doc/dinat.md b/docs/source/en/model_doc/dinat.md
index e6d3385003cb..89f0f5cb6572 100644
--- a/docs/source/en/model_doc/dinat.md
+++ b/docs/source/en/model_doc/dinat.md
@@ -65,6 +65,7 @@ DiNAT can be used as a *backbone*. When `output_hidden_states = True`,
it will output both `hidden_states` and `reshaped_hidden_states`. The `reshaped_hidden_states` have a shape of `(batch, num_channels, height, width)` rather than `(batch_size, height, width, num_channels)`.
Notes:
+
- DiNAT depends on [NATTEN](https://github.com/SHI-Labs/NATTEN/)'s implementation of Neighborhood Attention and Dilated Neighborhood Attention.
You can install it with pre-built wheels for Linux by referring to [shi-labs.com/natten](https://shi-labs.com/natten), or build on your system by running `pip install natten`.
Note that the latter will likely take time to compile. NATTEN does not support Windows devices yet.
diff --git a/docs/source/en/model_doc/dinov2.md b/docs/source/en/model_doc/dinov2.md
index 59256756acfd..0968641326af 100644
--- a/docs/source/en/model_doc/dinov2.md
+++ b/docs/source/en/model_doc/dinov2.md
@@ -19,7 +19,6 @@ specific language governing permissions and limitations under the License.
-
# DINOv2
[DINOv2](https://huggingface.co/papers/2304.07193) is a vision foundation model that uses [ViT](./vit) as a feature extractor for multiple downstream tasks like image classification and depth estimation. It focuses on stabilizing and accelerating training through techniques like a faster memory-efficient attention, sequence packing, improved stochastic depth, Fully Sharded Data Parallel (FSDP), and model distillation.
diff --git a/docs/source/en/model_doc/dinov2_with_registers.md b/docs/source/en/model_doc/dinov2_with_registers.md
index f89de76d2168..d6b9c08f2f8f 100644
--- a/docs/source/en/model_doc/dinov2_with_registers.md
+++ b/docs/source/en/model_doc/dinov2_with_registers.md
@@ -24,7 +24,8 @@ The [Vision Transformer](vit) (ViT) is a transformer encoder model (BERT-like) o
Next, people figured out ways to make ViT work really well on self-supervised image feature extraction (i.e. learning meaningful features, also called embeddings) on images without requiring any labels. Some example papers here include [DINOv2](dinov2) and [MAE](vit_mae).
-The authors of DINOv2 noticed that ViTs have artifacts in attention maps. It’s due to the model using some image patches as “registers”. The authors propose a fix: just add some new tokens (called "register" tokens), which you only use during pre-training (and throw away afterwards). This results in:
+The authors of DINOv2 noticed that ViTs have artifacts in attention maps. It's due to the model using some image patches as “registers”. The authors propose a fix: just add some new tokens (called "register" tokens), which you only use during pre-training (and throw away afterwards). This results in:
+
- no artifacts
- interpretable attention maps
- and improved performances.
@@ -45,7 +46,6 @@ Tips:
This model was contributed by [nielsr](https://huggingface.co/nielsr).
The original code can be found [here](https://github.com/facebookresearch/dinov2).
-
## Dinov2WithRegistersConfig
[[autodoc]] Dinov2WithRegistersConfig
@@ -58,4 +58,4 @@ The original code can be found [here](https://github.com/facebookresearch/dinov2
## Dinov2WithRegistersForImageClassification
[[autodoc]] Dinov2WithRegistersForImageClassification
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/dinov3.md b/docs/source/en/model_doc/dinov3.md
index a11a8fd10cca..94e531651566 100644
--- a/docs/source/en/model_doc/dinov3.md
+++ b/docs/source/en/model_doc/dinov3.md
@@ -19,7 +19,6 @@ specific language governing permissions and limitations under the License.
-
# DINOv3
[DINOv3](https://huggingface.co/papers/2508.10104) is a family of versatile vision foundation models that outperforms the specialized state of the art across a broad range of settings, without fine-tuning. DINOv3 produces high-quality dense features that achieve outstanding performance on various vision tasks, significantly surpassing previous self- and weakly-supervised foundation models.
diff --git a/docs/source/en/model_doc/dit.md b/docs/source/en/model_doc/dit.md
index 3027905fe38b..574ffe3ef11a 100644
--- a/docs/source/en/model_doc/dit.md
+++ b/docs/source/en/model_doc/dit.md
@@ -85,6 +85,7 @@ print(f"The predicted class label is: {predicted_class_label}")
## Notes
- The pretrained DiT weights can be loaded in a [BEiT] model with a modeling head to predict visual tokens.
+
```py
from transformers import BeitForMaskedImageModeling
diff --git a/docs/source/en/model_doc/doge.md b/docs/source/en/model_doc/doge.md
index 6221940d5d5a..b2e44356ddc4 100644
--- a/docs/source/en/model_doc/doge.md
+++ b/docs/source/en/model_doc/doge.md
@@ -17,7 +17,6 @@ rendered properly in your Markdown viewer.
# Doge
-
## Overview
Doge is a series of small language models based on the [Doge](https://github.com/SmallDoges/small-doge) architecture, aiming to combine the advantages of state-space and self-attention algorithms, calculate dynamic masks from cached value states using the zero-order hold method, and solve the problem of existing mainstream language models getting lost in context. It uses the `wsd_scheduler` scheduler to pre-train on the `smollm-corpus`, and can continue training on new datasets or add sparse activation feedforward networks from stable stage checkpoints.
@@ -28,7 +27,6 @@ As shown in the figure below, the sequence transformation part of the Doge archi
Checkout all Doge model checkpoints [here](https://huggingface.co/collections/SmallDoge/doge-slm-679cc991f027c4a3abbded4a).
-
## Usage
-
-
+
-
+
+
+
@@ -44,6 +42,7 @@ inputs = tokenizer("Hey how are you doing?", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.batch_decode(outputs))
```
+
@@ -82,6 +81,7 @@ outputs = model.generate(
streamer=steamer
)
```
+
## DogeConfig
@@ -101,4 +101,4 @@ outputs = model.generate(
## DogeForSequenceClassification
[[autodoc]] DogeForSequenceClassification
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/donut.md b/docs/source/en/model_doc/donut.md
index f06b6804d6e4..e582dab748ae 100644
--- a/docs/source/en/model_doc/donut.md
+++ b/docs/source/en/model_doc/donut.md
@@ -22,7 +22,7 @@ specific language governing permissions and limitations under the License. -->
# Donut
-[Donut (Document Understanding Transformer)](https://huggingface.co/papers/2111.15664) is a visual document understanding model that doesn't require an Optical Character Recognition (OCR) engine. Unlike traditional approaches that extract text using OCR before processing, Donut employs an end-to-end Transformer-based architecture to directly analyze document images. This eliminates OCR-related inefficiencies making it more accurate and adaptable to diverse languages and formats.
+[Donut (Document Understanding Transformer)](https://huggingface.co/papers/2111.15664) is a visual document understanding model that doesn't require an Optical Character Recognition (OCR) engine. Unlike traditional approaches that extract text using OCR before processing, Donut employs an end-to-end Transformer-based architecture to directly analyze document images. This eliminates OCR-related inefficiencies making it more accurate and adaptable to diverse languages and formats.
Donut features vision encoder ([Swin](./swin)) and a text decoder ([BART](./bart)). Swin converts document images into embeddings and BART processes them into meaningful text sequences.
diff --git a/docs/source/en/model_doc/dots1.md b/docs/source/en/model_doc/dots1.md
index 337cad8cb4c7..316ab3b1f5b9 100644
--- a/docs/source/en/model_doc/dots1.md
+++ b/docs/source/en/model_doc/dots1.md
@@ -25,7 +25,6 @@ The abstract from the report is the following:
*Mixture of Experts (MoE) models have emerged as a promising paradigm for scaling language models efficiently by activating only a subset of parameters for each input token. In this report, we present dots.llm1, a large-scale MoE model that activates 14B parameters out of a total of 142B parameters, delivering performance on par with state-of-the-art models while reducing training and inference costs. Leveraging our meticulously crafted and efficient data processing pipeline, dots.llm1 achieves performance comparable to Qwen2.5-72B after pretraining on high-quality corpus and post-training to fully unlock its capabilities. Notably, no synthetic data is used during pretraining. To foster further research, we open-source intermediate training checkpoints spanning the entire training process, providing valuable insights into the learning dynamics of large language models.*
-
## Dots1Config
[[autodoc]] Dots1Config
diff --git a/docs/source/en/model_doc/dpr.md b/docs/source/en/model_doc/dpr.md
index 5fe48bc47e7b..18b060cb111d 100644
--- a/docs/source/en/model_doc/dpr.md
+++ b/docs/source/en/model_doc/dpr.md
@@ -44,9 +44,9 @@ This model was contributed by [lhoestq](https://huggingface.co/lhoestq). The ori
- DPR consists in three models:
- * Question encoder: encode questions as vectors
- * Context encoder: encode contexts as vectors
- * Reader: extract the answer of the questions inside retrieved contexts, along with a relevance score (high if the inferred span actually answers the question).
+ * Question encoder: encode questions as vectors
+ * Context encoder: encode contexts as vectors
+ * Reader: extract the answer of the questions inside retrieved contexts, along with a relevance score (high if the inferred span actually answers the question).
## DPRConfig
diff --git a/docs/source/en/model_doc/edgetam.md b/docs/source/en/model_doc/edgetam.md
new file mode 100644
index 000000000000..780ccb3f70b3
--- /dev/null
+++ b/docs/source/en/model_doc/edgetam.md
@@ -0,0 +1,331 @@
+
+*This model was released on 2025-01-13 and added to Hugging Face Transformers on 2025-09-29.*
+
+
+
+# EdgeTAM
+
+## Overview
+
+The EdgeTAM model was proposed in [EdgeTAM: On-Device Track Anything Model](https://huggingface.co/papers/2501.07256) Chong Zhou, Chenchen Zhu, Yunyang Xiong, Saksham Suri, Fanyi Xiao, Lemeng Wu, Raghuraman Krishnamoorthi, Bo Dai, Chen Change Loy, Vikas Chandra, Bilge Soran.
+
+EdgeTAM is an efficient adaptation of SAM 2 that introduces a 2D Spatial Perceiver architecture to optimize memory attention mechanisms for real-time video segmentation on mobile devices.
+
+The abstract from the paper is the following:
+
+*On top of Segment Anything Model (SAM), SAM 2 further extends its capability from image to video inputs through a memory bank mechanism and obtains a remarkable performance compared with previous methods, making it a foundation model for video segmentation task. In this paper, we aim at making SAM 2 much more efficient so that it even runs on mobile devices while maintaining a comparable performance. Despite several works optimizing SAM for better efficiency, we find they are not sufficient for SAM 2 because they all focus on compressing the image encoder, while our benchmark shows that the newly introduced memory attention blocks are also the latency bottleneck. Given this observation, we propose EdgeTAM, which leverages a novel 2D Spatial Perceiver to reduce the computational cost. In particular, the proposed 2D Spatial Perceiver encodes the densely stored frame-level memories with a lightweight Transformer that contains a fixed set of learnable queries. Given that video segmentation is a dense prediction task, we find preserving the spatial structure of the memories is essential so that the queries are split into global-level and patch-level groups. We also propose a distillation pipeline that further improves the performance without inference overhead. As a result, EdgeTAM achieves 87.7, 70.0, 72.3, and 71.7 J&F on DAVIS 2017, MOSE, SA-V val, and SA-V test, while running at 16 FPS on iPhone 15 Pro Max.*
+
+This model was contributed by [yonigozlan](https://huggingface.co/yonigozlan).
+The original code can be found [here](https://github.com/facebookresearch/EdgeTAM).
+
+## Usage example
+
+### Automatic Mask Generation with Pipeline
+
+EdgeTAM can be used for automatic mask generation to segment all objects in an image using the `mask-generation` pipeline:
+
+```python
+>>> from transformers import pipeline
+
+>>> generator = pipeline("mask-generation", model="yonigozlan/edgetam-1", device=0)
+>>> image_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
+>>> outputs = generator(image_url, points_per_batch=64)
+
+>>> len(outputs["masks"]) # Number of masks generated
+39
+```
+
+### Basic Image Segmentation
+
+#### Single Point Click
+
+You can segment objects by providing a single point click on the object you want to segment:
+
+```python
+>>> from transformers import Sam2Processor, EdgeTamModel, infer_device
+>>> import torch
+>>> from PIL import Image
+>>> import requests
+
+>>> device = infer_device()
+
+>>> model = EdgeTamModel.from_pretrained("yonigozlan/edgetam-1").to(device)
+>>> processor = Sam2Processor.from_pretrained("yonigozlan/edgetam-1")
+
+>>> image_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
+>>> raw_image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
+
+>>> input_points = [[[[500, 375]]]] # Single point click, 4 dimensions (image_dim, object_dim, point_per_object_dim, coordinates)
+>>> input_labels = [[[1]]] # 1 for positive click, 0 for negative click, 3 dimensions (image_dim, object_dim, point_label)
+
+>>> inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs)
+
+>>> masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
+
+>>> # The model outputs multiple mask predictions ranked by quality score
+>>> print(f"Generated {masks.shape[1]} masks with shape {masks.shape}")
+Generated 3 masks with shape torch.Size([1, 3, 1200, 1800])
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.0463, 0.4859, 0.7616], device='cuda:0')
+```
+
+#### Multiple Points for Refinement
+
+You can provide multiple points to refine the segmentation:
+
+```python
+>>> # Add both positive and negative points to refine the mask
+>>> input_points = [[[[500, 375], [1125, 625]]]] # Multiple points for refinement
+>>> input_labels = [[[1, 1]]] # Both positive clicks
+
+>>> inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs)
+
+>>> masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.8362, 0.6900, 0.2120], device='cuda:0')
+```
+
+#### Bounding Box Input
+
+EdgeTAM also supports bounding box inputs for segmentation:
+
+```python
+>>> # Define bounding box as [x_min, y_min, x_max, y_max]
+>>> input_boxes = [[[75, 275, 1725, 850]]]
+
+>>> inputs = processor(images=raw_image, input_boxes=input_boxes, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs)
+
+>>> masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.9301, 0.9348, 0.6605], device='cuda:0')
+```
+
+#### Multiple Objects Segmentation
+
+You can segment multiple objects simultaneously:
+
+```python
+>>> # Define points for two different objects
+>>> input_points = [[[[500, 375]], [[650, 750]]]] # Points for two objects in same image
+>>> input_labels = [[[1], [1]]] # Positive clicks for both objects
+
+>>> inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs, multimask_output=False)
+
+>>> # Each object gets its own mask
+>>> masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
+>>> print(f"Generated masks for {masks.shape[0]} objects")
+Generated masks for 2 objects
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.7616, 0.9465], device='cuda:0')
+```
+
+### Batch Inference
+
+#### Batched Images
+
+Process multiple images simultaneously for improved efficiency:
+
+```python
+>>> from transformers import Sam2Processor, EdgeTamModel, infer_device
+>>> import torch
+>>> from PIL import Image
+>>> import requests
+
+>>> device = infer_device()
+
+>>> model = EdgeTamModel.from_pretrained("yonigozlan/edgetam-1").to(device)
+>>> processor = Sam2Processor.from_pretrained("yonigozlan/edgetam-1")
+
+>>> # Load multiple images
+>>> image_urls = [
+... "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg",
+... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/dog-sam.png"
+... ]
+>>> raw_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in image_urls]
+
+>>> # Single point per image
+>>> input_points = [[[[500, 375]]], [[[770, 200]]]] # One point for each image
+>>> input_labels = [[[1]], [[1]]] # Positive clicks for both images
+
+>>> inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs, multimask_output=False)
+
+>>> # Post-process masks for each image
+>>> all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
+>>> print(f"Processed {len(all_masks)} images, each with {all_masks[0].shape[0]} objects")
+Processed 2 images, each with 1 objects
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.7618, 0.7999], device='cuda:0')
+```
+
+#### Batched Objects per Image
+
+Segment multiple objects within each image using batch inference:
+
+```python
+>>> # Multiple objects per image - different numbers of objects per image
+>>> input_points = [
+... [[[500, 375]], [[650, 750]]], # Truck image: 2 objects
+... [[[770, 200]]] # Dog image: 1 object
+... ]
+>>> input_labels = [
+... [[1], [1]], # Truck image: positive clicks for both objects
+... [[1]] # Dog image: positive click for the object
+... ]
+
+>>> inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs, multimask_output=False)
+
+>>> all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
+```
+
+#### Batched Images with Batched Objects and Multiple Points
+
+Handle complex batch scenarios with multiple points per object:
+
+```python
+>>> # Add groceries image for more complex example
+>>> groceries_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/groceries.jpg"
+>>> groceries_image = Image.open(requests.get(groceries_url, stream=True).raw).convert("RGB")
+>>> raw_images = [raw_images[0], groceries_image] # Use truck and groceries images
+
+>>> # Complex batching: multiple images, multiple objects, multiple points per object
+>>> input_points = [
+... [[[500, 375]], [[650, 750]]], # Truck image: 2 objects with 1 point each
+... [[[400, 300]], [[630, 300], [550, 300]]] # Groceries image: obj1 has 1 point, obj2 has 2 points
+... ]
+>>> input_labels = [
+... [[1], [1]], # Truck image: positive clicks
+... [[1], [1, 1]] # Groceries image: positive clicks for refinement
+... ]
+
+>>> inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs, multimask_output=False)
+
+>>> all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
+```
+
+#### Batched Bounding Boxes
+
+Process multiple images with bounding box inputs:
+
+```python
+>>> # Multiple bounding boxes per image (using truck and groceries images)
+>>> input_boxes = [
+... [[75, 275, 1725, 850], [425, 600, 700, 875], [1375, 550, 1650, 800], [1240, 675, 1400, 750]], # Truck image: 4 boxes
+... [[450, 170, 520, 350], [350, 190, 450, 350], [500, 170, 580, 350], [580, 170, 640, 350]] # Groceries image: 4 boxes
+... ]
+
+>>> # Update images for this example
+>>> raw_images = [raw_images[0], groceries_image] # truck and groceries
+
+>>> inputs = processor(images=raw_images, input_boxes=input_boxes, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs, multimask_output=False)
+
+>>> all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
+>>> print(f"Processed {len(input_boxes)} images with {len(input_boxes[0])} and {len(input_boxes[1])} boxes respectively")
+Processed 2 images with 4 and 4 boxes respectively
+>>> print(f"IoU scores: {outputs.iou_scores.squeeze()}")
+IoU scores: tensor([0.9301, 0.9348, 0.6605, 0.9465], device='cuda:0')
+```
+
+### Using Previous Masks as Input
+
+EdgeTAM can use masks from previous predictions as input to refine segmentation:
+
+```python
+>>> # Get initial segmentation
+>>> input_points = [[[[500, 375]]]]
+>>> input_labels = [[[1]]]
+>>> inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+
+>>> with torch.no_grad():
+... outputs = model(**inputs)
+
+>>> # Use the best mask as input for refinement
+>>> mask_input = outputs.pred_masks[:, :, torch.argmax(outputs.iou_scores.squeeze())]
+
+>>> # Add additional points with the mask input
+>>> new_input_points = [[[[500, 375], [450, 300]]]]
+>>> new_input_labels = [[[1, 1]]]
+>>> inputs = processor(
+... input_points=new_input_points,
+... input_labels=new_input_labels,
+... original_sizes=inputs["original_sizes"],
+... return_tensors="pt",
+... ).to(device)
+
+>>> with torch.no_grad():
+... refined_outputs = model(
+... **inputs,
+... input_masks=mask_input,
+... image_embeddings=outputs.image_embeddings,
+... multimask_output=False,
+... )
+```
+
+
+## EdgeTamConfig
+
+[[autodoc]] EdgeTamConfig
+
+## EdgeTamVisionConfig
+
+[[autodoc]] EdgeTamVisionConfig
+
+## EdgeTamMaskDecoderConfig
+
+[[autodoc]] EdgeTamMaskDecoderConfig
+
+## EdgeTamPromptEncoderConfig
+
+[[autodoc]] EdgeTamPromptEncoderConfig
+
+## EdgeTamVisionModel
+
+[[autodoc]] EdgeTamVisionModel
+ - forward
+
+## EdgeTamModel
+
+[[autodoc]] EdgeTamModel
+ - forward
diff --git a/docs/source/en/model_doc/edgetam_video.md b/docs/source/en/model_doc/edgetam_video.md
new file mode 100644
index 000000000000..381bace4dbe0
--- /dev/null
+++ b/docs/source/en/model_doc/edgetam_video.md
@@ -0,0 +1,297 @@
+
+*This model was released on 2025-01-13 and added to Hugging Face Transformers on 2025-09-29.*
+
+
+
+
+
+
+
+
+
+
+
+
+
+# EdgeTAMVideo
+
+## Overview
+
+The EdgeTAM model was proposed in [EdgeTAM: On-Device Track Anything Model](https://huggingface.co/papers/2501.07256) Chong Zhou, Chenchen Zhu, Yunyang Xiong, Saksham Suri, Fanyi Xiao, Lemeng Wu, Raghuraman Krishnamoorthi, Bo Dai, Chen Change Loy, Vikas Chandra, Bilge Soran.
+
+EdgeTAM is an efficient adaptation of SAM 2 that introduces a 2D Spatial Perceiver architecture to optimize memory attention mechanisms for real-time video segmentation on mobile devices.
+
+The abstract from the paper is the following:
+
+*On top of Segment Anything Model (SAM), SAM 2 further extends its capability from image to video inputs through a memory bank mechanism and obtains a remarkable performance compared with previous methods, making it a foundation model for video segmentation task. In this paper, we aim at making SAM 2 much more efficient so that it even runs on mobile devices while maintaining a comparable performance. Despite several works optimizing SAM for better efficiency, we find they are not sufficient for SAM 2 because they all focus on compressing the image encoder, while our benchmark shows that the newly introduced memory attention blocks are also the latency bottleneck. Given this observation, we propose EdgeTAM, which leverages a novel 2D Spatial Perceiver to reduce the computational cost. In particular, the proposed 2D Spatial Perceiver encodes the densely stored frame-level memories with a lightweight Transformer that contains a fixed set of learnable queries. Given that video segmentation is a dense prediction task, we find preserving the spatial structure of the memories is essential so that the queries are split into global-level and patch-level groups. We also propose a distillation pipeline that further improves the performance without inference overhead. As a result, EdgeTAM achieves 87.7, 70.0, 72.3, and 71.7 J&F on DAVIS 2017, MOSE, SA-V val, and SA-V test, while running at 16 FPS on iPhone 15 Pro Max.*
+
+This model was contributed by [yonigozlan](https://huggingface.co/yonigozlan).
+The original code can be found [here](https://github.com/facebookresearch/EdgeTAM).
+
+## Usage example
+
+### Video Segmentation and Tracking
+
+EdgeTAM Video's key strength is its ability to track objects across video frames efficiently on mobile devices. Here's how to use it for video segmentation:
+
+#### Basic Video Tracking
+
+```python
+>>> from transformers import EdgeTamVideoModel, Sam2VideoProcessor, infer_device
+>>> import torch
+
+>>> device = infer_device()
+>>> model = EdgeTamVideoModel.from_pretrained("yonigozlan/edgetam-video-1").to(device, dtype=torch.bfloat16)
+>>> processor = Sam2VideoProcessor.from_pretrained("yonigozlan/edgetam-video-1")
+
+>>> # Load video frames (example assumes you have a list of PIL Images)
+>>> # video_frames = [Image.open(f"frame_{i:05d}.jpg") for i in range(num_frames)]
+
+>>> # For this example, we'll use the video loading utility
+>>> from transformers.video_utils import load_video
+>>> video_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/bedroom.mp4"
+>>> video_frames, _ = load_video(video_url)
+
+>>> # Initialize video inference session
+>>> inference_session = processor.init_video_session(
+... video=video_frames,
+... inference_device=device,
+... dtype=torch.bfloat16,
+... )
+
+>>> # Add click on first frame to select object
+>>> ann_frame_idx = 0
+>>> ann_obj_id = 1
+>>> points = [[[[210, 350]]]]
+>>> labels = [[[1]]]
+
+>>> processor.add_inputs_to_inference_session(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... obj_ids=ann_obj_id,
+... input_points=points,
+... input_labels=labels,
+... )
+
+>>> # Segment the object on the first frame
+>>> outputs = model(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... )
+>>> video_res_masks = processor.post_process_masks(
+... [outputs.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+>>> print(f"Segmentation shape: {video_res_masks.shape}")
+Segmentation shape: torch.Size([1, 1, 540, 960])
+
+>>> # Propagate through the entire video
+>>> video_segments = {}
+>>> for sam2_video_output in model.propagate_in_video_iterator(inference_session):
+... video_res_masks = processor.post_process_masks(
+... [sam2_video_output.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+... video_segments[sam2_video_output.frame_idx] = video_res_masks
+
+>>> print(f"Tracked object through {len(video_segments)} frames")
+Tracked object through 200 frames
+```
+
+#### Multi-Object Video Tracking
+
+Track multiple objects simultaneously across video frames:
+
+```python
+>>> # Reset for new tracking session
+>>> inference_session.reset_inference_session()
+
+>>> # Add multiple objects on the first frame
+>>> ann_frame_idx = 0
+>>> obj_ids = [2, 3]
+>>> input_points = [[[[200, 300]], [[400, 150]]]] # Points for two objects (batched)
+>>> input_labels = [[[1], [1]]]
+
+>>> processor.add_inputs_to_inference_session(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... obj_ids=obj_ids,
+... input_points=input_points,
+... input_labels=input_labels,
+... )
+
+>>> # Get masks for both objects on first frame
+>>> outputs = model(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... )
+
+>>> # Propagate both objects through video
+>>> video_segments = {}
+>>> for sam2_video_output in model.propagate_in_video_iterator(inference_session):
+... video_res_masks = processor.post_process_masks(
+... [sam2_video_output.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+... video_segments[sam2_video_output.frame_idx] = {
+... obj_id: video_res_masks[i]
+... for i, obj_id in enumerate(inference_session.obj_ids)
+... }
+
+>>> print(f"Tracked {len(inference_session.obj_ids)} objects through {len(video_segments)} frames")
+Tracked 2 objects through 200 frames
+```
+
+#### Refining Video Segmentation
+
+You can add additional clicks on any frame to refine the tracking:
+
+```python
+>>> # Add refinement click on a later frame
+>>> refine_frame_idx = 50
+>>> ann_obj_id = 2 # Refining first object
+>>> points = [[[[220, 280]]]] # Additional point
+>>> labels = [[[1]]] # Positive click
+
+>>> processor.add_inputs_to_inference_session(
+... inference_session=inference_session,
+... frame_idx=refine_frame_idx,
+... obj_ids=ann_obj_id,
+... input_points=points,
+... input_labels=labels,
+... )
+
+>>> # Re-propagate with the additional information
+>>> video_segments = {}
+>>> for sam2_video_output in model.propagate_in_video_iterator(inference_session):
+... video_res_masks = processor.post_process_masks(
+... [sam2_video_output.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+... video_segments[sam2_video_output.frame_idx] = video_res_masks
+```
+
+### Streaming Video Inference
+
+For real-time applications, EdgeTAM Video supports processing video frames as they arrive:
+
+```python
+>>> # Initialize session for streaming
+>>> inference_session = processor.init_video_session(
+... inference_device=device,
+... dtype=torch.bfloat16,
+... )
+
+>>> # Process frames one by one
+>>> for frame_idx, frame in enumerate(video_frames[:10]): # Process first 10 frames
+... inputs = processor(images=frame, device=device, return_tensors="pt")
+...
+... if frame_idx == 0:
+... # Add point input on first frame
+... processor.add_inputs_to_inference_session(
+... inference_session=inference_session,
+... frame_idx=0,
+... obj_ids=1,
+... input_points=[[[[210, 350], [250, 220]]]],
+... input_labels=[[[1, 1]]],
+... original_size=inputs.original_sizes[0], # need to be provided when using streaming video inference
+... )
+...
+... # Process current frame
+... sam2_video_output = model(inference_session=inference_session, frame=inputs.pixel_values[0])
+...
+... video_res_masks = processor.post_process_masks(
+... [sam2_video_output.pred_masks], original_sizes=inputs.original_sizes, binarize=False
+... )[0]
+... print(f"Frame {frame_idx}: mask shape {video_res_masks.shape}")
+
+Frame 0: mask shape torch.Size([1, 1, 540, 960])
+...
+```
+
+#### Video Batch Processing for Multiple Objects
+
+Track multiple objects simultaneously in video by adding them all at once:
+
+```python
+>>> # Initialize video session
+>>> inference_session = processor.init_video_session(
+... video=video_frames,
+... inference_device=device,
+... dtype=torch.bfloat16,
+... )
+
+>>> # Add multiple objects on the first frame using batch processing
+>>> ann_frame_idx = 0
+>>> obj_ids = [2, 3] # Track two different objects
+>>> input_points = [
+... [[[200, 300], [230, 250], [275, 175]], [[400, 150]]]
+... ] # Object 2: 3 points (2 positive, 1 negative); Object 3: 1 point
+>>> input_labels = [
+... [[1, 1, 0], [1]]
+... ] # Object 2: positive, positive, negative; Object 3: positive
+
+>>> processor.add_inputs_to_inference_session(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... obj_ids=obj_ids,
+... input_points=input_points,
+... input_labels=input_labels,
+... )
+
+>>> # Get masks for all objects on the first frame
+>>> outputs = model(
+... inference_session=inference_session,
+... frame_idx=ann_frame_idx,
+... )
+>>> video_res_masks = processor.post_process_masks(
+... [outputs.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+>>> print(f"Generated masks for {video_res_masks.shape[0]} objects")
+Generated masks for 2 objects
+
+>>> # Propagate all objects through the video
+>>> video_segments = {}
+>>> for sam2_video_output in model.propagate_in_video_iterator(inference_session):
+... video_res_masks = processor.post_process_masks(
+... [sam2_video_output.pred_masks], original_sizes=[[inference_session.video_height, inference_session.video_width]], binarize=False
+... )[0]
+... video_segments[sam2_video_output.frame_idx] = {
+... obj_id: video_res_masks[i]
+... for i, obj_id in enumerate(inference_session.obj_ids)
+... }
+
+>>> print(f"Tracked {len(inference_session.obj_ids)} objects through {len(video_segments)} frames")
+Tracked 2 objects through 200 frames
+```
+
+## EdgeTamVideoMaskDecoderConfig
+
+[[autodoc]] EdgeTamVideoMaskDecoderConfig
+
+## EdgeTamVideoPromptEncoderConfig
+
+[[autodoc]] EdgeTamVideoPromptEncoderConfig
+
+## EdgeTamVideoConfig
+
+[[autodoc]] EdgeTamVideoConfig
+
+## EdgeTamVideoInferenceSession
+
+[[autodoc]] EdgeTamVideoInferenceSession
+
+## EdgeTamVideoModel
+
+[[autodoc]] EdgeTamVideoModel
+ - forward
diff --git a/docs/source/en/model_doc/efficientloftr.md b/docs/source/en/model_doc/efficientloftr.md
index 2994ae83262d..4efd87502b67 100644
--- a/docs/source/en/model_doc/efficientloftr.md
+++ b/docs/source/en/model_doc/efficientloftr.md
@@ -45,6 +45,7 @@ results = keypoint_matcher([url_0, url_1], threshold=0.9)
print(results[0])
# {'keypoint_image_0': {'x': ..., 'y': ...}, 'keypoint_image_1': {'x': ..., 'y': ...}, 'score': ...}
```
+
+
+
+
+
+
+
+
+ 
-
## Usage Tips
### Generate text
@@ -84,7 +83,6 @@ generate_text = tokenizer.decode(output_ids, skip_special_tokens=True)
This model was contributed by [Anton Vlasjuk](https://huggingface.co/AntonV).
The original code can be found [here](https://github.com/PaddlePaddle/ERNIE).
-
## Ernie4_5Config
[[autodoc]] Ernie4_5Config
diff --git a/docs/source/en/model_doc/ernie4_5_moe.md b/docs/source/en/model_doc/ernie4_5_moe.md
index 20c4dcfd5435..fb6b8d791bec 100644
--- a/docs/source/en/model_doc/ernie4_5_moe.md
+++ b/docs/source/en/model_doc/ernie4_5_moe.md
@@ -40,7 +40,6 @@ Other models from the family can be found at [Ernie 4.5](./ernie4_5).
-
## Usage Tips
### Generate text
@@ -167,7 +166,6 @@ generate_text = tokenizer.decode(output_ids, skip_special_tokens=True)
This model was contributed by [Anton Vlasjuk](https://huggingface.co/AntonV).
The original code can be found [here](https://github.com/PaddlePaddle/ERNIE).
-
## Ernie4_5_MoeConfig
[[autodoc]] Ernie4_5_MoeConfig
diff --git a/docs/source/en/model_doc/ernie_m.md b/docs/source/en/model_doc/ernie_m.md
index 508fe2f596b2..e044614e7644 100644
--- a/docs/source/en/model_doc/ernie_m.md
+++ b/docs/source/en/model_doc/ernie_m.md
@@ -40,7 +40,6 @@ The abstract from the paper is the following:
*Recent studies have demonstrated that pre-trained cross-lingual models achieve impressive performance in downstream cross-lingual tasks. This improvement benefits from learning a large amount of monolingual and parallel corpora. Although it is generally acknowledged that parallel corpora are critical for improving the model performance, existing methods are often constrained by the size of parallel corpora, especially for lowresource languages. In this paper, we propose ERNIE-M, a new training method that encourages the model to align the representation of multiple languages with monolingual corpora, to overcome the constraint that the parallel corpus size places on the model performance. Our key insight is to integrate back-translation into the pre-training process. We generate pseudo-parallel sentence pairs on a monolingual corpus to enable the learning of semantic alignments between different languages, thereby enhancing the semantic modeling of cross-lingual models. Experimental results show that ERNIE-M outperforms existing cross-lingual models and delivers new state-of-the-art results in various cross-lingual downstream tasks.*
This model was contributed by [Susnato Dhar](https://huggingface.co/susnato). The original code can be found [here](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/paddlenlp/transformers/ernie_m).
-
## Usage tips
- Ernie-M is a BERT-like model so it is a stacked Transformer Encoder.
@@ -59,7 +58,6 @@ This model was contributed by [Susnato Dhar](https://huggingface.co/susnato). Th
[[autodoc]] ErnieMConfig
-
## ErnieMTokenizer
[[autodoc]] ErnieMTokenizer
@@ -68,7 +66,6 @@ This model was contributed by [Susnato Dhar](https://huggingface.co/susnato). Th
- create_token_type_ids_from_sequences
- save_vocabulary
-
## ErnieMModel
[[autodoc]] ErnieMModel
@@ -79,19 +76,16 @@ This model was contributed by [Susnato Dhar](https://huggingface.co/susnato). Th
[[autodoc]] ErnieMForSequenceClassification
- forward
-
## ErnieMForMultipleChoice
[[autodoc]] ErnieMForMultipleChoice
- forward
-
## ErnieMForTokenClassification
[[autodoc]] ErnieMForTokenClassification
- forward
-
## ErnieMForQuestionAnswering
[[autodoc]] ErnieMForQuestionAnswering
diff --git a/docs/source/en/model_doc/esm.md b/docs/source/en/model_doc/esm.md
index e83e2d5aa1da..a6190a71f020 100644
--- a/docs/source/en/model_doc/esm.md
+++ b/docs/source/en/model_doc/esm.md
@@ -44,12 +44,10 @@ sequence alignment (MSA) step at inference time, which means that ESMFold checkp
they do not require a database of known protein sequences and structures with associated external query tools
to make predictions, and are much faster as a result.
-
The abstract from
"Biological structure and function emerge from scaling unsupervised learning to 250
million protein sequences" is
-
*In the field of artificial intelligence, a combination of scale in data and model capacity enabled by unsupervised
learning has led to major advances in representation learning and statistical generation. In the life sciences, the
anticipated growth of sequencing promises unprecedented data on natural sequence diversity. Protein language modeling
@@ -63,7 +61,6 @@ can be identified by linear projections. Representation learning produces featur
applications, enabling state-of-the-art supervised prediction of mutational effect and secondary structure and
improving state-of-the-art features for long-range contact prediction.*
-
The abstract from
"Language models of protein sequences at the scale of evolution enable accurate structure prediction" is
diff --git a/docs/source/en/model_doc/evolla.md b/docs/source/en/model_doc/evolla.md
index a39103a06d12..ea8605050599 100644
--- a/docs/source/en/model_doc/evolla.md
+++ b/docs/source/en/model_doc/evolla.md
@@ -25,7 +25,7 @@ Evolla is an advanced 80-billion-parameter protein-language generative model des
The abstract from the paper is the following:
-*Proteins, nature’s intricate molecular machines, are the products of billions of years of evolution and play fundamental roles in sustaining life. Yet, deciphering their molecular language - that is, understanding how protein sequences and structures encode and determine biological functions - remains a corner-stone challenge in modern biology. Here, we introduce Evolla, an 80 billion frontier protein-language generative model designed to decode the molecular language of proteins. By integrating information from protein sequences, structures, and user queries, Evolla generates precise and contextually nuanced insights into protein function. A key innovation of Evolla lies in its training on an unprecedented AI-generated dataset: 546 million protein question-answer pairs and 150 billion word tokens, designed to reflect the immense complexity and functional diversity of proteins. Post-pretraining, Evolla integrates Direct Preference Optimization (DPO) to refine the model based on preference signals and Retrieval-Augmented Generation (RAG) for external knowledge incorporation, improving response quality and relevance. To evaluate its performance, we propose a novel framework, Instructional Response Space (IRS), demonstrating that Evolla delivers expert-level insights, advancing research in proteomics and functional genomics while shedding light on the molecular logic encoded in proteins. The online demo is available at http://www.chat-protein.com/.*
+*Proteins, nature's intricate molecular machines, are the products of billions of years of evolution and play fundamental roles in sustaining life. Yet, deciphering their molecular language - that is, understanding how protein sequences and structures encode and determine biological functions - remains a corner-stone challenge in modern biology. Here, we introduce Evolla, an 80 billion frontier protein-language generative model designed to decode the molecular language of proteins. By integrating information from protein sequences, structures, and user queries, Evolla generates precise and contextually nuanced insights into protein function. A key innovation of Evolla lies in its training on an unprecedented AI-generated dataset: 546 million protein question-answer pairs and 150 billion word tokens, designed to reflect the immense complexity and functional diversity of proteins. Post-pretraining, Evolla integrates Direct Preference Optimization (DPO) to refine the model based on preference signals and Retrieval-Augmented Generation (RAG) for external knowledge incorporation, improving response quality and relevance. To evaluate its performance, we propose a novel framework, Instructional Response Space (IRS), demonstrating that Evolla delivers expert-level insights, advancing research in proteomics and functional genomics while shedding light on the molecular logic encoded in proteins. The online demo is available at http://www.chat-protein.com/.*
Examples:
@@ -75,7 +75,6 @@ Tips:
- This model was contributed by [Xibin Bayes Zhou](https://huggingface.co/XibinBayesZhou).
- The original code can be found [here](https://github.com/westlake-repl/Evolla).
-
## EvollaConfig
[[autodoc]] EvollaConfig
diff --git a/docs/source/en/model_doc/exaone4.md b/docs/source/en/model_doc/exaone4.md
index 69d7ee0b2a81..9482f5be2c06 100644
--- a/docs/source/en/model_doc/exaone4.md
+++ b/docs/source/en/model_doc/exaone4.md
@@ -20,7 +20,7 @@ rendered properly in your Markdown viewer.
## Overview
**[EXAONE 4.0](https://github.com/LG-AI-EXAONE/EXAONE-4.0)** model is the language model, which integrates a **Non-reasoning mode** and **Reasoning mode** to achieve both the excellent usability of [EXAONE 3.5](https://github.com/LG-AI-EXAONE/EXAONE-3.5) and the advanced reasoning abilities of [EXAONE Deep](https://github.com/LG-AI-EXAONE/EXAONE-Deep). To pave the way for the agentic AI era, EXAONE 4.0 incorporates essential features such as agentic tool use, and its multilingual capabilities are extended
-to support Spanish in addition to English and Korean.
+to support Spanish in addition to English and Korean.
The EXAONE 4.0 model series consists of two sizes: a mid-size **32B** model optimized for high performance, and a small-size **1.2B** model designed for on-device applications.
@@ -33,7 +33,6 @@ For more details, please refer to our [technical report](https://huggingface.co/
All model weights including quantized versions are available at [Huggingface Collections](https://huggingface.co/collections/LGAI-EXAONE/exaone-40-686b2e0069800c835ed48375).
-
## Model Details
### Model Specifications
@@ -57,7 +56,6 @@ All model weights including quantized versions are available at [Huggingface Col
| Tied word embedding | False | True |
| Knowledge cut-off | Nov. 2024 | Nov. 2024 |
-
## Usage tips
### Non-reasoning mode
@@ -206,4 +204,4 @@ print(tokenizer.decode(output[0]))
## Exaone4ForQuestionAnswering
[[autodoc]] Exaone4ForQuestionAnswering
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/falcon3.md b/docs/source/en/model_doc/falcon3.md
index 368a5457ab6d..3d79a4e225dd 100644
--- a/docs/source/en/model_doc/falcon3.md
+++ b/docs/source/en/model_doc/falcon3.md
@@ -30,5 +30,6 @@ Depth up-scaling for improved reasoning: Building on recent studies on the effec
Knowledge distillation for better tiny models: To provide compact and efficient alternatives, we developed Falcon3-1B-Base and Falcon3-3B-Base by leveraging pruning and knowledge distillation techniques, using less than 100GT of curated high-quality data, thereby redefining pre-training efficiency.
## Resources
+
- [Blog post](https://huggingface.co/blog/falcon3)
- [Models on Huggingface](https://huggingface.co/collections/tiiuae/falcon3-67605ae03578be86e4e87026)
diff --git a/docs/source/en/model_doc/falcon_h1.md b/docs/source/en/model_doc/falcon_h1.md
index 981c00bd626b..48a647cd3797 100644
--- a/docs/source/en/model_doc/falcon_h1.md
+++ b/docs/source/en/model_doc/falcon_h1.md
@@ -21,7 +21,6 @@ The [FalconH1](https://huggingface.co/blog/tiiuae/falcon-h1) model was developed
This model was contributed by [DhiyaEddine](https://huggingface.co/DhiyaEddine), [ybelkada](https://huggingface.co/ybelkada), [JingweiZuo](https://huggingface.co/JingweiZuo), [IlyasChahed](https://huggingface.co/IChahed), and [MaksimVelikanov](https://huggingface.co/yellowvm).
The original code can be found [here](https://github.com/tiiuae/Falcon-H1).
-
## FalconH1Config
| Model | Depth | Dim | Attn Heads | KV | Mamba Heads | d_head | d_state | Ctx Len |
@@ -33,8 +32,6 @@ The original code can be found [here](https://github.com/tiiuae/Falcon-H1).
| H1 7B | 44 | 3072 | 12 | 2 | 24 | 128 / 128 | 256 | 256K |
| H1 34B | 72 | 5120 | 20 | 4 | 32 | 128 / 128 | 256 | 256K |
-
-
[[autodoc]] FalconH1Config
-*This model was released on 2025-07-09 and added to Hugging Face Transformers on 2025-09-15.*
+*This model was released on 2025-07-09 and added to Hugging Face Transformers on 2025-09-18.*
@@ -90,6 +89,7 @@ echo -e "Plants create energy through a process known as" | transformers-cli run
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [torchao](../quantization/torchao) to only quantize the weights to 4-bits.
+
```py
#pip install torchao
@@ -119,7 +119,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
```
-
## FlexOlmoConfig
[[autodoc]] FlexOlmoConfig
diff --git a/docs/source/en/model_doc/florence2.md b/docs/source/en/model_doc/florence2.md
index 77e8de10c31b..b7171e1faabd 100644
--- a/docs/source/en/model_doc/florence2.md
+++ b/docs/source/en/model_doc/florence2.md
@@ -138,21 +138,21 @@ print(parsed_answer)
## Notes
- Florence-2 is a prompt-based model. You need to provide a task prompt to tell the model what to do. Supported tasks are:
- - `
@@ -285,4 +284,3 @@ alt="drawing" width="600"/>
[[autodoc]] GotOcr2ForConditionalGeneration
- forward
-
diff --git a/docs/source/en/model_doc/gpt2.md b/docs/source/en/model_doc/gpt2.md
index 1645a92f6346..2740bfb33393 100644
--- a/docs/source/en/model_doc/gpt2.md
+++ b/docs/source/en/model_doc/gpt2.md
@@ -23,7 +23,6 @@ rendered properly in your Markdown viewer.
-
## GPTBigCodeConfig
[[autodoc]] GPTBigCodeConfig
diff --git a/docs/source/en/model_doc/gpt_neo.md b/docs/source/en/model_doc/gpt_neo.md
index f3de04d0e550..b0d13cf780b3 100644
--- a/docs/source/en/model_doc/gpt_neo.md
+++ b/docs/source/en/model_doc/gpt_neo.md
@@ -22,12 +22,10 @@ rendered properly in your Markdown viewer.
-
## GPT-Neo
[GPT-Neo](https://zenodo.org/records/5297715) is an open-source alternative to GPT-2 and GPT-3 models, built with Mesh TensorFlow for TPUs. GPT-Neo uses local attention in every other layer for more efficiency. It is trained on the [Pile](https://huggingface.co/datasets/EleutherAI/pile), a diverse dataset consisting of 22 smaller high-quality datasets. The original github repository can be found [here](https://github.com/EleutherAI/gpt-neo/tree/v1.1)
-
You can find all the original GPT-Neo checkpoints under the [EleutherAI](https://huggingface.co/EleutherAI?search_models=gpt-neo) organization.
> [!TIP]
@@ -45,6 +43,7 @@ from transformers import pipeline
pipeline = pipeline(task="text-generation", model="EleutherAI/gpt-neo-1.3B", dtype=torch.float16, device=0)
pipeline("Hello, I'm a language model")
```
+
-
## Using Scaled Dot Product Attention (SDPA)
PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
@@ -162,7 +160,6 @@ following speedups during training and inference.
| 4 | 1024 | 11.765 | 11.303 | 4.09 | 2558.96 | 2546.04 | 0.508 |
| 4 | 2048 | 19.568 | 17.735 | 10.33 | 4175.5 | 4165.26 | 0.246 |
-
## Resources
- [Causal language modeling task guide](../tasks/language_modeling)
diff --git a/docs/source/en/model_doc/gpt_neox_japanese.md b/docs/source/en/model_doc/gpt_neox_japanese.md
index 7b22484b9a76..bf786f7561d4 100644
--- a/docs/source/en/model_doc/gpt_neox_japanese.md
+++ b/docs/source/en/model_doc/gpt_neox_japanese.md
@@ -27,8 +27,6 @@ rendered properly in your Markdown viewer.
GPT-NeoX-Japanese, a Japanese language model based on [GPT-NeoX](./gpt_neox).
Japanese uses three types of characters (hiragana, katakana, kanji) and has a huge vocabulary. This model uses [BPEEncoder V2](https://github.com/tanreinama/Japanese-BPEEncoder_V2), a sub-word tokenizer to handle the different characters.
-
-
The model also removes some bias parameters for better performance.
You can find all the original GPT-NeoX-Japanese checkpoints under the [ABEJA](https://huggingface.co/abeja/models?search=gpt-neo-x) organization.
diff --git a/docs/source/en/model_doc/gpt_oss.md b/docs/source/en/model_doc/gpt_oss.md
index 136ebeb29570..60741d8473fa 100644
--- a/docs/source/en/model_doc/gpt_oss.md
+++ b/docs/source/en/model_doc/gpt_oss.md
@@ -35,13 +35,14 @@ The abstract from the paper is the following:
*
@@ -69,12 +68,14 @@ inputs = tokenizer("Explain quantum computing in simple terms", return_tensors="
outputs = model.generate(**inputs, max_length=50, cache_implementation="static")
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```
+
-
> Click on the I-JEPA models in the right sidebar for more examples of how to apply I-JEPA to different image representation and classification tasks.
The example below demonstrates how to extract image features with [`Pipeline`] or the [`AutoModel`] class.
@@ -88,10 +86,10 @@ embed_2 = infer(image_2)
similarity = cosine_similarity(embed_1, embed_2)
print(similarity)
```
+
-
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [bitsandbytes](../quantization/bitsandbytes) to only quantize the weights to 4-bits.
@@ -142,4 +140,3 @@ print(similarity)
[[autodoc]] IJepaForImageClassification
- forward
-
diff --git a/docs/source/en/model_doc/informer.md b/docs/source/en/model_doc/informer.md
index 7e79399cbc57..a9cea0f09cab 100644
--- a/docs/source/en/model_doc/informer.md
+++ b/docs/source/en/model_doc/informer.md
@@ -52,4 +52,4 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
## InformerForPrediction
[[autodoc]] InformerForPrediction
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/instructblip.md b/docs/source/en/model_doc/instructblip.md
index b0669f1c065f..ac84a71d887e 100644
--- a/docs/source/en/model_doc/instructblip.md
+++ b/docs/source/en/model_doc/instructblip.md
@@ -59,7 +59,6 @@ The attributes can be obtained from model config, as `model.config.num_query_tok
[[autodoc]] InstructBlipProcessor
-
## InstructBlipVisionModel
[[autodoc]] InstructBlipVisionModel
@@ -78,4 +77,4 @@ The attributes can be obtained from model config, as `model.config.num_query_tok
[[autodoc]] InstructBlipForConditionalGeneration
- forward
- - generate
\ No newline at end of file
+ - generate
diff --git a/docs/source/en/model_doc/instructblipvideo.md b/docs/source/en/model_doc/instructblipvideo.md
index e34b454a1237..d4d868b7f90e 100644
--- a/docs/source/en/model_doc/instructblipvideo.md
+++ b/docs/source/en/model_doc/instructblipvideo.md
@@ -59,7 +59,6 @@ The attributes can be obtained from model config, as `model.config.num_query_tok
[[autodoc]] InstructBlipVideoProcessor
-
## InstructBlipVideoVideoProcessor
[[autodoc]] InstructBlipVideoVideoProcessor
diff --git a/docs/source/en/model_doc/internvl.md b/docs/source/en/model_doc/internvl.md
index bf760fdbdd71..7e9fea7f4f20 100644
--- a/docs/source/en/model_doc/internvl.md
+++ b/docs/source/en/model_doc/internvl.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2025-04-14 and added to Hugging Face Transformers on 2025-04-18.*
-
@@ -32,19 +31,14 @@ The abstract from the paper is the following:
*We introduce InternVL3, a significant advancement in the InternVL series featuring a native multimodal pre-training paradigm. Rather than adapting a text-only large language model (LLM) into a multimodal large language model (MLLM) that supports visual inputs, InternVL3 jointly acquires multimodal and linguistic capabilities from both diverse multimodal data and pure-text corpora during a single pre-training stage. This unified training paradigm effectively addresses the complexities and alignment challenges commonly encountered in conventional post-hoc training pipelines for MLLMs. To further improve performance and scalability, InternVL3 incorporates variable visual position encoding (V2PE) to support extended multimodal contexts, employs advanced post-training techniques such as supervised fine-tuning (SFT) and mixed preference optimization (MPO), and adopts test-time scaling strategies alongside an optimized training infrastructure. Extensive empirical evaluations demonstrate that InternVL3 delivers superior performance across a wide range of multi-modal tasks. In particular, InternVL3-78B achieves a score of 72.2 on the MMMU benchmark, setting a new state-of-the-art among open-source MLLMs. Its capabilities remain highly competitive with leading proprietary models, including ChatGPT-4o, Claude 3.5 Sonnet, and Gemini 2.5 Pro, while also maintaining strong pure-language proficiency. In pursuit of open-science principles, we will publicly release both the training data and model weights to foster further research and development in next-generation MLLMs.*
-
Overview of InternVL3 models architecture, which is the same as InternVL2.5. Taken from the original checkpoint.
-
-
Comparison of InternVL3 performance on OpenCompass against other SOTA VLLMs. Taken from the original checkpoint.
-
-
This model was contributed by [yonigozlan](https://huggingface.co/yonigozlan).
The original code can be found [here](https://github.com/OpenGVLab/InternVL).
@@ -75,6 +69,7 @@ Here is how you can use the `image-text-to-text` pipeline to perform inference w
>>> outputs[0]["generated_text"]
'The image showcases a vibrant scene of nature, featuring several flowers and a bee. \n\n1. **Foreground Flowers**: \n - The primary focus is on a large, pink cosmos flower with a prominent yellow center. The petals are soft and slightly r'
```
+
### Inference on a single image
This example demonstrates how to perform inference on a single image with the InternVL models using chat templates.
@@ -112,7 +107,6 @@ This example demonstrates how to perform inference on a single image with the In
### Text-only generation
This example shows how to generate text using the InternVL model without providing any image input.
-
```python
>>> from transformers import AutoProcessor, AutoModelForImageTextToText
>>> import torch
diff --git a/docs/source/en/model_doc/jamba.md b/docs/source/en/model_doc/jamba.md
index 0aa06b16e90f..f85d08c5f64d 100644
--- a/docs/source/en/model_doc/jamba.md
+++ b/docs/source/en/model_doc/jamba.md
@@ -75,6 +75,7 @@ input_ids = tokenizer("Plants create energy through a process known as", return_
output = model.generate(**input_ids, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -19,7 +19,6 @@ specific language governing permissions and limitations under the License.
+
+
+
+# LFM2-VL
+
+## Overview
+
+[LFM2-VL](https://www.liquid.ai/blog/lfm2-vl-efficient-vision-language-models) first series of vision-language foundation models developed by [Liquid AI](https://liquid.ai/). These multimodal models are designed for low-latency and device-aware deployment. LFM2-VL extends the LFM2 family of open-weight Liquid Foundation Models (LFMs) into the vision-language space, supporting both text and image inputs with variable resolutions.
+
+## Architecture
+
+LFM2-VL consists of three main components: a language model backbone, a vision encoder, and a multimodal projector. LFM2-VL builds upon the LFM2 backbone, inheriting from either LFM2-1.2B (for LFM2-VL-1.6B) or LFM2-350M (for LFM2-VL-450M). For the vision tower, LFM2-VL uses SigLIP2 NaFlex encoders to convert input images into token sequences. Two variants are implemented:
+
+* Shape-optimized (400M) for more fine-grained vision capabilities for LFM2-VL-1.6B
+* Base (86M) for fast image processing for LFM2-VL-450M
+
+The encoder processes images at their native resolution up to 512×512 pixels, efficiently handling smaller images without upscaling and supporting non-standard aspect ratios without distortion. Larger images are split into non-overlapping square patches of 512×512 each, preserving detail. In LFM2-VL-1.6B, the model also receives a thumbnail (a small, downscaled version of the original image capturing the overall scene) to enhance global context understanding and alignment. Special tokens mark each patch’s position and indicate the thumbnail’s start. The multimodal connector is a 2-layer MLP connector with pixel unshuffle to reduce image token count.
+
+## Example
+
+The following example shows how to generate an answer using the `AutoModelForImageTextToText` class.
+
+```python
+from transformers import AutoProcessor, AutoModelForImageTextToText
+\
+# Load model and processor
+model_id = "LiquidAI/LFM2-VL-1.6B"
+model = AutoModelForImageTextToText.from_pretrained(
+ model_id,
+ device_map="auto",
+ dtype="bfloat16",
+)
+processor = AutoProcessor.from_pretrained(model_id)
+
+# Load image and create conversation
+conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "image": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ {"type": "text", "text": "What is in this image?"},
+ ],
+ },
+]
+
+# Generate snswer
+inputs = processor.apply_chat_template(
+ conversation,
+ add_generation_prompt=True,
+ return_tensors="pt",
+ return_dict=True,
+ tokenize=True,
+).to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=64)
+processor.batch_decode(outputs, skip_special_tokens=True)[0]
+
+```
+
+## Lfm2VlImageProcessorFast
+
+[[autodoc]] Lfm2VlImageProcessorFast
+
+## Lfm2VlProcessor
+
+[[autodoc]] Lfm2VlProcessor
+
+## Lfm2VlConfig
+
+[[autodoc]] Lfm2VlConfig
+
+## Lfm2VlModel
+
+[[autodoc]] Lfm2VlModel
+ - forward
+
+## Lfm2VlForConditionalGeneration
+
+[[autodoc]] Lfm2VlForConditionalGeneration
+ - forward
diff --git a/docs/source/en/model_doc/lightglue.md b/docs/source/en/model_doc/lightglue.md
index 13ac58a1b842..2a173a8e1422 100644
--- a/docs/source/en/model_doc/lightglue.md
+++ b/docs/source/en/model_doc/lightglue.md
@@ -143,10 +143,9 @@ processed_outputs = processor.post_process_keypoint_matching(outputs, image_size
## LightGlueImageProcessor
[[autodoc]] LightGlueImageProcessor
-
-- preprocess
-- post_process_keypoint_matching
-- visualize_keypoint_matching
+ - preprocess
+ - post_process_keypoint_matching
+ - visualize_keypoint_matching
+
@@ -28,9 +27,11 @@ rendered properly in your Markdown viewer.
[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/), developed by Meta, introduces a new auto-regressive Mixture-of-Experts (MoE) architecture.
This generation includes two models:
+
- The highly capable Llama 4 Maverick with 17B active parameters out of ~400B total, with 128 experts.
- The efficient Llama 4 Scout also has 17B active parameters out of ~109B total, using just 16 experts.
-
+
Both models leverage early fusion for native multimodality, enabling them to process text and image inputs.
Maverick and Scout are both trained on up to 40 trillion tokens on data encompassing 200 languages
(with specific fine-tuning support for 12 languages including Arabic, Spanish, German, and Hindi).
@@ -53,7 +54,6 @@ The examples below demonstrates how to generate with [`Pipeline`] or the [`AutoM
showcasing how to toggle the right attributes to enable very long-context generations, as some flavors of Llama 4
have context lengths going up to 10 million tokens.
-
@@ -29,7 +30,6 @@ rendered properly in your Markdown viewer.
This architecture turns out to coincide with Qwen2, with the main difference being the presence of biases in attention projections in Ministral.
-
You can find the Ministral checkpoints under the [Mistral AI](https://huggingface.co/mistralai) organization.
## Usage
@@ -83,4 +83,4 @@ The example below demonstrates how to use Ministral for text generation:
## MinistralForQuestionAnswering
[[autodoc]] MinistralForQuestionAnswering
-- forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/mistral.md b/docs/source/en/model_doc/mistral.md
index 3714f45e55a0..4c598fc79a71 100644
--- a/docs/source/en/model_doc/mistral.md
+++ b/docs/source/en/model_doc/mistral.md
@@ -86,7 +86,6 @@ echo -e "My favorite condiment is" | transformers chat mistralai/Mistral-7B-v0.3
-
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [bitsandbytes](../quantization/bitsandbytes) to only quantize the weights to 4-bits.
@@ -164,4 +163,4 @@ Use the [AttentionMaskVisualizer](https://github.com/huggingface/transformers/bl
## MistralForQuestionAnswering
[[autodoc]] MistralForQuestionAnswering
-- forward
+ - forward
diff --git a/docs/source/en/model_doc/mistral3.md b/docs/source/en/model_doc/mistral3.md
index 54af880ed467..4ac264ac9854 100644
--- a/docs/source/en/model_doc/mistral3.md
+++ b/docs/source/en/model_doc/mistral3.md
@@ -27,7 +27,6 @@ rendered properly in your Markdown viewer.
You can find the original Mistral 3 checkpoints under the [Mistral AI](https://huggingface.co/mistralai/models?search=mistral-small-3) organization.
-
> [!TIP]
> This model was contributed by [cyrilvallez](https://huggingface.co/cyrilvallez) and [yonigozlan](https://huggingface.co/yonigozlan).
> Click on the Mistral3 models in the right sidebar for more examples of how to apply Mistral3 to different tasks.
@@ -62,6 +61,7 @@ outputs = pipeline(text=messages, max_new_tokens=50, return_full_text=False)
outputs[0]["generated_text"]
'The image depicts a vibrant and lush garden scene featuring a variety of wildflowers and plants. The central focus is on a large, pinkish-purple flower, likely a Greater Celandine (Chelidonium majus), with a'
```
+
@@ -100,13 +100,15 @@ decoded_output = processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :
decoded_output
'The image depicts a vibrant and lush garden scene featuring a variety of wildflowers and plants. The central focus is on a large, pinkish-purple flower, likely a Greater Celandine (Chelidonium majus), with a'
```
+
-## Notes
+## Notes
+
+- Mistral 3 supports text-only generation.
-- Mistral 3 supports text-only generation.
-```py
+```py
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, infer_device
@@ -136,13 +138,16 @@ print(decoded_output)
5. Je me casse, à plus!
```
+
/\_/\
( o.o )
> ^ <
+
```"
````
-- Mistral 3 accepts batched image and text inputs.
+- Mistral 3 accepts batched image and text inputs.
+
```py
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, infer_device
@@ -184,7 +189,7 @@ messages = [
, "Describe this imageThe image depicts a vibrant street scene in what appears to be a Chinatown district. The focal point is a traditional Chinese"]
```
-- Mistral 3 also supported batched image and text inputs with a different number of images for each text. The example below quantizes the model with bitsandbytes.
+- Mistral 3 also supported batched image and text inputs with a different number of images for each text. The example below quantizes the model with bitsandbytes.
```py
import torch
diff --git a/docs/source/en/model_doc/mixtral.md b/docs/source/en/model_doc/mixtral.md
index ff501cd1a84d..1e9574145aa1 100644
--- a/docs/source/en/model_doc/mixtral.md
+++ b/docs/source/en/model_doc/mixtral.md
@@ -39,9 +39,10 @@ Mixtral-8x7B is the second large language model (LLM) released by [mistral.ai](h
Mixtral-8x7B is a decoder-only Transformer with the following architectural choices:
- Mixtral is a Mixture of Experts (MoE) model with 8 experts per MLP, with a total of 45 billion parameters. To learn more about mixture-of-experts, refer to the [blog post](https://huggingface.co/blog/moe).
-- Despite the model having 45 billion parameters, the compute required for a single forward pass is the same as that of a 14 billion parameter model. This is because even though each of the experts have to be loaded in RAM (70B like ram requirement) each token from the hidden states are dispatched twice (top 2 routing) and thus the compute (the operation required at each forward computation) is just 2 X sequence_length.
+- Despite the model having 45 billion parameters, the compute required for a single forward pass is the same as that of a 14 billion parameter model. This is because even though each of the experts have to be loaded in RAM (70B like ram requirement) each token from the hidden states are dispatched twice (top 2 routing) and thus the compute (the operation required at each forward computation) is just 2 X sequence_length.
The following implementation details are shared with Mistral AI's first model [Mistral-7B](mistral):
+
- Sliding Window Attention - Trained with 8k context length and fixed cache size, with a theoretical attention span of 128K tokens
- GQA (Grouped Query Attention) - allowing faster inference and lower cache size.
- Byte-fallback BPE tokenizer - ensures that characters are never mapped to out of vocabulary tokens.
@@ -55,6 +56,7 @@ For more details refer to the [release blog post](https://mistral.ai/news/mixtra
## Usage tips
The Mistral team has released 2 checkpoints:
+
- a base model, [Mixtral-8x7B-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1), which has been pre-trained to predict the next token on internet-scale data.
- an instruction tuned model, [Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1), which is the base model optimized for chat purposes using supervised fine-tuning (SFT) and direct preference optimization (DPO).
@@ -138,8 +140,8 @@ Below is a expected speedup diagram that compares pure inference time between th
### Sliding window Attention
-The current implementation supports the sliding window attention mechanism and memory efficient cache management.
-To enable sliding window attention, just make sure to have a `flash-attn` version that is compatible with sliding window attention (`>=2.3.0`).
+The current implementation supports the sliding window attention mechanism and memory efficient cache management.
+To enable sliding window attention, just make sure to have a `flash-attn` version that is compatible with sliding window attention (`>=2.3.0`).
The Flash Attention-2 model uses also a more memory efficient cache slicing mechanism - as recommended per the official implementation of Mistral model that use rolling cache mechanism we keep the cache size fixed (`self.config.sliding_window`), support batched generation only for `padding_side="left"` and use the absolute position of the current token to compute the positional embedding.
diff --git a/docs/source/en/model_doc/mlcd.md b/docs/source/en/model_doc/mlcd.md
index 1ce785ee76bb..7ff2fb434da0 100644
--- a/docs/source/en/model_doc/mlcd.md
+++ b/docs/source/en/model_doc/mlcd.md
@@ -32,9 +32,9 @@ Tips:
- We adopted the official [LLaVA-NeXT](https://github.com/LLaVA-VL/LLaVA-NeXT) and the official training dataset [LLaVA-NeXT-Data](https://huggingface.co/datasets/lmms-lab/LLaVA-NeXT-Data) for evaluating the foundational visual models.
-- The language model is [Qwen2.5-7B](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct).
+- The language model is [Qwen2.5-7B](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct).
-Result:
+Result:
| Vision Tower | RoPE2D | ChartQA | DocVQA | InfoVQA | OCRBench | MMMU |
| :-------------------------------------------------------------------------------------------- | :----: | :-------- | :-------- | :-------- | :--------- | :-------- |
@@ -45,7 +45,6 @@ Result:
| **[MLCD (ViT-bigG-14-336px)](https://huggingface.co/DeepGlint-AI/mlcd-vit-bigG-patch14-336)** | √ | 71.07 | 79.63 | 44.38 | 572.00 | 46.78 |
| **[MLCD (ViT-bigG-14-448px)](https://huggingface.co/DeepGlint-AI/mlcd-vit-bigG-patch14-448)** | √ | **73.80** | **83.34** | **46.59** | **582.00** | 46.00 |
-
## Usage
```python
diff --git a/docs/source/en/model_doc/mllama.md b/docs/source/en/model_doc/mllama.md
index 1ea7f172bb3a..a0fc5db41cfe 100644
--- a/docs/source/en/model_doc/mllama.md
+++ b/docs/source/en/model_doc/mllama.md
@@ -35,15 +35,12 @@ The [Llama 3.2-Vision](https://ai.meta.com/blog/llama-3-2-connect-2024-vision-ed
- The text passed to the processor should have the `"<|image|>"` tokens where the images should be inserted.
- The processor has its own `apply_chat_template` method to convert chat messages to text that can then be passed as text to the processor. If you're using `transformers>=4.49.0`, you can also get a vectorized output from `apply_chat_template`. See the **Usage Examples** below for more details on how to use it.
-
-
Mllama has an extra token used as a placeholder for image positions in the text. It means that input ids and an input embedding layer will have an extra token. But since the weights for input and output embeddings are not tied, the `lm_head` layer has one less token and will fail if you want to calculate loss on image tokens or apply some logit processors. In case you are training, make sure to mask out special `"<|image|>"` tokens in the `labels` as the model should not be trained on predicting them.
Otherwise if you see CUDA-side index errors when generating, use the below code to expand the `lm_head` by one more token.
-
```python
old_embeddings = model.get_output_embeddings()
@@ -52,12 +49,13 @@ resized_embeddings = model._get_resized_lm_head(old_embeddings, new_num_tokens=n
resized_embeddings.requires_grad_(old_embeddings.weight.requires_grad)
model.set_output_embeddings(resized_embeddings)
```
-
+
## Usage Example
#### Instruct model
+
```python
import torch
from transformers import MllamaForConditionalGeneration, AutoProcessor
@@ -83,6 +81,7 @@ print(processor.decode(output[0]))
```
#### Base model
+
```python
import requests
import torch
@@ -102,7 +101,6 @@ output = model.generate(**inputs, do_sample=False, max_new_tokens=25)
print(processor.decode(output[0], skip_special_tokens=True))
```
-
## MllamaConfig
[[autodoc]] MllamaConfig
@@ -111,7 +109,6 @@ print(processor.decode(output[0], skip_special_tokens=True))
[[autodoc]] MllamaProcessor
-
## MllamaImageProcessor
[[autodoc]] MllamaImageProcessor
diff --git a/docs/source/en/model_doc/mm-grounding-dino.md b/docs/source/en/model_doc/mm-grounding-dino.md
index e411ef5defb6..0d628c3b31de 100644
--- a/docs/source/en/model_doc/mm-grounding-dino.md
+++ b/docs/source/en/model_doc/mm-grounding-dino.md
@@ -100,7 +100,6 @@ for box, score, labels in zip(result["boxes"], result["scores"], result["labels"
| [mm_grounding_dino_tiny_o365v1_goldg_v3det](https://huggingface.co/openmmlab-community/mm_grounding_dino_tiny_o365v1_goldg_v3det) | O365,GoldG,V3Det | 33.0 | 36.0 | 45.9 | 40.5(+11.7) | 21.5 | 25.5 | 40.2 | 30.6(+10.5) |
| [mm_grounding_dino_tiny_o365v1_goldg_grit_v3det](https://huggingface.co/openmmlab-community/mm_grounding_dino_tiny_o365v1_goldg_grit_v3det) | O365,GoldG,GRIT,V3Det | 34.2 | 37.4 | 46.2 | 41.4(+12.6) | 23.6 | 27.6 | 40.5 | 31.9(+11.8) |
-
- This implementation also supports inference for [LLMDet](https://github.com/iSEE-Laboratory/LLMDet). Here's a table of LLMDet models and their performance on LVIS (results from [official repo](https://github.com/iSEE-Laboratory/LLMDet)):
| Model | Pre-Train Data | MiniVal APr | MiniVal APc | MiniVal APf | MiniVal AP | Val1.0 APr | Val1.0 APc | Val1.0 APf | Val1.0 AP |
@@ -109,7 +108,6 @@ for box, score, labels in zip(result["boxes"], result["scores"], result["labels"
| [llmdet_base](https://huggingface.co/iSEE-Laboratory/llmdet_base) | (O365,GoldG,V3Det) + GroundingCap-1M | 48.3 | 40.8 | 43.1 | 54.3 | 38.5 | 28.2 | 34.3 | 47.8 |
| [llmdet_large](https://huggingface.co/iSEE-Laboratory/llmdet_large) | (O365V2,OpenImageV6,GoldG) + GroundingCap-1M | 51.1 | 45.1 | 46.1 | 56.6 | 42.0 | 31.6 | 38.8 | 50.2 |
-
## MMGroundingDinoConfig
[[autodoc]] MMGroundingDinoConfig
diff --git a/docs/source/en/model_doc/mms.md b/docs/source/en/model_doc/mms.md
index 3ac351d0ddcb..171beaf440d1 100644
--- a/docs/source/en/model_doc/mms.md
+++ b/docs/source/en/model_doc/mms.md
@@ -376,6 +376,7 @@ detected_lang = model.config.id2label[lang_id]
```
To see all the supported languages of a checkpoint, you can print out the language ids as follows:
+
```py
processor.id2label.values()
```
diff --git a/docs/source/en/model_doc/mobilebert.md b/docs/source/en/model_doc/mobilebert.md
index 4e3cc2e5d647..08486ace56eb 100644
--- a/docs/source/en/model_doc/mobilebert.md
+++ b/docs/source/en/model_doc/mobilebert.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2020-04-06 and added to Hugging Face Transformers on 2020-11-16.*
-
@@ -47,6 +46,7 @@ pipeline = pipeline(
)
pipeline("The capital of France is [MASK].")
```
+
@@ -85,7 +85,6 @@ echo -e "The capital of France is [MASK]." | transformers run --task fill-mask -
-
## Notes
- Inputs should be padded on the right because BERT uses absolute position embeddings.
diff --git a/docs/source/en/model_doc/mobilenet_v1.md b/docs/source/en/model_doc/mobilenet_v1.md
index c77bef730423..eea159bdd738 100644
--- a/docs/source/en/model_doc/mobilenet_v1.md
+++ b/docs/source/en/model_doc/mobilenet_v1.md
@@ -32,7 +32,6 @@ You can all the original MobileNet checkpoints under the [Google](https://huggin
The example below demonstrates how to classify an image with [`Pipeline`] or the [`AutoModel`] class.
-
@@ -84,23 +83,24 @@ print(f"The predicted class label is: {predicted_class_label}")
-
## Notes
-- Checkpoint names follow the pattern `mobilenet_v1_{depth_multiplier}_{resolution}`, like `mobilenet_v1_1.0_224`. `1.0` is the depth multiplier and `224` is the image resolution.
-- While trained on images of a specific sizes, the model architecture works with images of different sizes (minimum 32x32). The [`MobileNetV1ImageProcessor`] handles the necessary preprocessing.
-- MobileNet is pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k), a dataset with 1000 classes. However, the model actually predicts 1001 classes. The additional class is an extra "background" class (index 0).
-- The original TensorFlow checkpoints determines the padding amount at inference because it depends on the input image size. To use the native PyTorch padding behavior, set `tf_padding=False` in [`MobileNetV1Config`].
+- Checkpoint names follow the pattern `mobilenet_v1_{depth_multiplier}_{resolution}`, like `mobilenet_v1_1.0_224`. `1.0` is the depth multiplier and `224` is the image resolution.
+- While trained on images of a specific sizes, the model architecture works with images of different sizes (minimum 32x32). The [`MobileNetV1ImageProcessor`] handles the necessary preprocessing.
+- MobileNet is pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k), a dataset with 1000 classes. However, the model actually predicts 1001 classes. The additional class is an extra "background" class (index 0).
+- The original TensorFlow checkpoints determines the padding amount at inference because it depends on the input image size. To use the native PyTorch padding behavior, set `tf_padding=False` in [`MobileNetV1Config`].
+
```python
from transformers import MobileNetV1Config
config = MobileNetV1Config.from_pretrained("google/mobilenet_v1_1.0_224", tf_padding=True)
```
-- The Transformers implementation does not support the following features.
- - Uses global average pooling instead of the optional 7x7 average pooling with stride 2. For larger inputs, this gives a pooled output that is larger than a 1x1 pixel.
- - Does not support other `output_stride` values (fixed at 32). For smaller `output_strides`, the original implementation uses dilated convolution to prevent spatial resolution from being reduced further. (which would require dilated convolutions).
- - `output_hidden_states=True` returns *all* intermediate hidden states. It is not possible to extract the output from specific layers for other downstream purposes.
- - Does not include the quantized models from the original checkpoints because they include "FakeQuantization" operations to unquantize the weights.
+
+- The Transformers implementation does not support the following features.
+ - Uses global average pooling instead of the optional 7x7 average pooling with stride 2. For larger inputs, this gives a pooled output that is larger than a 1x1 pixel.
+ - Does not support other `output_stride` values (fixed at 32). For smaller `output_strides`, the original implementation uses dilated convolution to prevent spatial resolution from being reduced further. (which would require dilated convolutions).
+ - `output_hidden_states=True` returns *all* intermediate hidden states. It is not possible to extract the output from specific layers for other downstream purposes.
+ - Does not include the quantized models from the original checkpoints because they include "FakeQuantization" operations to unquantize the weights.
## MobileNetV1Config
diff --git a/docs/source/en/model_doc/mobilenet_v2.md b/docs/source/en/model_doc/mobilenet_v2.md
index 3e1379e3f079..bf94454e438d 100644
--- a/docs/source/en/model_doc/mobilenet_v2.md
+++ b/docs/source/en/model_doc/mobilenet_v2.md
@@ -30,10 +30,8 @@ You can all the original MobileNet checkpoints under the [Google](https://huggin
> [!TIP]
> Click on the MobileNet V2 models in the right sidebar for more examples of how to apply MobileNet to different vision tasks.
-
The examples below demonstrate how to classify an image with [`Pipeline`] or the [`AutoModel`] class.
-
@@ -82,24 +80,25 @@ print(f"The predicted class label is: {predicted_class_label}")
-
## Notes
-- Classification checkpoint names follow the pattern `mobilenet_v2_{depth_multiplier}_{resolution}`, like `mobilenet_v2_1.4_224`. `1.4` is the depth multiplier and `224` is the image resolution. Segmentation checkpoint names follow the pattern `deeplabv3_mobilenet_v2_{depth_multiplier}_{resolution}`.
-- While trained on images of a specific sizes, the model architecture works with images of different sizes (minimum 32x32). The [`MobileNetV2ImageProcessor`] handles the necessary preprocessing.
-- MobileNet is pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k), a dataset with 1000 classes. However, the model actually predicts 1001 classes. The additional class is an extra "background" class (index 0).
-- The segmentation models use a [DeepLabV3+](https://huggingface.co/papers/1802.02611) head which is often pretrained on datasets like [PASCAL VOC](https://huggingface.co/datasets/merve/pascal-voc).
-- The original TensorFlow checkpoints determines the padding amount at inference because it depends on the input image size. To use the native PyTorch padding behavior, set `tf_padding=False` in [`MobileNetV2Config`].
+- Classification checkpoint names follow the pattern `mobilenet_v2_{depth_multiplier}_{resolution}`, like `mobilenet_v2_1.4_224`. `1.4` is the depth multiplier and `224` is the image resolution. Segmentation checkpoint names follow the pattern `deeplabv3_mobilenet_v2_{depth_multiplier}_{resolution}`.
+- While trained on images of a specific sizes, the model architecture works with images of different sizes (minimum 32x32). The [`MobileNetV2ImageProcessor`] handles the necessary preprocessing.
+- MobileNet is pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k), a dataset with 1000 classes. However, the model actually predicts 1001 classes. The additional class is an extra "background" class (index 0).
+- The segmentation models use a [DeepLabV3+](https://huggingface.co/papers/1802.02611) head which is often pretrained on datasets like [PASCAL VOC](https://huggingface.co/datasets/merve/pascal-voc).
+- The original TensorFlow checkpoints determines the padding amount at inference because it depends on the input image size. To use the native PyTorch padding behavior, set `tf_padding=False` in [`MobileNetV2Config`].
+
```python
from transformers import MobileNetV2Config
config = MobileNetV2Config.from_pretrained("google/mobilenet_v2_1.4_224", tf_padding=True)
```
-- The Transformers implementation does not support the following features.
- - Uses global average pooling instead of the optional 7x7 average pooling with stride 2. For larger inputs, this gives a pooled output that is larger than a 1x1 pixel.
- - `output_hidden_states=True` returns *all* intermediate hidden states. It is not possible to extract the output from specific layers for other downstream purposes.
- - Does not include the quantized models from the original checkpoints because they include "FakeQuantization" operations to unquantize the weights.
- - For segmentation models, the final convolution layer of the backbone is computed even though the DeepLabV3+ head doesn't use it.
+
+- The Transformers implementation does not support the following features.
+ - Uses global average pooling instead of the optional 7x7 average pooling with stride 2. For larger inputs, this gives a pooled output that is larger than a 1x1 pixel.
+ - `output_hidden_states=True` returns *all* intermediate hidden states. It is not possible to extract the output from specific layers for other downstream purposes.
+ - Does not include the quantized models from the original checkpoints because they include "FakeQuantization" operations to unquantize the weights.
+ - For segmentation models, the final convolution layer of the backbone is computed even though the DeepLabV3+ head doesn't use it.
## MobileNetV2Config
diff --git a/docs/source/en/model_doc/mobilevit.md b/docs/source/en/model_doc/mobilevit.md
index b4a51bd200f2..ca0a35f6ece8 100644
--- a/docs/source/en/model_doc/mobilevit.md
+++ b/docs/source/en/model_doc/mobilevit.md
@@ -11,11 +11,8 @@ Unless required by applicable law or agreed to in writing, software distributed
-->
*This model was released on 2021-10-05 and added to Hugging Face Transformers on 2022-06-29.*
-
-
# MobileViT
-
@@ -24,21 +21,18 @@ Unless required by applicable law or agreed to in writing, software distributed
[MobileViT](https://huggingface.co/papers/2110.02178) is a lightweight vision transformer for mobile devices that merges CNNs's efficiency and inductive biases with transformers global context modeling. It treats transformers as convolutions, enabling global information processing without the heavy computational cost of standard ViTs.
-
-
You can find all the original MobileViT checkpoints under the [Apple](https://huggingface.co/apple/models?search=mobilevit) organization.
-
> [!TIP]
+>
> - This model was contributed by [matthijs](https://huggingface.co/Matthijs).
>
> Click on the MobileViT models in the right sidebar for more examples of how to apply MobileViT to different vision tasks.
-
The example below demonstrates how to do [Image Classification] with [`Pipeline`] and the [`AutoModel`] class.
@@ -92,7 +86,6 @@ print(f"The predicted class label is:{predicted_class_label}")
-
## Notes
- Does **not** operate on sequential data, it's purely designed for image tasks.
@@ -102,8 +95,6 @@ print(f"The predicted class label is:{predicted_class_label}")
- The classification models are pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k).
- The segmentation models use a [DeepLabV3](https://huggingface.co/papers/1706.05587) head and are pretrained on [PASCAL VOC](http://host.robots.ox.ac.uk/pascal/VOC/).
-
-
## MobileViTConfig
[[autodoc]] MobileViTConfig
diff --git a/docs/source/en/model_doc/modernbert-decoder.md b/docs/source/en/model_doc/modernbert-decoder.md
index 013b9d24b5f4..1ab96700659b 100644
--- a/docs/source/en/model_doc/modernbert-decoder.md
+++ b/docs/source/en/model_doc/modernbert-decoder.md
@@ -36,7 +36,7 @@ You can find all the original ModernBERT Decoder checkpoints under the [jhu-clsp
>
> Click on the ModernBERT Decoder models in the right sidebar for more examples of how to apply ModernBERT Decoder to different text generation tasks.
-The example below demonstrates how to use ModernBERT Decoder for text generation with [`Pipeline`], [`AutoModel`] (with and without quantization), and from the command line.
+The example below demonstrates how to use ModernBERT Decoder for text generation with [`Pipeline`], [`AutoModel`] (with and without quantization), and from the command line.
@@ -119,7 +119,7 @@ print(f"Prediction probabilities: {predictions}")
-```
+```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
@@ -151,6 +151,7 @@ with torch.no_grad():
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Generated text: {generated_text}")
```
+
@@ -162,7 +163,6 @@ echo "The future of artificial intelligence is" | transformers run --task text-g
-
## ModernBertDecoderConfig
[[autodoc]] ModernBertDecoderConfig
diff --git a/docs/source/en/model_doc/moonshine.md b/docs/source/en/model_doc/moonshine.md
index 7abe123b88e2..b85a174a86fb 100644
--- a/docs/source/en/model_doc/moonshine.md
+++ b/docs/source/en/model_doc/moonshine.md
@@ -83,6 +83,7 @@ predicted_ids = model.generate(**input_features, cache_implementation="static")
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
transcription[0]
```
+
@@ -101,4 +102,3 @@ transcription[0]
[[autodoc]] MoonshineForConditionalGeneration
- forward
- generate
-
diff --git a/docs/source/en/model_doc/moshi.md b/docs/source/en/model_doc/moshi.md
index e17a1b7b8b14..885623b26e52 100644
--- a/docs/source/en/model_doc/moshi.md
+++ b/docs/source/en/model_doc/moshi.md
@@ -35,9 +35,10 @@ Moshi is a speech-text foundation model that casts spoken dialogue as speech-to-
The abstract from the paper is the following:
-*We introduce Moshi, a speech-text foundation model and full-duplex spoken dialogue framework. Current systems for spoken dialogue rely on pipelines of independent components, namely voice activity detection, speech recognition, textual dialogue and text-to-speech. Such frameworks cannot emulate the experience of real conversations. First, their complexity induces a latency of several seconds between interactions. Second, text being the intermediate modality for dialogue, non-linguistic information that modifies meaning— such as emotion or non-speech sounds— is lost in the interaction. Finally, they rely on a segmentation into speaker turns, which does not take into account overlapping speech, interruptions and interjections. Moshi solves these independent issues altogether by casting spoken dialogue as speech-to-speech generation. Starting from a text language model backbone, Moshi generates speech as tokens from the residual quantizer of a neural audio codec, while modeling separately its own speech and that of the user into parallel streams. This allows for the removal of explicit speaker turns, and the modeling of arbitrary conversational dynamics. We moreover extend the hierarchical semantic-to-acoustic token generation of previous work to first predict time-aligned text tokens as a prefix to audio tokens. Not only this “Inner Monologue” method significantly improves the linguistic quality of generated speech, but we also illustrate how it can provide streaming speech recognition and text-to-speech. Our resulting model is the first real-time full-duplex spoken large language model, with a theoretical latency of 160ms, 200ms in practice, and is available at github.com/kyutai-labs/moshi.*
+*We introduce Moshi, a speech-text foundation model and full-duplex spoken dialogue framework. Current systems for spoken dialogue rely on pipelines of independent components, namely voice activity detection, speech recognition, textual dialogue and text-to-speech. Such frameworks cannot emulate the experience of real conversations. First, their complexity induces a latency of several seconds between interactions. Second, text being the intermediate modality for dialogue, non-linguistic information that modifies meaning— such as emotion or non-speech sounds— is lost in the interaction. Finally, they rely on a segmentation into speaker turns, which does not take into account overlapping speech, interruptions and interjections. Moshi solves these independent issues altogether by casting spoken dialogue as speech-to-speech generation. Starting from a text language model backbone, Moshi generates speech as tokens from the residual quantizer of a neural audio codec, while modeling separately its own speech and that of the user into parallel streams. This allows for the removal of explicit speaker turns, and the modeling of arbitrary conversational dynamics. We moreover extend the hierarchical semantic-to-acoustic token generation of previous work to first predict time-aligned text tokens as a prefix to audio tokens. Not only this “Inner Monologue” method significantly improves the linguistic quality of generated speech, but we also illustrate how it can provide streaming speech recognition and text-to-speech. Our resulting model is the first real-time full-duplex spoken large language model, with a theoretical latency of 160ms, 200ms in practice, and is available at github.com/kyutai-labs/moshi.*
Moshi deals with 3 streams of information:
+
1. The user's audio
2. Moshi's audio
3. Moshi's textual output
@@ -49,7 +50,7 @@ Moshi's made of 3 components:
**1. The main decoder (Helium in the paper)**
-It corresponds to [`MoshiForCausalLM`]. It is strictly a classic text LLM, that uses an architecture similar to [` ~GemmaForCausalLM`]. In other words, it takes text tokens, embeds them, pass them through the decoder and a language head, to get text logits.
+It corresponds to [`MoshiForCausalLM`]. It is strictly a classic text LLM, that uses an architecture similar to [`~GemmaForCausalLM`]. In other words, it takes text tokens, embeds them, pass them through the decoder and a language head, to get text logits.
**2. The depth decoder**
@@ -63,15 +64,14 @@ Note that each timestamp - i.e each codebook - gets its own set of Linear Layers
It's the audio encoder from Kyutai, that has recently been integrated to transformers, which is used to "tokenize" audio. It has the same use that [`~EncodecModel`] has in [`~MusicgenModel`].
-
## Tips:
-The original checkpoints can be converted using the conversion script `src/transformers/models/moshi/convert_moshi_transformers.py`
-
+The original checkpoints can be converted using the conversion script `src/transformers/models/moshi/convert_moshi_transformers.py`
### How to use the model:
This implementation has two main aims:
+
1. quickly test model generation by simplifying the original API
2. simplify training. A training guide will come soon, but user contributions are welcomed!
@@ -86,6 +86,7 @@ It is designed for intermediate use. We strongly recommend using the original [i
Moshi is a streaming auto-regressive model with two streams of audio. To put it differently, one audio stream corresponds to what the model said/will say and the other audio stream corresponds to what the user said/will say.
[`MoshiForConditionalGeneration.generate`] thus needs 3 inputs:
+
1. `input_ids` - corresponding to the text token history
2. `moshi_input_values` or `moshi_audio_codes`- corresponding to the model audio history
3. `user_input_values` or `user_audio_codes` - corresponding to the user audio history
@@ -93,6 +94,7 @@ Moshi is a streaming auto-regressive model with two streams of audio. To put it
These three inputs must be synchronized. Meaning that their lengths must correspond to the same number of tokens.
You can dynamically use the 3 inputs depending on what you want to test:
+
1. Simply check the model response to an user prompt - in that case, `input_ids` can be filled with pad tokens and `user_input_values` can be a zero tensor of the same shape than the user prompt.
2. Test more complex behaviour - in that case, you must be careful about how the input tokens are synchronized with the audios.
@@ -108,12 +110,9 @@ To follow the example of the following image, `"Hello, I'm Moshi"` could be tran
-
[`MoshiForConditionalGeneration.generate`] then auto-regressively feeds to itself its own audio stream, but since it doesn't have access to the user input stream while using `transformers`, it will thus **assume that the user is producing blank audio**.
-
-
-```python
+```python
>>> from datasets import load_dataset, Audio
>>> import torch, math
>>> from transformers import MoshiForConditionalGeneration, AutoFeatureExtractor, AutoTokenizer, infer_device
@@ -149,7 +148,7 @@ To follow the example of the following image, `"Hello, I'm Moshi"` could be tran
Most of the work has to be done during data creation/pre-processing, because of the need to align/synchronize streams.
Once it's done, you can simply forward `text_labels` and `audio_labels` to [`MoshiForConditionalGeneration.forward`], alongside the usual inputs, to get the model loss.
-
+
A training guide will come soon, but user contributions are welcomed!
### How does the model forward the inputs / generate:
@@ -162,13 +161,10 @@ A training guide will come soon, but user contributions are welcomed!
3. The depth decoder switches the dimension on which we forward / generate (codebooks instead of time). It uses the token generated from `text logits` and the `temporal context` to auto-regressively generate audio codebooks.
-
This model was contributed by [Yoach Lacombe (ylacombe)](https://huggingface.co/ylacombe).
The original code can be found [here](https://github.com/kyutai-labs/moshi).
-
-
## MoshiConfig
[[autodoc]] MoshiConfig
diff --git a/docs/source/en/model_doc/mpt.md b/docs/source/en/model_doc/mpt.md
index 9482e6a91958..60d14641177c 100644
--- a/docs/source/en/model_doc/mpt.md
+++ b/docs/source/en/model_doc/mpt.md
@@ -23,11 +23,11 @@ rendered properly in your Markdown viewer.
## Overview
-The MPT model was proposed by the [MosaicML](https://www.mosaicml.com/) team and released with multiple sizes and finetuned variants. The MPT models are a series of open source and commercially usable LLMs pre-trained on 1T tokens.
+The MPT model was proposed by the [MosaicML](https://www.mosaicml.com/) team and released with multiple sizes and finetuned variants. The MPT models are a series of open source and commercially usable LLMs pre-trained on 1T tokens.
-MPT models are GPT-style decoder-only transformers with several improvements: performance-optimized layer implementations, architecture changes that provide greater training stability, and the elimination of context length limits by replacing positional embeddings with ALiBi.
+MPT models are GPT-style decoder-only transformers with several improvements: performance-optimized layer implementations, architecture changes that provide greater training stability, and the elimination of context length limits by replacing positional embeddings with ALiBi.
-- MPT base: MPT base pre-trained models on next token prediction
+- MPT base: MPT base pre-trained models on next token prediction
- MPT instruct: MPT base models fine-tuned on instruction based tasks
- MPT storywriter: MPT base models fine-tuned for 2500 steps on 65k-token excerpts of fiction books contained in the books3 corpus, this enables the model to handle very long sequences
diff --git a/docs/source/en/model_doc/mra.md b/docs/source/en/model_doc/mra.md
index ed11d1d9e04f..422ed3cec515 100644
--- a/docs/source/en/model_doc/mra.md
+++ b/docs/source/en/model_doc/mra.md
@@ -64,4 +64,4 @@ The original code can be found [here](https://github.com/mlpen/mra-attention).
## MraForQuestionAnswering
[[autodoc]] MraForQuestionAnswering
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/mt5.md b/docs/source/en/model_doc/mt5.md
index fa02ee4c3c08..4e652458e1b3 100644
--- a/docs/source/en/model_doc/mt5.md
+++ b/docs/source/en/model_doc/mt5.md
@@ -133,7 +133,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
See [`T5Tokenizer`] for all details.
-
## MT5TokenizerFast
[[autodoc]] MT5TokenizerFast
diff --git a/docs/source/en/model_doc/musicgen.md b/docs/source/en/model_doc/musicgen.md
index 7e91b2265fe3..1b0e8868ac82 100644
--- a/docs/source/en/model_doc/musicgen.md
+++ b/docs/source/en/model_doc/musicgen.md
@@ -77,9 +77,9 @@ Generation is limited by the sinusoidal positional embeddings to 30 second input
than 30 seconds of audio (1503 tokens), and input audio passed by Audio-Prompted Generation contributes to this limit so,
given an input of 20 seconds of audio, MusicGen cannot generate more than 10 seconds of additional audio.
-Transformers supports both mono (1-channel) and stereo (2-channel) variants of MusicGen. The mono channel versions
-generate a single set of codebooks. The stereo versions generate 2 sets of codebooks, 1 for each channel (left/right),
-and each set of codebooks is decoded independently through the audio compression model. The audio streams for each
+Transformers supports both mono (1-channel) and stereo (2-channel) variants of MusicGen. The mono channel versions
+generate a single set of codebooks. The stereo versions generate 2 sets of codebooks, 1 for each channel (left/right),
+and each set of codebooks is decoded independently through the audio compression model. The audio streams for each
channel are combined to give the final stereo output.
### Unconditional Generation
@@ -208,7 +208,7 @@ For batched audio-prompted generation, the generated `audio_values` can be post-
### Generation Configuration
-The default parameters that control the generation process, such as sampling, guidance scale and number of generated
+The default parameters that control the generation process, such as sampling, guidance scale and number of generated
tokens, can be found in the model's generation config, and updated as desired:
```python
@@ -226,20 +226,21 @@ tokens, can be found in the model's generation config, and updated as desired:
>>> model.generation_config.max_length = 256
```
-Note that any arguments passed to the generate method will **supersede** those in the generation config, so setting
-`do_sample=False` in the call to generate will supersede the setting of `model.generation_config.do_sample` in the
+Note that any arguments passed to the generate method will **supersede** those in the generation config, so setting
+`do_sample=False` in the call to generate will supersede the setting of `model.generation_config.do_sample` in the
generation config.
## Model Structure
The MusicGen model can be de-composed into three distinct stages:
+
1. Text encoder: maps the text inputs to a sequence of hidden-state representations. The pre-trained MusicGen models use a frozen text encoder from either T5 or Flan-T5
2. MusicGen decoder: a language model (LM) that auto-regressively generates audio tokens (or codes) conditional on the encoder hidden-state representations
3. Audio encoder/decoder: used to encode an audio prompt to use as prompt tokens, and recover the audio waveform from the audio tokens predicted by the decoder
Thus, the MusicGen model can either be used as a standalone decoder model, corresponding to the class [`MusicgenForCausalLM`],
or as a composite model that includes the text encoder and audio encoder/decoder, corresponding to the class
-[`MusicgenForConditionalGeneration`]. If only the decoder needs to be loaded from the pre-trained checkpoint, it can be loaded by first
+[`MusicgenForConditionalGeneration`]. If only the decoder needs to be loaded from the pre-trained checkpoint, it can be loaded by first
specifying the correct config, or be accessed through the `.decoder` attribute of the composite model:
```python
@@ -259,6 +260,7 @@ be combined with the frozen text encoder and audio encoder/decoders to recover t
model.
Tips:
+
* MusicGen is trained on the 32kHz checkpoint of Encodec. You should ensure you use a compatible version of the Encodec model.
* Sampling mode tends to deliver better results than greedy - you can toggle sampling with the variable `do_sample` in the call to [`MusicgenForConditionalGeneration.generate`]
diff --git a/docs/source/en/model_doc/musicgen_melody.md b/docs/source/en/model_doc/musicgen_melody.md
index d2cd51bbcf2c..baf21adaab21 100644
--- a/docs/source/en/model_doc/musicgen_melody.md
+++ b/docs/source/en/model_doc/musicgen_melody.md
@@ -35,13 +35,12 @@ The abstract from the paper is the following:
*We tackle the task of conditional music generation. We introduce MusicGen, a single Language Model (LM) that operates over several streams of compressed discrete music representation, i.e., tokens. Unlike prior work, MusicGen is comprised of a single-stage transformer LM together with efficient token interleaving patterns, which eliminates the need for cascading several models, e.g., hierarchically or upsampling. Following this approach, we demonstrate how MusicGen can generate high-quality samples, while being conditioned on textual description or melodic features, allowing better controls over the generated output. We conduct extensive empirical evaluation, considering both automatic and human studies, showing the proposed approach is superior to the evaluated baselines on a standard text-to-music benchmark. Through ablation studies, we shed light over the importance of each of the components comprising MusicGen.*
-
This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The original code can be found [here](https://github.com/facebookresearch/audiocraft). The pre-trained checkpoints can be found on the [Hugging Face Hub](https://huggingface.co/models?sort=downloads&search=facebook%2Fmusicgen).
-
## Difference with [MusicGen](https://huggingface.co/docs/transformers/main/en/model_doc/musicgen)
There are two key differences with MusicGen:
+
1. The audio prompt is used here as a conditional signal for the generated audio sample, whereas it's used for audio continuation in [MusicGen](https://huggingface.co/docs/transformers/main/en/model_doc/musicgen).
2. Conditional text and audio signals are concatenated to the decoder's hidden states instead of being used as a cross-attention signal, as in MusicGen.
@@ -54,19 +53,19 @@ MusicGen Melody is compatible with two generation modes: greedy and sampling. In
Transformers supports both mono (1-channel) and stereo (2-channel) variants of MusicGen Melody. The mono channel versions generate a single set of codebooks. The stereo versions generate 2 sets of codebooks, 1 for each channel (left/right), and each set of codebooks is decoded independently through the audio compression model. The audio streams for each channel are combined to give the final stereo output.
-
#### Audio Conditional Generation
The model can generate an audio sample conditioned on a text and an audio prompt through use of the [`MusicgenMelodyProcessor`] to pre-process the inputs.
In the following examples, we load an audio file using the 🤗 Datasets library, which can be pip installed through the command below:
-```
+```bash
pip install --upgrade pip
pip install datasets[audio]
```
The audio file we are about to use is loaded as follows:
+
```python
>>> from datasets import load_dataset
@@ -147,10 +146,9 @@ Or save them as a `.wav` file using a third-party library, e.g. `soundfile`:
>>> sf.write("musicgen_out.wav", audio_values[0].T.numpy(), sampling_rate)
```
-
### Text-only Conditional Generation
-The same [`MusicgenMelodyProcessor`] can be used to pre-process a text-only prompt.
+The same [`MusicgenMelodyProcessor`] can be used to pre-process a text-only prompt.
```python
>>> from transformers import AutoProcessor, MusicgenMelodyForConditionalGeneration
@@ -168,7 +166,6 @@ The same [`MusicgenMelodyProcessor`] can be used to pre-process a text-only prom
The `guidance_scale` is used in classifier free guidance (CFG), setting the weighting between the conditional logits (which are predicted from the text prompts) and the unconditional logits (which are predicted from an unconditional or 'null' prompt). Higher guidance scale encourages the model to generate samples that are more closely linked to the input prompt, usually at the expense of poorer audio quality. CFG is enabled by setting `guidance_scale > 1`. For best results, use `guidance_scale=3` (default).
-
You can also generate in batch:
```python
@@ -231,6 +228,7 @@ Note that any arguments passed to the generate method will **supersede** those i
## Model Structure
The MusicGen model can be de-composed into three distinct stages:
+
1. Text encoder: maps the text inputs to a sequence of hidden-state representations. The pre-trained MusicGen models use a frozen text encoder from either T5 or Flan-T5.
2. MusicGen Melody decoder: a language model (LM) that auto-regressively generates audio tokens (or codes) conditional on the encoder hidden-state representations
3. Audio decoder: used to recover the audio waveform from the audio tokens predicted by the decoder.
@@ -260,10 +258,10 @@ python src/transformers/models/musicgen_melody/convert_musicgen_melody_transform
```
Tips:
+
* MusicGen is trained on the 32kHz checkpoint of Encodec. You should ensure you use a compatible version of the Encodec model.
* Sampling mode tends to deliver better results than greedy - you can toggle sampling with the variable `do_sample` in the call to [`MusicgenMelodyForConditionalGeneration.generate`]
-
## MusicgenMelodyDecoderConfig
[[autodoc]] MusicgenMelodyDecoderConfig
@@ -294,4 +292,4 @@ Tips:
## MusicgenMelodyForConditionalGeneration
[[autodoc]] MusicgenMelodyForConditionalGeneration
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/mvp.md b/docs/source/en/model_doc/mvp.md
index 2cce9bd6cac1..26aa2f29b76d 100644
--- a/docs/source/en/model_doc/mvp.md
+++ b/docs/source/en/model_doc/mvp.md
@@ -25,7 +25,6 @@ rendered properly in your Markdown viewer.
The MVP model was proposed in [MVP: Multi-task Supervised Pre-training for Natural Language Generation](https://huggingface.co/papers/2206.12131) by Tianyi Tang, Junyi Li, Wayne Xin Zhao and Ji-Rong Wen.
-
According to the abstract,
- MVP follows a standard Transformer encoder-decoder architecture.
@@ -67,6 +66,7 @@ For summarization, it is an example to use MVP and MVP with summarization-specif
```
For data-to-text generation, it is an example to use MVP and multi-task pre-trained variants.
+
```python
>>> from transformers import MvpTokenizerFast, MvpForConditionalGeneration
diff --git a/docs/source/en/model_doc/myt5.md b/docs/source/en/model_doc/myt5.md
index 409735751252..35ab716a8e71 100644
--- a/docs/source/en/model_doc/myt5.md
+++ b/docs/source/en/model_doc/myt5.md
@@ -44,4 +44,3 @@ The original code can be found [here](https://github.com/tomlimi/MYTE).
## MyT5Tokenizer
[[autodoc]] MyT5Tokenizer
-
diff --git a/docs/source/en/model_doc/nat.md b/docs/source/en/model_doc/nat.md
index dadcae6f17f0..36662173f2f4 100644
--- a/docs/source/en/model_doc/nat.md
+++ b/docs/source/en/model_doc/nat.md
@@ -68,6 +68,7 @@ The `reshaped_hidden_states` have a shape of `(batch, num_channels, height, widt
`(batch_size, height, width, num_channels)`.
Notes:
+
- NAT depends on [NATTEN](https://github.com/SHI-Labs/NATTEN/)'s implementation of Neighborhood Attention.
You can install it with pre-built wheels for Linux by referring to [shi-labs.com/natten](https://shi-labs.com/natten),
or build on your system by running `pip install natten`.
diff --git a/docs/source/en/model_doc/nemotron.md b/docs/source/en/model_doc/nemotron.md
index 360a6ba22267..50f6f99eae2f 100644
--- a/docs/source/en/model_doc/nemotron.md
+++ b/docs/source/en/model_doc/nemotron.md
@@ -97,7 +97,6 @@ Minitron is released under the [NVIDIA Open Model License Agreement](https://dev
| :------------- | :------------- | :------------- | :------------- | :------------- |
| 75.0 | 74.0 | 24.1 | 50.9 | 29.5
-
*Code generation performance*. Evaluated using [HumanEval](https://github.com/openai/human-eval):
| p@1, 0-Shot |
@@ -109,7 +108,8 @@ Please refer to our [paper](https://huggingface.co/papers/2407.14679) for the fu
### Citation
If you find our work helpful, please consider citing our paper:
-```
+
+```bibtex
@article{minitron2024,
title={Compact Language Models via Pruning and Knowledge Distillation},
author={Saurav Muralidharan and Sharath Turuvekere Sreenivas and Raviraj Joshi and Marcin Chochowski and Mostofa Patwary and Mohammad Shoeybi and Bryan Catanzaro and Jan Kautz and Pavlo Molchanov},
@@ -123,13 +123,11 @@ If you find our work helpful, please consider citing our paper:
[[autodoc]] NemotronConfig
-
## NemotronModel
[[autodoc]] NemotronModel
- forward
-
## NemotronForCausalLM
[[autodoc]] NemotronForCausalLM
@@ -140,13 +138,11 @@ If you find our work helpful, please consider citing our paper:
[[autodoc]] NemotronForSequenceClassification
- forward
-
## NemotronForQuestionAnswering
[[autodoc]] NemotronForQuestionAnswering
- forward
-
## NemotronForTokenClassification
[[autodoc]] NemotronForTokenClassification
diff --git a/docs/source/en/model_doc/nllb-moe.md b/docs/source/en/model_doc/nllb-moe.md
index f1456ee402dd..d8c44a5fc0f8 100644
--- a/docs/source/en/model_doc/nllb-moe.md
+++ b/docs/source/en/model_doc/nllb-moe.md
@@ -110,7 +110,6 @@ See example below for a translation from romanian to german:
- [Translation task guide](../tasks/translation)
- [Summarization task guide](../tasks/summarization)
-
## NllbMoeConfig
[[autodoc]] NllbMoeConfig
@@ -135,4 +134,3 @@ See example below for a translation from romanian to german:
[[autodoc]] NllbMoeForConditionalGeneration
- forward
-
diff --git a/docs/source/en/model_doc/nllb.md b/docs/source/en/model_doc/nllb.md
index 6f12a3aa746b..f44c03dcfdd3 100644
--- a/docs/source/en/model_doc/nllb.md
+++ b/docs/source/en/model_doc/nllb.md
@@ -29,7 +29,6 @@ rendered properly in your Markdown viewer.
[NLLB: No Language Left Behind](https://huggingface.co/papers/2207.04672) is a multilingual translation model. It's trained on data using data mining techniques tailored for low-resource languages and supports over 200 languages. NLLB features a conditional compute architecture using a Sparsely Gated Mixture of Experts.
-
You can find all the original NLLB checkpoints under the [AI at Meta](https://huggingface.co/facebook/models?search=nllb) organization.
> [!TIP]
@@ -129,9 +128,10 @@ visualizer("UN Chief says there is no military solution in Syria")
>>> tokenizer = NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M", legacy_behaviour=True)
```
- - For non-English languages, specify the language's [BCP-47](https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200) code with the `src_lang` keyword as shown below.
+- For non-English languages, specify the language's [BCP-47](https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200) code with the `src_lang` keyword as shown below.
+
+- See example below for a translation from Romanian to German.
- - See example below for a translation from Romanian to German.
```python
>>> from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
diff --git a/docs/source/en/model_doc/olmo2.md b/docs/source/en/model_doc/olmo2.md
index 158909c085c3..ba2c93d3ab26 100644
--- a/docs/source/en/model_doc/olmo2.md
+++ b/docs/source/en/model_doc/olmo2.md
@@ -87,6 +87,7 @@ echo -e "Plants create energy through a process known as" | transformers-cli run
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [torchao](../quantization/torchao) to only quantize the weights to 4-bits.
+
```py
#pip install torchao
@@ -116,7 +117,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
```
-
## Notes
- OLMo2 uses RMSNorm instead of standard layer norm. The RMSNorm is applied to attention queries and keys, and it is applied after the attention and feedforward layers rather than before.
@@ -129,7 +129,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
model = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B", revision="stage1-step140000-tokens294B")
```
-
## Olmo2Config
[[autodoc]] Olmo2Config
diff --git a/docs/source/en/model_doc/olmo3.md b/docs/source/en/model_doc/olmo3.md
index e320181925ca..07a3cc3ebed9 100644
--- a/docs/source/en/model_doc/olmo3.md
+++ b/docs/source/en/model_doc/olmo3.md
@@ -12,11 +12,11 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-
⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer.
-->
-*This model was released on {release_date} and added to Hugging Face Transformers on 2025-09-08.*
+*This model was released on {release_date} and added to Hugging Face Transformers on 2025-09-16.*
+
@@ -46,7 +46,7 @@ pipe = pipeline(
dtype=torch.bfloat16,
device=0,
)
-
+
result = pipe("Plants create energy through a process known as")
print(result)
```
@@ -87,6 +87,7 @@ echo -e "Plants create energy through a process known as" | transformers-cli run
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [torchao](../quantization/torchao) to only quantize the weights to 4-bits.
+
```py
#pip install torchao
@@ -116,18 +117,16 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
```
-
## Notes
-- Load specific intermediate checkpoints by adding the `revision` parameter to [`~PreTrainedModel.from_pretrained`].
+- Load specific intermediate checkpoints by adding the `revision` parameter to [`~PreTrainedModel.from_pretrained`].
```py
from transformers import AutoModelForCausalLM
-
+
model = AutoModelForCausalLM.from_pretrained("allenai/TBA", revision="stage1-step140000-tokens294B")
```
-
## Olmo3Config
[[autodoc]] Olmo3Config
@@ -144,4 +143,4 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
## Olmo3PreTrainedModel
[[autodoc]] Olmo3PreTrainedModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/oneformer.md b/docs/source/en/model_doc/oneformer.md
index c4b3bd142fe0..7f5d32bc55a8 100644
--- a/docs/source/en/model_doc/oneformer.md
+++ b/docs/source/en/model_doc/oneformer.md
@@ -39,7 +39,7 @@ This model was contributed by [Jitesh Jain](https://huggingface.co/praeclarumjj3
## Usage tips
-- OneFormer requires two inputs during inference: *image* and *task token*.
+- OneFormer requires two inputs during inference: *image* and *task token*.
- During training, OneFormer only uses panoptic annotations.
- If you want to train the model in a distributed environment across multiple nodes, then one should update the
`get_num_masks` function inside in the `OneFormerLoss` class of `modeling_oneformer.py`. When training on multiple nodes, this should be
diff --git a/docs/source/en/model_doc/openai-gpt.md b/docs/source/en/model_doc/openai-gpt.md
index b45b205e2592..04d37d89cc49 100644
--- a/docs/source/en/model_doc/openai-gpt.md
+++ b/docs/source/en/model_doc/openai-gpt.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2018-06-11 and added to Hugging Face Transformers on 2023-06-20.*
-
@@ -24,8 +23,6 @@ rendered properly in your Markdown viewer.
-
-
# GPT
[GPT (Generative Pre-trained Transformer)](https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf) ([blog post](https://openai.com/index/language-unsupervised/)) focuses on effectively learning text representations and transferring them to tasks. This model trains the Transformer decoder to predict the next word, and then fine-tuned on labeled data.
@@ -39,12 +36,9 @@ You can find all the original GPT checkpoints under the [OpenAI community](https
The example below demonstrates how to generate text with [`Pipeline`], [`AutoModel`], and from the command line.
-
-
-
```python
import torch
from transformers import pipeline
@@ -75,6 +69,7 @@ print(tokenizer.decode(outputs[0], skip_special_tokens=True))
echo -e "The future of AI is" | transformers run --task text-generation --model openai-community/openai-gpt --device 0
```
+
@@ -89,22 +84,22 @@ echo -e "The future of AI is" | transformers run --task text-generation --model
## OpenAIGPTModel
[[autodoc]] OpenAIGPTModel
-- forward
+ - forward
## OpenAIGPTLMHeadModel
[[autodoc]] OpenAIGPTLMHeadModel
-- forward
+ - forward
## OpenAIGPTDoubleHeadsModel
[[autodoc]] OpenAIGPTDoubleHeadsModel
-- forward
+ - forward
## OpenAIGPTForSequenceClassification
[[autodoc]] OpenAIGPTForSequenceClassification
-- forward
+ - forward
## OpenAIGPTTokenizer
diff --git a/docs/source/en/model_doc/opt.md b/docs/source/en/model_doc/opt.md
index e645956f1ece..7c65689594e4 100644
--- a/docs/source/en/model_doc/opt.md
+++ b/docs/source/en/model_doc/opt.md
@@ -36,7 +36,6 @@ You can find all the original OPT checkpoints under the [OPT](https://huggingfac
The example below demonstrates how to generate text with [`Pipeline`], [`AutoModel`], and from the command line.
-
@@ -65,12 +64,14 @@ model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
generated_ids = model.generate(**model_inputs, max_new_tokens=30, do_sample=False)
tokenizer.batch_decode(generated_ids)[0]
```
+
```py
echo -e "Plants create energy through a process known as" | transformers run --task text-generation --model facebook/opt-125m --device 0
```
+
diff --git a/docs/source/en/model_doc/ovis2.md b/docs/source/en/model_doc/ovis2.md
index ab1d761f19ed..731ebbb83f08 100644
--- a/docs/source/en/model_doc/ovis2.md
+++ b/docs/source/en/model_doc/ovis2.md
@@ -13,12 +13,13 @@ specific language governing permissions and limitations under the License.
rendered properly in your Markdown viewer.
-->
+*This model was released on 2024-05-31 and added to Hugging Face Transformers on 2025-08-18.*
# Ovis2
## Overview
-The [Ovis2](https://github.com/AIDC-AI/Ovis) is an updated version of the [Ovis](https://huggingface.co/papers/2405.20797) model developed by the AIDC-AI team at Alibaba International Digital Commerce Group.
+The [Ovis2](https://github.com/AIDC-AI/Ovis) is an updated version of the [Ovis](https://huggingface.co/papers/2405.20797) model developed by the AIDC-AI team at Alibaba International Digital Commerce Group.
Ovis2 is the latest advancement in multi-modal large language models (MLLMs), succeeding Ovis1.6. It retains the architectural design of the Ovis series, which focuses on aligning visual and textual embeddings, and introduces major improvements in data curation and training methods.
diff --git a/docs/source/en/model_doc/paligemma.md b/docs/source/en/model_doc/paligemma.md
index 58aa622a0d37..fa7c193da453 100644
--- a/docs/source/en/model_doc/paligemma.md
+++ b/docs/source/en/model_doc/paligemma.md
@@ -140,6 +140,7 @@ visualizer("![]()
What is in this image?")
answer = "a pallas cat"
inputs = processor(images=image, text=prompt, suffix=answer, return_tensors="pt")
```
+
- PaliGemma can support multiple input images if it is fine-tuned to accept multiple images. For example, the [NLVR2](https://huggingface.co/google/paligemma-3b-ft-nlvr2-448) checkpoint supports multiple images. Pass the images as a list to the processor.
```py
diff --git a/docs/source/en/model_doc/parakeet.md b/docs/source/en/model_doc/parakeet.md
new file mode 100644
index 000000000000..4cb72e7e4585
--- /dev/null
+++ b/docs/source/en/model_doc/parakeet.md
@@ -0,0 +1,221 @@
+
+*This model was released on {release_date} and added to Hugging Face Transformers on 2025-09-25.*
+
+
+
+
+```py
+from transformers import pipeline
+
+pipe = pipeline("automatic-speech-recognition", model="nvidia/parakeet-ctc-1.1b")
+out = pipe("https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/bcn_weather.mp3")
+print(out)
+```
+
+
+
+
+```py
+from transformers import AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, Audio
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+processor = AutoProcessor.from_pretrained("nvidia/parakeet-ctc-1.1b")
+model = AutoModelForCTC.from_pretrained("nvidia/parakeet-ctc-1.1b", dtype="auto", device_map=device)
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el['array'] for el in ds["audio"][:5]]
+
+inputs = processor(speech_samples, sampling_rate=processor.feature_extractor.sampling_rate)
+inputs.to(model.device, dtype=model.dtype)
+outputs = model.generate(**inputs)
+print(processor.batch_decode(outputs))
+```
+
+
+
+
+### Making The Model Go Brrr
+
+Parakeet supports full-graph compilation with CUDA graphs! This optimization is most effective when you know the maximum audio length you want to transcribe. The key idea is using static input shapes to avoid recompilation. For example, if you know your audio will be under 30 seconds, you can use the processor to pad all inputs to 30 seconds, preparing consistent input features and attention masks. See the example below!
+
+```python
+from transformers import AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, Audio
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+processor = AutoProcessor.from_pretrained("nvidia/parakeet-ctc-1.1b")
+model = AutoModelForCTC.from_pretrained("nvidia/parakeet-ctc-1.1b", dtype="auto", device_map=device)
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el['array'] for el in ds["audio"][:5]]
+
+# Compile the generate method with fullgraph and CUDA graphs
+model.generate = torch.compile(model.generate, fullgraph=True, mode="reduce-overhead")
+
+# let's define processor kwargs to pad to 30 seconds
+processor_kwargs = {
+ "padding": "max_length",
+ "max_length": 30 * processor.feature_extractor.sampling_rate,
+}
+
+# Define a timing context using CUDA events
+class TimerContext:
+ def __init__(self, name="Execution"):
+ self.name = name
+ self.start_event = None
+ self.end_event = None
+
+ def __enter__(self):
+ # Use CUDA events for more accurate GPU timing
+ self.start_event = torch.cuda.Event(enable_timing=True)
+ self.end_event = torch.cuda.Event(enable_timing=True)
+ self.start_event.record()
+ return self
+
+ def __exit__(self, *args):
+ self.end_event.record()
+ torch.cuda.synchronize()
+ elapsed_time = self.start_event.elapsed_time(self.end_event) / 1000.0
+ print(f"{self.name} time: {elapsed_time:.4f} seconds")
+
+
+inputs = processor(speech_samples[0], **processor_kwargs)
+inputs.to(device, dtype=model.dtype)
+print("\n" + "="*50)
+print("First generation - compiling...")
+# Generate with the compiled model
+with TimerContext("First generation"):
+ outputs = model.generate(**inputs)
+print(processor.batch_decode(outputs))
+
+inputs = processor(speech_samples[1], **processor_kwargs)
+inputs.to(device, dtype=model.dtype)
+print("\n" + "="*50)
+print("Second generation - recording CUDA graphs...")
+with TimerContext("Second generation"):
+ outputs = model.generate(**inputs)
+print(processor.batch_decode(outputs))
+
+inputs = processor(speech_samples[2], **processor_kwargs)
+inputs.to(device, dtype=model.dtype)
+print("\n" + "="*50)
+print("Third generation - fast !!!")
+with TimerContext("Third generation"):
+ outputs = model.generate(**inputs)
+print(processor.batch_decode(outputs))
+
+inputs = processor(speech_samples[3], **processor_kwargs)
+inputs.to(device, dtype=model.dtype)
+print("\n" + "="*50)
+print("Fourth generation - still fast !!!")
+with TimerContext("Fourth generation"):
+ outputs = model.generate(**inputs)
+print(processor.batch_decode(outputs))
+```
+
+### Training
+
+```python
+from transformers import AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, Audio
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+processor = AutoProcessor.from_pretrained("nvidia/parakeet-ctc-1.1b")
+model = AutoModelForCTC.from_pretrained("nvidia/parakeet-ctc-1.1b", dtype="auto", device_map=device)
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el['array'] for el in ds["audio"][:5]]
+text_samples = [el for el in ds["text"][:5]]
+
+# passing `text` to the processor will prepare inputs' `labels` key
+inputs = processor(audio=speech_samples, text=text_samples, sampling_rate=processor.feature_extractor.sampling_rate)
+inputs.to(device, dtype=model.dtype)
+
+outputs = model(**inputs)
+outputs.loss.backward()
+```
+
+## ParakeetTokenizerFast
+
+[[autodoc]] ParakeetTokenizerFast
+
+## ParakeetFeatureExtractor
+
+[[autodoc]] ParakeetFeatureExtractor
+ - __call__
+
+## ParakeetProcessor
+
+[[autodoc]] ParakeetProcessor
+ - __call__
+ - batch_decode
+ - decode
+
+## ParakeetEncoderConfig
+
+[[autodoc]] ParakeetEncoderConfig
+
+## ParakeetCTCConfig
+
+[[autodoc]] ParakeetCTCConfig
+
+## ParakeetEncoder
+
+[[autodoc]] ParakeetEncoder
+
+## ParakeetForCTC
+
+[[autodoc]] ParakeetForCTC
diff --git a/docs/source/en/model_doc/patchtsmixer.md b/docs/source/en/model_doc/patchtsmixer.md
index 5541f4d80936..4a9ddef46416 100644
--- a/docs/source/en/model_doc/patchtsmixer.md
+++ b/docs/source/en/model_doc/patchtsmixer.md
@@ -25,15 +25,13 @@ rendered properly in your Markdown viewer.
The PatchTSMixer model was proposed in [TSMixer: Lightweight MLP-Mixer Model for Multivariate Time Series Forecasting](https://huggingface.co/papers/2306.09364) by Vijay Ekambaram, Arindam Jati, Nam Nguyen, Phanwadee Sinthong and Jayant Kalagnanam.
-
PatchTSMixer is a lightweight time-series modeling approach based on the MLP-Mixer architecture. In this HuggingFace implementation, we provide PatchTSMixer's capabilities to effortlessly facilitate lightweight mixing across patches, channels, and hidden features for effective multivariate time-series modeling. It also supports various attention mechanisms starting from simple gated attention to more complex self-attention blocks that can be customized accordingly. The model can be pretrained and subsequently used for various downstream tasks such as forecasting, classification and regression.
-
The abstract from the paper is the following:
*TSMixer is a lightweight neural architecture exclusively composed of multi-layer perceptron (MLP) modules designed for multivariate forecasting and representation learning on patched time series. Our model draws inspiration from the success of MLP-Mixer models in computer vision. We demonstrate the challenges involved in adapting Vision MLP-Mixer for time series and introduce empirically validated components to enhance accuracy. This includes a novel design paradigm of attaching online reconciliation heads to the MLP-Mixer backbone, for explicitly modeling the time-series properties such as hierarchy and channel-correlations. We also propose a Hybrid channel modeling approach to effectively handle noisy channel interactions and generalization across diverse datasets, a common challenge in existing patch channel-mixing methods. Additionally, a simple gated attention mechanism is introduced in the backbone to prioritize important features. By incorporating these lightweight components, we significantly enhance the learning capability of simple MLP structures, outperforming complex Transformer models with minimal computing usage. Moreover, TSMixer's modular design enables compatibility with both supervised and masked self-supervised learning methods, making it a promising building block for time-series Foundation Models. TSMixer outperforms state-of-the-art MLP and Transformer models in forecasting by a considerable margin of 8-60%. It also outperforms the latest strong benchmarks of Patch-Transformer models (by 1-2%) with a significant reduction in memory and runtime (2-3X).*
-This model was contributed by [ajati](https://huggingface.co/ajati), [vijaye12](https://huggingface.co/vijaye12),
+This model was contributed by [ajati](https://huggingface.co/ajati), [vijaye12](https://huggingface.co/vijaye12),
[gsinthong](https://huggingface.co/gsinthong), [namctin](https://huggingface.co/namctin),
[wmgifford](https://huggingface.co/wmgifford), [kashif](https://huggingface.co/kashif).
@@ -68,32 +66,27 @@ The model can also be used for time series classification and time series regres
[[autodoc]] PatchTSMixerConfig
-
## PatchTSMixerModel
[[autodoc]] PatchTSMixerModel
- forward
-
## PatchTSMixerForPrediction
[[autodoc]] PatchTSMixerForPrediction
- forward
-
## PatchTSMixerForTimeSeriesClassification
[[autodoc]] PatchTSMixerForTimeSeriesClassification
- forward
-
## PatchTSMixerForPretraining
[[autodoc]] PatchTSMixerForPretraining
- forward
-
## PatchTSMixerForRegression
[[autodoc]] PatchTSMixerForRegression
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/pegasus_x.md b/docs/source/en/model_doc/pegasus_x.md
index 791618c67d30..4f42b787b925 100644
--- a/docs/source/en/model_doc/pegasus_x.md
+++ b/docs/source/en/model_doc/pegasus_x.md
@@ -53,6 +53,7 @@ Through photosynthesis, plants capture energy from sunlight using a green pigmen
These ingredients are then transformed into glucose, a type of sugar that serves as a source of chemical energy, and oxygen, which is released as a byproduct into the atmosphere. The glucose produced during photosynthesis is not just used immediately; plants also store it as starch or convert it into other organic compounds like cellulose, which is essential for building their cellular structure.
This energy reserve allows them to grow, develop leaves, produce flowers, bear fruit, and carry out various physiological processes throughout their lifecycle.""")
```
+
@@ -78,12 +79,14 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = model.generate(**input_ids, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
```bash
echo -e "Plants are among the most remarkable and essential life forms on Earth, possessing a unique ability to produce their own food through a process known as photosynthesis. This complex biochemical process is fundamental not only to plant life but to virtually all life on the planet. Through photosynthesis, plants capture energy from sunlight using a green pigment called chlorophyll, which is located in specialized cell structures called chloroplasts." | transformers-cli run --task summarization --model google/pegasus-x-large --device 0
```
+
diff --git a/docs/source/en/model_doc/perception_lm.md b/docs/source/en/model_doc/perception_lm.md
index ee6b63fce6fd..7d3d608253fc 100644
--- a/docs/source/en/model_doc/perception_lm.md
+++ b/docs/source/en/model_doc/perception_lm.md
@@ -38,11 +38,9 @@ video captions. Additionally, we introduce PLM–VideoBench, a suite for evaluat
understanding tasks focusing on the ability to reason about “what”, “where”, “when”, and “how” of a
video. We make our work fully reproducible by providing data, training recipes, code & models.*
-
This model was contributed by [shumingh](https://huggingface.co/shumingh).
The original code can be found [here](https://github.com/facebookresearch/perception_models).
-
## PerceptionLMConfig
[[autodoc]] PerceptionLMConfig
diff --git a/docs/source/en/model_doc/persimmon.md b/docs/source/en/model_doc/persimmon.md
index 764c959879ad..854eaee835df 100644
--- a/docs/source/en/model_doc/persimmon.md
+++ b/docs/source/en/model_doc/persimmon.md
@@ -39,7 +39,7 @@ The original code can be found [here](https://github.com/persimmon-ai-labs/adept
The `Persimmon` models were trained using `bfloat16`, but the original inference uses `float16` The checkpoints uploaded on the hub use `dtype = 'float16'` which will be
-used by the `AutoModel` API to cast the checkpoints from `torch.float32` to `torch.float16`.
+used by the `AutoModel` API to cast the checkpoints from `torch.float32` to `torch.float16`.
The `dtype` of the online weights is mostly irrelevant, unless you are using `dtype="auto"` when initializing a model using `model = AutoModelForCausalLM.from_pretrained("path", dtype = "auto")`. The reason is that the model will first be downloaded ( using the `dtype` of the checkpoints online) then it will be cast to the default `dtype` of `torch` (becomes `torch.float32`). Users should specify the `dtype` they want, and if they don't it will be `torch.float32`.
@@ -47,7 +47,6 @@ Finetuning the model in `float16` is not recommended and known to produce `nan`,
-
Tips:
- To convert the model, you need to clone the original repository using `git clone https://github.com/persimmon-ai-labs/adept-inference`, then get the checkpoints:
@@ -62,6 +61,7 @@ python src/transformers/models/persimmon/convert_persimmon_weights_to_hf.py --i
```
For the chat model:
+
```bash
wget https://axtkn4xl5cip.objectstorage.us-phoenix-1.oci.customer-oci.com/n/axtkn4xl5cip/b/adept-public-data/o/8b_chat_model_release.tar
tar -xvf 8b_base_model_release.tar
@@ -76,13 +76,11 @@ model = PersimmonForCausalLM.from_pretrained("/output/path")
tokenizer = PersimmonTokenizer.from_pretrained("/output/path")
```
-
- Perismmon uses a `sentencepiece` based tokenizer, with a `Unigram` model. It supports bytefallback, which is only available in `tokenizers==0.14.0` for the fast tokenizer.
The `LlamaTokenizer` is used as it is a standard wrapper around sentencepiece. The `chat` template will be updated with the templating functions in a follow up PR!
- The authors suggest to use the following prompt format for the chat mode: `f"human: {prompt}\n\nadept:"`
-
## PersimmonConfig
[[autodoc]] PersimmonConfig
diff --git a/docs/source/en/model_doc/phimoe.md b/docs/source/en/model_doc/phimoe.md
index 319cbc470b91..64a12e3820ae 100644
--- a/docs/source/en/model_doc/phimoe.md
+++ b/docs/source/en/model_doc/phimoe.md
@@ -45,12 +45,14 @@ The original code for PhiMoE can be found [here](https://huggingface.co/microsof
Phi-3.5-MoE-instruct has been integrated in the development version (4.44.2.dev) of `transformers`. Until the official version is released through `pip`, ensure that you are doing the following:
+
* When loading the model, ensure that `trust_remote_code=True` is passed as an argument of the `from_pretrained()` function.
The current `transformers` version can be verified with: `pip list | grep transformers`.
Examples of required packages:
-```
+
+```bash
flash_attn==2.5.8
torch==2.3.1
accelerate==0.31.0
diff --git a/docs/source/en/model_doc/pix2struct.md b/docs/source/en/model_doc/pix2struct.md
index c43c9b3b92ed..412d2c2fef95 100644
--- a/docs/source/en/model_doc/pix2struct.md
+++ b/docs/source/en/model_doc/pix2struct.md
@@ -79,4 +79,4 @@ The original code can be found [here](https://github.com/google-research/pix2str
## Pix2StructForConditionalGeneration
[[autodoc]] Pix2StructForConditionalGeneration
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/pixtral.md b/docs/source/en/model_doc/pixtral.md
index 55ba09084292..bb175973bd23 100644
--- a/docs/source/en/model_doc/pixtral.md
+++ b/docs/source/en/model_doc/pixtral.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2024-09-17 and added to Hugging Face Transformers on 2024-09-14.*
-
diff --git a/docs/source/en/model_doc/plbart.md b/docs/source/en/model_doc/plbart.md
index d8ce330cb0f7..b3459299437e 100644
--- a/docs/source/en/model_doc/plbart.md
+++ b/docs/source/en/model_doc/plbart.md
@@ -120,4 +120,4 @@ it's passed with the `text_target` keyword argument.
## PLBartForCausalLM
[[autodoc]] PLBartForCausalLM
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/pop2piano.md b/docs/source/en/model_doc/pop2piano.md
index 5f68b1805000..c934d8789037 100644
--- a/docs/source/en/model_doc/pop2piano.md
+++ b/docs/source/en/model_doc/pop2piano.md
@@ -21,14 +21,14 @@ specific language governing permissions and limitations under the License.
The Pop2Piano model was proposed in [Pop2Piano : Pop Audio-based Piano Cover Generation](https://huggingface.co/papers/2211.00895) by Jongho Choi and Kyogu Lee.
-Piano covers of pop music are widely enjoyed, but generating them from music is not a trivial task. It requires great
-expertise with playing piano as well as knowing different characteristics and melodies of a song. With Pop2Piano you
-can directly generate a cover from a song's audio waveform. It is the first model to directly generate a piano cover
-from pop audio without melody and chord extraction modules.
-
-Pop2Piano is an encoder-decoder Transformer model based on [T5](https://huggingface.co/papers/1910.10683). The input audio
-is transformed to its waveform and passed to the encoder, which transforms it to a latent representation. The decoder
-uses these latent representations to generate token ids in an autoregressive way. Each token id corresponds to one of four
+Piano covers of pop music are widely enjoyed, but generating them from music is not a trivial task. It requires great
+expertise with playing piano as well as knowing different characteristics and melodies of a song. With Pop2Piano you
+can directly generate a cover from a song's audio waveform. It is the first model to directly generate a piano cover
+from pop audio without melody and chord extraction modules.
+
+Pop2Piano is an encoder-decoder Transformer model based on [T5](https://huggingface.co/papers/1910.10683). The input audio
+is transformed to its waveform and passed to the encoder, which transforms it to a latent representation. The decoder
+uses these latent representations to generate token ids in an autoregressive way. Each token id corresponds to one of four
different token types: time, velocity, note and 'special'. The token ids are then decoded to their equivalent MIDI file.
The abstract from the paper is the following:
@@ -53,10 +53,13 @@ The original code can be found [here](https://github.com/sweetcocoa/pop2piano).
## Usage tips
* To use Pop2Piano, you will need to install the 🤗 Transformers library, as well as the following third party modules:
+
```bash
pip install pretty-midi==0.2.9 essentia==2.1b6.dev1034 librosa scipy
```
+
Please note that you may need to restart your runtime after installation.
+
* Pop2Piano is an Encoder-Decoder based model like T5.
* Pop2Piano can be used to generate midi-audio files for a given audio sequence.
* Choosing different composers in `Pop2PianoForConditionalGeneration.generate()` can lead to variety of different results.
@@ -131,7 +134,6 @@ Please note that you may need to restart your runtime after installation.
>>> tokenizer_output[1].write("./Outputs/midi_output2.mid")
```
-
- Example of processing multiple audio files in batch (Using `Pop2PianoFeatureExtractor` and `Pop2PianoTokenizer`):
```python
@@ -166,7 +168,6 @@ Please note that you may need to restart your runtime after installation.
>>> tokenizer_output[1].write("./Outputs/midi_output2.mid")
```
-
## Pop2PianoConfig
[[autodoc]] Pop2PianoConfig
diff --git a/docs/source/en/model_doc/prompt_depth_anything.md b/docs/source/en/model_doc/prompt_depth_anything.md
index 5af13c5d630e..d4b6f4cc2598 100644
--- a/docs/source/en/model_doc/prompt_depth_anything.md
+++ b/docs/source/en/model_doc/prompt_depth_anything.md
@@ -19,8 +19,7 @@ rendered properly in your Markdown viewer.
## Overview
-The Prompt Depth Anything model was introduced in [Prompting Depth Anything for 4K Resolution Accurate Metric Depth Estimation](https://huggingface.co/papers/2412.14015) by Haotong Lin, Sida Peng, Jingxiao Chen, Songyou Peng, Jiaming Sun, Minghuan Liu, Hujun Bao, Jiashi Feng, Xiaowei Zhou, Bingyi Kang.
-
+The Prompt Depth Anything model was introduced in [Prompting Depth Anything for 4K Resolution Accurate Metric Depth Estimation](https://huggingface.co/papers/2412.14015) by Haotong Lin, Sida Peng, Jingxiao Chen, Songyou Peng, Jiaming Sun, Minghuan Liu, Hujun Bao, Jiashi Feng, Xiaowei Zhou, Bingyi Kang.
The abstract from the paper is as follows:
@@ -100,4 +99,4 @@ If you are interested in submitting a resource to be included here, please feel
[[autodoc]] PromptDepthAnythingImageProcessorFast
- preprocess
- - post_process_depth_estimation
\ No newline at end of file
+ - post_process_depth_estimation
diff --git a/docs/source/en/model_doc/pvt.md b/docs/source/en/model_doc/pvt.md
index e7902affe5f4..38858db55529 100644
--- a/docs/source/en/model_doc/pvt.md
+++ b/docs/source/en/model_doc/pvt.md
@@ -29,23 +29,22 @@ is used to further reduce the resource consumption when learning high-resolution
The abstract from the paper is the following:
-*Although convolutional neural networks (CNNs) have achieved great success in computer vision, this work investigates a
-simpler, convolution-free backbone network useful for many dense prediction tasks. Unlike the recently proposed Vision
-Transformer (ViT) that was designed for image classification specifically, we introduce the Pyramid Vision Transformer
-(PVT), which overcomes the difficulties of porting Transformer to various dense prediction tasks. PVT has several
-merits compared to current state of the arts. Different from ViT that typically yields low resolution outputs and
-incurs high computational and memory costs, PVT not only can be trained on dense partitions of an image to achieve high
-output resolution, which is important for dense prediction, but also uses a progressive shrinking pyramid to reduce the
-computations of large feature maps. PVT inherits the advantages of both CNN and Transformer, making it a unified
-backbone for various vision tasks without convolutions, where it can be used as a direct replacement for CNN backbones.
+*Although convolutional neural networks (CNNs) have achieved great success in computer vision, this work investigates a
+simpler, convolution-free backbone network useful for many dense prediction tasks. Unlike the recently proposed Vision
+Transformer (ViT) that was designed for image classification specifically, we introduce the Pyramid Vision Transformer
+(PVT), which overcomes the difficulties of porting Transformer to various dense prediction tasks. PVT has several
+merits compared to current state of the arts. Different from ViT that typically yields low resolution outputs and
+incurs high computational and memory costs, PVT not only can be trained on dense partitions of an image to achieve high
+output resolution, which is important for dense prediction, but also uses a progressive shrinking pyramid to reduce the
+computations of large feature maps. PVT inherits the advantages of both CNN and Transformer, making it a unified
+backbone for various vision tasks without convolutions, where it can be used as a direct replacement for CNN backbones.
We validate PVT through extensive experiments, showing that it boosts the performance of many downstream tasks, including
-object detection, instance and semantic segmentation. For example, with a comparable number of parameters, PVT+RetinaNet
-achieves 40.4 AP on the COCO dataset, surpassing ResNet50+RetinNet (36.3 AP) by 4.1 absolute AP (see Figure 2). We hope
+object detection, instance and semantic segmentation. For example, with a comparable number of parameters, PVT+RetinaNet
+achieves 40.4 AP on the COCO dataset, surpassing ResNet50+RetinNet (36.3 AP) by 4.1 absolute AP (see Figure 2). We hope
that PVT could serve as an alternative and useful backbone for pixel-level predictions and facilitate future research.*
This model was contributed by [Xrenya](https://huggingface.co/Xrenya). The original code can be found [here](https://github.com/whai362/PVT).
-
- PVTv1 on ImageNet-1K
| **Model variant** |**Size** |**Acc@1**|**Params (M)**|
@@ -55,7 +54,6 @@ This model was contributed by [Xrenya](https://huggingface.co/Xrenya). The origi
| PVT-Medium | 224 | 81.2 | 44.2 |
| PVT-Large | 224 | 81.7 | 61.4 |
-
## PvtConfig
[[autodoc]] PvtConfig
diff --git a/docs/source/en/model_doc/pvt_v2.md b/docs/source/en/model_doc/pvt_v2.md
index 0d0ee3cca751..5be8998f4cc2 100644
--- a/docs/source/en/model_doc/pvt_v2.md
+++ b/docs/source/en/model_doc/pvt_v2.md
@@ -26,7 +26,7 @@ The PVTv2 encoder structure has been successfully deployed to achieve state-of-t
PVTv2 belongs to a family of models called [hierarchical transformers](https://natecibik.medium.com/the-rise-of-vision-transformers-f623c980419f) , which make adaptations to transformer layers in order to generate multi-scale feature maps. Unlike the columnal structure of Vision Transformer ([ViT](https://huggingface.co/papers/2010.11929)) which loses fine-grained detail, multi-scale feature maps are known preserve this detail and aid performance in dense prediction tasks. In the case of PVTv2, this is achieved by generating image patch tokens using 2D convolution with overlapping kernels in each encoder layer.
-The multi-scale features of hierarchical transformers allow them to be easily swapped in for traditional workhorse computer vision backbone models like ResNet in larger architectures. Both Segformer and Panoptic Segformer demonstrated that configurations using PVTv2 for a backbone consistently outperformed those with similarly sized ResNet backbones.
+The multi-scale features of hierarchical transformers allow them to be easily swapped in for traditional workhorse computer vision backbone models like ResNet in larger architectures. Both Segformer and Panoptic Segformer demonstrated that configurations using PVTv2 for a backbone consistently outperformed those with similarly sized ResNet backbones.
Another powerful feature of the PVTv2 is the complexity reduction in the self-attention layers called Spatial Reduction Attention (SRA), which uses 2D convolution layers to project hidden states to a smaller resolution before attending to them with the queries, improving the $O(n^2)$ complexity of self-attention to $O(n^2/R)$, with $R$ being the spatial reduction ratio (`sr_ratio`, aka kernel size and stride in the 2D convolution).
@@ -48,6 +48,7 @@ This model was contributed by [FoamoftheSea](https://huggingface.co/FoamoftheSea
- ImageNet pretrained weights for all model sizes can be found on the [hub](https://huggingface.co/models?other=pvt_v2).
The best way to get started with the PVTv2 is to load the pretrained checkpoint with the size of your choosing using `AutoModelForImageClassification`:
+
```python
import requests
import torch
@@ -99,7 +100,6 @@ outputs = model(torch.tensor(processed["pixel_values"]))
| PVT-V2-B4 | 224 | 83.6 | 62.6 |
| PVT-V2-B5 | 224 | 83.8 | 82.0 |
-
## PvtV2Config
[[autodoc]] PvtV2Config
diff --git a/docs/source/en/model_doc/qdqbert.md b/docs/source/en/model_doc/qdqbert.md
index 4c934d92d5fc..b791b4b2afe6 100644
--- a/docs/source/en/model_doc/qdqbert.md
+++ b/docs/source/en/model_doc/qdqbert.md
@@ -115,7 +115,7 @@ tensors. After setting up the tensor quantizers, one can use the following examp
The goal of exporting to ONNX is to deploy inference by [TensorRT](https://developer.nvidia.com/tensorrt). Fake
quantization will be broken into a pair of QuantizeLinear/DequantizeLinear ONNX ops. After setting static member of
-TensorQuantizer to use Pytorch’s own fake quantization functions, fake quantized model can be exported to ONNX, follow
+TensorQuantizer to use Pytorch's own fake quantization functions, fake quantized model can be exported to ONNX, follow
the instructions in [torch.onnx](https://pytorch.org/docs/stable/onnx.html). Example:
```python
diff --git a/docs/source/en/model_doc/qwen2.md b/docs/source/en/model_doc/qwen2.md
index 3f872302cc27..feeb69959b21 100644
--- a/docs/source/en/model_doc/qwen2.md
+++ b/docs/source/en/model_doc/qwen2.md
@@ -142,7 +142,6 @@ outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```
-
## Notes
- Ensure your Transformers library version is up-to-date. Qwen2 requires Transformers>=4.37.0 for full support.
diff --git a/docs/source/en/model_doc/qwen2_5_omni.md b/docs/source/en/model_doc/qwen2_5_omni.md
index e124f7cdb421..e2e0dc348a1c 100644
--- a/docs/source/en/model_doc/qwen2_5_omni.md
+++ b/docs/source/en/model_doc/qwen2_5_omni.md
@@ -29,9 +29,7 @@ The [Qwen2.5-Omni](https://qwenlm.github.io/blog/qwen2.5-omni/) model is a unifi
The abstract from the technical report is the following:
-*We present Qwen2.5-Omni, an end-to-end multimodal model designed to perceive diverse modalities, including text, images, audio, and video, while simultaneously generating text and natural speech responses in a streaming manner. To enable the streaming of multimodal information inputs, both audio and visual encoders utilize a block-wise processing approach. This strategy effectively decouples the handling of long sequences of multimodal data, assigning the perceptual responsibilities to the multimodal encoder and entrusting the modeling of extended sequences to a large language model. Such a division of labor enhances the fusion of different modalities via the shared attention mechanism. To synchronize the timestamps of video inputs with audio, we organized the audio and video sequentially in an interleaved manner and propose a novel position embedding approach, named TMRoPE (Time-aligned Multimodal RoPE). To concurrently generate text and speech while avoiding interference between the two modalities, we propose Thinker-Talker architecture. In this framework, Thinker functions as a large language model tasked with text generation, while Talker is a dual-track autoregressive model that directly utilizes the hidden representations from the Thinker to produce audio tokens as output. Both the Thinker and Talker models are designed to be trained and inferred in an end-to-end manner. For decoding audio tokens in a streaming manner, we introduce a sliding-window DiT that restricts the receptive field, aiming to reduce the initial package delay. Qwen2.5-Omni outperforms the similarly sized Qwen2-VL and Qwen2-Audio in both image and audio capabilities. Furthermore, Qwen2.5-Omni achieves state-of-the-art performance on multimodal benchmarks like Omni-Bench. Notably, Qwen2.5-Omni is the first open-source model to achieve a level of performance in end-to-end speech instruction following that is comparable to its capabilities with text inputs, as evidenced by benchmarks such as MMLU and GSM8K. As for speech generation, Qwen2.5-Omni’s streaming Talker outperform most existing streaming and non-streaming alternatives in robustness and naturalness.*
-
-
+*We present Qwen2.5-Omni, an end-to-end multimodal model designed to perceive diverse modalities, including text, images, audio, and video, while simultaneously generating text and natural speech responses in a streaming manner. To enable the streaming of multimodal information inputs, both audio and visual encoders utilize a block-wise processing approach. This strategy effectively decouples the handling of long sequences of multimodal data, assigning the perceptual responsibilities to the multimodal encoder and entrusting the modeling of extended sequences to a large language model. Such a division of labor enhances the fusion of different modalities via the shared attention mechanism. To synchronize the timestamps of video inputs with audio, we organized the audio and video sequentially in an interleaved manner and propose a novel position embedding approach, named TMRoPE (Time-aligned Multimodal RoPE). To concurrently generate text and speech while avoiding interference between the two modalities, we propose Thinker-Talker architecture. In this framework, Thinker functions as a large language model tasked with text generation, while Talker is a dual-track autoregressive model that directly utilizes the hidden representations from the Thinker to produce audio tokens as output. Both the Thinker and Talker models are designed to be trained and inferred in an end-to-end manner. For decoding audio tokens in a streaming manner, we introduce a sliding-window DiT that restricts the receptive field, aiming to reduce the initial package delay. Qwen2.5-Omni outperforms the similarly sized Qwen2-VL and Qwen2-Audio in both image and audio capabilities. Furthermore, Qwen2.5-Omni achieves state-of-the-art performance on multimodal benchmarks like Omni-Bench. Notably, Qwen2.5-Omni is the first open-source model to achieve a level of performance in end-to-end speech instruction following that is comparable to its capabilities with text inputs, as evidenced by benchmarks such as MMLU and GSM8K. As for speech generation, Qwen2.5-Omni's streaming Talker outperform most existing streaming and non-streaming alternatives in robustness and naturalness.*
## Notes
@@ -40,7 +38,6 @@ The abstract from the technical report is the following:
- In case out out-of-memory errors hwen working with video input, decrease `processor.max_pixels`. By default the maximum is set to a very arge value and high resolution visuals will not be resized, unless resolution exceeds `processor.max_pixels`.
- The processor has its own [`~ProcessorMixin.apply_chat_template`] method to convert chat messages to model inputs.
-
## Usage example
`Qwen2.5-Omni` can be found on the [Huggingface Hub](https://huggingface.co/Qwen).
@@ -275,7 +272,8 @@ processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B", min_pixels=min
#### Prompt for audio output
If users need audio output, the system prompt must be set as "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.", otherwise the audio output may not work as expected.
-```
+
+```python
{
"role": "system",
"content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.",
@@ -285,6 +283,7 @@ If users need audio output, the system prompt must be set as "You are Qwen, a vi
#### Use audio output or not
The model supports both text and audio outputs, if users do not need audio outputs, they can set `enable_audio_output` in the `from_pretrained` function. This option will save about `~2GB` of GPU memory but the `return_audio` option for `generate` function will only allow to be set at `False`.
+
```python
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
"Qwen/Qwen2.5-Omni-7B",
@@ -341,8 +340,6 @@ model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
)
```
-
-
## Qwen2_5OmniConfig
[[autodoc]] Qwen2_5OmniConfig
diff --git a/docs/source/en/model_doc/qwen2_5_vl.md b/docs/source/en/model_doc/qwen2_5_vl.md
index 62527ea4963a..7f682bf80201 100644
--- a/docs/source/en/model_doc/qwen2_5_vl.md
+++ b/docs/source/en/model_doc/qwen2_5_vl.md
@@ -26,7 +26,6 @@ rendered properly in your Markdown viewer.
[Qwen2.5-VL](https://huggingface.co/papers/2502.13923) is a multimodal vision-language model, available in 3B, 7B, and 72B parameters, pretrained on 4.1T tokens. The model introduces window attention in the ViT encoder to accelerate training and inference, dynamic FPS sampling on the spatial and temporal dimensions for better video understanding across different sampling rates, and an upgraded MRoPE (multi-resolutional rotary positional encoding) mechanism to better capture and learn temporal dynamics.
-
You can find all the original Qwen2.5-VL checkpoints under the [Qwen2.5-VL](https://huggingface.co/collections/Qwen/qwen25-vl-6795ffac22b334a837c0f9a5) collection.
> [!TIP]
@@ -61,6 +60,7 @@ messages = [
pipe(text=messages,max_new_tokens=20, return_full_text=False)
```
+
@@ -110,6 +110,7 @@ output_text = processor.batch_decode(
)
print(output_text)
```
+
@@ -130,9 +131,11 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
)
```
+
### Notes
- Use Qwen2.5-VL for video inputs by setting `"type": "video"` as shown below.
+
```python
conversation = [
{
@@ -159,8 +162,10 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
print(output_text)
```
+
- Use Qwen2.5-VL for a mixed batch of inputs (images, videos, text). Add labels when handling multiple images or videos for better reference
as show below.
+
```python
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
@@ -221,14 +226,15 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
max_pixels = 2048*2048
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
```
-
+
Higher resolution can require more compute whereas reducing the resolution can save memory as follows:
-
+
```python
min_pixels = 256*28*28
max_pixels = 1024*28*28
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
```
+
## Qwen2_5_VLConfig
[[autodoc]] Qwen2_5_VLConfig
diff --git a/docs/source/en/model_doc/qwen2_audio.md b/docs/source/en/model_doc/qwen2_audio.md
index 7cdcd52119c0..9b9dd43a919d 100644
--- a/docs/source/en/model_doc/qwen2_audio.md
+++ b/docs/source/en/model_doc/qwen2_audio.md
@@ -36,7 +36,6 @@ The abstract from the paper is the following:
*We introduce the latest progress of Qwen-Audio, a large-scale audio-language model called Qwen2-Audio, which is capable of accepting various audio signal inputs and performing audio analysis or direct textual responses with regard to speech instructions. In contrast to complex hierarchical tags, we have simplified the pre-training process by utilizing natural language prompts for different data and tasks, and have further expanded the data volume. We have boosted the instruction-following capability of Qwen2-Audio and implemented two distinct audio interaction modes for voice chat and audio analysis. In the voice chat mode, users can freely engage in voice interactions with Qwen2-Audio without text input. In the audio analysis mode, users could provide audio and text instructions for analysis during the interaction. Note that we do not use any system prompts to switch between voice chat and audio analysis modes. Qwen2-Audio is capable of intelligently comprehending the content within audio and following voice commands to respond appropriately. For instance, in an audio segment that simultaneously contains sounds, multi-speaker conversations, and a voice command, Qwen2-Audio can directly understand the command and provide an interpretation and response to the audio. Additionally, DPO has optimized the model's performance in terms of factuality and adherence to desired behavior. According to the evaluation results from AIR-Bench, Qwen2-Audio outperformed previous SOTAs, such as Gemini-1.5-pro, in tests focused on audio-centric instruction-following capabilities. Qwen2-Audio is open-sourced with the aim of fostering the advancement of the multi-modal language community. *
-
## Usage tips
`Qwen2-Audio-7B` and `Qwen2-Audio-7B-Instruct` can be found on the [Huggingface Hub](https://huggingface.co/Qwen)
@@ -79,6 +78,7 @@ In the following, we demonstrate how to use `Qwen2-Audio-7B-Instruct` for the in
### Voice Chat Inference
In the voice chat mode, users can freely engage in voice interactions with Qwen2-Audio without text input:
+
```python
from io import BytesIO
from urllib.request import urlopen
@@ -119,6 +119,7 @@ response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_
### Audio Analysis Inference
In the audio analysis, users could provide both audio and text instructions for analysis:
+
```python
from io import BytesIO
from urllib.request import urlopen
@@ -167,6 +168,7 @@ response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_
### Batch Inference
We also support batch inference:
+
```python
from io import BytesIO
from urllib.request import urlopen
diff --git a/docs/source/en/model_doc/qwen2_moe.md b/docs/source/en/model_doc/qwen2_moe.md
index b8a3fe65d310..9d55de63e16d 100644
--- a/docs/source/en/model_doc/qwen2_moe.md
+++ b/docs/source/en/model_doc/qwen2_moe.md
@@ -24,7 +24,6 @@ rendered properly in your Markdown viewer.
# Qwen2MoE
-
[Qwen2MoE](https://huggingface.co/papers/2407.10671) is a Mixture-of-Experts (MoE) variant of [Qwen2](./qwen2), available as a base model and an aligned chat model. It uses SwiGLU activation, group query attention and a mixture of sliding window attention and full attention. The tokenizer can also be adapted to multiple languages and codes.
The MoE architecture uses upcyled models from the dense language models. For example, Qwen1.5-MoE-A2.7B is upcycled from Qwen-1.8B. It has 14.3B parameters but only 2.7B parameters are activated during runtime.
@@ -57,6 +56,7 @@ messages = [
outputs = pipe(messages, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
print(outputs[0]["generated_text"][-1]['content'])
```
+
@@ -100,14 +100,14 @@ generated_ids = [
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
```
-
+
+
```bash
transformers chat Qwen/Qwen1.5-MoE-A2.7B-Chat --dtype auto --attn_implementation flash_attention_2
```
-
-
+
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
diff --git a/docs/source/en/model_doc/qwen2_vl.md b/docs/source/en/model_doc/qwen2_vl.md
index 8ff09ca57238..59dc25b5e085 100644
--- a/docs/source/en/model_doc/qwen2_vl.md
+++ b/docs/source/en/model_doc/qwen2_vl.md
@@ -25,7 +25,7 @@ rendered properly in your Markdown viewer.
## Overview
-The [Qwen2-VL](https://huggingface.co/papers/2409.12191) ([blog post](https://qwenlm.github.io/blog/qwen2-vl/)) model is a major update to [Qwen-VL](https://huggingface.co/papers/2308.12966) from the Qwen team at Alibaba Research.
+The [Qwen2-VL](https://huggingface.co/papers/2409.12191) ([blog post](https://qwenlm.github.io/blog/qwen2-vl/)) model is a major update to [Qwen-VL](https://huggingface.co/papers/2308.12966) from the Qwen team at Alibaba Research.
The abstract from the blog is the following:
@@ -203,8 +203,8 @@ min_pixels = 256*28*28
max_pixels = 1024*28*28
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
```
-This ensures each image gets encoded using a number between 256-1024 tokens. The 28 comes from the fact that the model uses a patch size of 14 and a temporal patch size of 2 (14 x 2 = 28).
+This ensures each image gets encoded using a number between 256-1024 tokens. The 28 comes from the fact that the model uses a patch size of 14 and a temporal patch size of 2 (14 x 2 = 28).
#### Multiple Image Inputs
@@ -307,7 +307,7 @@ model = Qwen2VLForConditionalGeneration.from_pretrained(
[[autodoc]] Qwen2VLTextModel
- forward
-
+
## Qwen2VLModel
[[autodoc]] Qwen2VLModel
diff --git a/docs/source/en/model_doc/qwen3.md b/docs/source/en/model_doc/qwen3.md
index 87e6ba500f96..0141388fb97f 100644
--- a/docs/source/en/model_doc/qwen3.md
+++ b/docs/source/en/model_doc/qwen3.md
@@ -25,7 +25,6 @@ rendered properly in your Markdown viewer.
To be released with the official model launch.
-
## Usage tips
To be released with the official model launch.
diff --git a/docs/source/en/model_doc/qwen3_next.md b/docs/source/en/model_doc/qwen3_next.md
index f2e003182ee7..62b52e3d6d5e 100644
--- a/docs/source/en/model_doc/qwen3_next.md
+++ b/docs/source/en/model_doc/qwen3_next.md
@@ -13,18 +13,21 @@ specific language governing permissions and limitations under the License.
rendered properly in your Markdown viewer.
-->
+*This model was released on {release_date} and added to Hugging Face Transformers on 2025-09-10.*
+
## Overview
-The Qwen3-Next series represents our next-generation foundation models, optimized for extreme context length and large-scale parameter efficiency.
+The Qwen3-Next series represents our next-generation foundation models, optimized for extreme context length and large-scale parameter efficiency.
The series introduces a suite of architectural innovations designed to maximize performance while minimizing computational cost:
-- **Hybrid Attention**: Replaces standard attention with the combination of **Gated DeltaNet** and **Gated Attention**, enabling efficient context modeling.
+
+- **Hybrid Attention**: Replaces standard attention with the combination of **Gated DeltaNet** and **Gated Attention**, enabling efficient context modeling.
- **High-Sparsity MoE**: Achieves an extreme low activation ratio as 1:50 in MoE layers — drastically reducing FLOPs per token while preserving model capacity.
- **Multi-Token Prediction(MTP)**: Boosts pretraining model performance, and accelerates inference.
-- **Other Optimizations**: Includes techniques such as **zero-centered and weight-decayed layernorm**, **Gated Attention**, and other stabilizing enhancements for robust training.
+- **Other Optimizations**: Includes techniques such as **zero-centered and weight-decayed layernorm**, **Gated Attention**, and other stabilizing enhancements for robust training.
Built on this architecture, we trained and open-sourced Qwen3-Next-80B-A3B — 80B total parameters, only 3B active — achieving extreme sparsity and efficiency.
-Despite its ultra-efficiency, it outperforms Qwen3-32B on downstream tasks — while requiring **less than 1/10 of the training cost**.
+Despite its ultra-efficiency, it outperforms Qwen3-32B on downstream tasks — while requiring **less than 1/10 of the training cost**.
Moreover, it delivers over **10x higher inference throughput** than Qwen3-32B when handling contexts longer than 32K tokens.
For more details, please visit our blog [Qwen3-Next](qwen3_next) ([blog post](https://qwenlm.github.io/blog/qwen3_next/)).
@@ -60,7 +63,7 @@ generated_ids = model.generate(
**model_inputs,
max_new_tokens=512
)
-output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
content = tokenizer.decode(output_ids, skip_special_tokens=True)
diff --git a/docs/source/en/model_doc/qwen3_omni_moe.md b/docs/source/en/model_doc/qwen3_omni_moe.md
new file mode 100644
index 000000000000..9b7fa18d3812
--- /dev/null
+++ b/docs/source/en/model_doc/qwen3_omni_moe.md
@@ -0,0 +1,409 @@
+
+*This model was released on 2025-03-26 and added to Hugging Face Transformers on 2025-09-21.*
+
+# Qwen3-Omni-MOE
+
+
@@ -47,6 +46,7 @@ pipeline = pipeline(
)
pipeline("The capital of France is [MASK].")
```
+
@@ -24,21 +21,18 @@ Unless required by applicable law or agreed to in writing, software distributed
[MobileViT](https://huggingface.co/papers/2110.02178) is a lightweight vision transformer for mobile devices that merges CNNs's efficiency and inductive biases with transformers global context modeling. It treats transformers as convolutions, enabling global information processing without the heavy computational cost of standard ViTs.
-
@@ -46,7 +46,7 @@ pipe = pipeline(
dtype=torch.bfloat16,
device=0,
)
-
+
result = pipe("Plants create energy through a process known as")
print(result)
```
@@ -87,6 +87,7 @@ echo -e "Plants create energy through a process known as" | transformers-cli run
Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
The example below uses [torchao](../quantization/torchao) to only quantize the weights to 4-bits.
+
```py
#pip install torchao
@@ -116,18 +117,16 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
```
-
## Notes
-- Load specific intermediate checkpoints by adding the `revision` parameter to [`~PreTrainedModel.from_pretrained`].
+- Load specific intermediate checkpoints by adding the `revision` parameter to [`~PreTrainedModel.from_pretrained`].
```py
from transformers import AutoModelForCausalLM
-
+
model = AutoModelForCausalLM.from_pretrained("allenai/TBA", revision="stage1-step140000-tokens294B")
```
-
## Olmo3Config
[[autodoc]] Olmo3Config
@@ -144,4 +143,4 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
## Olmo3PreTrainedModel
[[autodoc]] Olmo3PreTrainedModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/oneformer.md b/docs/source/en/model_doc/oneformer.md
index c4b3bd142fe0..7f5d32bc55a8 100644
--- a/docs/source/en/model_doc/oneformer.md
+++ b/docs/source/en/model_doc/oneformer.md
@@ -39,7 +39,7 @@ This model was contributed by [Jitesh Jain](https://huggingface.co/praeclarumjj3
## Usage tips
-- OneFormer requires two inputs during inference: *image* and *task token*.
+- OneFormer requires two inputs during inference: *image* and *task token*.
- During training, OneFormer only uses panoptic annotations.
- If you want to train the model in a distributed environment across multiple nodes, then one should update the
`get_num_masks` function inside in the `OneFormerLoss` class of `modeling_oneformer.py`. When training on multiple nodes, this should be
diff --git a/docs/source/en/model_doc/openai-gpt.md b/docs/source/en/model_doc/openai-gpt.md
index b45b205e2592..04d37d89cc49 100644
--- a/docs/source/en/model_doc/openai-gpt.md
+++ b/docs/source/en/model_doc/openai-gpt.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2018-06-11 and added to Hugging Face Transformers on 2023-06-20.*
-
@@ -24,8 +23,6 @@ rendered properly in your Markdown viewer.
+
+
+
+
+# Parakeet
+
+## Overview
+
+Parakeet models, [introduced by NVIDIA NeMo](https://developer.nvidia.com/blog/pushing-the-boundaries-of-speech-recognition-with-nemo-parakeet-asr-models/), are models that combine a [Fast Conformer](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/models.html#fast-conformer) encoder with connectionist temporal classification (CTC), recurrent neural network transducer (RNNT) or token and duration transducer (TDT) decoder for automatic speech recognition.
+
+**Model Architecture**
+
+- **Fast Conformer Encoder**: A linearly scalable Conformer architecture that processes mel-spectrogram features and reduces sequence length through subsampling. This is more efficient version of the Conformer Encoder found in [FastSpeech2Conformer](./fastspeech2_conformer.md) (see [`ParakeetEncoder`] for the encoder implementation and details).
+- [**ParakeetForCTC**](#parakeetforctc): a Fast Conformer Encoder + a CTC decoder
+ - **CTC Decoder**: Simple but effective decoder consisting of:
+ - 1D convolution projection from encoder hidden size to vocabulary size (for optimal NeMo compatibility).
+ - CTC loss computation for training.
+ - Greedy CTC decoding for inference.
+
+The original implementation can be found in [NVIDIA NeMo](https://github.com/NVIDIA/NeMo).
+Model checkpoints are to be found under [the NVIDIA organization](https://huggingface.co/nvidia/models?search=parakeet).
+
+This model was contributed by [Nithin Rao Koluguri](https://huggingface.co/nithinraok), [Eustache Le Bihan](https://huggingface.co/eustlb) and [Eric Bezzam](https://huggingface.co/bezzam).
+
+## Usage
+
+### Basic usage
+
+
+
+
diff --git a/docs/source/en/model_doc/plbart.md b/docs/source/en/model_doc/plbart.md
index d8ce330cb0f7..b3459299437e 100644
--- a/docs/source/en/model_doc/plbart.md
+++ b/docs/source/en/model_doc/plbart.md
@@ -120,4 +120,4 @@ it's passed with the `text_target` keyword argument.
## PLBartForCausalLM
[[autodoc]] PLBartForCausalLM
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/pop2piano.md b/docs/source/en/model_doc/pop2piano.md
index 5f68b1805000..c934d8789037 100644
--- a/docs/source/en/model_doc/pop2piano.md
+++ b/docs/source/en/model_doc/pop2piano.md
@@ -21,14 +21,14 @@ specific language governing permissions and limitations under the License.
The Pop2Piano model was proposed in [Pop2Piano : Pop Audio-based Piano Cover Generation](https://huggingface.co/papers/2211.00895) by Jongho Choi and Kyogu Lee.
-Piano covers of pop music are widely enjoyed, but generating them from music is not a trivial task. It requires great
-expertise with playing piano as well as knowing different characteristics and melodies of a song. With Pop2Piano you
-can directly generate a cover from a song's audio waveform. It is the first model to directly generate a piano cover
-from pop audio without melody and chord extraction modules.
-
-Pop2Piano is an encoder-decoder Transformer model based on [T5](https://huggingface.co/papers/1910.10683). The input audio
-is transformed to its waveform and passed to the encoder, which transforms it to a latent representation. The decoder
-uses these latent representations to generate token ids in an autoregressive way. Each token id corresponds to one of four
+Piano covers of pop music are widely enjoyed, but generating them from music is not a trivial task. It requires great
+expertise with playing piano as well as knowing different characteristics and melodies of a song. With Pop2Piano you
+can directly generate a cover from a song's audio waveform. It is the first model to directly generate a piano cover
+from pop audio without melody and chord extraction modules.
+
+Pop2Piano is an encoder-decoder Transformer model based on [T5](https://huggingface.co/papers/1910.10683). The input audio
+is transformed to its waveform and passed to the encoder, which transforms it to a latent representation. The decoder
+uses these latent representations to generate token ids in an autoregressive way. Each token id corresponds to one of four
different token types: time, velocity, note and 'special'. The token ids are then decoded to their equivalent MIDI file.
The abstract from the paper is the following:
@@ -53,10 +53,13 @@ The original code can be found [here](https://github.com/sweetcocoa/pop2piano).
## Usage tips
* To use Pop2Piano, you will need to install the 🤗 Transformers library, as well as the following third party modules:
+
```bash
pip install pretty-midi==0.2.9 essentia==2.1b6.dev1034 librosa scipy
```
+
Please note that you may need to restart your runtime after installation.
+
* Pop2Piano is an Encoder-Decoder based model like T5.
* Pop2Piano can be used to generate midi-audio files for a given audio sequence.
* Choosing different composers in `Pop2PianoForConditionalGeneration.generate()` can lead to variety of different results.
@@ -131,7 +134,6 @@ Please note that you may need to restart your runtime after installation.
>>> tokenizer_output[1].write("./Outputs/midi_output2.mid")
```
-
- Example of processing multiple audio files in batch (Using `Pop2PianoFeatureExtractor` and `Pop2PianoTokenizer`):
```python
@@ -166,7 +168,6 @@ Please note that you may need to restart your runtime after installation.
>>> tokenizer_output[1].write("./Outputs/midi_output2.mid")
```
-
## Pop2PianoConfig
[[autodoc]] Pop2PianoConfig
diff --git a/docs/source/en/model_doc/prompt_depth_anything.md b/docs/source/en/model_doc/prompt_depth_anything.md
index 5af13c5d630e..d4b6f4cc2598 100644
--- a/docs/source/en/model_doc/prompt_depth_anything.md
+++ b/docs/source/en/model_doc/prompt_depth_anything.md
@@ -19,8 +19,7 @@ rendered properly in your Markdown viewer.
## Overview
-The Prompt Depth Anything model was introduced in [Prompting Depth Anything for 4K Resolution Accurate Metric Depth Estimation](https://huggingface.co/papers/2412.14015) by Haotong Lin, Sida Peng, Jingxiao Chen, Songyou Peng, Jiaming Sun, Minghuan Liu, Hujun Bao, Jiashi Feng, Xiaowei Zhou, Bingyi Kang.
-
+The Prompt Depth Anything model was introduced in [Prompting Depth Anything for 4K Resolution Accurate Metric Depth Estimation](https://huggingface.co/papers/2412.14015) by Haotong Lin, Sida Peng, Jingxiao Chen, Songyou Peng, Jiaming Sun, Minghuan Liu, Hujun Bao, Jiashi Feng, Xiaowei Zhou, Bingyi Kang.
The abstract from the paper is as follows:
@@ -100,4 +99,4 @@ If you are interested in submitting a resource to be included here, please feel
[[autodoc]] PromptDepthAnythingImageProcessorFast
- preprocess
- - post_process_depth_estimation
\ No newline at end of file
+ - post_process_depth_estimation
diff --git a/docs/source/en/model_doc/pvt.md b/docs/source/en/model_doc/pvt.md
index e7902affe5f4..38858db55529 100644
--- a/docs/source/en/model_doc/pvt.md
+++ b/docs/source/en/model_doc/pvt.md
@@ -29,23 +29,22 @@ is used to further reduce the resource consumption when learning high-resolution
The abstract from the paper is the following:
-*Although convolutional neural networks (CNNs) have achieved great success in computer vision, this work investigates a
-simpler, convolution-free backbone network useful for many dense prediction tasks. Unlike the recently proposed Vision
-Transformer (ViT) that was designed for image classification specifically, we introduce the Pyramid Vision Transformer
-(PVT), which overcomes the difficulties of porting Transformer to various dense prediction tasks. PVT has several
-merits compared to current state of the arts. Different from ViT that typically yields low resolution outputs and
-incurs high computational and memory costs, PVT not only can be trained on dense partitions of an image to achieve high
-output resolution, which is important for dense prediction, but also uses a progressive shrinking pyramid to reduce the
-computations of large feature maps. PVT inherits the advantages of both CNN and Transformer, making it a unified
-backbone for various vision tasks without convolutions, where it can be used as a direct replacement for CNN backbones.
+*Although convolutional neural networks (CNNs) have achieved great success in computer vision, this work investigates a
+simpler, convolution-free backbone network useful for many dense prediction tasks. Unlike the recently proposed Vision
+Transformer (ViT) that was designed for image classification specifically, we introduce the Pyramid Vision Transformer
+(PVT), which overcomes the difficulties of porting Transformer to various dense prediction tasks. PVT has several
+merits compared to current state of the arts. Different from ViT that typically yields low resolution outputs and
+incurs high computational and memory costs, PVT not only can be trained on dense partitions of an image to achieve high
+output resolution, which is important for dense prediction, but also uses a progressive shrinking pyramid to reduce the
+computations of large feature maps. PVT inherits the advantages of both CNN and Transformer, making it a unified
+backbone for various vision tasks without convolutions, where it can be used as a direct replacement for CNN backbones.
We validate PVT through extensive experiments, showing that it boosts the performance of many downstream tasks, including
-object detection, instance and semantic segmentation. For example, with a comparable number of parameters, PVT+RetinaNet
-achieves 40.4 AP on the COCO dataset, surpassing ResNet50+RetinNet (36.3 AP) by 4.1 absolute AP (see Figure 2). We hope
+object detection, instance and semantic segmentation. For example, with a comparable number of parameters, PVT+RetinaNet
+achieves 40.4 AP on the COCO dataset, surpassing ResNet50+RetinNet (36.3 AP) by 4.1 absolute AP (see Figure 2). We hope
that PVT could serve as an alternative and useful backbone for pixel-level predictions and facilitate future research.*
This model was contributed by [Xrenya](https://huggingface.co/Xrenya). The original code can be found [here](https://github.com/whai362/PVT).
-
- PVTv1 on ImageNet-1K
| **Model variant** |**Size** |**Acc@1**|**Params (M)**|
@@ -55,7 +54,6 @@ This model was contributed by [Xrenya](https://huggingface.co/Xrenya). The origi
| PVT-Medium | 224 | 81.2 | 44.2 |
| PVT-Large | 224 | 81.7 | 61.4 |
-
## PvtConfig
[[autodoc]] PvtConfig
diff --git a/docs/source/en/model_doc/pvt_v2.md b/docs/source/en/model_doc/pvt_v2.md
index 0d0ee3cca751..5be8998f4cc2 100644
--- a/docs/source/en/model_doc/pvt_v2.md
+++ b/docs/source/en/model_doc/pvt_v2.md
@@ -26,7 +26,7 @@ The PVTv2 encoder structure has been successfully deployed to achieve state-of-t
PVTv2 belongs to a family of models called [hierarchical transformers](https://natecibik.medium.com/the-rise-of-vision-transformers-f623c980419f) , which make adaptations to transformer layers in order to generate multi-scale feature maps. Unlike the columnal structure of Vision Transformer ([ViT](https://huggingface.co/papers/2010.11929)) which loses fine-grained detail, multi-scale feature maps are known preserve this detail and aid performance in dense prediction tasks. In the case of PVTv2, this is achieved by generating image patch tokens using 2D convolution with overlapping kernels in each encoder layer.
-The multi-scale features of hierarchical transformers allow them to be easily swapped in for traditional workhorse computer vision backbone models like ResNet in larger architectures. Both Segformer and Panoptic Segformer demonstrated that configurations using PVTv2 for a backbone consistently outperformed those with similarly sized ResNet backbones.
+The multi-scale features of hierarchical transformers allow them to be easily swapped in for traditional workhorse computer vision backbone models like ResNet in larger architectures. Both Segformer and Panoptic Segformer demonstrated that configurations using PVTv2 for a backbone consistently outperformed those with similarly sized ResNet backbones.
Another powerful feature of the PVTv2 is the complexity reduction in the self-attention layers called Spatial Reduction Attention (SRA), which uses 2D convolution layers to project hidden states to a smaller resolution before attending to them with the queries, improving the $O(n^2)$ complexity of self-attention to $O(n^2/R)$, with $R$ being the spatial reduction ratio (`sr_ratio`, aka kernel size and stride in the 2D convolution).
@@ -48,6 +48,7 @@ This model was contributed by [FoamoftheSea](https://huggingface.co/FoamoftheSea
- ImageNet pretrained weights for all model sizes can be found on the [hub](https://huggingface.co/models?other=pvt_v2).
The best way to get started with the PVTv2 is to load the pretrained checkpoint with the size of your choosing using `AutoModelForImageClassification`:
+
```python
import requests
import torch
@@ -99,7 +100,6 @@ outputs = model(torch.tensor(processed["pixel_values"]))
| PVT-V2-B4 | 224 | 83.6 | 62.6 |
| PVT-V2-B5 | 224 | 83.8 | 82.0 |
-
## PvtV2Config
[[autodoc]] PvtV2Config
diff --git a/docs/source/en/model_doc/qdqbert.md b/docs/source/en/model_doc/qdqbert.md
index 4c934d92d5fc..b791b4b2afe6 100644
--- a/docs/source/en/model_doc/qdqbert.md
+++ b/docs/source/en/model_doc/qdqbert.md
@@ -115,7 +115,7 @@ tensors. After setting up the tensor quantizers, one can use the following examp
The goal of exporting to ONNX is to deploy inference by [TensorRT](https://developer.nvidia.com/tensorrt). Fake
quantization will be broken into a pair of QuantizeLinear/DequantizeLinear ONNX ops. After setting static member of
-TensorQuantizer to use Pytorch’s own fake quantization functions, fake quantized model can be exported to ONNX, follow
+TensorQuantizer to use Pytorch's own fake quantization functions, fake quantized model can be exported to ONNX, follow
the instructions in [torch.onnx](https://pytorch.org/docs/stable/onnx.html). Example:
```python
diff --git a/docs/source/en/model_doc/qwen2.md b/docs/source/en/model_doc/qwen2.md
index 3f872302cc27..feeb69959b21 100644
--- a/docs/source/en/model_doc/qwen2.md
+++ b/docs/source/en/model_doc/qwen2.md
@@ -142,7 +142,6 @@ outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```
-
## Notes
- Ensure your Transformers library version is up-to-date. Qwen2 requires Transformers>=4.37.0 for full support.
diff --git a/docs/source/en/model_doc/qwen2_5_omni.md b/docs/source/en/model_doc/qwen2_5_omni.md
index e124f7cdb421..e2e0dc348a1c 100644
--- a/docs/source/en/model_doc/qwen2_5_omni.md
+++ b/docs/source/en/model_doc/qwen2_5_omni.md
@@ -29,9 +29,7 @@ The [Qwen2.5-Omni](https://qwenlm.github.io/blog/qwen2.5-omni/) model is a unifi
The abstract from the technical report is the following:
-*We present Qwen2.5-Omni, an end-to-end multimodal model designed to perceive diverse modalities, including text, images, audio, and video, while simultaneously generating text and natural speech responses in a streaming manner. To enable the streaming of multimodal information inputs, both audio and visual encoders utilize a block-wise processing approach. This strategy effectively decouples the handling of long sequences of multimodal data, assigning the perceptual responsibilities to the multimodal encoder and entrusting the modeling of extended sequences to a large language model. Such a division of labor enhances the fusion of different modalities via the shared attention mechanism. To synchronize the timestamps of video inputs with audio, we organized the audio and video sequentially in an interleaved manner and propose a novel position embedding approach, named TMRoPE (Time-aligned Multimodal RoPE). To concurrently generate text and speech while avoiding interference between the two modalities, we propose Thinker-Talker architecture. In this framework, Thinker functions as a large language model tasked with text generation, while Talker is a dual-track autoregressive model that directly utilizes the hidden representations from the Thinker to produce audio tokens as output. Both the Thinker and Talker models are designed to be trained and inferred in an end-to-end manner. For decoding audio tokens in a streaming manner, we introduce a sliding-window DiT that restricts the receptive field, aiming to reduce the initial package delay. Qwen2.5-Omni outperforms the similarly sized Qwen2-VL and Qwen2-Audio in both image and audio capabilities. Furthermore, Qwen2.5-Omni achieves state-of-the-art performance on multimodal benchmarks like Omni-Bench. Notably, Qwen2.5-Omni is the first open-source model to achieve a level of performance in end-to-end speech instruction following that is comparable to its capabilities with text inputs, as evidenced by benchmarks such as MMLU and GSM8K. As for speech generation, Qwen2.5-Omni’s streaming Talker outperform most existing streaming and non-streaming alternatives in robustness and naturalness.*
-
-
+*We present Qwen2.5-Omni, an end-to-end multimodal model designed to perceive diverse modalities, including text, images, audio, and video, while simultaneously generating text and natural speech responses in a streaming manner. To enable the streaming of multimodal information inputs, both audio and visual encoders utilize a block-wise processing approach. This strategy effectively decouples the handling of long sequences of multimodal data, assigning the perceptual responsibilities to the multimodal encoder and entrusting the modeling of extended sequences to a large language model. Such a division of labor enhances the fusion of different modalities via the shared attention mechanism. To synchronize the timestamps of video inputs with audio, we organized the audio and video sequentially in an interleaved manner and propose a novel position embedding approach, named TMRoPE (Time-aligned Multimodal RoPE). To concurrently generate text and speech while avoiding interference between the two modalities, we propose Thinker-Talker architecture. In this framework, Thinker functions as a large language model tasked with text generation, while Talker is a dual-track autoregressive model that directly utilizes the hidden representations from the Thinker to produce audio tokens as output. Both the Thinker and Talker models are designed to be trained and inferred in an end-to-end manner. For decoding audio tokens in a streaming manner, we introduce a sliding-window DiT that restricts the receptive field, aiming to reduce the initial package delay. Qwen2.5-Omni outperforms the similarly sized Qwen2-VL and Qwen2-Audio in both image and audio capabilities. Furthermore, Qwen2.5-Omni achieves state-of-the-art performance on multimodal benchmarks like Omni-Bench. Notably, Qwen2.5-Omni is the first open-source model to achieve a level of performance in end-to-end speech instruction following that is comparable to its capabilities with text inputs, as evidenced by benchmarks such as MMLU and GSM8K. As for speech generation, Qwen2.5-Omni's streaming Talker outperform most existing streaming and non-streaming alternatives in robustness and naturalness.*
## Notes
@@ -40,7 +38,6 @@ The abstract from the technical report is the following:
- In case out out-of-memory errors hwen working with video input, decrease `processor.max_pixels`. By default the maximum is set to a very arge value and high resolution visuals will not be resized, unless resolution exceeds `processor.max_pixels`.
- The processor has its own [`~ProcessorMixin.apply_chat_template`] method to convert chat messages to model inputs.
-
## Usage example
`Qwen2.5-Omni` can be found on the [Huggingface Hub](https://huggingface.co/Qwen).
@@ -275,7 +272,8 @@ processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B", min_pixels=min
#### Prompt for audio output
If users need audio output, the system prompt must be set as "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.", otherwise the audio output may not work as expected.
-```
+
+```python
{
"role": "system",
"content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.",
@@ -285,6 +283,7 @@ If users need audio output, the system prompt must be set as "You are Qwen, a vi
#### Use audio output or not
The model supports both text and audio outputs, if users do not need audio outputs, they can set `enable_audio_output` in the `from_pretrained` function. This option will save about `~2GB` of GPU memory but the `return_audio` option for `generate` function will only allow to be set at `False`.
+
```python
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
"Qwen/Qwen2.5-Omni-7B",
@@ -341,8 +340,6 @@ model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
)
```
-
-
## Qwen2_5OmniConfig
[[autodoc]] Qwen2_5OmniConfig
diff --git a/docs/source/en/model_doc/qwen2_5_vl.md b/docs/source/en/model_doc/qwen2_5_vl.md
index 62527ea4963a..7f682bf80201 100644
--- a/docs/source/en/model_doc/qwen2_5_vl.md
+++ b/docs/source/en/model_doc/qwen2_5_vl.md
@@ -26,7 +26,6 @@ rendered properly in your Markdown viewer.
[Qwen2.5-VL](https://huggingface.co/papers/2502.13923) is a multimodal vision-language model, available in 3B, 7B, and 72B parameters, pretrained on 4.1T tokens. The model introduces window attention in the ViT encoder to accelerate training and inference, dynamic FPS sampling on the spatial and temporal dimensions for better video understanding across different sampling rates, and an upgraded MRoPE (multi-resolutional rotary positional encoding) mechanism to better capture and learn temporal dynamics.
-
You can find all the original Qwen2.5-VL checkpoints under the [Qwen2.5-VL](https://huggingface.co/collections/Qwen/qwen25-vl-6795ffac22b334a837c0f9a5) collection.
> [!TIP]
@@ -61,6 +60,7 @@ messages = [
pipe(text=messages,max_new_tokens=20, return_full_text=False)
```
+
+
+
+
+
+
+## Overview
+
+The Qwen3-Omni-MOE model is a unified multiple modalities model proposed in [Qwen3-Omni Technical Report](https://huggingface.co/papers/2509.17765) from Qwen team, Alibaba Group.
+
+The abstract from the technical report is the following:
+
+*We present Qwen3-Omni, a single multimodal model that, for the first time, maintains state-of-the-art performance across text, image, audio, and video without any degradation relative to single-modal counterparts. Qwen3-Omni matches the performance of same-sized single-modal models within the Qwen series and excels particularly on audio tasks. Across 36 audio and audio-visual benchmarks, Qwen3-Omni achieves open-source SOTA on 32 benchmarks and overall SOTA on 22, outperforming strong closed-source models such as Gemini-2.5-Pro, Seed-ASR, and GPT-4o-Transcribe. Qwen3-Omni adopts a Thinker-Talker MoE architecture that unifies perception and generation across text, images, audio, and video, yielding fluent text and natural real-time speech. It supports text interaction in 119 languages, speech understanding in 19 languages, and speech generation in 10 languages. To reduce first-packet latency in streaming synthesis, Talker autoregressively predicts discrete speech codecs using a multi-codebook scheme. Leveraging the representational capacity of these codebooks, we replace computationally intensive block-wise diffusion with a lightweight causal ConvNet, enabling streaming from the first codec frame. In cold-start settings, Qwen3-Omni achieves a theoretical end-to-end first-packet latency of 234 ms. To further strengthen multimodal reasoning, we introduce a Thinking model that explicitly reasons over inputs from any modality. Since the research community currently lacks a general-purpose audio captioning model, we fine-tuned Qwen3-Omni-30B-A3B to obtain Qwen3-Omni-30B-A3B-Captioner, which produces detailed, low-hallucination captions for arbitrary audio inputs. Qwen3-Omni-30B-A3B, Qwen3-Omni-30B-A3B-Thinking, and Qwen3-Omni-30B-A3B-Captioner are publicly released under the Apache 2.0 license.
+
+## Notes
+
+- Use [`Qwen3OmniMoeForConditionalGeneration`] to generate audio and text output. To generate only one output type, use [`Qwen3OmniMoeThinkerForConditionalGeneration`] for text-only and [`Qwen3OmniMoeTalkerForConditionalGeneration`] for audio-only outputs.
+- Audio generation with [`Qwen3OmniMoeForConditionalGeneration`] supports only single batch size at the moment.
+- In case out out-of-memory errors hwen working with video input, decrease `processor.max_pixels`. By default the maximum is set to a very arge value and high resolution visuals will not be resized, unless resolution exceeds `processor.max_pixels`.
+- The processor has its own [`~ProcessorMixin.apply_chat_template`] method to convert chat messages to model inputs.
+
+## Usage example
+
+`Qwen3-Omni` can be found on the [Huggingface Hub](https://huggingface.co/Qwen).
+
+### Single Media inference
+
+The model can accept text, images, audio and videos as input. Here's an example code for inference.
+
+```python
+import soundfile as sf
+from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor
+
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ dtype="auto",
+ device_map="auto"
+)
+processor = Qwen3OmniMoeProcessor.from_pretrained("Qwen/Qwen3-Omni-30B-A3B-Instruct")
+
+conversations = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "video", "video": "/path/to/video.mp4"},
+ {"type": "text", "text": "What cant you hear and see in this video?"},
+ ],
+ },
+]
+
+inputs = processor.apply_chat_template(
+ conversations,
+ load_audio_from_video=True,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=True,
+ return_tensors="pt",
+ video_fps=1,
+
+ # kwargs to be passed to `Qwen3OmniMoeProcessor`
+ padding=True,
+ use_audio_in_video=True,
+).to(model.device)
+
+# Generation params for audio or text can be different and have to be prefixed with `thinker_` or `talker_`
+text_ids, audio = model.generate(**inputs, use_audio_in_video=True, thinker_do_sample=False, talker_do_sample=True)
+text = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+
+sf.write(
+ "output.wav",
+ audio.reshape(-1).detach().cpu().numpy(),
+ samplerate=24000,
+)
+print(text)
+```
+
+### Text-only generation
+
+To generate only text output and save compute by not loading the audio generation model, we can use `Qwen3OmniMoeThinkerForConditionalGeneration` model.
+
+```python
+from transformers import Qwen3OmniMoeThinkerForConditionalGeneration, Qwen3OmniMoeProcessor
+
+model = Qwen3OmniMoeThinkerForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ dtype="auto",
+ device_map="auto",
+)
+processor = Qwen3OmniMoeProcessor.from_pretrained("Qwen/Qwen3-Omni-30B-A3B-Instruct")
+
+conversations = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "video", "video": "/path/to/video.mp4"},
+ {"type": "text", "text": "What cant you hear and see in this video?"},
+ ],
+ },
+]
+
+inputs = processor.apply_chat_template(
+ conversations,
+ load_audio_from_video=True,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=True,
+ return_tensors="pt",
+ video_fps=1,
+
+ # kwargs to be passed to `Qwen3OmniMoeProcessor`
+ padding=True,
+ use_audio_in_video=True,
+).to(model.device)
+
+
+text_ids = model.generate(**inputs, use_audio_in_video=True)
+text = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+
+sf.write(
+ "output.wav",
+ audio.reshape(-1).detach().cpu().numpy(),
+ samplerate=24000,
+)
+print(text)
+```
+
+### Batch Mixed Media Inference
+
+The model can batch inputs composed of mixed samples of various types such as text, images, audio and videos as input when using `Qwen3OmniMoeThinkerForConditionalGeneration` model. Here is an example.
+
+```python
+import soundfile as sf
+from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor
+
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ dtype="auto",
+ device_map="auto"
+)
+processor = Qwen3OmniMoeProcessor.from_pretrained("Qwen/Qwen3-Omni-30B-A3B-Instruct")
+
+# Conversation with video only
+conversation1 = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "video", "path": "/path/to/video.mp4"},
+ ]
+ }
+]
+
+# Conversation with audio only
+conversation2 = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "audio", "path": "/path/to/audio.wav"},
+ ]
+ }
+]
+
+# Conversation with pure text
+conversation3 = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": "who are you?"}],
+ }
+]
+
+
+# Conversation with mixed media
+conversation4 = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "path": "/path/to/image.jpg"},
+ {"type": "video", "path": "/path/to/video.mp4"},
+ {"type": "audio", "path": "/path/to/audio.wav"},
+ {"type": "text", "text": "What are the elements can you see and hear in these medias?"},
+ ],
+ }
+]
+
+conversations = [conversation1, conversation2, conversation3, conversation4]
+
+inputs = processor.apply_chat_template(
+ conversations,
+ load_audio_from_video=True,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=True,
+ return_tensors="pt",
+ video_fps=1,
+
+ # kwargs to be passed to `Qwen3OmniMoeProcessor`
+ padding=True,
+ use_audio_in_video=True,
+).to(model.thinker.device)
+
+text_ids = model.generate(**inputs, use_audio_in_video=True)
+text = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+
+print(text)
+```
+
+### Usage Tips
+
+#### Image Resolution trade-off
+
+The model supports a wide range of resolution inputs. By default, it uses the native resolution for input, but higher resolutions can enhance performance at the cost of more computation. Users can set the minimum and maximum number of pixels to achieve an optimal configuration for their needs.
+
+```python
+min_pixels = 128*28*28
+max_pixels = 768*28*28
+processor = AutoProcessor.from_pretrained("Qwen/Qwen3-Omni-30B-A3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+```
+
+#### Prompt for audio output
+If users need audio output, the system prompt must be set as "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.", otherwise the audio output may not work as expected.
+
+```json
+{
+ "role": "system",
+ "content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.",
+}
+```
+
+#### Use audio output or not
+
+The model supports both text and audio outputs, if users do not need audio outputs, they can set `enable_audio_output` in the `from_pretrained` function. This option will save about `~2GB` of GPU memory but the `return_audio` option for `generate` function will only allow to be set at `False`.
+
+```python
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ dtype="auto",
+ device_map="auto",
+ enable_audio_output=False,
+)
+```
+
+In order to obtain a flexible experience, we recommend that users set `enable_audio_output` at `True` when initializing the model through `from_pretrained` function, and then decide whether to return audio when `generate` function is called. When `return_audio` is set to `False`, the model will only return text outputs to get text responses faster.
+
+```python
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ dtype="auto",
+ device_map="auto",
+ enable_audio_output=True,
+)
+...
+text_ids = model.generate(**inputs, return_audio=False)
+```
+
+#### Change voice type of output audio
+Qwen3-Omni-MOE supports the ability to change the voice of the output audio. Users can use the `spk` parameter of `generate` function to specify the voice type. The `"Qwen/Qwen3-Omni-30B-A3B-Instruct"` checkpoint support two voice types: `Chelsie` and `Ethan`, while `Chelsie` is a female voice and `Ethan` is a male voice. By default, if `spk` is not specified, the default voice type is `Chelsie`.
+
+```python
+text_ids, audio = model.generate(**inputs, spk="Chelsie")
+```
+
+```python
+text_ids, audio = model.generate(**inputs, spk="Ethan")
+```
+
+#### Flash-Attention 2 to speed up generation
+
+First, make sure to install the latest version of Flash Attention 2:
+
+```bash
+pip install -U flash-attn --no-build-isolation
+```
+
+Also, you should have hardware that is compatible with FlashAttention 2. Read more about it in the official documentation of the [flash attention repository](https://github.com/Dao-AILab/flash-attention). FlashAttention-2 can only be used when a model is loaded in `torch.float16` or `torch.bfloat16`.
+
+To load and run a model using FlashAttention-2, add `attn_implementation="flash_attention_2"` when loading the model:
+
+```python
+from transformers import Qwen3OmniMoeForConditionalGeneration
+
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ device_map="auto",
+ dtype=torch.bfloat16,
+ attn_implementation="flash_attention_2",
+)
+```
+
+## Qwen3OmniMoeConfig
+
+[[autodoc]] Qwen3OmniMoeConfig
+
+## Qwen3OmniMoeThinkerConfig
+
+[[autodoc]] Qwen3OmniMoeThinkerConfig
+
+## Qwen3OmniMoeTalkerConfig
+
+[[autodoc]] Qwen3OmniMoeTalkerConfig
+
+## Qwen3OmniMoeForConditionalGeneration
+
+[[autodoc]] Qwen3OmniMoeForConditionalGeneration
+
+## Qwen3OmniMoeThinkerTextModel
+
+[[autodoc]] Qwen3OmniMoeThinkerTextModel
+
+## Qwen3OmniMoeThinkerForConditionalGeneration
+
+[[autodoc]] Qwen3OmniMoeThinkerForConditionalGeneration
+
+## Qwen3OmniMoeTalkerForConditionalGeneration
+
+[[autodoc]] Qwen3OmniMoeTalkerForConditionalGeneration
+
+## Qwen3OmniMoePreTrainedModel
+
+[[autodoc]] Qwen3OmniMoePreTrainedModel
+
+## Qwen3OmniMoePreTrainedModelForConditionalGeneration
+
+[[autodoc]] Qwen3OmniMoePreTrainedModelForConditionalGeneration
+
+## Qwen3OmniMoeTalkerModel
+
+[[autodoc]] Qwen3OmniMoeTalkerModel
+
+## Qwen3OmniMoeThinkerTextPreTrainedModel
+
+[[autodoc]] Qwen3OmniMoeThinkerTextPreTrainedModel
+
+## Qwen3OmniMoeProcessor
+
+[[autodoc]] Qwen3OmniMoeProcessor
+
+## Qwen3OmniMoeCode2Wav
+
+[[autodoc]] Qwen3OmniMoeCode2Wav
+
+## Qwen3OmniMoeCode2WavDecoderBlock
+
+[[autodoc]] Qwen3OmniMoeCode2WavDecoderBlock
+
+## Qwen3OmniMoeCode2WavTransformerModel
+
+[[autodoc]] Qwen3OmniMoeCode2WavTransformerModel
+
+## Qwen3OmniMoeTalkerCodePredictorModel
+
+[[autodoc]] Qwen3OmniMoeTalkerCodePredictorModel
+
+## Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration
+
+[[autodoc]] Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration
diff --git a/docs/source/en/model_doc/qwen3_vl.md b/docs/source/en/model_doc/qwen3_vl.md
index 9e90363a1eba..33c8c7e96aee 100644
--- a/docs/source/en/model_doc/qwen3_vl.md
+++ b/docs/source/en/model_doc/qwen3_vl.md
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
rendered properly in your Markdown viewer.
-->
-*This model was released on None and added to Hugging Face Transformers on 2025-08-16.*
+*This model was released on 2025-09-23 and added to Hugging Face Transformers on 2025-09-15.*
+
+
+
@@ -77,6 +77,7 @@ output_text = processor.batch_decode(
)
print(output_text)
```
+
diff --git a/docs/source/en/model_doc/qwen3_vl_moe.md b/docs/source/en/model_doc/qwen3_vl_moe.md
index 76d046efff2d..771f6d411cf2 100644
--- a/docs/source/en/model_doc/qwen3_vl_moe.md
+++ b/docs/source/en/model_doc/qwen3_vl_moe.md
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
rendered properly in your Markdown viewer.
-->
-*This model was released on None and added to Hugging Face Transformers on 2025-08-17.*
+*This model was released on 2025-02-19 and added to Hugging Face Transformers on 2025-09-15.*
@@ -77,6 +77,7 @@ output_text = processor.batch_decode(
)
print(output_text)
```
+
diff --git a/docs/source/en/model_doc/recurrent_gemma.md b/docs/source/en/model_doc/recurrent_gemma.md
index 1cd4e784a5bd..2d7c940e00a9 100644
--- a/docs/source/en/model_doc/recurrent_gemma.md
+++ b/docs/source/en/model_doc/recurrent_gemma.md
@@ -31,16 +31,14 @@ The abstract from the paper is the following:
Tips:
-- The original checkpoints can be converted using the conversion script [`src/transformers/models/recurrent_gemma/convert_recurrent_gemma_weights_to_hf.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/recurrent_gemma/convert_recurrent_gemma_to_hf.py).
+- The original checkpoints can be converted using the conversion script [`src/transformers/models/recurrent_gemma/convert_recurrent_gemma_weights_to_hf.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/recurrent_gemma/convert_recurrent_gemma_to_hf.py).
This model was contributed by [Arthur Zucker](https://huggingface.co/ArthurZ). The original code can be found [here](https://github.com/google-deepmind/recurrentgemma).
-
## RecurrentGemmaConfig
[[autodoc]] RecurrentGemmaConfig
-
## RecurrentGemmaModel
[[autodoc]] RecurrentGemmaModel
@@ -50,4 +48,3 @@ This model was contributed by [Arthur Zucker](https://huggingface.co/ArthurZ). T
[[autodoc]] RecurrentGemmaForCausalLM
- forward
-
diff --git a/docs/source/en/model_doc/reformer.md b/docs/source/en/model_doc/reformer.md
index f94134609d2b..c556e01ba13c 100644
--- a/docs/source/en/model_doc/reformer.md
+++ b/docs/source/en/model_doc/reformer.md
@@ -41,8 +41,8 @@ found [here](https://github.com/google/trax/tree/master/trax/models/reformer).
## Usage tips
- Reformer does **not** work with *torch.nn.DataParallel* due to a bug in PyTorch, see [issue #36035](https://github.com/pytorch/pytorch/issues/36035).
-- Use Axial position encoding (see below for more details). It’s a mechanism to avoid having a huge positional encoding matrix (when the sequence length is very big) by factorizing it into smaller matrices.
-- Replace traditional attention by LSH (local-sensitive hashing) attention (see below for more details). It’s a technique to avoid computing the full product query-key in the attention layers.
+- Use Axial position encoding (see below for more details). It's a mechanism to avoid having a huge positional encoding matrix (when the sequence length is very big) by factorizing it into smaller matrices.
+- Replace traditional attention by LSH (local-sensitive hashing) attention (see below for more details). It's a technique to avoid computing the full product query-key in the attention layers.
- Avoid storing the intermediate results of each layer by using reversible transformer layers to obtain them during the backward pass (subtracting the residuals from the input of the next layer gives them back) or recomputing them for results inside a given layer (less efficient than storing them but saves memory).
- Compute the feedforward operations by chunks and not on the whole batch.
@@ -89,7 +89,6 @@ equal to `config.hidden_size` and `config.axial_pos_shape` is set to a tuple \\(
product has to be equal to `config.max_embedding_size`, which during training has to be equal to the *sequence
length* of the `input_ids`.
-
### LSH Self Attention
In Locality sensitive hashing (LSH) self attention the key and query projection weights are tied. Therefore, the key
@@ -122,7 +121,6 @@ Using LSH self attention, the memory and time complexity of the query-key matmul
\\(\mathcal{O}(n_s \times n_s)\\) to \\(\mathcal{O}(n_s \times \log(n_s))\\), which usually represents the memory
and time bottleneck in a transformer model, with \\(n_s\\) being the sequence length.
-
### Local Self Attention
Local self attention is essentially a "normal" self attention layer with key, query and value projections, but is
@@ -134,7 +132,6 @@ Using Local self attention, the memory and time complexity of the query-key matm
\\(\mathcal{O}(n_s \times n_s)\\) to \\(\mathcal{O}(n_s \times \log(n_s))\\), which usually represents the memory
and time bottleneck in a transformer model, with \\(n_s\\) being the sequence length.
-
### Training
During training, we must ensure that the sequence length is set to a value that can be divided by the least common
diff --git a/docs/source/en/model_doc/retribert.md b/docs/source/en/model_doc/retribert.md
index 871bdc6e8c86..829fed24215f 100644
--- a/docs/source/en/model_doc/retribert.md
+++ b/docs/source/en/model_doc/retribert.md
@@ -39,7 +39,6 @@ pair of BERT encoders with lower-dimension projection for dense semantic indexin
This model was contributed by [yjernite](https://huggingface.co/yjernite). Code to train and use the model can be
found [here](https://github.com/huggingface/transformers/tree/main/examples/research-projects/distillation).
-
## RetriBertConfig
[[autodoc]] RetriBertConfig
diff --git a/docs/source/en/model_doc/roberta.md b/docs/source/en/model_doc/roberta.md
index da393646442a..43414fac4c88 100644
--- a/docs/source/en/model_doc/roberta.md
+++ b/docs/source/en/model_doc/roberta.md
@@ -28,7 +28,6 @@ rendered properly in your Markdown viewer.
You can find all the original RoBERTa checkpoints under the [Facebook AI](https://huggingface.co/FacebookAI) organization.
-
> [!TIP]
> Click on the RoBERTa models in the right sidebar for more examples of how to apply RoBERTa to different language tasks.
diff --git a/docs/source/en/model_doc/rt_detr.md b/docs/source/en/model_doc/rt_detr.md
index 02accfd6d9f7..d4c85f63fc37 100644
--- a/docs/source/en/model_doc/rt_detr.md
+++ b/docs/source/en/model_doc/rt_detr.md
@@ -23,7 +23,6 @@ rendered properly in your Markdown viewer.
## Overview
-
The RT-DETR model was proposed in [DETRs Beat YOLOs on Real-time Object Detection](https://huggingface.co/papers/2304.08069) by Wenyu Lv, Yian Zhao, Shangliang Xu, Jinman Wei, Guanzhong Wang, Cheng Cui, Yuning Du, Qingqing Dang, Yi Liu.
RT-DETR is an object detection model that stands for "Real-Time DEtection Transformer." This model is designed to perform object detection tasks with a focus on achieving real-time performance while maintaining high accuracy. Leveraging the transformer architecture, which has gained significant popularity in various fields of deep learning, RT-DETR processes images to identify and locate multiple objects within them.
@@ -39,7 +38,6 @@ alt="drawing" width="600"/>
The model version was contributed by [rafaelpadilla](https://huggingface.co/rafaelpadilla) and [sangbumchoi](https://github.com/SangbumChoi). The original code can be found [here](https://github.com/lyuwenyu/RT-DETR/).
-
## Usage tips
Initially, an image is processed using a pre-trained convolutional neural network, specifically a Resnet-D variant as referenced in the original code. This network extracts features from the final three layers of the architecture. Following this, a hybrid encoder is employed to convert the multi-scale features into a sequential array of image features. Then, a decoder, equipped with auxiliary prediction heads is used to refine the object queries. This process facilitates the direct generation of bounding boxes, eliminating the need for any additional post-processing to acquire the logits and coordinates for the bounding boxes.
diff --git a/docs/source/en/model_doc/rt_detr_v2.md b/docs/source/en/model_doc/rt_detr_v2.md
index f5eb54625c84..3f814ce0d649 100644
--- a/docs/source/en/model_doc/rt_detr_v2.md
+++ b/docs/source/en/model_doc/rt_detr_v2.md
@@ -34,9 +34,9 @@ The abstract from the paper is the following:
This model was contributed by [jadechoghari](https://huggingface.co/jadechoghari).
The original code can be found [here](https://github.com/lyuwenyu/RT-DETR).
-## Usage tips
+## Usage tips
-This second version of RT-DETR improves how the decoder finds objects in an image.
+This second version of RT-DETR improves how the decoder finds objects in an image.
- **better sampling** – adjusts offsets so the model looks at the right areas
- **flexible attention** – can use smooth (bilinear) or fixed (discrete) sampling
@@ -85,17 +85,15 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- See also: [Object detection task guide](../tasks/object_detection).
- Notebooks for [inference](https://github.com/qubvel/transformers-notebooks/blob/main/notebooks/RT_DETR_v2_inference.ipynb) and [fine-tuning](https://github.com/qubvel/transformers-notebooks/blob/main/notebooks/RT_DETR_v2_finetune_on_a_custom_dataset.ipynb) RT-DETRv2 on a custom dataset (🌎).
-
## RTDetrV2Config
[[autodoc]] RTDetrV2Config
-
## RTDetrV2Model
[[autodoc]] RTDetrV2Model
- forward
-
+
## RTDetrV2ForObjectDetection
[[autodoc]] RTDetrV2ForObjectDetection
diff --git a/docs/source/en/model_doc/rwkv.md b/docs/source/en/model_doc/rwkv.md
index 4d9d6bbb8860..9b5d64fedbb7 100644
--- a/docs/source/en/model_doc/rwkv.md
+++ b/docs/source/en/model_doc/rwkv.md
@@ -58,7 +58,7 @@ torch.allclose(torch.cat([output_one, output_two], dim=1), output_whole, atol=1e
If you want to make sure the model stops generating when `'\n\n'` is detected, we recommend using the following stopping criteria:
-```python
+```python
from transformers import StoppingCriteria
class RwkvStoppingCriteria(StoppingCriteria):
@@ -152,4 +152,4 @@ $$D_{i} = e^{u + K_{i} - q} + e^{M_{i}} \tilde{D}_{i} \hbox{ where } q = \max(
which finally gives us
-$$O_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}$$
\ No newline at end of file
+$$O_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}$$
diff --git a/docs/source/en/model_doc/sam.md b/docs/source/en/model_doc/sam.md
index 49a58254630a..65286eb8428d 100644
--- a/docs/source/en/model_doc/sam.md
+++ b/docs/source/en/model_doc/sam.md
@@ -41,7 +41,6 @@ Tips:
- Fine-tuning the model is not supported yet
- According to the paper, textual input should be also supported. However, at this time of writing this seems not to be supported according to [the official repository](https://github.com/facebookresearch/segment-anything/issues/4#issuecomment-1497626844).
-
This model was contributed by [ybelkada](https://huggingface.co/ybelkada) and [ArthurZ](https://huggingface.co/ArthurZ).
The original code can be found [here](https://github.com/facebookresearch/segment-anything).
@@ -98,6 +97,7 @@ masks = processor.image_processor.post_process_masks(
)
scores = outputs.iou_scores
```
+
## Resources
A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with SAM.
diff --git a/docs/source/en/model_doc/sam_hq.md b/docs/source/en/model_doc/sam_hq.md
index 2bd14229c37c..9dea1de7a77e 100644
--- a/docs/source/en/model_doc/sam_hq.md
+++ b/docs/source/en/model_doc/sam_hq.md
@@ -25,7 +25,6 @@ The model is an enhancement to the original SAM model that produces significantl
-
SAM-HQ introduces several key improvements over the original SAM model:
1. High-Quality Output Token: A learnable token injected into SAM's mask decoder for higher quality mask prediction
@@ -105,7 +104,6 @@ masks = processor.image_processor.post_process_masks(
scores = outputs.iou_scores
```
-
## Resources
A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with SAM-HQ:
@@ -137,7 +135,6 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
[[autodoc]] SamHQVisionModel
-
## SamHQModel
[[autodoc]] SamHQModel
diff --git a/docs/source/en/model_doc/seamless_m4t.md b/docs/source/en/model_doc/seamless_m4t.md
index c6f3a56f9ba1..e7fc00d047c3 100644
--- a/docs/source/en/model_doc/seamless_m4t.md
+++ b/docs/source/en/model_doc/seamless_m4t.md
@@ -67,7 +67,6 @@ Here is how to use the processor to process text and audio:
>>> text_inputs = processor(text = "Hello, my dog is cute", src_lang="eng", return_tensors="pt")
```
-
### Speech
[`SeamlessM4TModel`] can *seamlessly* generate text or speech with few or no changes. Let's target Russian voice translation:
@@ -84,7 +83,7 @@ With basically the same code, I've translated English text and Arabic speech to
Similarly, you can generate translated text from audio files or from text with the same model. You only have to pass `generate_speech=False` to [`SeamlessM4TModel.generate`].
This time, let's translate to French.
-```python
+```python
>>> # from audio
>>> output_tokens = model.generate(**audio_inputs, tgt_lang="fra", generate_speech=False)
>>> translated_text_from_audio = processor.decode(output_tokens[0].tolist()[0], skip_special_tokens=True)
@@ -96,11 +95,10 @@ This time, let's translate to French.
### Tips
-
#### 1. Use dedicated models
[`SeamlessM4TModel`] is transformers top level model to generate speech and text, but you can also use dedicated models that perform the task without additional components, thus reducing the memory footprint.
-For example, you can replace the audio-to-audio generation snippet with the model dedicated to the S2ST task, the rest is exactly the same code:
+For example, you can replace the audio-to-audio generation snippet with the model dedicated to the S2ST task, the rest is exactly the same code:
```python
>>> from transformers import SeamlessM4TForSpeechToSpeech
@@ -130,7 +128,6 @@ Use `return_intermediate_token_ids=True` with [`SeamlessM4TModel`] to return bot
## Model architecture
-
SeamlessM4T features a versatile architecture that smoothly handles the sequential generation of text and speech. This setup comprises two sequence-to-sequence (seq2seq) models. The first model translates the input modality into translated text, while the second model generates speech tokens, known as "unit tokens," from the translated text.
Each modality has its own dedicated encoder with a unique architecture. Additionally, for speech output, a vocoder inspired by the [HiFi-GAN](https://huggingface.co/papers/2010.05646) architecture is placed on top of the second seq2seq model.
@@ -142,7 +139,6 @@ Here's how the generation process works:
- If speech generation is required, the second seq2seq model, following a standard encoder-decoder structure, generates unit tokens.
- These unit tokens are then passed through the final vocoder to produce the actual speech.
-
This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The original code can be found [here](https://github.com/facebookresearch/seamless_communication).
## SeamlessM4TModel
@@ -150,19 +146,16 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
[[autodoc]] SeamlessM4TModel
- generate
-
## SeamlessM4TForTextToSpeech
[[autodoc]] SeamlessM4TForTextToSpeech
- generate
-
## SeamlessM4TForSpeechToSpeech
[[autodoc]] SeamlessM4TForSpeechToSpeech
- generate
-
## SeamlessM4TForTextToText
[[autodoc]] transformers.SeamlessM4TForTextToText
@@ -179,7 +172,6 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
[[autodoc]] SeamlessM4TConfig
-
## SeamlessM4TTokenizer
[[autodoc]] SeamlessM4TTokenizer
@@ -189,7 +181,6 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
- create_token_type_ids_from_sequences
- save_vocabulary
-
## SeamlessM4TTokenizerFast
[[autodoc]] SeamlessM4TTokenizerFast
@@ -209,7 +200,6 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
[[autodoc]] SeamlessM4TCodeHifiGan
-
## SeamlessM4THifiGan
[[autodoc]] SeamlessM4THifiGan
@@ -221,5 +211,3 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
## SeamlessM4TTextToUnitForConditionalGeneration
[[autodoc]] SeamlessM4TTextToUnitForConditionalGeneration
-
-
diff --git a/docs/source/en/model_doc/seamless_m4t_v2.md b/docs/source/en/model_doc/seamless_m4t_v2.md
index 8a4ab82d2e98..4a32199243ab 100644
--- a/docs/source/en/model_doc/seamless_m4t_v2.md
+++ b/docs/source/en/model_doc/seamless_m4t_v2.md
@@ -35,7 +35,7 @@ SeamlessM4T-v2 enables multiple tasks without relying on separate models:
The abstract from the paper is the following:
-*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one’s voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
+*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one's voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
## Usage
@@ -67,7 +67,6 @@ Here is how to use the processor to process text and audio:
>>> text_inputs = processor(text = "Hello, my dog is cute", src_lang="eng", return_tensors="pt")
```
-
### Speech
[`SeamlessM4Tv2Model`] can *seamlessly* generate text or speech with few or no changes. Let's target Russian voice translation:
@@ -84,7 +83,7 @@ With basically the same code, I've translated English text and Arabic speech to
Similarly, you can generate translated text from audio files or from text with the same model. You only have to pass `generate_speech=False` to [`SeamlessM4Tv2Model.generate`].
This time, let's translate to French.
-```python
+```python
>>> # from audio
>>> output_tokens = model.generate(**audio_inputs, tgt_lang="fra", generate_speech=False)
>>> translated_text_from_audio = processor.decode(output_tokens[0].tolist()[0], skip_special_tokens=True)
@@ -96,11 +95,10 @@ This time, let's translate to French.
### Tips
-
#### 1. Use dedicated models
[`SeamlessM4Tv2Model`] is transformers top level model to generate speech and text, but you can also use dedicated models that perform the task without additional components, thus reducing the memory footprint.
-For example, you can replace the audio-to-audio generation snippet with the model dedicated to the S2ST task, the rest is exactly the same code:
+For example, you can replace the audio-to-audio generation snippet with the model dedicated to the S2ST task, the rest is exactly the same code:
```python
>>> from transformers import SeamlessM4Tv2ForSpeechToSpeech
@@ -141,6 +139,7 @@ The architecture of this new version differs from the first in a few aspects:
#### Improvements on the second-pass model
The second seq2seq model, named text-to-unit model, is now non-auto regressive, meaning that it computes units in a **single forward pass**. This achievement is made possible by:
+
- the use of **character-level embeddings**, meaning that each character of the predicted translated text has its own embeddings, which are then used to predict the unit tokens.
- the use of an intermediate duration predictor, that predicts speech duration at the **character-level** on the predicted translated text.
- the use of a new text-to-unit decoder mixing convolutions and self-attention to handle longer context.
@@ -148,6 +147,7 @@ The second seq2seq model, named text-to-unit model, is now non-auto regressive,
#### Difference in the speech encoder
The speech encoder, which is used during the first-pass generation process to predict the translated text, differs mainly from the previous speech encoder through these mechanisms:
+
- the use of chunked attention mask to prevent attention across chunks, ensuring that each position attends only to positions within its own chunk and a fixed number of previous chunks.
- the use of relative position embeddings which only considers distance between sequence elements rather than absolute positions. Please refer to [Self-Attentionwith Relative Position Representations (Shaw et al.)](https://huggingface.co/papers/1803.02155) for more details.
- the use of a causal depth-wise convolution instead of a non-causal one.
@@ -161,7 +161,6 @@ Here's how the generation process works:
- If speech generation is required, the second seq2seq model, generates unit tokens in an non auto-regressive way.
- These unit tokens are then passed through the final vocoder to produce the actual speech.
-
This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The original code can be found [here](https://github.com/facebookresearch/seamless_communication).
## SeamlessM4Tv2Model
@@ -169,19 +168,16 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
[[autodoc]] SeamlessM4Tv2Model
- generate
-
## SeamlessM4Tv2ForTextToSpeech
[[autodoc]] SeamlessM4Tv2ForTextToSpeech
- generate
-
## SeamlessM4Tv2ForSpeechToSpeech
[[autodoc]] SeamlessM4Tv2ForSpeechToSpeech
- generate
-
## SeamlessM4Tv2ForTextToText
[[autodoc]] transformers.SeamlessM4Tv2ForTextToText
diff --git a/docs/source/en/model_doc/seed_oss.md b/docs/source/en/model_doc/seed_oss.md
index 0f0dacb2be90..dbcddcb5f2c7 100644
--- a/docs/source/en/model_doc/seed_oss.md
+++ b/docs/source/en/model_doc/seed_oss.md
@@ -1,17 +1,20 @@
-
+
+*This model was released on {release_date} and added to Hugging Face Transformers on 2025-08-22.*
# SeedOss
@@ -54,4 +57,4 @@ To be released with the official model launch.
## SeedOssForQuestionAnswering
[[autodoc]] SeedOssForQuestionAnswering
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/segformer.md b/docs/source/en/model_doc/segformer.md
index 756c98d45f08..a6b407e58793 100644
--- a/docs/source/en/model_doc/segformer.md
+++ b/docs/source/en/model_doc/segformer.md
@@ -71,8 +71,6 @@ logits = outputs.logits # shape [batch, num_labels, height, width]
-
-
## Notes
- SegFormer works with **any input size**, padding inputs to be divisible by `config.patch_sizes`.
diff --git a/docs/source/en/model_doc/seggpt.md b/docs/source/en/model_doc/seggpt.md
index 9e8c08cf2d2e..356b0f7abcf6 100644
--- a/docs/source/en/model_doc/seggpt.md
+++ b/docs/source/en/model_doc/seggpt.md
@@ -30,6 +30,7 @@ The abstract from the paper is the following:
*We present SegGPT, a generalist model for segmenting everything in context. We unify various segmentation tasks into a generalist in-context learning framework that accommodates different kinds of segmentation data by transforming them into the same format of images. The training of SegGPT is formulated as an in-context coloring problem with random color mapping for each data sample. The objective is to accomplish diverse tasks according to the context, rather than relying on specific colors. After training, SegGPT can perform arbitrary segmentation tasks in images or videos via in-context inference, such as object instance, stuff, part, contour, and text. SegGPT is evaluated on a broad range of tasks, including few-shot semantic segmentation, video object segmentation, semantic segmentation, and panoptic segmentation. Our results show strong capabilities in segmenting in-domain and out-of*
Tips:
+
- One can use [`SegGptImageProcessor`] to prepare image input, prompt and mask to the model.
- One can either use segmentation maps or RGB images as prompt masks. If using the latter make sure to set `do_convert_rgb=False` in the `preprocess` method.
- It's highly advisable to pass `num_labels` when using `segmentation_maps` (not considering background) during preprocessing and postprocessing with [`SegGptImageProcessor`] for your use case.
@@ -74,7 +75,6 @@ mask = image_processor.post_process_semantic_segmentation(outputs, target_sizes,
This model was contributed by [EduardoPacheco](https://huggingface.co/EduardoPacheco).
The original code can be found [here]([(https://github.com/baaivision/Painter/tree/main)).
-
## SegGptConfig
[[autodoc]] SegGptConfig
diff --git a/docs/source/en/model_doc/shieldgemma2.md b/docs/source/en/model_doc/shieldgemma2.md
index 99ffde6288ff..6a67c2d61b5a 100644
--- a/docs/source/en/model_doc/shieldgemma2.md
+++ b/docs/source/en/model_doc/shieldgemma2.md
@@ -22,9 +22,9 @@ rendered properly in your Markdown viewer.
The ShieldGemma 2 model was proposed in a [technical report](https://huggingface.co/papers/2504.01081) by Google. ShieldGemma 2, built on [Gemma 3](https://ai.google.dev/gemma/docs/core/model_card_3), is a 4 billion (4B) parameter model that checks the safety of both synthetic and natural images against key categories to help you build robust datasets and models. With this addition to the Gemma family of models, researchers and developers can now easily minimize the risk of harmful content in their models across key areas of harm as defined below:
-- No Sexually Explicit content: The image shall not contain content that depicts explicit or graphic sexual acts (e.g., pornography, erotic nudity, depictions of rape or sexual assault).
-- No Dangerous Content: The image shall not contain content that facilitates or encourages activities that could cause real-world harm (e.g., building firearms and explosive devices, promotion of terrorism, instructions for suicide).
-- No Violence/Gore content: The image shall not contain content that depicts shocking, sensational, or gratuitous violence (e.g., excessive blood and gore, gratuitous violence against animals, extreme injury or moment of death).
+- No Sexually Explicit content: The image shall not contain content that depicts explicit or graphic sexual acts (e.g., pornography, erotic nudity, depictions of rape or sexual assault).
+- No Dangerous Content: The image shall not contain content that facilitates or encourages activities that could cause real-world harm (e.g., building firearms and explosive devices, promotion of terrorism, instructions for suicide).
+- No Violence/Gore content: The image shall not contain content that depicts shocking, sensational, or gratuitous violence (e.g., excessive blood and gore, gratuitous violence against animals, extreme injury or moment of death).
We recommend using ShieldGemma 2 as an input filter to vision language models, or as an output filter of image generation systems. To train a robust image safety model, we curated training datasets of natural and synthetic images and instruction-tuned Gemma 3 to demonstrate strong performance.
@@ -86,7 +86,6 @@ output = model(**inputs)
print(output.probabilities)
```
-
## ShieldGemma2Processor
[[autodoc]] ShieldGemma2Processor
diff --git a/docs/source/en/model_doc/siglip.md b/docs/source/en/model_doc/siglip.md
index c0eb9a8ac6b5..bf9c0a460348 100644
--- a/docs/source/en/model_doc/siglip.md
+++ b/docs/source/en/model_doc/siglip.md
@@ -31,7 +31,6 @@ Unlike CLIP, SigLIP employs a pairwise sigmoid loss on image-text pairs during t
You can find all the original SigLIP checkpoints under the [SigLIP](https://huggingface.co/collections/google/siglip-659d5e62f0ae1a57ae0e83ba) collection.
-
> [!TIP]
> Click on the SigLIP models in the right sidebar for more examples of how to apply SigLIP to different image and text tasks.
@@ -107,12 +106,14 @@ logits_per_image = outputs.logits_per_image
probs = torch.sigmoid(logits_per_image)
print(f"{probs[0][0]:.1%} that image 0 is '{candidate_labels[0]}'")
```
+
## Notes
- Training is supported for DDP and FSDP on single-node multi-GPU setups. However, it does not use [torch.distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) utilities which may limit the scalability of batch size.
- When using the standalone [`SiglipTokenizer`] or [`SiglipProcessor`], make sure to pass `padding="max_length"` because that is how the model was trained.
- To get the same results as the [`Pipeline`], a prompt template of `"This is a photo of {label}."` should be passed to the processor.
- Toggle the `attn_implementation` parameter to either `"sdpa"` or `"flash_attention_2"` to use a more memory-efficient attention.
+
```py
# pip install -U flash-attn --no-build-isolation
@@ -126,7 +127,6 @@ print(f"{probs[0][0]:.1%} that image 0 is '{candidate_labels[0]}'")
)
```
-
## SiglipConfig
[[autodoc]] SiglipConfig
@@ -179,7 +179,6 @@ print(f"{probs[0][0]:.1%} that image 0 is '{candidate_labels[0]}'")
[[autodoc]] SiglipVisionModel
- forward
-
## SiglipForImageClassification
[[autodoc]] SiglipForImageClassification
diff --git a/docs/source/en/model_doc/siglip2.md b/docs/source/en/model_doc/siglip2.md
index f2684c6defcf..6a058f8907a4 100644
--- a/docs/source/en/model_doc/siglip2.md
+++ b/docs/source/en/model_doc/siglip2.md
@@ -32,7 +32,6 @@ rendered properly in your Markdown viewer.
- NaFlex supports different resolutions and maintains the native image aspect ratio
- FixRes supports fixed resolutions and is backwards compatible with [SigLIP](./siglip)
-
You can find all the original SigLIP2 checkpoints under the [SigLIP2](https://huggingface.co/collections/google/siglip2-67b5dcef38c175486e240107) collection.
> [!TIP]
@@ -157,6 +156,7 @@ print(f"{probs[0][0]:.1%} that image 0 is '{candidate_labels[0]}'")
NaFlex resizes the input image so the height and width are multiples of the patch size after resizing. It keeps the aspect ratio distortion as low as possible and produces a sequence length of at most the desired target sequence length (`max_num_patches`). After resizing, the image is split into a sequence of patches and a mask with padding information is added.
- Toggle the `attn_implementation` parameter to either `"sdpa"` or `"flash_attention_2"` to use a more memory-efficient attention.
+
```py
# pip install -U flash-attn --no-build-isolation
@@ -169,6 +169,7 @@ print(f"{probs[0][0]:.1%} that image 0 is '{candidate_labels[0]}'")
device_map=device,
)
```
+
## Siglip2Config
[[autodoc]] Siglip2Config
diff --git a/docs/source/en/model_doc/smollm3.md b/docs/source/en/model_doc/smollm3.md
index da98a15e33b5..db2ddd336013 100644
--- a/docs/source/en/model_doc/smollm3.md
+++ b/docs/source/en/model_doc/smollm3.md
@@ -139,7 +139,6 @@ outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```
-
## Notes
- Ensure your Transformers library version is up-to-date. SmolLM3 requires Transformers>=4.53.0 for full support.
diff --git a/docs/source/en/model_doc/smolvlm.md b/docs/source/en/model_doc/smolvlm.md
index c9a886ac8769..61400bac177b 100644
--- a/docs/source/en/model_doc/smolvlm.md
+++ b/docs/source/en/model_doc/smolvlm.md
@@ -38,7 +38,8 @@ Videos should not be upsampled.
If `do_resize` is set to `True`, the model resizes images so that the longest edge is 4*512 pixels by default.
The default resizing behavior can be customized by passing a dictionary to the `size` parameter. For example, `{"longest_edge": 4 * 512}` is the default, but you can change it to a different value if needed.
-Here’s how to control resizing and set a custom size:
+Here's how to control resizing and set a custom size:
+
```python
image_processor = SmolVLMImageProcessor(do_resize=True, size={"longest_edge": 2 * 512}, max_image_size=512)
```
@@ -47,8 +48,6 @@ Additionally, the `max_image_size` parameter, which controls the size of each sq
This model was contributed by [orrzohar](https://huggingface.co/orrzohar).
-
-
## Usage example
### Single Media inference
diff --git a/docs/source/en/model_doc/stablelm.md b/docs/source/en/model_doc/stablelm.md
index 29f32a0004e2..e47598a8f852 100644
--- a/docs/source/en/model_doc/stablelm.md
+++ b/docs/source/en/model_doc/stablelm.md
@@ -92,7 +92,6 @@ Now, to run the model with Flash Attention 2, refer to the snippet below:
['The weather is always wonderful in Costa Rica, which makes it a prime destination for retirees. That’s where the Pensionado program comes in, offering']
```
-
## StableLmConfig
[[autodoc]] StableLmConfig
diff --git a/docs/source/en/model_doc/starcoder2.md b/docs/source/en/model_doc/starcoder2.md
index 2d27aed399cd..b67e5dedd2cc 100644
--- a/docs/source/en/model_doc/starcoder2.md
+++ b/docs/source/en/model_doc/starcoder2.md
@@ -34,7 +34,7 @@ The abstract of the paper is the following:
## License
The models are licensed under the [BigCode OpenRAIL-M v1 license agreement](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement).
-
+
## Usage tips
The StarCoder2 models can be found in the [HuggingFace hub](https://huggingface.co/collections/bigcode/starcoder2-65de6da6e87db3383572be1a). You can find some examples for inference and fine-tuning in StarCoder2's [GitHub repo](https://github.com/bigcode-project/starcoder2).
diff --git a/docs/source/en/model_doc/superglue.md b/docs/source/en/model_doc/superglue.md
index 3e42b002ec6a..061f3ec2b9fb 100644
--- a/docs/source/en/model_doc/superglue.md
+++ b/docs/source/en/model_doc/superglue.md
@@ -143,10 +143,9 @@ processed_outputs = processor.post_process_keypoint_matching(outputs, image_size
## SuperGlueImageProcessor
[[autodoc]] SuperGlueImageProcessor
-
-- preprocess
-- post_process_keypoint_matching
-- visualize_keypoint_matching
+ - preprocess
+ - post_process_keypoint_matching
+ - visualize_keypoint_matching
@@ -157,4 +156,4 @@ processed_outputs = processor.post_process_keypoint_matching(outputs, image_size
- forward
-
\ No newline at end of file
+
diff --git a/docs/source/en/model_doc/superpoint.md b/docs/source/en/model_doc/superpoint.md
index b86f7fd4aa77..3efd5ecf90f2 100644
--- a/docs/source/en/model_doc/superpoint.md
+++ b/docs/source/en/model_doc/superpoint.md
@@ -33,8 +33,6 @@ You can find all the original SuperPoint checkpoints under the [Magic Leap Commu
>
> Click on the SuperPoint models in the right sidebar for more examples of how to apply SuperPoint to different computer vision tasks.
-
-
The example below demonstrates how to detect interest points in an image with the [`AutoModel`] class.
@@ -101,6 +99,7 @@ processed_outputs = processor.post_process_keypoint_detection(outputs, [image_si
```
- You can then print the keypoints on the image of your choice to visualize the result:
+
```py
import matplotlib.pyplot as plt
plt.axis("off")
@@ -130,16 +129,15 @@ processed_outputs = processor.post_process_keypoint_detection(outputs, [image_si
## SuperPointImageProcessor
[[autodoc]] SuperPointImageProcessor
-
-- preprocess
+ - preprocess
## SuperPointImageProcessorFast
[[autodoc]] SuperPointImageProcessorFast
-- preprocess
-- post_process_keypoint_detection
+ - preprocess
+ - post_process_keypoint_detection
## SuperPointForKeypointDetection
[[autodoc]] SuperPointForKeypointDetection
-- forward
+ - forward
diff --git a/docs/source/en/model_doc/swin.md b/docs/source/en/model_doc/swin.md
index f6a994ef69bc..81142f6c4111 100644
--- a/docs/source/en/model_doc/swin.md
+++ b/docs/source/en/model_doc/swin.md
@@ -47,6 +47,7 @@ pipeline = pipeline(
)
pipeline("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg")
```
+
@@ -79,6 +80,7 @@ class_labels = model.config.id2label
predicted_class_label = class_labels[predicted_class_id]
print(f"The predicted class label is: {predicted_class_label}")
```
+
diff --git a/docs/source/en/model_doc/swinv2.md b/docs/source/en/model_doc/swinv2.md
index 507b79fc7cf1..0dc008767ac3 100644
--- a/docs/source/en/model_doc/swinv2.md
+++ b/docs/source/en/model_doc/swinv2.md
@@ -81,7 +81,7 @@ print(f"The predicted class label is: {predicted_class_label}")
## Notes
-- Swin Transformer V2 can pad the inputs for any input height and width divisible by `32`.
+- Swin Transformer V2 can pad the inputs for any input height and width divisible by `32`.
- Swin Transformer V2 can be used as a [backbone](../backbones). When `output_hidden_states = True`, it outputs both `hidden_states` and `reshaped_hidden_states`. The `reshaped_hidden_states` have a shape of `(batch, num_channels, height, width)` rather than `(batch_size, sequence_length, num_channels)`.
## Swinv2Config
diff --git a/docs/source/en/model_doc/switch_transformers.md b/docs/source/en/model_doc/switch_transformers.md
index efa6bd499dbc..5eb27a9e7d8c 100644
--- a/docs/source/en/model_doc/switch_transformers.md
+++ b/docs/source/en/model_doc/switch_transformers.md
@@ -27,7 +27,6 @@ rendered properly in your Markdown viewer.
You can find all the original Switch Transformers checkpoints under the [Switch Transformer](https://huggingface.co/collections/google/switch-transformers-release-6548c35c6507968374b56d1f) collection.
-
> [!TIP]
> This model was contributed by [ybelkada](https://huggingface.co/ybelkada) and [ArthurZ](https://huggingface.co/ArthurZ).
>
@@ -99,7 +98,6 @@ outputs = model.generate(input_ids)
print(tokenizer.decode(outputs[0]))
```
-
## SwitchTransformersConfig
[[autodoc]] SwitchTransformersConfig
diff --git a/docs/source/en/model_doc/t5gemma.md b/docs/source/en/model_doc/t5gemma.md
index aa8d3b7880ed..80880cf6559d 100644
--- a/docs/source/en/model_doc/t5gemma.md
+++ b/docs/source/en/model_doc/t5gemma.md
@@ -39,7 +39,6 @@ The example below demonstrates how to chat with the model with [`Pipeline`] or t
-
```python
import torch
from transformers import pipeline
@@ -86,9 +85,10 @@ print(tokenizer.decode(outputs[0]))
-```
+```bash
echo -e "Write me a poem about Machine Learning. Answer:" | transformers run --task text2text-generation --model google/t5gemma-2b-2b-prefixlm --device 0
```
+
diff --git a/docs/source/en/model_doc/t5v1.1.md b/docs/source/en/model_doc/t5v1.1.md
index 4ad072addcc0..62787d5f9d62 100644
--- a/docs/source/en/model_doc/t5v1.1.md
+++ b/docs/source/en/model_doc/t5v1.1.md
@@ -68,7 +68,6 @@ Google has released the following variants:
- [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl).
-
Refer to [T5's documentation page](t5) for all API reference, tips, code examples and notebooks.
diff --git a/docs/source/en/model_doc/table-transformer.md b/docs/source/en/model_doc/table-transformer.md
index b35df2aec311..c982d3059072 100644
--- a/docs/source/en/model_doc/table-transformer.md
+++ b/docs/source/en/model_doc/table-transformer.md
@@ -43,8 +43,8 @@ alt="drawing" width="600"/>
Table detection and table structure recognition clarified. Taken from the original paper.
-The authors released 2 models, one for [table detection](https://huggingface.co/microsoft/table-transformer-detection) in
-documents, one for [table structure recognition](https://huggingface.co/microsoft/table-transformer-structure-recognition)
+The authors released 2 models, one for [table detection](https://huggingface.co/microsoft/table-transformer-detection) in
+documents, one for [table structure recognition](https://huggingface.co/microsoft/table-transformer-structure-recognition)
(the task of recognizing the individual rows, columns etc. in a table).
This model was contributed by [nielsr](https://huggingface.co/nielsr). The original code can be
diff --git a/docs/source/en/model_doc/tapas.md b/docs/source/en/model_doc/tapas.md
index 4dfac5edce37..09c624c7fb7e 100644
--- a/docs/source/en/model_doc/tapas.md
+++ b/docs/source/en/model_doc/tapas.md
@@ -30,6 +30,7 @@ token types that encode tabular structure. TAPAS is pre-trained on the masked la
millions of tables from English Wikipedia and corresponding texts.
For question answering, TAPAS has 2 heads on top: a cell selection head and an aggregation head, for (optionally) performing aggregations (such as counting or summing) among selected cells. TAPAS has been fine-tuned on several datasets:
+
- [SQA](https://www.microsoft.com/en-us/download/details.aspx?id=54253) (Sequential Question Answering by Microsoft)
- [WTQ](https://github.com/ppasupat/WikiTableQuestions) (Wiki Table Questions by Stanford University)
- [WikiSQL](https://github.com/salesforce/WikiSQL) (by Salesforce).
@@ -76,7 +77,6 @@ To summarize:
| Weak supervision for aggregation | WTQ | Questions might involve aggregation, and the model must learn this given only the answer as supervision |
| Strong supervision for aggregation | WikiSQL-supervised | Questions might involve aggregation, and the model must learn this given the gold aggregation operator |
-
Initializing a model with a pre-trained base and randomly initialized classification heads from the hub can be done as shown below.
```py
@@ -105,7 +105,6 @@ Of course, you don't necessarily have to follow one of these three ways in which
>>> model = TapasForQuestionAnswering.from_pretrained("google/tapas-base", config=config)
```
-
What you can also do is start from an already fine-tuned checkpoint. A note here is that the already fine-tuned checkpoint on WTQ has some issues due to the L2-loss which is somewhat brittle. See [here](https://github.com/google-research/tapas/issues/91#issuecomment-735719340) for more info.
For a list of all pre-trained and fine-tuned TAPAS checkpoints available on HuggingFace's hub, see [here](https://huggingface.co/models?search=tapas).
@@ -128,7 +127,6 @@ The tables themselves should be present in a folder, each table being a separate
**STEP 3: Convert your data into tensors using TapasTokenizer**
-
Third, given that you've prepared your data in this TSV/CSV format (and corresponding CSV files containing the tabular data), you can then use [`TapasTokenizer`] to convert table-question pairs into `input_ids`, `attention_mask`, `token_type_ids` and so on. Again, based on which of the three cases you picked above, [`TapasForQuestionAnswering`] requires different
inputs to be fine-tuned:
@@ -214,13 +212,11 @@ Of course, this only shows how to encode a single training example. It is advise
>>> train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32)
```
-
Note that here, we encode each table-question pair independently. This is fine as long as your dataset is **not conversational**. In case your dataset involves conversational questions (such as in SQA), then you should first group together the `queries`, `answer_coordinates` and `answer_text` per table (in the order of their `position`
index) and batch encode each table with its questions. This will make sure that the `prev_labels` token types (see docs of [`TapasTokenizer`]) are set correctly. See [this notebook](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/TAPAS/Fine_tuning_TapasForQuestionAnswering_on_SQA.ipynb) for more info.
**STEP 4: Train (fine-tune) the model
-
You can then fine-tune [`TapasForQuestionAnswering`] as follows (shown here for the weak supervision for aggregation case):
```py
@@ -272,10 +268,8 @@ You can then fine-tune [`TapasForQuestionAnswering`] as follows (shown here for
... optimizer.step()
```
-
## Usage: inference
-
Here we explain how you can use [`TapasForQuestionAnswering`] for inference (i.e. making predictions on new data). For inference, only `input_ids`, `attention_mask` and `token_type_ids` (which you can obtain using [`TapasTokenizer`]) have to be provided to the model to obtain the logits. Next, you can use the handy [`~models.tapas.tokenization_tapas.convert_logits_to_predictions`] method to convert these into predicted coordinates and optional aggregation indices.
However, note that inference is **different** depending on whether or not the setup is conversational. In a non-conversational set-up, inference can be done in parallel on all table-question pairs of a batch. Here's an example of that:
@@ -333,7 +327,6 @@ What is the total number of movies?
Predicted answer: SUM > 87, 53, 69
```
-
In case of a conversational set-up, then each table-question pair must be provided **sequentially** to the model, such that the `prev_labels` token types can be overwritten by the predicted `labels` of the previous table-question pair. Again, more info can be found in [this notebook](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/TAPAS/Fine_tuning_TapasForQuestionAnswering_on_SQA.ipynb).
## Resources
diff --git a/docs/source/en/model_doc/tapex.md b/docs/source/en/model_doc/tapex.md
index 0a10826ee1af..606d8940c4ed 100644
--- a/docs/source/en/model_doc/tapex.md
+++ b/docs/source/en/model_doc/tapex.md
@@ -37,6 +37,7 @@ Bei Chen, Jiaqi Guo, Morteza Ziyadi, Zeqi Lin, Weizhu Chen, Jian-Guang Lou. TAPE
which it can be fine-tuned to answer natural language questions related to tabular data, as well as performing table fact checking.
TAPEX has been fine-tuned on several datasets:
+
- [SQA](https://www.microsoft.com/en-us/download/details.aspx?id=54253) (Sequential Question Answering by Microsoft)
- [WTQ](https://github.com/ppasupat/WikiTableQuestions) (Wiki Table Questions by Stanford University)
- [WikiSQL](https://github.com/salesforce/WikiSQL) (by Salesforce)
diff --git a/docs/source/en/model_doc/textnet.md b/docs/source/en/model_doc/textnet.md
index 9c29a8b16bee..c986b17dbff0 100644
--- a/docs/source/en/model_doc/textnet.md
+++ b/docs/source/en/model_doc/textnet.md
@@ -34,7 +34,7 @@ This model was contributed by [Raghavan](https://huggingface.co/Raghavan), [jade
## Usage tips
-TextNet is mainly used as a backbone network for the architecture search of text detection. Each stage of the backbone network is comprised of a stride-2 convolution and searchable blocks.
+TextNet is mainly used as a backbone network for the architecture search of text detection. Each stage of the backbone network is comprised of a stride-2 convolution and searchable blocks.
Specifically, we present a layer-level candidate set, defined as {conv3×3, conv1×3, conv3×1, identity}. As the 1×3 and 3×1 convolutions have asymmetric kernels and oriented structure priors, they may help to capture the features of extreme aspect-ratio and rotated text lines.
TextNet is the backbone for Fast, but can also be used as an efficient text/image classification, we add a `TextNetForImageClassification` as is it would allow people to train an image classifier on top of the pre-trained textnet weights
@@ -62,4 +62,3 @@ TextNet is the backbone for Fast, but can also be used as an efficient text/imag
[[autodoc]] TextNetForImageClassification
- forward
-
diff --git a/docs/source/en/model_doc/time_series_transformer.md b/docs/source/en/model_doc/time_series_transformer.md
index c38671f00fb3..36a68af80ca8 100644
--- a/docs/source/en/model_doc/time_series_transformer.md
+++ b/docs/source/en/model_doc/time_series_transformer.md
@@ -35,16 +35,16 @@ point forecasting model. This means that the model learns a distribution, from w
and a decoder, which predicts a `prediction_length` of time series values into the future (called `future_values`). During training, one needs to provide
pairs of (`past_values` and `future_values`) to the model.
- In addition to the raw (`past_values` and `future_values`), one typically provides additional features to the model. These can be the following:
- - `past_time_features`: temporal features which the model will add to `past_values`. These serve as "positional encodings" for the Transformer encoder.
+ - `past_time_features`: temporal features which the model will add to `past_values`. These serve as "positional encodings" for the Transformer encoder.
Examples are "day of the month", "month of the year", etc. as scalar values (and then stacked together as a vector).
e.g. if a given time-series value was obtained on the 11th of August, then one could have [11, 8] as time feature vector (11 being "day of the month", 8 being "month of the year").
- - `future_time_features`: temporal features which the model will add to `future_values`. These serve as "positional encodings" for the Transformer decoder.
+ - `future_time_features`: temporal features which the model will add to `future_values`. These serve as "positional encodings" for the Transformer decoder.
Examples are "day of the month", "month of the year", etc. as scalar values (and then stacked together as a vector).
e.g. if a given time-series value was obtained on the 11th of August, then one could have [11, 8] as time feature vector (11 being "day of the month", 8 being "month of the year").
- - `static_categorical_features`: categorical features which are static over time (i.e., have the same value for all `past_values` and `future_values`).
+ - `static_categorical_features`: categorical features which are static over time (i.e., have the same value for all `past_values` and `future_values`).
An example here is the store ID or region ID that identifies a given time-series.
Note that these features need to be known for ALL data points (also those in the future).
- - `static_real_features`: real-valued features which are static over time (i.e., have the same value for all `past_values` and `future_values`).
+ - `static_real_features`: real-valued features which are static over time (i.e., have the same value for all `past_values` and `future_values`).
An example here is the image representation of the product for which you have the time-series values (like the [ResNet](resnet) embedding of a "shoe" picture,
if your time-series is about the sales of shoes).
Note that these features need to be known for ALL data points (also those in the future).
@@ -61,7 +61,6 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h
- Check out the Time Series Transformer blog-post in HuggingFace blog: [Probabilistic Time Series Forecasting with 🤗 Transformers](https://huggingface.co/blog/time-series-transformers)
-
## TimeSeriesTransformerConfig
[[autodoc]] TimeSeriesTransformerConfig
diff --git a/docs/source/en/model_doc/timesfm.md b/docs/source/en/model_doc/timesfm.md
index 83dee48e71be..e8938202ee9e 100644
--- a/docs/source/en/model_doc/timesfm.md
+++ b/docs/source/en/model_doc/timesfm.md
@@ -25,16 +25,13 @@ rendered properly in your Markdown viewer.
TimesFM (Time Series Foundation Model) is a pretrained time-series foundation model proposed in [A decoder-only foundation model for time-series forecasting](https://huggingface.co/papers/2310.10688) by Abhimanyu Das, Weihao Kong, Rajat Sen, and Yichen Zhou. It is a decoder only model that uses non-overlapping patches of time-series data as input and outputs some output patch length prediction in an autoregressive fashion.
-
The abstract from the paper is the following:
*Motivated by recent advances in large language models for Natural Language Processing (NLP), we design a time-series foundation model for forecasting whose out-of-the-box zero-shot performance on a variety of public datasets comes close to the accuracy of state-of-the-art supervised forecasting models for each individual dataset. Our model is based on pretraining a patched-decoder style attention model on a large time-series corpus, and can work well across different forecasting history lengths, prediction lengths and temporal granularities.*
-
This model was contributed by [kashif](https://huggingface.co/kashif).
The original code can be found [here](https://github.com/google-research/timesfm).
-
To use the model:
```python
diff --git a/docs/source/en/model_doc/timesformer.md b/docs/source/en/model_doc/timesformer.md
index 59e9ee71817d..1d87158d72e1 100644
--- a/docs/source/en/model_doc/timesformer.md
+++ b/docs/source/en/model_doc/timesformer.md
@@ -54,4 +54,4 @@ the number of input frames per clip changes based on the model size so you shoul
## TimesformerForVideoClassification
[[autodoc]] TimesformerForVideoClassification
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/transfo-xl.md b/docs/source/en/model_doc/transfo-xl.md
index 5d9b92f7946f..0bd1b0f57e1d 100644
--- a/docs/source/en/model_doc/transfo-xl.md
+++ b/docs/source/en/model_doc/transfo-xl.md
@@ -90,7 +90,6 @@ This model was contributed by [thomwolf](https://huggingface.co/thomwolf). The o
- Basically, the hidden states of the previous segment are concatenated to the current input to compute the attention scores. This allows the model to pay attention to information that was in the previous segment as well as the current one. By stacking multiple attention layers, the receptive field can be increased to multiple previous segments.
- This changes the positional embeddings to positional relative embeddings (as the regular positional embeddings would give the same results in the current input and the current hidden state at a given position) and needs to make some adjustments in the way attention scores are computed.
-
TransformerXL does **not** work with *torch.nn.DataParallel* due to a bug in PyTorch, see [issue #36035](https://github.com/pytorch/pytorch/issues/36035)
diff --git a/docs/source/en/model_doc/trocr.md b/docs/source/en/model_doc/trocr.md
index 6346977dafa1..da5c71edde36 100644
--- a/docs/source/en/model_doc/trocr.md
+++ b/docs/source/en/model_doc/trocr.md
@@ -14,8 +14,6 @@ rendered properly in your Markdown viewer.
specific language governing permissions and limitations under the License. -->
*This model was released on 2021-09-21 and added to Hugging Face Transformers on 2021-10-13.*
-
-
@@ -32,13 +30,11 @@ You can find all the original TrOCR checkpoints under the [Microsoft](https://hu
alt="drawing" width="600"/>
TrOCR architecture. Taken from the original paper.
-
> [!TIP]
> This model was contributed by [nielsr](https://huggingface.co/nielsr).
>
> Click on the TrOCR models in the right sidebar for more examples of how to apply TrOCR to different image and text tasks.
-
The example below demonstrates how to perform optical character recognition (OCR) with the [`AutoModel`] class.
@@ -113,7 +109,6 @@ print(generated_text)
- A notebook on [inference with TrOCR](https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/TrOCR/Inference_with_TrOCR_%2B_Gradio_demo.ipynb) and Gradio demo.
- A notebook on [evaluating TrOCR on the IAM test set](https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/TrOCR/Evaluating_TrOCR_base_handwritten_on_the_IAM_test_set.ipynb).
-
## TrOCRConfig
[[autodoc]] TrOCRConfig
diff --git a/docs/source/en/model_doc/tvp.md b/docs/source/en/model_doc/tvp.md
index 49a538ffa8c4..2df4da02555a 100644
--- a/docs/source/en/model_doc/tvp.md
+++ b/docs/source/en/model_doc/tvp.md
@@ -47,6 +47,7 @@ The [`TvpProcessor`] wraps [`BertTokenizer`] and [`TvpImageProcessor`] into a si
encode the text and prepare the images respectively.
The following example shows how to run temporal video grounding using [`TvpProcessor`] and [`TvpForVideoGrounding`].
+
```python
import av
import cv2
@@ -165,7 +166,6 @@ Tips:
- Checkpoints for pre-trained [tvp-base](https://huggingface.co/Intel/tvp-base) is released.
- Please refer to [Table 2](https://huggingface.co/papers/2303.04995) for TVP's performance on Temporal Video Grounding task.
-
## TvpConfig
[[autodoc]] TvpConfig
diff --git a/docs/source/en/model_doc/udop.md b/docs/source/en/model_doc/udop.md
index eb400cc39d5f..cc370accf3e3 100644
--- a/docs/source/en/model_doc/udop.md
+++ b/docs/source/en/model_doc/udop.md
@@ -115,4 +115,4 @@ to fine-tune UDOP on a custom dataset as well as inference. 🌎
## UdopEncoderModel
[[autodoc]] UdopEncoderModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/umt5.md b/docs/source/en/model_doc/umt5.md
index 349dcecf03cc..784cc9974df1 100644
--- a/docs/source/en/model_doc/umt5.md
+++ b/docs/source/en/model_doc/umt5.md
@@ -39,7 +39,7 @@ Google has released the following variants:
This model was contributed by [agemagician](https://huggingface.co/agemagician) and [stefan-it](https://huggingface.co/stefan-it). The original code can be
found [here](https://github.com/google-research/t5x).
-## Usage tips
+## Usage tips
- UMT5 was only pre-trained on [mC4](https://huggingface.co/datasets/mc4) excluding any supervised training.
Therefore, this model has to be fine-tuned before it is usable on a downstream task, unlike the original T5 model.
@@ -67,7 +67,7 @@ The conversion script is also different because the model was saved in t5x's lat
['nyone who drink a alcohol A A. This I']
```
-
+
Refer to [T5's documentation page](t5) for more tips, code examples and notebooks.
@@ -105,4 +105,3 @@ Refer to [T5's documentation page](t5) for more tips, code examples and notebook
[[autodoc]] UMT5ForQuestionAnswering
- forward
-
diff --git a/docs/source/en/model_doc/univnet.md b/docs/source/en/model_doc/univnet.md
index e20bc5c405e8..4329846ab7f9 100644
--- a/docs/source/en/model_doc/univnet.md
+++ b/docs/source/en/model_doc/univnet.md
@@ -69,7 +69,6 @@ write("sample_audio.wav", feature_extractor.sampling_rate, audio)
This model was contributed by [dg845](https://huggingface.co/dg845).
To the best of my knowledge, there is no official code release, but an unofficial implementation can be found at [maum-ai/univnet](https://github.com/maum-ai/univnet) with pretrained checkpoints [here](https://github.com/maum-ai/univnet#pre-trained-model).
-
## UnivNetConfig
[[autodoc]] UnivNetConfig
@@ -82,4 +81,4 @@ To the best of my knowledge, there is no official code release, but an unofficia
## UnivNetModel
[[autodoc]] UnivNetModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/upernet.md b/docs/source/en/model_doc/upernet.md
index 2c2e50fc560d..900b5635fc16 100644
--- a/docs/source/en/model_doc/upernet.md
+++ b/docs/source/en/model_doc/upernet.md
@@ -81,4 +81,4 @@ If you're interested in submitting a resource to be included here, please feel f
## UperNetForSemanticSegmentation
[[autodoc]] UperNetForSemanticSegmentation
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/van.md b/docs/source/en/model_doc/van.md
index 0e07e314bee9..0a4ded430211 100644
--- a/docs/source/en/model_doc/van.md
+++ b/docs/source/en/model_doc/van.md
@@ -74,4 +74,3 @@ If you're interested in submitting a resource to be included here, please feel f
[[autodoc]] VanForImageClassification
- forward
-
diff --git a/docs/source/en/model_doc/vaultgemma.md b/docs/source/en/model_doc/vaultgemma.md
index c9eb36124fca..deada15dc0f7 100644
--- a/docs/source/en/model_doc/vaultgemma.md
+++ b/docs/source/en/model_doc/vaultgemma.md
@@ -12,10 +12,10 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-
⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer.
-->
+*This model was released on 2016-07-01 and added to Hugging Face Transformers on 2025-09-12.*
# VaultGemma
@@ -30,7 +30,7 @@ sequence length.
VaultGemma was trained from scratch with sequence-level differential privacy (DP). Its training data includes the same
mixture as the [Gemma 2 models](https://huggingface.co/collections/google/gemma-2-release-667d6600fd5220e7b967f315),
consisting of a number of documents of varying lengths. Additionally, it is trained using
-[DP stochastic gradient descent (DP-SGD)](https://arxiv.org/abs/1607.00133) and provides a
+[DP stochastic gradient descent (DP-SGD)](https://huggingface.co/papers/1607.00133) and provides a
(ε ≤ 2.0, δ ≤ 1.1e-10)-sequence-level DP guarantee, where a sequence consists of 1024 consecutive tokens extracted from
heterogeneous data sources. Specifically, the privacy unit of the guarantee is for the sequences after sampling and
packing of the mixture.
@@ -44,7 +44,6 @@ command line.
-
```python
from transformers import pipeline
@@ -82,7 +81,7 @@ print(tokenizer.decode(outputs[0]))
-```
+```bash
echo -e "Write me a poem about Machine Learning. Answer:" | transformers run --task text2text-generation --model google/vaultgemma-1b-pt --device 0
```
diff --git a/docs/source/en/model_doc/video_llava.md b/docs/source/en/model_doc/video_llava.md
index 6b09367f37c8..2e1bf19abdc6 100644
--- a/docs/source/en/model_doc/video_llava.md
+++ b/docs/source/en/model_doc/video_llava.md
@@ -27,7 +27,6 @@ rendered properly in your Markdown viewer.
Video-LLaVa is an open-source multimodal LLM trained by fine-tuning LlamA/Vicuna on multimodal instruction-following data generated by Llava1.5 and VideChat. It is an auto-regressive language model, based on the transformer architecture. Video-LLaVa unifies visual representations to the language feature space, and enables an LLM to perform visual reasoning capabilities on both images and videos simultaneously.
-
The Video-LLaVA model was proposed in [Video-LLaVA: Learning United Visual Representation by Alignment Before Projection](https://huggingface.co/papers/2311.10122) by Bin Lin, Yang Ye, Bin Zhu, Jiaxi Cui, Munang Ning, Peng Jin, Li Yuan.
The abstract from the paper is the following:
@@ -55,18 +54,16 @@ for the LLM*
- Note the model has not been explicitly trained to process multiple images/videos in the same prompt, although this is technically possible, you may experience inaccurate results.
-- Note that the video inputs should have exactly 8 frames at the input, since the models were trained in that setting.
+- Note that the video inputs should have exactly 8 frames at the input, since the models were trained in that setting.
This model was contributed by [RaushanTurganbay](https://huggingface.co/RaushanTurganbay).
The original code can be found [here](https://github.com/PKU-YuanGroup/Video-LLaVA).
-
> [!NOTE]
-> LLaVA models after release v4.46 will raise warnings about adding `processor.patch_size = {{patch_size}}`, `processor.num_additional_image_tokens = {{num_additional_image_tokens}}` and processor.vision_feature_select_strategy = {{vision_feature_select_strategy}}`. It is strongly recommended to add the attributes to the processor if you own the model checkpoint, or open a PR if it is not owned by you.
+> LLaVA models after release v4.46 will raise warnings about adding `processor.patch_size = {{patch_size}}`, `processor.num_additional_image_tokens = {{num_additional_image_tokens}}` and `processor.vision_feature_select_strategy = {{vision_feature_select_strategy}}`. It is strongly recommended to add the attributes to the processor if you own the model checkpoint, or open a PR if it is not owned by you.
Adding these attributes means that LLaVA will try to infer the number of image tokens required per image and expand the text with as many `` placeholders as there will be tokens. Usually it is around 500 tokens per image, so make sure that the text is not truncated as otherwise there will be failure when merging the embeddings.
The attributes can be obtained from model config, as `model.config.vision_config.patch_size` or `model.config.vision_feature_select_strategy`. The `num_additional_image_tokens` should be `1` if the vision backbone adds a CLS token or `0` if nothing extra is added to the vision patches.
-
## Usage example
### Single Media Mode
@@ -126,7 +123,7 @@ For multiple turns conversation change the prompt format to:
### Mixed Media Mode
-The model can also generate from an interleaved image-video inputs. However note, that it was not trained in interleaved image-video setting which might affect the performance. Below is an example usage for mixed media input, add the following lines to the above code snippet:
+The model can also generate from an interleaved image-video inputs. However note, that it was not trained in interleaved image-video setting which might affect the performance. Below is an example usage for mixed media input, add the following lines to the above code snippet:
```python
from PIL import Image
@@ -150,7 +147,7 @@ processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokeniza
### Quantization using Bitsandbytes for memory efficiency
-The model can be loaded in lower bits, significantly reducing memory burden while maintaining the performance of the original model. his allows for efficient deployment on resource-constrained cases.
+The model can be loaded in lower bits, significantly reducing memory burden while maintaining the performance of the original model. his allows for efficient deployment on resource-constrained cases.
First make sure to install bitsandbytes by running `pip install bitsandbytes` and to have access to a GPU/accelerator that is supported by the library.
@@ -164,7 +161,6 @@ We value your feedback to help identify bugs before the full release! Check out
Load the quantized model by simply adding [`BitsAndBytesConfig`](../main_classes/quantization#transformers.BitsAndBytesConfig) as shown below:
-
```python
from transformers import VideoLlavaForConditionalGeneration, BitsAndBytesConfig
@@ -178,7 +174,6 @@ quantization_config = BitsAndBytesConfig(
model = VideoLlavaForConditionalGeneration.from_pretrained("LanguageBind/Video-LLaVA-7B-hf", quantization_config=quantization_config, device_map="auto")
```
-
### Flash-Attention 2 to speed-up generation
Additionally, we can greatly speed-up model inference by using [Flash Attention](../perf_train_gpu_one#flash-attention-2), which is a faster implementation of the attention mechanism used inside the model.
@@ -203,7 +198,6 @@ model = VideoLlavaForConditionalGeneration.from_pretrained(
).to(0)
```
-
## VideoLlavaConfig
[[autodoc]] VideoLlavaConfig
@@ -212,7 +206,6 @@ model = VideoLlavaForConditionalGeneration.from_pretrained(
[[autodoc]] VideoLlavaImageProcessor
-
## VideoLlavaVideoProcessor
[[autodoc]] VideoLlavaVideoProcessor
diff --git a/docs/source/en/model_doc/videomae.md b/docs/source/en/model_doc/videomae.md
index e0ebbaa42885..eb02fc48bb40 100644
--- a/docs/source/en/model_doc/videomae.md
+++ b/docs/source/en/model_doc/videomae.md
@@ -42,16 +42,16 @@ The original code can be found [here](https://github.com/MCG-NJU/VideoMAE).
## Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import VideoMAEForVideoClassification
model = VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics", attn_implementation="sdpa", dtype=torch.float16)
...
@@ -75,6 +75,7 @@ you're interested in submitting a resource to be included here, please feel free
review it! The resource should ideally demonstrate something new instead of duplicating an existing resource.
**Video classification**
+
- [A notebook](https://github.com/huggingface/notebooks/blob/main/examples/video_classification.ipynb) that shows how
to fine-tune a VideoMAE model on a custom dataset.
- [Video classification task guide](../tasks/video_classification)
diff --git a/docs/source/en/model_doc/vipllava.md b/docs/source/en/model_doc/vipllava.md
index 0d0a209c27a6..a6554c91b57c 100644
--- a/docs/source/en/model_doc/vipllava.md
+++ b/docs/source/en/model_doc/vipllava.md
@@ -37,7 +37,6 @@ The original code can be found [here](https://github.com/mu-cai/ViP-LLaVA).
This model was contributed by [Younes Belkada](https://huggingface.co/ybelkada)
-
## Usage tips:
- The architecture is similar than llava architecture except that the multi-modal projector takes a set of concatenated vision hidden states and has an additional layernorm layer on that module.
@@ -47,11 +46,10 @@ This model was contributed by [Younes Belkada](https://huggingface.co/ybelkada)
- Note the model has not been explicitly trained to process multiple images in the same prompt, although this is technically possible, you may experience inaccurate results.
> [!NOTE]
-> LLaVA models after release v4.46 will raise warnings about adding `processor.patch_size = {{patch_size}}`, `processor.num_additional_image_tokens = {{num_additional_image_tokens}}` and processor.vision_feature_select_strategy = {{vision_feature_select_strategy}}`. It is strongly recommended to add the attributes to the processor if you own the model checkpoint, or open a PR if it is not owned by you.
+> LLaVA models after release v4.46 will raise warnings about adding `processor.patch_size = {{patch_size}}`, `processor.num_additional_image_tokens = {{num_additional_image_tokens}}` and `processor.vision_feature_select_strategy = {{vision_feature_select_strategy}}`. It is strongly recommended to add the attributes to the processor if you own the model checkpoint, or open a PR if it is not owned by you.
Adding these attributes means that LLaVA will try to infer the number of image tokens required per image and expand the text with as many `` placeholders as there will be tokens. Usually it is around 500 tokens per image, so make sure that the text is not truncated as otherwise there will be failure when merging the embeddings.
The attributes can be obtained from model config, as `model.config.vision_config.patch_size` or `model.config.vision_feature_select_strategy`. The `num_additional_image_tokens` should be `1` if the vision backbone adds a CLS token or `0` if nothing extra is added to the vision patches.
-
- For better results, we recommend users to use the processor's `apply_chat_template()` method to format your prompt correctly. For that you need to construct a conversation history, passing in a plain string will not format your prompt. Each message in the conversation history for chat templates is a dictionary with keys "role" and "content". The "content" should be a list of dictionaries, for "text" and "image" modalities, as follows:
```python
@@ -88,16 +86,17 @@ print(text_prompt)
```
- If you want to construct a chat prompt yourself, below is a list of prompt formats accepted by VipLLaVa checkpoints:
+
```bash
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:
```
For multiple turns conversation:
+
```bash
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: ###Human: ###Assistant:
```
-
## VipLlavaConfig
[[autodoc]] VipLlavaConfig
diff --git a/docs/source/en/model_doc/visual_bert.md b/docs/source/en/model_doc/visual_bert.md
index 7a7ac24e4dbf..a9912144c4f9 100644
--- a/docs/source/en/model_doc/visual_bert.md
+++ b/docs/source/en/model_doc/visual_bert.md
@@ -27,7 +27,6 @@ rendered properly in your Markdown viewer.
You can find all the original VisualBERT checkpoints under the [UCLA NLP](https://huggingface.co/uclanlp/models?search=visualbert) organization.
-
> [!TIP]
> This model was contributed by [gchhablani](https://huggingface.co/gchhablani).
> Click on the VisualBERT models in the right sidebar for more examples of how to apply VisualBERT to different image and language tasks.
diff --git a/docs/source/en/model_doc/vit_hybrid.md b/docs/source/en/model_doc/vit_hybrid.md
index 86c2c7229f58..c10d1c489b76 100644
--- a/docs/source/en/model_doc/vit_hybrid.md
+++ b/docs/source/en/model_doc/vit_hybrid.md
@@ -55,16 +55,16 @@ found [here](https://github.com/google-research/vision_transformer).
## Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import ViTHybridForImageClassification
model = ViTHybridForImageClassification.from_pretrained("google/vit-hybrid-base-bit-384", attn_implementation="sdpa", dtype=torch.float16)
...
diff --git a/docs/source/en/model_doc/vit_mae.md b/docs/source/en/model_doc/vit_mae.md
index b8b9867e8812..0547594ae118 100644
--- a/docs/source/en/model_doc/vit_mae.md
+++ b/docs/source/en/model_doc/vit_mae.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2021-11-11 and added to Hugging Face Transformers on 2022-01-18.*
-
@@ -67,6 +66,7 @@ reconstruction = outputs.logits
## Notes
+
- ViTMAE is typically used in two stages. Self-supervised pretraining with [`ViTMAEForPreTraining`], and then discarding the decoder and fine-tuning the encoder. After fine-tuning, the weights can be plugged into a model like [`ViTForImageClassification`].
- Use [`ViTImageProcessor`] for input preparation.
diff --git a/docs/source/en/model_doc/vit_msn.md b/docs/source/en/model_doc/vit_msn.md
index 5b727f34256c..d7a8172a18f3 100644
--- a/docs/source/en/model_doc/vit_msn.md
+++ b/docs/source/en/model_doc/vit_msn.md
@@ -40,11 +40,11 @@ while producing representations of a high semantic level that perform competitiv
on ImageNet-1K, with only 5,000 annotated images, our base MSN model achieves 72.4% top-1 accuracy,
and with 1% of ImageNet-1K labels, we achieve 75.7% top-1 accuracy, setting a new state-of-the-art for self-supervised learning on this benchmark.*
-
+
MSN architecture. Taken from the original paper.
-This model was contributed by [sayakpaul](https://huggingface.co/sayakpaul). The original code can be found [here](https://github.com/facebookresearch/msn).
+This model was contributed by [sayakpaul](https://huggingface.co/sayakpaul). The original code can be found [here](https://github.com/facebookresearch/msn).
## Usage tips
@@ -58,16 +58,16 @@ labels when fine-tuned.
### Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import ViTMSNForImageClassification
model = ViTMSNForImageClassification.from_pretrained("facebook/vit-msn-base", attn_implementation="sdpa", dtype=torch.float16)
...
diff --git a/docs/source/en/model_doc/vitdet.md b/docs/source/en/model_doc/vitdet.md
index 539ae5e376c8..a1250f1bb909 100644
--- a/docs/source/en/model_doc/vitdet.md
+++ b/docs/source/en/model_doc/vitdet.md
@@ -40,4 +40,4 @@ Tips:
## VitDetModel
[[autodoc]] VitDetModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/vitmatte.md b/docs/source/en/model_doc/vitmatte.md
index 519a2dd74d66..0584df8e67a5 100644
--- a/docs/source/en/model_doc/vitmatte.md
+++ b/docs/source/en/model_doc/vitmatte.md
@@ -62,4 +62,4 @@ The model expects both the image and trimap (concatenated) as input. Use [`ViTMa
## VitMatteForImageMatting
[[autodoc]] VitMatteForImageMatting
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/vits.md b/docs/source/en/model_doc/vits.md
index 2c1777b77f18..96dc93892470 100644
--- a/docs/source/en/model_doc/vits.md
+++ b/docs/source/en/model_doc/vits.md
@@ -149,11 +149,10 @@ Audio(waveform, rate=model.config.sampling_rate)
## VitsTokenizer
[[autodoc]] VitsTokenizer
-- __call__
-- save_vocabulary
+ - __call__
+ - save_vocabulary
## VitsModel
[[autodoc]] VitsModel
-- forward
-
+ - forward
diff --git a/docs/source/en/model_doc/vivit.md b/docs/source/en/model_doc/vivit.md
index 041f80f61ae6..fc127fa6f595 100644
--- a/docs/source/en/model_doc/vivit.md
+++ b/docs/source/en/model_doc/vivit.md
@@ -32,16 +32,16 @@ This model was contributed by [jegormeister](https://huggingface.co/jegormeister
### Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import VivitModel
model = VivitModel.from_pretrained("google/vivit-b-16x2-kinetics400", attn_implementation="sdpa", dtype=torch.float16)
...
@@ -56,8 +56,6 @@ On a local benchmark (A100-40GB, PyTorch 2.3.0, OS Ubuntu 22.04) with `float32`
|---------------------:|-------------:|----------:|--------------:|----------------------:|---------------------:|-----------------:|
| 100 | 1 | True | 7.122 | 2575.28 | 5932.54 | 130.364 |
-
-
### Inference
| num_batches | batch_size | is cuda | is half | Speedup (%) | Mem eager (MB) | Mem BT (MB) | Mem saved (%) |
|---------------|--------------|-----------|-----------|---------------|------------------|---------------|-----------------|
@@ -65,7 +63,6 @@ On a local benchmark (A100-40GB, PyTorch 2.3.0, OS Ubuntu 22.04) with `float32`
| 20 | 2 | True | False | 17.146 | 1234.75 | 447.175 | 176.122 |
| 20 | 4 | True | False | 18.093 | 2275.82 | 709.864 | 220.6 |
| 20 | 8 | True | False | 19.284 | 4358.19 | 1233.24 | 253.393 |
-
## VivitConfig
diff --git a/docs/source/en/model_doc/vjepa2.md b/docs/source/en/model_doc/vjepa2.md
index 93960f051893..049c7ff98f21 100644
--- a/docs/source/en/model_doc/vjepa2.md
+++ b/docs/source/en/model_doc/vjepa2.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2025-06-11 and added to Hugging Face Transformers on 2025-06-11.*
-
@@ -34,7 +33,6 @@ rendered properly in your Markdown viewer.
You can find all original V-JEPA2 checkpoints under the [V-JEPA 2](https://huggingface.co/collections/facebook/v-jepa-2-6841bad8413014e185b497a6) collection.
-
This model was contributed by [koustuvs](https://huggingface.co/koustuvs), [yonigozlan](https://huggingface.co/yonigozlan) and [qubvel](https://huggingface.co/qubvel-hf). The original code can be found [here](https://github.com/facebookresearch/vjepa2).
## Usage example
diff --git a/docs/source/en/model_doc/voxtral.md b/docs/source/en/model_doc/voxtral.md
index 71f0661c8276..3dd2fc9e0d31 100644
--- a/docs/source/en/model_doc/voxtral.md
+++ b/docs/source/en/model_doc/voxtral.md
@@ -22,6 +22,7 @@ Voxtral is an upgrade of [Ministral 3B and Mistral Small 3B](https://mistral.ai/
You can read more in Mistral's [realease blog post](https://mistral.ai/news/voxtral).
The model is available in two checkpoints:
+
- 3B: [mistralai/Voxtral-Mini-3B-2507](https://huggingface.co/mistralai/Voxtral-Mini-3B-2507)
- 24B: [mistralai/Voxtral-Small-24B-2507](https://huggingface.co/mistralai/Voxtral-Small-24B-2507)
@@ -43,6 +44,7 @@ Voxtral builds on Ministral-3B by adding audio processing capabilities:
The model supports audio-text instructions, including multi-turn and multi-audio interactions, all processed in batches.
➡️ audio + text instruction
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -78,7 +80,8 @@ print(decoded_outputs[0])
print("=" * 80)
```
-➡️ multi-audio + text instruction
+➡️ multi-audio + text instruction
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -119,6 +122,7 @@ print("=" * 80)
```
➡️ multi-turn:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -173,6 +177,7 @@ print("=" * 80)
```
➡️ text only:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -208,6 +213,7 @@ print("=" * 80)
```
➡️ audio only:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -243,6 +249,7 @@ print("=" * 80)
```
➡️ batched inference!
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device()
diff --git a/docs/source/en/model_doc/wav2vec2-bert.md b/docs/source/en/model_doc/wav2vec2-bert.md
index 4edb67498aaa..23409b0898c3 100644
--- a/docs/source/en/model_doc/wav2vec2-bert.md
+++ b/docs/source/en/model_doc/wav2vec2-bert.md
@@ -31,7 +31,7 @@ The official results of the model can be found in Section 3.2.1 of the paper.
The abstract from the paper is the following:
-*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one’s voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
+*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one's voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The original code can be found [here](https://github.com/facebookresearch/seamless_communication).
@@ -54,7 +54,6 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
- [`Wav2Vec2BertForSequenceClassification`] can be used by adapting this [example script](https://github.com/huggingface/transformers/tree/main/examples/pytorch/audio-classification).
- See also: [Audio classification task guide](../tasks/audio_classification)
-
## Wav2Vec2BertConfig
[[autodoc]] Wav2Vec2BertConfig
diff --git a/docs/source/en/model_doc/wav2vec2-conformer.md b/docs/source/en/model_doc/wav2vec2-conformer.md
index e2a56b450df3..663b6163011b 100644
--- a/docs/source/en/model_doc/wav2vec2-conformer.md
+++ b/docs/source/en/model_doc/wav2vec2-conformer.md
@@ -38,7 +38,7 @@ Note: Meta (FAIR) released a new version of [Wav2Vec2-BERT 2.0](https://huggingf
- Wav2Vec2-Conformer follows the same architecture as Wav2Vec2, but replaces the *Attention*-block with a *Conformer*-block
as introduced in [Conformer: Convolution-augmented Transformer for Speech Recognition](https://huggingface.co/papers/2005.08100).
-- For the same number of layers, Wav2Vec2-Conformer requires more parameters than Wav2Vec2, but also yields
+- For the same number of layers, Wav2Vec2-Conformer requires more parameters than Wav2Vec2, but also yields
an improved word error rate.
- Wav2Vec2-Conformer uses the same tokenizer and feature extractor as Wav2Vec2.
- Wav2Vec2-Conformer can use either no relative position embeddings, Transformer-XL-like position embeddings, or
diff --git a/docs/source/en/model_doc/wav2vec2.md b/docs/source/en/model_doc/wav2vec2.md
index 6c4772f90bc8..1f5f4a905767 100644
--- a/docs/source/en/model_doc/wav2vec2.md
+++ b/docs/source/en/model_doc/wav2vec2.md
@@ -80,13 +80,10 @@ model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-large-960h-lv60-self",
Below is an expected speedup diagram comparing the pure inference time between the native implementation in transformers of the `facebook/wav2vec2-large-960h-lv60-self` model and the flash-attention-2 and sdpa (scale-dot-product-attention) versions. . We show the average speedup obtained on the `librispeech_asr` `clean` validation split:
-
-
-
## Resources
A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with Wav2Vec2. If you're interested in submitting a resource to be included here, please feel free to open a Pull Request and we'll review it! The resource should ideally demonstrate something new instead of duplicating an existing resource.
diff --git a/docs/source/en/model_doc/wav2vec2_phoneme.md b/docs/source/en/model_doc/wav2vec2_phoneme.md
index fe989def3bdd..206ea048c023 100644
--- a/docs/source/en/model_doc/wav2vec2_phoneme.md
+++ b/docs/source/en/model_doc/wav2vec2_phoneme.md
@@ -53,7 +53,6 @@ The original code can be found [here](https://github.com/pytorch/fairseq/tree/ma
- By default, the model outputs a sequence of phonemes. In order to transform the phonemes to a sequence of words one
should make use of a dictionary and language model.
-
Wav2Vec2Phoneme's architecture is based on the Wav2Vec2 model, for API reference, check out [`Wav2Vec2`](wav2vec2)'s documentation page
@@ -64,7 +63,7 @@ except for the tokenizer.
## Wav2Vec2PhonemeCTCTokenizer
[[autodoc]] Wav2Vec2PhonemeCTCTokenizer
- - __call__
- - batch_decode
- - decode
- - phonemize
+ - __call__
+ - batch_decode
+ - decode
+ - phonemize
diff --git a/docs/source/en/model_doc/whisper.md b/docs/source/en/model_doc/whisper.md
index 673085ac3e7d..5e19e870bddc 100644
--- a/docs/source/en/model_doc/whisper.md
+++ b/docs/source/en/model_doc/whisper.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2022-12-06 and added to Hugging Face Transformers on 2022-10-05.*
-
diff --git a/docs/source/en/model_doc/xcodec.md b/docs/source/en/model_doc/xcodec.md
index c4a0b92a26f6..957a74093484 100644
--- a/docs/source/en/model_doc/xcodec.md
+++ b/docs/source/en/model_doc/xcodec.md
@@ -33,9 +33,10 @@ The X-Codec model is a neural audio codec that integrates semantic information f
The abstract of the paper states the following:
-*Recent advancements in audio generation have been significantly propelled by the capabilities of Large Language Models (LLMs). The existing research on audio LLM has primarily focused on enhancing the architecture and scale of audio language models, as well as leveraging larger datasets, and generally, acoustic codecs, such as EnCodec, are used for audio tokenization. However, these codecs were originally designed for audio compression, which may lead to suboptimal performance in the context of audio LLM. Our research aims to address the shortcomings of current audio LLM codecs, particularly their challenges in maintaining semantic integrity in generated audio. For instance, existing methods like VALL-E, which condition acoustic token generation on text transcriptions, often suffer from content inaccuracies and elevated word error rates (WER) due to semantic misinterpretations of acoustic tokens, resulting in word skipping and errors. To overcome these issues, we propose a straightforward yet effective approach called X-Codec. X-Codec incorporates semantic features from a pre-trained semantic encoder before the Residual Vector Quantization (RVQ) stage and introduces a semantic reconstruction loss after RVQ. By enhancing the semantic ability of the codec, X-Codec significantly reduces WER in speech synthesis tasks and extends these benefits to non-speech applications, including music and sound generation. Our experiments in text-to-speech, music continuation, and text-to-sound tasks demonstrate that integrating semantic information substantially improves the overall performance of language models in audio generation.*
+*Recent advancements in audio generation have been significantly propelled by the capabilities of Large Language Models (LLMs). The existing research on audio LLM has primarily focused on enhancing the architecture and scale of audio language models, as well as leveraging larger datasets, and generally, acoustic codecs, such as EnCodec, are used for audio tokenization. However, these codecs were originally designed for audio compression, which may lead to suboptimal performance in the context of audio LLM. Our research aims to address the shortcomings of current audio LLM codecs, particularly their challenges in maintaining semantic integrity in generated audio. For instance, existing methods like VALL-E, which condition acoustic token generation on text transcriptions, often suffer from content inaccuracies and elevated word error rates (WER) due to semantic misinterpretations of acoustic tokens, resulting in word skipping and errors. To overcome these issues, we propose a straightforward yet effective approach called X-Codec. X-Codec incorporates semantic features from a pre-trained semantic encoder before the Residual Vector Quantization (RVQ) stage and introduces a semantic reconstruction loss after RVQ. By enhancing the semantic ability of the codec, X-Codec significantly reduces WER in speech synthesis tasks and extends these benefits to non-speech applications, including music and sound generation. Our experiments in text-to-speech, music continuation, and text-to-sound tasks demonstrate that integrating semantic information substantially improves the overall performance of language models in audio generation.*
Model cards:
+
- [xcodec-hubert-librispeech](https://huggingface.co/hf-audio/xcodec-hubert-librispeech) (for speech)
- [xcodec-wavlm-mls](https://huggingface.co/hf-audio/xcodec-wavlm-mls) (for speech)
- [xcodec-wavlm-more-data](https://huggingface.co/hf-audio/xcodec-wavlm-more-data) (for speech)
@@ -46,12 +47,11 @@ This model was contributed by [Manal El Aidouni](https://huggingface.co/Manel).
Demos can be found on this [page](https://x-codec-audio.github.io/).
-
-## Usage example
+## Usage example
Here is a quick example of how to encode and decode an audio using this model:
-```python
+```python
from datasets import load_dataset, Audio
from transformers import XcodecModel, AutoFeatureExtractor
dummy_dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
@@ -75,6 +75,7 @@ audio_values = decoder_outputs.audio_values
audio_values = model(inputs["input_values"]).audio_values
```
+
To listen to the original and reconstructed audio, run the snippet below and then open the generated `original.wav` and `reconstruction.wav` files in your music player to compare.
```python
@@ -88,15 +89,13 @@ sf.write("original.wav", original, sampling_rate)
sf.write("reconstruction.wav", reconstruction.T, sampling_rate)
```
-
## XcodecConfig
[[autodoc]] XcodecConfig
-
## XcodecModel
[[autodoc]] XcodecModel
- decode
- encode
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/xglm.md b/docs/source/en/model_doc/xglm.md
index d82bba7d23f9..9372b52af1f7 100644
--- a/docs/source/en/model_doc/xglm.md
+++ b/docs/source/en/model_doc/xglm.md
@@ -44,7 +44,6 @@ showing in particular that it enables cross-lingual in-context learning on some
on surface form robustness and adaptation to tasks that do not have a natural cloze form. Finally, we evaluate our models
in social value tasks such as hate speech detection in five languages and find it has limitations similar to comparable sized GPT-3 models.*
-
This model was contributed by [Suraj](https://huggingface.co/valhalla). The original code can be found [here](https://github.com/pytorch/fairseq/tree/main/examples/xglm).
## Resources
diff --git a/docs/source/en/model_doc/xlm-prophetnet.md b/docs/source/en/model_doc/xlm-prophetnet.md
index 4dad4c0afa78..fbf47d8c422a 100644
--- a/docs/source/en/model_doc/xlm-prophetnet.md
+++ b/docs/source/en/model_doc/xlm-prophetnet.md
@@ -41,7 +41,6 @@ You can do so by running the following command: `pip install -U transformers==4.
**DISCLAIMER:** If you see something strange, file a [Github Issue](https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title) and assign
@patrickvonplaten
-
## Overview
The XLM-ProphetNet model was proposed in [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training,](https://huggingface.co/papers/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei
diff --git a/docs/source/en/model_doc/xlm-roberta-xl.md b/docs/source/en/model_doc/xlm-roberta-xl.md
index 988107fdacc6..97dc6f1a7445 100644
--- a/docs/source/en/model_doc/xlm-roberta-xl.md
+++ b/docs/source/en/model_doc/xlm-roberta-xl.md
@@ -77,6 +77,7 @@ predicted_token = tokenizer.decode(predicted_token_id)
print(f"The predicted token is: {predicted_token}")
```
+
@@ -84,6 +85,7 @@ print(f"The predicted token is: {predicted_token}")
```bash
echo -e "Plants create through a process known as photosynthesis." | transformers-cli run --task fill-mask --model facebook/xlm-roberta-xl --device 0
```
+
diff --git a/docs/source/en/model_doc/xlm-roberta.md b/docs/source/en/model_doc/xlm-roberta.md
index a662742c2674..3a4b8e682603 100644
--- a/docs/source/en/model_doc/xlm-roberta.md
+++ b/docs/source/en/model_doc/xlm-roberta.md
@@ -87,6 +87,7 @@ print(f"The predicted token is: {predicted_token}")
```bash
echo -e "Plants create through a process known as photosynthesis." | transformers-cli run --task fill-mask --model FacebookAI/xlm-roberta-base --device 0
```
+
diff --git a/docs/source/en/model_doc/xlm.md b/docs/source/en/model_doc/xlm.md
index dc51fa4be4cd..11c00f4ec8ed 100644
--- a/docs/source/en/model_doc/xlm.md
+++ b/docs/source/en/model_doc/xlm.md
@@ -79,6 +79,7 @@ print(f"Predicted token: {predicted_token}")
```bash
echo -e "Plants create through a process known as photosynthesis." | transformers-cli run --task fill-mask --model FacebookAI/xlm-mlm-en-2048 --device 0
```
+
diff --git a/docs/source/en/model_doc/xlstm.md b/docs/source/en/model_doc/xlstm.md
index b239d631fbbc..e1ba3195eccf 100644
--- a/docs/source/en/model_doc/xlstm.md
+++ b/docs/source/en/model_doc/xlstm.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2024-05-07 and added to Hugging Face Transformers on 2025-07-25.*
-
# xLSTM
## Overview
@@ -32,7 +31,6 @@ The abstract from the paper is the following:
This model was contributed by [NX-AI](https://huggingface.co/NX-AI).
The original code can be found [here](https://github.com/NX-AI/xlstm).
-
## xLSTMConfig
[[autodoc]] xLSTMConfig
diff --git a/docs/source/en/model_doc/xmod.md b/docs/source/en/model_doc/xmod.md
index 0593e9940bd6..624b7ebb2d23 100644
--- a/docs/source/en/model_doc/xmod.md
+++ b/docs/source/en/model_doc/xmod.md
@@ -36,6 +36,7 @@ The original code can be found [here](https://github.com/facebookresearch/fairse
## Usage tips
Tips:
+
- X-MOD is similar to [XLM-R](xlm-roberta), but a difference is that the input language needs to be specified so that the correct language adapter can be activated.
- The main models – base and large – have adapters for 81 languages.
@@ -44,6 +45,7 @@ Tips:
### Input language
There are two ways to specify the input language:
+
1. By setting a default language before using the model:
```python
diff --git a/docs/source/en/model_doc/yolos.md b/docs/source/en/model_doc/yolos.md
index 5c31b539e59c..4a75b2ed020f 100644
--- a/docs/source/en/model_doc/yolos.md
+++ b/docs/source/en/model_doc/yolos.md
@@ -26,14 +26,12 @@ rendered properly in your Markdown viewer.
[YOLOS](https://huggingface.co/papers/2106.00666) uses a [Vision Transformer (ViT)](./vit) for object detection with minimal modifications and region priors. It can achieve performance comparable to specialized object detection models and frameworks with knowledge about 2D spatial structures.
-
You can find all the original YOLOS checkpoints under the [HUST Vision Lab](https://huggingface.co/hustvl/models?search=yolos) organization.
YOLOS architecture. Taken from the original paper.
-
> [!TIP]
> This model wasa contributed by [nielsr](https://huggingface.co/nielsr).
> Click on the YOLOS models in the right sidebar for more examples of how to apply YOLOS to different object detection tasks.
@@ -98,8 +96,8 @@ for score, label, box in zip(filtered_scores, filtered_labels, pixel_boxes):
-
## Notes
+
- Use [`YolosImageProcessor`] for preparing images (and optional targets) for the model. Contrary to [DETR](./detr), YOLOS doesn't require a `pixel_mask`.
## Resources
diff --git a/docs/source/en/model_doc/yoso.md b/docs/source/en/model_doc/yoso.md
index f07e5aba0827..211b0dcf8091 100644
--- a/docs/source/en/model_doc/yoso.md
+++ b/docs/source/en/model_doc/yoso.md
@@ -26,20 +26,20 @@ rendered properly in your Markdown viewer.
The YOSO model was proposed in [You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling](https://huggingface.co/papers/2111.09714)
by Zhanpeng Zeng, Yunyang Xiong, Sathya N. Ravi, Shailesh Acharya, Glenn Fung, Vikas Singh. YOSO approximates standard softmax self-attention
via a Bernoulli sampling scheme based on Locality Sensitive Hashing (LSH). In principle, all the Bernoulli random variables can be sampled with
-a single hash.
+a single hash.
The abstract from the paper is the following:
-*Transformer-based models are widely used in natural language processing (NLP). Central to the transformer model is
-the self-attention mechanism, which captures the interactions of token pairs in the input sequences and depends quadratically
-on the sequence length. Training such models on longer sequences is expensive. In this paper, we show that a Bernoulli sampling
-attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity of such models to linear.
-We bypass the quadratic cost by considering self-attention as a sum of individual tokens associated with Bernoulli random
-variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant).
-This leads to an efficient sampling scheme to estimate self-attention which relies on specific modifications of
-LSH (to enable deployment on GPU architectures). We evaluate our algorithm on the GLUE benchmark with standard 512 sequence
-length where we see favorable performance relative to a standard pretrained Transformer. On the Long Range Arena (LRA) benchmark,
-for evaluating performance on long sequences, our method achieves results consistent with softmax self-attention but with sizable
+*Transformer-based models are widely used in natural language processing (NLP). Central to the transformer model is
+the self-attention mechanism, which captures the interactions of token pairs in the input sequences and depends quadratically
+on the sequence length. Training such models on longer sequences is expensive. In this paper, we show that a Bernoulli sampling
+attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity of such models to linear.
+We bypass the quadratic cost by considering self-attention as a sum of individual tokens associated with Bernoulli random
+variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant).
+This leads to an efficient sampling scheme to estimate self-attention which relies on specific modifications of
+LSH (to enable deployment on GPU architectures). We evaluate our algorithm on the GLUE benchmark with standard 512 sequence
+length where we see favorable performance relative to a standard pretrained Transformer. On the Long Range Arena (LRA) benchmark,
+for evaluating performance on long sequences, our method achieves results consistent with softmax self-attention but with sizable
speed-ups and memory savings and often outperforms other efficient self-attention methods. Our code is available at this https URL*
This model was contributed by [novice03](https://huggingface.co/novice03). The original code can be found [here](https://github.com/mlpen/YOSO).
@@ -50,12 +50,12 @@ This model was contributed by [novice03](https://huggingface.co/novice03). The o
in parallel on a GPU.
- The kernels provide a `fast_hash` function, which approximates the random projections of the queries and keys using the Fast Hadamard Transform. Using these
hash codes, the `lsh_cumulation` function approximates self-attention via LSH-based Bernoulli sampling.
-- To use the custom kernels, the user should set `config.use_expectation = False`. To ensure that the kernels are compiled successfully,
-the user must install the correct version of PyTorch and cudatoolkit. By default, `config.use_expectation = True`, which uses YOSO-E and
+- To use the custom kernels, the user should set `config.use_expectation = False`. To ensure that the kernels are compiled successfully,
+the user must install the correct version of PyTorch and cudatoolkit. By default, `config.use_expectation = True`, which uses YOSO-E and
does not require compiling CUDA kernels.
+alt="drawing" width="600"/>
YOSO Attention Algorithm. Taken from the original paper.
@@ -99,4 +99,4 @@ alt="drawing" width="600"/>
## YosoForQuestionAnswering
[[autodoc]] YosoForQuestionAnswering
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/zamba.md b/docs/source/en/model_doc/zamba.md
index bb9740807703..847f0532e2a7 100644
--- a/docs/source/en/model_doc/zamba.md
+++ b/docs/source/en/model_doc/zamba.md
@@ -24,7 +24,6 @@ rendered properly in your Markdown viewer.
This model was contributed by [pglo](https://huggingface.co/pglo).
-
## Model details
Zamba-7B-v1 is a hybrid between state-space models (Specifically [Mamba](https://github.com/state-spaces/mamba)) and transformer, and was trained using next-token prediction. Zamba uses a shared transformer layer after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba-7B-v1 was pre-trained on 1T tokens of text and code data.
@@ -33,23 +32,24 @@ Zamba-7B-v1 is a hybrid between state-space models (Specifically [Mamba](https:/
## Quick start
-
### Presequities
Zamba requires you use `transformers` version 4.46.0 or higher:
+
```bash
pip install transformers>=4.45.0
```
In order to run optimized Mamba implementations, you first need to install `mamba-ssm` and `causal-conv1d`:
+
```bash
pip install mamba-ssm causal-conv1d>=1.2.0
```
+
You also have to have the model on a CUDA device.
You can run the model not using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify `use_mamba_kernels=False` when loading the model.
-
## Inference
```python
@@ -66,39 +66,33 @@ outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))
```
-
## Model card
The model cards can be found at:
-* [Zamba-7B](https://huggingface.co/Zyphra/Zamba-7B-v1)
+* [Zamba-7B](https://huggingface.co/Zyphra/Zamba-7B-v1)
## Issues
For issues with model output, or community discussion, please use the Hugging Face community [forum](https://huggingface.co/Zyphra/Zamba-7B-v1/discussions)
-
## License
The model weights are open-sourced via an Apache 2.0 license.
-
## ZambaConfig
[[autodoc]] ZambaConfig
-
## ZambaModel
[[autodoc]] ZambaModel
- forward
-
## ZambaForCausalLM
[[autodoc]] ZambaForCausalLM
- forward
-
## ZambaForSequenceClassification
[[autodoc]] transformers.ZambaForSequenceClassification
diff --git a/docs/source/en/model_doc/zamba2.md b/docs/source/en/model_doc/zamba2.md
index 1d911a59c277..c9d3d3d1de75 100644
--- a/docs/source/en/model_doc/zamba2.md
+++ b/docs/source/en/model_doc/zamba2.md
@@ -26,19 +26,18 @@ rendered properly in your Markdown viewer.
This model was contributed by [pglo](https://huggingface.co/pglo).
-
## Model details
-[Zamba2-1.2B](https://www.zyphra.com/post/zamba2-mini), [Zamba2-2.7B](https://www.zyphra.com/post/zamba2-small) and [Zamba2-7B](https://www.zyphra.com/post/zamba2-7b) are hybrid models combining state-space models (Specifically [Mamba](https://github.com/state-spaces/mamba)) and transformer, and were trained using next-token prediction. Zamba2 uses shared transformer layers after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba2-1.2B, Zamba2-2.7B and Zamba2-7B were pre-trained on 2T and 3T tokens, respectively.
+[Zamba2-1.2B](https://www.zyphra.com/post/zamba2-mini), [Zamba2-2.7B](https://www.zyphra.com/post/zamba2-small) and [Zamba2-7B](https://www.zyphra.com/post/zamba2-7b) are hybrid models combining state-space models (Specifically [Mamba2](https://github.com/state-spaces/mamba)) and transformer, and were trained using next-token prediction. Zamba2 uses shared transformer layers after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba2-1.2B, Zamba2-2.7B and Zamba2-7B were pre-trained on 2T and 3T tokens, respectively.
## Quick start
-
### Presequities
Zamba2 requires you use `transformers` version 4.48.0 or higher:
+
```bash
pip install transformers>=4.48.0
```
@@ -59,41 +58,35 @@ outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))
```
-
## Model card
The model cards can be found at:
+
* [Zamba2-1.2B](https://huggingface.co/Zyphra/Zamba2-1.2B)
* [Zamba2-2.7B](https://huggingface.co/Zyphra/Zamba2-2.7B)
* [Zamba2-7B](https://huggingface.co/Zyphra/Zamba2-7B)
-
## Issues
For issues with model output, or community discussion, please use the Hugging Face community [forum](https://huggingface.co/Zyphra/Zamba2-7B/discussions)
-
## License
The model weights are open-sourced via an Apache 2.0 license.
-
## Zamba2Config
[[autodoc]] Zamba2Config
-
## Zamba2Model
[[autodoc]] Zamba2Model
- forward
-
## Zamba2ForCausalLM
[[autodoc]] Zamba2ForCausalLM
- forward
-
## Zamba2ForSequenceClassification
[[autodoc]] transformers.Zamba2ForSequenceClassification
diff --git a/docs/source/en/model_doc/zoedepth.md b/docs/source/en/model_doc/zoedepth.md
index 367c630a3224..92840a770462 100644
--- a/docs/source/en/model_doc/zoedepth.md
+++ b/docs/source/en/model_doc/zoedepth.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2023-02-23 and added to Hugging Face Transformers on 2024-07-08.*
-
@@ -97,6 +96,7 @@ Image.fromarray(depth.astype("uint8"))
## Notes
- In the [original implementation](https://github.com/isl-org/ZoeDepth/blob/edb6daf45458569e24f50250ef1ed08c015f17a7/zoedepth/models/depth_model.py#L131) ZoeDepth performs inference on both the original and flipped images and averages the results. The `post_process_depth_estimation` function handles this by passing the flipped outputs to the optional `outputs_flipped` argument as shown below.
+
```py
with torch.no_grad():
outputs = model(pixel_values)
@@ -107,8 +107,9 @@ Image.fromarray(depth.astype("uint8"))
outputs_flipped=outputs_flipped,
)
```
-
+
## Resources
+
- Refer to this [notebook](https://github.com/NielsRogge/Transformers-Tutorials/tree/master/ZoeDepth) for an inference example.
## ZoeDepthConfig
diff --git a/docs/source/en/model_memory_anatomy.md b/docs/source/en/model_memory_anatomy.md
index 7ef53f40566e..f0a215b05c1b 100644
--- a/docs/source/en/model_memory_anatomy.md
+++ b/docs/source/en/model_memory_anatomy.md
@@ -16,24 +16,23 @@ limitations under the License.
# Model training anatomy
-To understand performance optimization techniques that one can apply to improve efficiency of model training
-speed and memory utilization, it's helpful to get familiar with how GPU is utilized during training, and how compute
+To understand performance optimization techniques that one can apply to improve efficiency of model training
+speed and memory utilization, it's helpful to get familiar with how GPU is utilized during training, and how compute
intensity varies depending on an operation performed.
-Let's start by exploring a motivating example of GPU utilization and the training run of a model. For the demonstration,
-we'll need to install a few libraries:
+Let's start by exploring a motivating example of GPU utilization and the training run of a model. For the demonstration,
+we'll need to install a few libraries:
```bash
-pip install transformers datasets accelerate nvidia-ml-py3
+pip install transformers datasets accelerate nvidia-ml-py
```
-The `nvidia-ml-py3` library allows us to monitor the memory usage of the models from within Python. You might be familiar
+The `nvidia-ml-py` library allows us to monitor the memory usage of the models from within Python. You might be familiar
with the `nvidia-smi` command in the terminal - this library allows to access the same information in Python directly.
-Then, we create some dummy data: random token IDs between 100 and 30000 and binary labels for a classifier.
+Then, we create some dummy data: random token IDs between 100 and 30000 and binary labels for a classifier.
In total, we get 512 sequences each with length 512 and store them in a [`~datasets.Dataset`] with PyTorch format.
-
```py
>>> import numpy as np
>>> from datasets import Dataset
@@ -74,9 +73,9 @@ Let's verify that we start with a free GPU memory:
GPU memory occupied: 0 MB.
```
-That looks good: the GPU memory is not occupied as we would expect before we load any models. If that's not the case on
-your machine make sure to stop all processes that are using GPU memory. However, not all free GPU memory can be used by
-the user. When a model is loaded to the GPU the kernels are also loaded, which can take up 1-2GB of memory. To see how
+That looks good: the GPU memory is not occupied as we would expect before we load any models. If that's not the case on
+your machine make sure to stop all processes that are using GPU memory. However, not all free GPU memory can be used by
+the user. When a model is loaded to the GPU the kernels are also loaded, which can take up 1-2GB of memory. To see how
much it is we load a tiny tensor into the GPU which triggers the kernels to be loaded as well.
```py
@@ -92,10 +91,9 @@ We see that the kernels alone take up 1.3GB of GPU memory. Now let's see how muc
## Load Model
-First, we load the `google-bert/bert-large-uncased` model. We load the model weights directly to the GPU so that we can check
+First, we load the `google-bert/bert-large-uncased` model. We load the model weights directly to the GPU so that we can check
how much space just the weights use.
-
```py
>>> from transformers import AutoModelForSequenceClassification
@@ -105,17 +103,16 @@ how much space just the weights use.
GPU memory occupied: 2631 MB.
```
-We can see that the model weights alone take up 1.3 GB of GPU memory. The exact number depends on the specific
-GPU you are using. Note that on newer GPUs a model can sometimes take up more space since the weights are loaded in an
-optimized fashion that speeds up the usage of the model. Now we can also quickly check if we get the same result
+We can see that the model weights alone take up 1.3 GB of GPU memory. The exact number depends on the specific
+GPU you are using. Note that on newer GPUs a model can sometimes take up more space since the weights are loaded in an
+optimized fashion that speeds up the usage of the model. Now we can also quickly check if we get the same result
as with `nvidia-smi` CLI:
-
```bash
nvidia-smi
```
-```bash
+```text
Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03 Driver Version: 460.91.03 CUDA Version: 11.2 |
@@ -138,8 +135,8 @@ Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
```
-We get the same number as before and you can also see that we are using a V100 GPU with 16GB of memory. So now we can
-start training the model and see how the GPU memory consumption changes. First, we set up a few standard training
+We get the same number as before and you can also see that we are using a V100 GPU with 16GB of memory. So now we can
+start training the model and see how the GPU memory consumption changes. First, we set up a few standard training
arguments:
```py
@@ -154,7 +151,7 @@ default_args = {
- If you plan to run multiple experiments, in order to properly clear the memory between experiments, restart the Python
+ If you plan to run multiple experiments, in order to properly clear the memory between experiments, restart the Python
kernel between experiments.
@@ -175,15 +172,15 @@ Let's use the [`Trainer`] and train the model without using any GPU performance
>>> print_summary(result)
```
-```
+```text
Time: 57.82
Samples/second: 8.86
GPU memory occupied: 14949 MB.
```
-We see that already a relatively small batch size almost fills up our GPU's entire memory. However, a larger batch size
+We see that already a relatively small batch size almost fills up our GPU's entire memory. However, a larger batch size
can often result in faster model convergence or better end performance. So ideally we want to tune the batch size to our
-model's needs and not to the GPU limitations. What's interesting is that we use much more memory than the size of the model.
+model's needs and not to the GPU limitations. What's interesting is that we use much more memory than the size of the model.
To understand a bit better why this is the case let's have a look at a model's operations and memory needs.
## Anatomy of Model's Operations
@@ -206,10 +203,9 @@ This knowledge can be helpful to know when analyzing performance bottlenecks.
This summary is derived from [Data Movement Is All You Need: A Case Study on Optimizing Transformers 2020](https://huggingface.co/papers/2007.00072)
-
## Anatomy of Model's Memory
-We've seen that training the model uses much more memory than just putting the model on the GPU. This is because there
+We've seen that training the model uses much more memory than just putting the model on the GPU. This is because there
are many components during training that use GPU memory. The components on GPU memory are the following:
1. model weights
@@ -219,8 +215,8 @@ are many components during training that use GPU memory. The components on GPU m
5. temporary buffers
6. functionality-specific memory
-A typical model trained in mixed precision with AdamW requires 18 bytes per model parameter plus activation memory. For
-inference there are no optimizer states and gradients, so we can subtract those. And thus we end up with 6 bytes per
+A typical model trained in mixed precision with AdamW requires 18 bytes per model parameter plus activation memory. For
+inference there are no optimizer states and gradients, so we can subtract those. And thus we end up with 6 bytes per
model parameter for mixed precision inference, plus activation memory.
Let's look at the details.
@@ -244,29 +240,29 @@ Let's look at the details.
- size depends on many factors, the key ones being sequence length, hidden size and batch size.
-There are the input and output that are being passed and returned by the forward and the backward functions and the
+There are the input and output that are being passed and returned by the forward and the backward functions and the
forward activations saved for gradient computation.
**Temporary Memory**
-Additionally, there are all kinds of temporary variables which get released once the calculation is done, but in the
-moment these could require additional memory and could push to OOM. Therefore, when coding it's crucial to think
+Additionally, there are all kinds of temporary variables which get released once the calculation is done, but in the
+moment these could require additional memory and could push to OOM. Therefore, when coding it's crucial to think
strategically about such temporary variables and sometimes to explicitly free those as soon as they are no longer needed.
**Functionality-specific memory**
-Then, your software could have special memory needs. For example, when generating text using beam search, the software
+Then, your software could have special memory needs. For example, when generating text using beam search, the software
needs to maintain multiple copies of inputs and outputs.
**`forward` vs `backward` Execution Speed**
-For convolutions and linear layers there are 2x flops in the backward compared to the forward, which generally translates
-into ~2x slower (sometimes more, because sizes in the backward tend to be more awkward). Activations are usually
-bandwidth-limited, and it’s typical for an activation to have to read more data in the backward than in the forward
-(e.g. activation forward reads once, writes once, activation backward reads twice, gradOutput and output of the forward,
+For convolutions and linear layers there are 2x flops in the backward compared to the forward, which generally translates
+into ~2x slower (sometimes more, because sizes in the backward tend to be more awkward). Activations are usually
+bandwidth-limited, and it's typical for an activation to have to read more data in the backward than in the forward
+(e.g. activation forward reads once, writes once, activation backward reads twice, gradOutput and output of the forward,
and writes once, gradInput).
-As you can see, there are potentially a few places where we could save GPU memory or speed up operations.
-Now that you understand what affects GPU utilization and computation speed, refer to
-the [Methods and tools for efficient training on a single GPU](perf_train_gpu_one) documentation page to learn about
-performance optimization techniques.
+As you can see, there are potentially a few places where we could save GPU memory or speed up operations.
+Now that you understand what affects GPU utilization and computation speed, refer to
+the [Methods and tools for efficient training on a single GPU](perf_train_gpu_one) documentation page to learn about
+performance optimization techniques.
diff --git a/docs/source/en/models.md b/docs/source/en/models.md
index fdfcfba6585a..ae5572c0c77a 100644
--- a/docs/source/en/models.md
+++ b/docs/source/en/models.md
@@ -45,7 +45,6 @@ There are two general types of models you can load:
1. A barebones model, like [`AutoModel`] or [`LlamaModel`], that outputs hidden states.
2. A model with a specific *head* attached, like [`AutoModelForCausalLM`] or [`LlamaForCausalLM`], for performing specific tasks.
-
## Model classes
To get a pretrained model, you need to load the weights into the model. This is done by calling [`~PreTrainedModel.from_pretrained`] which accepts weights from the Hugging Face Hub or a local directory.
@@ -111,7 +110,6 @@ You need enough memory to hold two copies of the model weights (random and pretr
Transformers reduces some of these memory-related challenges with fast initialization, sharded checkpoints, Accelerate's [Big Model Inference](https://hf.co/docs/accelerate/usage_guides/big_modeling) feature, and supporting lower bit data types.
-
### Sharded checkpoints
The [`~PreTrainedModel.save_pretrained`] method automatically shards checkpoints larger than 10GB.
diff --git a/docs/source/en/models_timeline.md b/docs/source/en/models_timeline.md
new file mode 100644
index 000000000000..61514d08ea47
--- /dev/null
+++ b/docs/source/en/models_timeline.md
@@ -0,0 +1,28 @@
+
+
+# Models Timeline
+
+The [Models Timeline](https://huggingface.co/spaces/yonigozlan/Transformers-Timeline) is an interactive chart of how architectures in Transformers have changed over time. You can scroll through models in order, spanning text, vision, audio, video, and multimodal use cases.
+
+Use the filters to narrow models by modality or task. Set custom date ranges to focus on models added during specific periods. Click a model card to see its capabilities, supported tasks, and documentation.
+
+
diff --git a/docs/source/en/modular_transformers.md b/docs/source/en/modular_transformers.md
index 39d29f8a6cd4..17001cc81ee9 100644
--- a/docs/source/en/modular_transformers.md
+++ b/docs/source/en/modular_transformers.md
@@ -82,7 +82,7 @@ class RobertaForMaskedLM(BertForMaskedLM):
If you don't use the defined dependency, you'll receive the following error.
-```
+```text
ValueError: You defined `RobertaEmbeddings` in the modular_roberta.py, it should be used when you define `BertModel`, as it is one of it's direct dependencies. Make sure you use it in the `__init__` function.
```
diff --git a/docs/source/en/open_webui.md b/docs/source/en/open_webui.md
index 9042131631e7..2946fc95f145 100644
--- a/docs/source/en/open_webui.md
+++ b/docs/source/en/open_webui.md
@@ -9,6 +9,7 @@ transformers serve --enable-cors
```
Before you can speak into Open WebUI, you need to update its settings to use your server for speech to text (STT) tasks. Launch Open WebUI, and navigate to the audio tab inside the admin settings. If you're using Open WebUI with the default ports, [this link (default)](http://localhost:3000/admin/settings/audio) or [this link (python deployment)](http://localhost:8080/admin/settings/audio) will take you there. Do the following changes there:
+
1. Change the type of "Speech-to-Text Engine" to "OpenAI";
2. Update the address to your server's address -- `http://localhost:8000/v1` by default;
3. Type your model of choice into the "STT Model" field, e.g. `openai/whisper-large-v3` ([available models](https://huggingface.co/models?pipeline_tag=automatic-speech-recognition&sort=trending)).
diff --git a/docs/source/en/pad_truncation.md b/docs/source/en/pad_truncation.md
index 345f86283d12..45b2509e86de 100644
--- a/docs/source/en/pad_truncation.md
+++ b/docs/source/en/pad_truncation.md
@@ -22,25 +22,25 @@ In most cases, padding your batch to the length of the longest sequence and trun
The `padding` argument controls padding. It can be a boolean or a string:
- - `True` or `'longest'`: pad to the longest sequence in the batch (no padding is applied if you only provide
+- `True` or `'longest'`: pad to the longest sequence in the batch (no padding is applied if you only provide
a single sequence).
- - `'max_length'`: pad to a length specified by the `max_length` argument or the maximum length accepted
+- `'max_length'`: pad to a length specified by the `max_length` argument or the maximum length accepted
by the model if no `max_length` is provided (`max_length=None`). Padding will still be applied if you only provide a single sequence.
- - `False` or `'do_not_pad'`: no padding is applied. This is the default behavior.
+- `False` or `'do_not_pad'`: no padding is applied. This is the default behavior.
The `truncation` argument controls truncation. It can be a boolean or a string:
- - `True` or `'longest_first'`: truncate to a maximum length specified by the `max_length` argument or
+- `True` or `'longest_first'`: truncate to a maximum length specified by the `max_length` argument or
the maximum length accepted by the model if no `max_length` is provided (`max_length=None`). This will
truncate token by token, removing a token from the longest sequence in the pair until the proper length is
reached.
- - `'only_second'`: truncate to a maximum length specified by the `max_length` argument or the maximum
+- `'only_second'`: truncate to a maximum length specified by the `max_length` argument or the maximum
length accepted by the model if no `max_length` is provided (`max_length=None`). This will only truncate
the second sentence of a pair if a pair of sequences (or a batch of pairs of sequences) is provided.
- - `'only_first'`: truncate to a maximum length specified by the `max_length` argument or the maximum
+- `'only_first'`: truncate to a maximum length specified by the `max_length` argument or the maximum
length accepted by the model if no `max_length` is provided (`max_length=None`). This will only truncate
the first sentence of a pair if a pair of sequences (or a batch of pairs of sequences) is provided.
- - `False` or `'do_not_truncate'`: no truncation is applied. This is the default behavior.
+- `False` or `'do_not_truncate'`: no truncation is applied. This is the default behavior.
The `max_length` argument controls the length of the padding and truncation. It can be an integer or `None`, in which case it will default to the maximum length the model can accept. If the model has no specific maximum input length, truncation or padding to `max_length` is deactivated.
diff --git a/docs/source/en/perf_infer_gpu_multi.md b/docs/source/en/perf_infer_gpu_multi.md
index 01823dd5b200..21d1817e302b 100644
--- a/docs/source/en/perf_infer_gpu_multi.md
+++ b/docs/source/en/perf_infer_gpu_multi.md
@@ -45,13 +45,7 @@ This guide shows how to enable tensor parallelism with Transformers and differen
## Partitioning a model
-Transformers supports tensor parallelism if a model has a `tp_plan`. There are two plans to partition a model.
-
-- The `auto` tensor parallelism plan partitions a model (see the supported models above) based on a predefined configuration.
-- You can also manually specify your own partitioning plan and pass it to the `tp_plan` parameter in [`~PreTrainedModel.from_pretrained`].
-
-
-
+Transformers supports tensor parallelism if a model has a `tp_plan`. Set `tp_plan="auto"` to automatically use a tensor parallelism plan based on a model's predefined configuration.
```py
import os
@@ -78,32 +72,6 @@ Launch the inference script above on [torchrun](https://pytorch.org/docs/stable/
torchrun --nproc-per-node 4 demo.py
```
-
-
-
-Define a tensor parallel plan for each layer in `tp_plan` and pass it to [`~PreTrainedModel.from_pretrained`]. The example below uses a combination of column and row partitioning. Refer to the [Partitioning strategies](#partitioning-strategies) section to learn about other supported partitioning strategies.
-
-> [!WARNING]
-> Manually specifying your own partitioning plan requires a good understanding of the model architecture and how the partitioning strategies interact together. If you are not sure about the partitioning strategies, the resulting model can be very slow, even failing or incorrect. Refer to the [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) to learn more.
-
-```py
-from transformers import AutoModelForCausalLM
-
-tp_plan = {
- "model.layers.*.self_attn.q_proj": "colwise",
- "model.layers.*.self_attn.k_proj": "colwise",
- "model.layers.*.self_attn.v_proj": "colwise",
- "model.layers.*.self_attn.o_proj": "rowwise",
- ...
-}
-
-model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, tp_plan=tp_plan)
-print(model._tp_plan)
-```
-
-
-
-
## Partitioning strategies
All partitioning strategies are defined in the [`ParallelInterface`] class which maps a string to the strategy implementation. You don't need to interact with this class directly since all the strategies are set with `tp_plan` in [`~PreTrainedModel.from_pretrained`], but it is useful for checking what strategies are available.
diff --git a/docs/source/en/perf_infer_gpu_one.md b/docs/source/en/perf_infer_gpu_one.md
index 33fe9358fe7d..ed6c2b4a8d1a 100644
--- a/docs/source/en/perf_infer_gpu_one.md
+++ b/docs/source/en/perf_infer_gpu_one.md
@@ -69,7 +69,7 @@ Learn in more detail the concepts underlying 8-bit quantization in the [Gentle I
Set up a [`BitsAndBytesConfig`] and set `load_in_4bit=True` to load a model in 4-bit precision. The [`BitsAndBytesConfig`] is passed to the `quantization_config` parameter in [`~PreTrainedModel.from_pretrained`].
-Allow Accelerate to automatically distribute the model across your available hardware by setting `device_map=“auto”`.
+Allow Accelerate to automatically distribute the model across your available hardware by setting `device_map="auto"`.
Place all inputs on the same device as the model.
diff --git a/docs/source/en/perf_train_gaudi.md b/docs/source/en/perf_train_gaudi.md
index 2ba792d484a3..0e5140d731ec 100644
--- a/docs/source/en/perf_train_gaudi.md
+++ b/docs/source/en/perf_train_gaudi.md
@@ -20,14 +20,17 @@ The Intel Gaudi AI accelerator family includes [Intel Gaudi 1](https://habana.ai
[`TrainingArguments`], [`Trainer`] and [`Pipeline`] detect and set the backend device to `hpu` if an Intel Gaudi device is available. No additional changes are required to enable training and inference on your device.
Some modeling code in Transformers is not optimized for HPU lazy mode. If you encounter any errors, set the environment variable below to use eager mode:
-```
-PT_HPU_LAZY_MODE=0
+
+```bash
+export PT_HPU_LAZY_MODE=0
```
In some cases, you'll also need to enable int64 support to avoid casting issues with long integers:
+
+```bash
+export PT_ENABLE_INT64_SUPPORT=1
```
-PT_ENABLE_INT64_SUPPORT=1
-```
+
Refer to the [Gaudi docs](https://docs.habana.ai/en/latest/index.html) for more details.
> [!TIP]
diff --git a/docs/source/en/philosophy.md b/docs/source/en/philosophy.md
index 7cfa46458b75..e98b1fa57bd9 100644
--- a/docs/source/en/philosophy.md
+++ b/docs/source/en/philosophy.md
@@ -26,24 +26,24 @@ The library was designed with two strong goals in mind:
1. Be as easy and fast to use as possible:
- - We strongly limited the number of user-facing abstractions to learn, in fact, there are almost no abstractions,
+- We strongly limited the number of user-facing abstractions to learn, in fact, there are almost no abstractions,
just three standard classes required to use each model: [configuration](main_classes/configuration),
[models](main_classes/model), and a preprocessing class ([tokenizer](main_classes/tokenizer) for NLP, [image processor](main_classes/image_processor) for vision, [feature extractor](main_classes/feature_extractor) for audio, and [processor](main_classes/processors) for multimodal inputs).
- - All of these classes can be initialized in a simple and unified way from pretrained instances by using a common
+- All of these classes can be initialized in a simple and unified way from pretrained instances by using a common
`from_pretrained()` method which downloads (if needed), caches and
loads the related class instance and associated data (configurations' hyperparameters, tokenizers' vocabulary,
and models' weights) from a pretrained checkpoint provided on [Hugging Face Hub](https://huggingface.co/models) or your own saved checkpoint.
- - On top of those three base classes, the library provides two APIs: [`pipeline`] for quickly
+- On top of those three base classes, the library provides two APIs: [`pipeline`] for quickly
using a model for inference on a given task and [`Trainer`] to quickly train or fine-tune a PyTorch model.
- - As a consequence, this library is NOT a modular toolbox of building blocks for neural nets. If you want to
+- As a consequence, this library is NOT a modular toolbox of building blocks for neural nets. If you want to
extend or build upon the library, just use regular Python or PyTorch and inherit from the base
classes of the library to reuse functionalities like model loading and saving. If you'd like to learn more about our coding philosophy for models, check out our [Repeat Yourself](https://huggingface.co/blog/transformers-design-philosophy) blog post.
2. Provide state-of-the-art models with performances as close as possible to the original models:
- - We provide at least one example for each architecture which reproduces a result provided by the official authors
+- We provide at least one example for each architecture which reproduces a result provided by the official authors
of said architecture.
- - The code is usually as close to the original code base as possible which means some PyTorch code may be not as
+- The code is usually as close to the original code base as possible which means some PyTorch code may be not as
*pytorchic* as it could be as a result of being converted from other Deep Learning frameworks.
A few other goals:
diff --git a/docs/source/en/pipeline_gradio.md b/docs/source/en/pipeline_gradio.md
index 0cd65665d33d..b53bcc8bd184 100644
--- a/docs/source/en/pipeline_gradio.md
+++ b/docs/source/en/pipeline_gradio.md
@@ -45,8 +45,8 @@ gr.Interface.from_pipeline(pipeline).launch(share=True)
The Space below is created with the code above and hosted on Spaces.
diff --git a/docs/source/en/pipeline_webserver.md b/docs/source/en/pipeline_webserver.md
index 0112d116c47d..37d245483b94 100644
--- a/docs/source/en/pipeline_webserver.md
+++ b/docs/source/en/pipeline_webserver.md
@@ -82,6 +82,7 @@ Query the server with a POST request.
```bash
curl -X POST -d "Paris is the [MASK] of France." http://localhost:8000/
```
+
This should return the output below.
```bash
diff --git a/docs/source/en/pr_checks.md b/docs/source/en/pr_checks.md
index a5634c29ee49..5fdbbbab05bc 100644
--- a/docs/source/en/pr_checks.md
+++ b/docs/source/en/pr_checks.md
@@ -21,6 +21,7 @@ rendered properly in your Markdown viewer.
# Checks on a Pull Request
When you open a pull request on 🤗 Transformers, a fair number of checks will be run to make sure the patch you are adding is not breaking anything existing. Those checks are of four types:
+
- regular tests
- documentation build
- code and documentation style
@@ -52,7 +53,6 @@ or for an editable install:
pip install -e .[quality]
```
-
## Tests
All the jobs that begin with `ci/circleci: run_tests_` run parts of the Transformers testing suite. Each of those jobs focuses on a part of the library in a certain environment: for instance `ci/circleci: run_tests_pipelines` runs the pipeline tests in an environment where all pipeline-related requirements are installed.
@@ -195,6 +195,7 @@ Another way when the patterns are just different casings of the same replacement
```
In this case, the code is copied from `BertForSequenceClassification` by replacing:
+
- `Bert` by `MobileBert` (for instance when using `MobileBertModel` in the init)
- `bert` by `mobilebert` (for instance when defining `self.mobilebert`)
- `BERT` by `MOBILEBERT` (in the constant `MOBILEBERT_INPUTS_DOCSTRING`)
diff --git a/docs/source/en/quantization/auto_round.md b/docs/source/en/quantization/auto_round.md
index 15abf9faa846..7526597ee86f 100644
--- a/docs/source/en/quantization/auto_round.md
+++ b/docs/source/en/quantization/auto_round.md
@@ -11,18 +11,17 @@ rendered properly in your Markdown viewer.
# AutoRound
-[AutoRound](https://github.com/intel/auto-round) is an advanced quantization algorithm that delivers strong accuracy, even at 2-bit precision.
-It leverages sign gradient descent to fine-tune both rounding values and min-max clipping thresholds in just 200 steps. Designed for broad compatibility, it seamlessly supports a wide range of LLMs and is actively expanding to cover more VLMs as well.
+[AutoRound](https://github.com/intel/auto-round) is an advanced quantization algorithm that delivers strong accuracy, even at 2-bit precision.
+It leverages sign gradient descent to fine-tune both rounding values and min-max clipping thresholds in just 200 steps. Designed for broad compatibility, it seamlessly supports a wide range of LLMs and is actively expanding to cover more VLMs as well.
It also supports quantization and inference across multiple hardware platforms, including CPU, XPU, and CUDA.
-AutoRound also offers a variety of useful features, including mixed-bit tuning and inference, lm-head quantization, support for exporting to formats like GPTQ/AWQ/GGUF, and flexible tuning recipes.
+AutoRound also offers a variety of useful features, including mixed-bit tuning and inference, lm-head quantization, support for exporting to formats like GPTQ/AWQ/GGUF, and flexible tuning recipes.
For a comprehensive overview and the latest updates, check out the AutoRound [README](https://github.com/intel/auto-round).
-AutoRound was originally developed as part of the [Intel Neural Compressor](https://github.com/intel/neural-compressor), serving as a general-purpose model compression library for deep learning.
-It has since evolved into a standalone library focused specifically on low-precision optimization for large language models (LLMs).
+AutoRound was originally developed as part of the [Intel Neural Compressor](https://github.com/intel/neural-compressor), serving as a general-purpose model compression library for deep learning.
+It has since evolved into a standalone library focused specifically on low-precision optimization for large language models (LLMs).
AutoRound remains fully integrated with the Intel Neural Compressor, and you can explore the repository for more details.
-
## Installation
```bash
@@ -51,6 +50,7 @@ Currently, only offline mode is supported to generate quantized models.
### Command Line Usage
+
```bash
auto-round \
--model facebook/opt-125m \
@@ -59,7 +59,7 @@ auto-round \
--output_dir ./tmp_autoround
```
-AutoRound also offer another two recipes, `auto-round-best` and `auto-round-light`, designed for optimal accuracy and improved speed, respectively.
+AutoRound also offer another two recipes, `auto-round-best` and `auto-round-light`, designed for optimal accuracy and improved speed, respectively.
For 2 bits, we recommend using `auto-round-best` or `auto-round`.
@@ -99,6 +99,7 @@ autoround.quantize_and_save(output_dir, format='auto_round')
### AutoRoundBest recipe
This setting provides the best accuracy in most scenarios but is 4–5× slower than the standard AutoRound recipe. It is especially recommended for 2-bit quantization and is a good choice if sufficient resources are available.
+
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from auto_round import AutoRound
@@ -121,6 +122,7 @@ autoround = AutoRound(
output_dir = "./tmp_autoround"
autoround.quantize_and_save(output_dir, format='auto_round')
```
+
@@ -230,7 +232,7 @@ print(tokenizer.decode(model.generate(**inputs, max_new_tokens=50, do_sample=Fal
AutoRound automatically selects the backend for each layer based on compatibility. In general, the priority order is Marlin > ExLLaMAV2 > Triton, but the final choice depends on factors such as group size, bit width, packing format, hardware device, and other implementation details. For more details, please refer to [backends](https://github.com/intel/auto-round?tab=readme-ov-file#specify-backend),
-The backend may not always be the most suitable for certain devices.
+The backend may not always be the most suitable for certain devices.
You can specify your preferred backend such as "ipex" for CPU, "ipex/triton" for XPU, "marlin/exllamav2/triton" for CUDA, according to your needs or hardware compatibility. Please note that additional corresponding libraries may be required.
```python
@@ -247,7 +249,6 @@ print(tokenizer.decode(model.generate(**inputs, max_new_tokens=50, do_sample=Fal
-
### Convert GPTQ/AWQ to AutoRound
@@ -277,7 +278,6 @@ the [transformers](https://github.com/huggingface/transformers/issues) repositor
If you encounter any issues with auto-round, please open an issue on
the [AutoRound](https://github.com/intel/auto-round/issues) repository.
-
## Acknowledgement
Special thanks to open-source low precision libraries such as AutoGPTQ, AutoAWQ, GPTQModel, Triton, Marlin, and ExLLaMAV2 for providing low-precision CUDA kernels, which are leveraged in AutoRound.
diff --git a/docs/source/en/quantization/awq.md b/docs/source/en/quantization/awq.md
index b6437e2588a8..b2cf4b9ecdf6 100644
--- a/docs/source/en/quantization/awq.md
+++ b/docs/source/en/quantization/awq.md
@@ -25,6 +25,7 @@ Run the command below to install autoawq
```bash
pip install autoawq
```
+
> [!WARNING]
> AutoAWQ downgrades Transformers to version 4.47.1. If you want to do inference with AutoAWQ, you may need to reinstall your Transformers' version after installing AutoAWQ.
diff --git a/docs/source/en/quantization/bitnet.md b/docs/source/en/quantization/bitnet.md
index 922210b2137b..31474e1d3213 100644
--- a/docs/source/en/quantization/bitnet.md
+++ b/docs/source/en/quantization/bitnet.md
@@ -41,7 +41,7 @@ model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")
## Kernels
-`@torch.compile` is used to unpack the weights and perform the forward pass. It’s very straightforward to implement and delivers significant speed improvements. Additional optimized kernels will be integrated in future versions.
+`@torch.compile` is used to unpack the weights and perform the forward pass. It's very straightforward to implement and delivers significant speed improvements. Additional optimized kernels will be integrated in future versions.
## Resources
diff --git a/docs/source/en/quantization/bitsandbytes.md b/docs/source/en/quantization/bitsandbytes.md
index 60c3c2dfebf9..81238c0707e7 100644
--- a/docs/source/en/quantization/bitsandbytes.md
+++ b/docs/source/en/quantization/bitsandbytes.md
@@ -16,7 +16,7 @@ rendered properly in your Markdown viewer.
# Bitsandbytes
-The [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes) library provides quantization tools for LLMs through a lightweight Python wrapper around CUDA functions. It enables working with large models using limited computational resources by reducing their memory footprint.
+The [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes) library provides quantization tools for LLMs through a lightweight Python wrapper around hardware accelerator functions. It enables working with large models using limited computational resources by reducing their memory footprint.
At its core, bitsandbytes provides:
@@ -32,36 +32,38 @@ bitsandbytes offers two main quantization features:
> **Note:** For a user-friendly quantization experience, you can use the `bitsandbytes` [community space](https://huggingface.co/spaces/bnb-community/bnb-my-repo).
-
Run the command below to install bitsandbytes.
```bash
pip install --upgrade transformers accelerate bitsandbytes
```
+
To compile from source, follow the instructions in the [bitsandbytes installation guide](https://huggingface.co/docs/bitsandbytes/main/en/installation).
## Hardware Compatibility
-bitsandbytes is currently only supported on CUDA GPUs for CUDA versions 11.0 - 12.8. However, there's an ongoing multi-backend effort under development, which is currently in alpha. If you're interested in providing feedback or testing, check out the [bitsandbytes repository](https://github.com/bitsandbytes-foundation/bitsandbytes) for more information.
+bitsandbytes is supported on NVIDIA GPUs for CUDA versions 11.8 - 13.0, Intel XPU, Intel Gaudi (HPU), and CPU. There is an ongoing effort to support additional platforms. If you're interested in providing feedback or testing, check out the [bitsandbytes repository](https://github.com/bitsandbytes-foundation/bitsandbytes) for more information.
-### CUDA
+### NVIDIA GPUs (CUDA)
+
+This backend is supported on Linux x86-64, Linux aarch64, and Windows platforms.
| Feature | Minimum Hardware Requirement |
|---------|-------------------------------|
-| 8-bit optimizers | NVIDIA Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs * |
-| LLM.int8() | NVIDIA Turing (RTX 20 series, T4) or newer GPUs |
-| NF4/FP4 quantization | NVIDIA Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs * |
+| 8-bit optimizers | NVIDIA Pascal (GTX 10X0 series, P100) or newer GPUs * |
+| LLM.int8() | NVIDIA Turing (RTX 20X0 series, T4) or newer GPUs |
+| NF4/FP4 quantization | NVIDIA Pascal (GTX 10X0 series, P100) or newer GPUs * |
+
+### Intel GPUs (XPU)
-### Multi-backend
+This backend is supported on Linux x86-64 and Windows x86-64 platforms.
-| Backend | Supported Versions | Python versions | Architecture Support | Status |
-|---------|-------------------|----------------|---------------------|---------|
-| AMD ROCm | 6.1+ | 3.10+ | minimum CDNA - gfx90a, RDNA - gfx1100 | Alpha |
-| Apple Silicon (MPS) | WIP | 3.10+ | M1/M2 chips | Planned |
-| Intel CPU | v2.4.0+ (ipex) | 3.10+ | Intel CPU | Alpha |
-| Intel GPU | v2.4.0+ (ipex) | 3.10+ | Intel GPU | Experimental |
-| Ascend NPU | 2.1.0+ (torch_npu) | 3.10+ | Ascend NPU | Experimental |
+### Intel Gaudi (HPU)
-> **Note:** Bitsandbytes is moving away from the multi-backend approach towards using [Pytorch Custom Operators](https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html), as the main mechanism for supporting new hardware, and dispatching to the correct backend.
+This backend is supported on Linux x86-64 for Gaudi2 and Gaudi3.
+
+### CPU
+
+This backend is supported on Linux x86-64, Linux aarch64, and Windows x86-64 platforms.
## Quantization Examples
@@ -116,6 +118,7 @@ model = AutoModelForCausalLM.from_pretrained(
model.push_to_hub("bloom-560m-8bit")
```
+
@@ -166,6 +169,7 @@ model = AutoModelForCausalLM.from_pretrained(
model.push_to_hub("bloom-560m-4bit")
```
+
diff --git a/docs/source/en/quantization/compressed_tensors.md b/docs/source/en/quantization/compressed_tensors.md
index a3b01a1b4489..4f55f008aa8d 100644
--- a/docs/source/en/quantization/compressed_tensors.md
+++ b/docs/source/en/quantization/compressed_tensors.md
@@ -65,11 +65,11 @@ print(f"{mem_params/2**30:.4f} GB")
## Model checkpoint
-compressed-tensor models are defined through its configuration entry. The following example is taken from the [nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf](https://huggingface.co/nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf/blob/main/config.json) `config.json` file.
+Compressed-tensor models are defined through its configuration entry. The following example is taken from the [nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf](https://huggingface.co/nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf/blob/main/config.json) `config.json` file.
There are a lot of entries to allow for flexible expression both during and after compression, but the entries for loading and inference can be simplified to focus on just a few key entries.
-```yaml
+```json
"quantization_config": {
"config_groups": {
"group_0": {
@@ -97,31 +97,31 @@ The config file specifies the quantization of a config group (`group_0`), which
For a more detailed look at the model weights, use the [safetensors viewer](https://huggingface.co/nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf?show_file_info=model.safetensors.index.json) on the model card to see the quantized weights, input scale, and weight scale for all [nn.Linear](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) modules.
-| Tensors | Shape | Precision |
+| Tensors | Shape | Precision |
| ------- | ----- | --------- |
-model.layers.0.input_layernorm.weight | [4 096] | BF16
-model.layers.0.mlp.down_proj.input_scale | [1] | BF16
-model.layers.0.mlp.down_proj.weight | [4 096, 14 336] | F8_E4M3
-model.layers.0.mlp.down_proj.weight_scale | [1] | BF16
-model.layers.0.mlp.gate_proj.input_scale | [1] | BF16
-model.layers.0.mlp.gate_proj.weight | [14 336, 4 096] | F8_E4M3
-model.layers.0.mlp.gate_proj.weight_scale | [1] | BF16
-model.layers.0.mlp.up_proj.input_scale| [1] |BF16
-model.layers.0.mlp.up_proj.weight | [14 336, 4 096] | F8_E4M3
-model.layers.0.mlp.up_proj.weight_scale | [1] | BF16
-model.layers.0.post_attention_layernorm.weight | [4 096] |BF16
-model.layers.0.self_attn.k_proj.input_scale | [1] | BF16
-model.layers.0.self_attn.k_proj.weight | [1 024, 4 096]| F8_E4M3
-model.layers.0.self_attn.k_proj.weight_scale |[1] | BF16
-model.layers.0.self_attn.o_proj.input_scale | [1] | BF16
-model.layers.0.self_attn.o_proj.weight | [4 096, 4 096] | F8_E4M3
-model.layers.0.self_attn.o_proj.weight_scale | [1] | BF16
-model.layers.0.self_attn.q_proj.input_scale | [1] | BF16
-model.layers.0.self_attn.q_proj.weight | [4 096, 4 096] | F8_E4M3
-model.layers.0.self_attn.q_proj.weight_scale | [1] | BF16
-model.layers.0.self_attn.v_proj.input_scale | [1] | BF16
-model.layers.0.self_attn.v_proj.weight | [1 024, 4 096] | F8_E4M3
-model.layers.0.self_attn.v_proj.weight_scale | [1] | BF16
+|model.layers.0.input_layernorm.weight | [4 096] | BF16|
+|model.layers.0.mlp.down_proj.input_scale | [1] | BF16|
+|model.layers.0.mlp.down_proj.weight | [4 096, 14 336] | F8_E4M3|
+|model.layers.0.mlp.down_proj.weight_scale | [1] | BF16|
+|model.layers.0.mlp.gate_proj.input_scale | [1] | BF16|
+|model.layers.0.mlp.gate_proj.weight | [14 336, 4 096] | F8_E4M3|
+|model.layers.0.mlp.gate_proj.weight_scale | [1] | BF16|
+|model.layers.0.mlp.up_proj.input_scale| [1] |BF16|
+|model.layers.0.mlp.up_proj.weight | [14 336, 4 096] | F8_E4M3|
+|model.layers.0.mlp.up_proj.weight_scale | [1] | BF16|
+|model.layers.0.post_attention_layernorm.weight | [4 096] |BF16|
+|model.layers.0.self_attn.k_proj.input_scale | [1] | BF16|
+|model.layers.0.self_attn.k_proj.weight | [1 024, 4 096]| F8_E4M3|
+|model.layers.0.self_attn.k_proj.weight_scale |[1] | BF16|
+|model.layers.0.self_attn.o_proj.input_scale | [1] | BF16|
+|model.layers.0.self_attn.o_proj.weight | [4 096, 4 096] | F8_E4M3|
+|model.layers.0.self_attn.o_proj.weight_scale | [1] | BF16|
+|model.layers.0.self_attn.q_proj.input_scale | [1] | BF16|
+|model.layers.0.self_attn.q_proj.weight | [4 096, 4 096] | F8_E4M3|
+|model.layers.0.self_attn.q_proj.weight_scale | [1] | BF16|
+|model.layers.0.self_attn.v_proj.input_scale | [1] | BF16|
+|model.layers.0.self_attn.v_proj.weight | [1 024, 4 096] | F8_E4M3|
+|model.layers.0.self_attn.v_proj.weight_scale | [1] | BF16|
When loading a compressed-tensors model with the [`~quantizers.HFQuantizer`] integration, all the [nn.Linear](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) modules specified in the quantization config are replaced by [CompressedLinear](https://github.com/neuralmagic/compressed-tensors/blob/975cb223b19fcac2b98a4271d17668462d4d6e1d/src/compressed_tensors/linear/compressed_linear.py#L30) modules that manage the compressed weights and forward pass for inference. The `lm_head` module is still kept as an unquantized nn.Linear module.
diff --git a/docs/source/en/quantization/concept_guide.md b/docs/source/en/quantization/concept_guide.md
index ff300b9d48a5..df3a2bdc6f2a 100644
--- a/docs/source/en/quantization/concept_guide.md
+++ b/docs/source/en/quantization/concept_guide.md
@@ -18,12 +18,11 @@ rendered properly in your Markdown viewer.
Quantization reduces the memory footprint and computational cost of large machine learning models like those found in the Transformers library. It achieves this by representing the model's weights and or activations with lower-precision data types (like 8-bit integers or int8) instead of the standard 32-bit floating-point (float32).
-
Reducing a model's precision offers several significant benefits:
-- Smaller model size: Lower-precision data types require less storage space. An int8 model, for example, is roughly 4 times smaller than its float32 counterpart.
-- Faster inference: Operations on lower-precision data types, especially integers, can be significantly faster on compatible hardware (CPUs and GPUs often have specialized instructions for int8 operations). This leads to lower latency.
-- Reduced energy consumption: Faster computations and smaller memory transfers often translate to lower power usage.
+- Smaller model size: Lower-precision data types require less storage space. An int8 model, for example, is roughly 4 times smaller than its float32 counterpart.
+- Faster inference: Operations on lower-precision data types, especially integers, can be significantly faster on compatible hardware (CPUs and GPUs often have specialized instructions for int8 operations). This leads to lower latency.
+- Reduced energy consumption: Faster computations and smaller memory transfers often translate to lower power usage.
The primary trade-off in quantization is *efficiency* vs. *accuracy*. Reducing precision saves resources but inevitably introduces small errors (quantization noise). The goal is to minimize this error using appropriate schemes (affine/symmetric), granularity (per-tensor/channel), and techniques (PTQ/QAT) so that the model's performance on its target task degrades as little as possible.
@@ -46,8 +45,7 @@ The most common method is *affine quantization*. For a given float32 tensor (lik
There are two main ways to perform this mapping, *symmetric* and *asymmetric*. The choice between symmetric and asymmetric quantization determines how the float32 range is mapped to the int8 range.
- Symmetric: This method assumes the original float32 range is symmetric around zero ( \\([ -a, a ]\\) ). This range is mapped symmetrically to the int8 range, for example, \\([-127, 127]\\). A key characteristic is that the float32 value \\(0.0\\) maps directly to the int8 value \\(0\\). This only requires one parameter, the **scale ( \\(S\\) )**, to define the mapping. It can simplify computations, but it might be less accurate if the original data distribution isn't naturally centered around zero.
-- Asymmetric (Affine): This method does not assume the data is centered around zero. It maps the exact range \\([val_{min}, val_{max}]\\) from float32 to the full int8 range, like \\([-128, 127]\\). This requires two parameters, a **scale ( \\(S\\) )** and a **zero-point ( \\(Z\\) )**.
-
+- Asymmetric (Affine): This method does not assume the data is centered around zero. It maps the exact range \\([val_{min}, val_{max}]\\) from float32 to the full int8 range, like \\([-128, 127]\\). This requires two parameters, a **scale ( \\(S\\) )** and a **zero-point ( \\(Z\\) )**.
scale ( \\(S\\) ): A positive float32 number representing the ratio between the float32 and the int8 range.
@@ -134,8 +132,7 @@ There are two main types of quantization techniques.
## Quantization in Transformers
-Transformers integrates several quantization backends such as bitsandbytes, torchao, compressed-tensors, and more (refer to the quantization [overview](./overview) for more backends).
-
+Transformers integrates several quantization backends such as bitsandbytes, torchao, compressed-tensors, and more (refer to the quantization [overview](./overview) for more backends).
All backends are unified under the [`HfQuantizer`] API and associated [`QuantizationConfig`] classes. You can integrate your own custom quantization backends by implementing a custom [`HfQuantizer`] and [`QuantizationConfig`], as shown in the [Contribution](./contribute) guide.
@@ -165,7 +162,6 @@ model = AutoModelForCausalLM.from_pretrained(
)
```
-
## Resources
To explore quantization and related performance optimization concepts more deeply, check out the following resources.
@@ -175,4 +171,4 @@ To explore quantization and related performance optimization concepts more deepl
- [Introduction to Quantization cooked in 🤗 with 💗🧑🍳](https://huggingface.co/blog/merve/quantization)
- [EfficientML.ai Lecture 5 - Quantization Part I](https://www.youtube.com/watch?v=RP23-dRVDWM)
- [Making Deep Learning Go Brrrr From First Principles](https://horace.io/brrr_intro.html)
-- [Accelerating Generative AI with PyTorch Part 2: LLM Optimizations](https://pytorch.org/blog/accelerating-generative-ai-2/)
\ No newline at end of file
+- [Accelerating Generative AI with PyTorch Part 2: LLM Optimizations](https://pytorch.org/blog/accelerating-generative-ai-2/)
diff --git a/docs/source/en/quantization/finegrained_fp8.md b/docs/source/en/quantization/finegrained_fp8.md
index bbf273d8d933..1afd1505029b 100644
--- a/docs/source/en/quantization/finegrained_fp8.md
+++ b/docs/source/en/quantization/finegrained_fp8.md
@@ -59,4 +59,4 @@ Use [`~PreTrainedModel.save_pretrained`] to save the quantized model and reload
quant_path = "/path/to/save/quantized/model"
model.save_pretrained(quant_path)
model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto")
-```
\ No newline at end of file
+```
diff --git a/docs/source/en/quantization/fp_quant.md b/docs/source/en/quantization/fp_quant.md
index 7c12fb870531..4888795a6d77 100644
--- a/docs/source/en/quantization/fp_quant.md
+++ b/docs/source/en/quantization/fp_quant.md
@@ -18,7 +18,9 @@ rendered properly in your Markdown viewer.
[FP-Quant](https://github.com/IST-DASLab/FP-Quant) is a family of quantization algorithms tailored for the Blackwell generation of Nvidia GPUs. The goal is to allow for efficient post-training quantization (PTQ) and quantization-aware training (QAT) of LLMs in the [MXFP4 and NVFP4 data-types](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
-Currently, only PTQ with MXFP4 is supported. Models can either be quantized on the fly with `quantization_config=FPQuantConfig()`:
+This integration accompanies the pre-print of the [**Bridging the Gap Between Promise and Performance for Microscaling FP4 Quantization**](https://arxiv.org/abs/2509.23202) pre-print.
+
+Currently, only QAT is only supported with `pseudoquantization=True`. Models can either be quantized on the fly with `quantization_config=FPQuantConfig()`:
```python
from transformers import AutoModelForCausalLM, AutoTokenizer, FPQuantConfig
@@ -34,6 +36,8 @@ model = AutoModelForCausalLM.from_pretrained(
or pre-processed with GPTQ for better quality (see [FP Format Quantization Harness](https://github.com/IST-DASLab/FP-Quant)).
+You can choose between MXFP4 and NVFP4 with `FPQuantConfig(forward_dtype="mxfp4")`. NVFP4 provides better quality but uses a little more memory.
+
A **Blackwell-generation GPU is required** to run the kernels. Runtime support for FP-Quant is implemented through the [QuTLASS](https://github.com/IST-DASLab/qutlass) library and a lightweight PyTorch interface lib [`fp_quant`](https://github.com/IST-DASLab/FP-Quant/tree/master/inference_lib). We recommend installing the former **from source** and the latter with `pip install fp_quant`.
Users **without a Blackwell-generation GPU** , can use the method with `quantization_config=FPQuantConfig(pseudoquant=True)` without having to install [QuTLASS](https://github.com/IST-DASLab/qutlass). This would provide no speedups but would fully emulate the effect of quantization.
diff --git a/docs/source/en/quantization/mxfp4.md b/docs/source/en/quantization/mxfp4.md
index a2b9f7634c8d..dd313c5555ed 100644
--- a/docs/source/en/quantization/mxfp4.md
+++ b/docs/source/en/quantization/mxfp4.md
@@ -16,7 +16,7 @@ rendered properly in your Markdown viewer.
# MXFP4
-Note: MXFP4 quantisation currently only works for OpenAI GPT-OSS 120b and 20b.
+Note: MXFP4 quantisation currently only works for OpenAI GPT-OSS 120b and 20b.
MXFP4 is a 4-bit floating point format that dramatically reduces the memory requirements of large models. Large models (GPT-OSS-120B) can fit on a single 80GB GPU and smaller models (GPT-OSS-20B) only require 16GB of memory. It uses blockwise scaling to preserve it's range and accuracy, which typically becomes degraded at lower precisions.
@@ -25,7 +25,6 @@ To use MXPF4, make sure your hardware meets the following requirements.
- Install Accelerate, kernels, and Triton ≥ 3.4. Only manually install Triton ≥ 3.4 if you're using PyTorch 2.7 because it is already supported in PyTorch 2.8.
- NVIDIA GPU Compute Capability ≥ 7.5 which includes Tesla GPUs and newer. Use [get_device_capability](https://docs.pytorch.org/docs/stable/generated/torch.cuda.get_device_capability.html) to check Compute Capability.
-
```python
from torch import cuda
cuda.get_device_capability()
@@ -54,7 +53,6 @@ print(cfg.quantization_config)
# }
```
-
## MXFP4 kernels
Transformers automatically pulls the MXFP4-aware Triton kernels from the community repository when you load a model that needs them. The kernels are stored in your local cache and used during the forward pass.
@@ -67,7 +65,6 @@ You can use [hf cache scan](https://huggingface.co/docs/huggingface_hub/en/guide
hf cache scan
```
-
```shell
REPO ID REPO TYPE SIZE ON DISK
-------------------------------- --------- ------------
diff --git a/docs/source/en/quantization/overview.md b/docs/source/en/quantization/overview.md
index ceab195b2b59..0a8dee1e33ae 100644
--- a/docs/source/en/quantization/overview.md
+++ b/docs/source/en/quantization/overview.md
@@ -27,14 +27,14 @@ Use the Space below to help you pick a quantization method depending on your har
| [AQLM](./aqlm) | 🔴 | 🟢 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 1/2 | 🟢 | 🟢 | 🟢 | https://github.com/Vahe1994/AQLM |
| [AutoRound](./auto_round) | 🔴 | 🟢 | 🟢 | 🔴 | 🔴 | 🟢 | 🔴 | 2/3/4/8 | 🔴 | 🟢 | 🟢 | https://github.com/intel/auto-round |
| [AWQ](./awq) | 🔴 | 🟢 | 🟢 | 🟢 | 🔴 | 🟢 | ? | 4 | 🟢 | 🟢 | 🟢 | https://github.com/casper-hansen/AutoAWQ |
-| [bitsandbytes](./bitsandbytes) | 🟢 | 🟡 | 🟢 | 🟡 | 🔴 | 🟡 | 🟢 | 4/8 | 🟢 | 🟢 | 🟢 | https://github.com/bitsandbytes-foundation/bitsandbytes |
+| [bitsandbytes](./bitsandbytes) | 🟢 | 🟢 | 🟢 | 🟡 | 🟡 | 🟢 | 🟢 | 4/8 | 🟢 | 🟢 | 🟢 | https://github.com/bitsandbytes-foundation/bitsandbytes |
| [compressed-tensors](./compressed_tensors) | 🔴 | 🟢 | 🟢 | 🟢 | 🔴 | 🔴 | 🔴 | 1/8 | 🟢 | 🟢 | 🟢 | https://github.com/neuralmagic/compressed-tensors |
| [EETQ](./eetq) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | ? | 8 | 🟢 | 🟢 | 🟢 | https://github.com/NetEase-FuXi/EETQ |
| [FP-Quant](./fp_quant) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 4 | 🔴 | 🟢 | 🟢 | https://github.com/IST-DASLab/FP-Quant |
| [GGUF / GGML (llama.cpp)](../gguf) | 🟢 | 🟢 | 🟢 | 🔴 | 🟢 | 🟢 | 🔴 | 1/8 | 🔴 | [See Notes](../gguf) | [See Notes](../gguf) | https://github.com/ggerganov/llama.cpp |
| [GPTQModel](./gptq) | 🔴 | 🟢 | 🟢 | 🟢 | 🟢 | 🟢 | 🔴 | 2/3/4/8 | 🟢 | 🟢 | 🟢 | https://github.com/ModelCloud/GPTQModel |
| [AutoGPTQ](./gptq) | 🔴 | 🔴 | 🟢 | 🟢 | 🔴 | 🔴 | 🔴 | 2/3/4/8 | 🟢 | 🟢 | 🟢 | https://github.com/AutoGPTQ/AutoGPTQ |
-| [HIGGS](./higgs) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 2/4 | 🔴 | 🟢 | 🟢 | https://github.com/HanGuo97/flute |
+| [HIGGS](./higgs) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 2/4 | 🔴 | 🟢 | 🟢 | https://github.com/HanGuo97/flute |
| [HQQ](./hqq) | 🟢 | 🟢 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 1/8 | 🟢 | 🔴 | 🟢 | https://github.com/mobiusml/hqq/ |
| [optimum-quanto](./quanto) | 🟢 | 🟢 | 🟢 | 🔴 | 🟢 | 🟢 | 🟢 | 2/4/8 | 🔴 | 🔴 | 🟢 | https://github.com/huggingface/optimum-quanto |
| [FBGEMM_FP8](./fbgemm_fp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🔴 | 8 | 🔴 | 🟢 | 🟢 | https://github.com/pytorch/FBGEMM |
@@ -53,7 +53,7 @@ If you are new to quantization, we recommend checking out these beginner-friendl
## User-Friendly Quantization Tools
-If you are looking for a user-friendly quantization experience, you can use the following community spaces and notebooks:
+If you are looking for a user-friendly quantization experience, you can use the following community spaces and notebooks:
* [Bitsandbytes Space](https://huggingface.co/spaces/bnb-community/bnb-my-repo)
* [GGUF Space](https://huggingface.co/spaces/ggml-org/gguf-my-repo)
diff --git a/docs/source/en/quantization/quanto.md b/docs/source/en/quantization/quanto.md
index b3cf58b5b6ad..f58f93025f45 100644
--- a/docs/source/en/quantization/quanto.md
+++ b/docs/source/en/quantization/quanto.md
@@ -66,4 +66,4 @@ model = torch.compile(model)
Read the [Quanto: a PyTorch quantization backend for Optimum](https://huggingface.co/blog/quanto-introduction) blog post to learn more about the library design and benchmarks.
-For more hands-on examples, take a look at the Quanto [notebook](https://colab.research.google.com/drive/16CXfVmtdQvciSh9BopZUDYcmXCDpvgrT?usp=sharing).
\ No newline at end of file
+For more hands-on examples, take a look at the Quanto [notebook](https://colab.research.google.com/drive/16CXfVmtdQvciSh9BopZUDYcmXCDpvgrT?usp=sharing).
diff --git a/docs/source/en/quantization/selecting.md b/docs/source/en/quantization/selecting.md
index 7653e946dd80..e2c7bdf27076 100644
--- a/docs/source/en/quantization/selecting.md
+++ b/docs/source/en/quantization/selecting.md
@@ -26,7 +26,7 @@ Consider the quantization methods below for inference.
| quantization method | use case |
|---|---|
-| bitsandbytes | ease of use and QLoRA fine-tuning on NVIDIA GPUs |
+| bitsandbytes | ease of use and QLoRA fine-tuning on NVIDIA and Intel GPUs |
| compressed-tensors | loading specific quantized formats (FP8, Sparse) |
| GPTQModel or AWQ | good 4-bit accuracy with upfront calibration |
| HQQ | fast on the fly quantization without calibration |
@@ -112,22 +112,22 @@ Consider the quantization method below during fine-tuning to save memory.
### bitsandbytes[[training]]
-* **Description:** The standard method for QLoRA fine-tuning via PEFT.
-* **Pros:** Enables fine-tuning large models on consumer GPUs; widely supported and documented for PEFT.
-* **Cons:** Primarily for NVIDIA GPUs.
+* **Description:** The standard method for QLoRA fine-tuning via PEFT.
+* **Pros:** Enables fine-tuning large models on consumer GPUs; widely supported and documented for PEFT.
+* **Cons:** Primarily for NVIDIA GPUs.
Other methods offer PEFT compatibility, though bitsandbytes is the most established and straightforward path for QLoRA.
-See the [bitsandbytes documentation](./bitsandbytes#qlora) and [PEFT Docs](https://huggingface.co/docs/peft/developer_guides/quantization#aqlm-quantization) for more details.
+See the [bitsandbytes documentation](./bitsandbytes#qlora) and [PEFT Docs](https://huggingface.co/docs/peft/developer_guides/quantization#aqlm-quantization) for more details.
## Research
Methods like [AQLM](./aqlm), [SpQR](./spqr), [VPTQ](./vptq), [HIGGS](./higgs), etc., push the boundaries of compression (< 2-bit) or explore novel techniques.
-* Consider these if:
- * You need extreme compression (sub-4-bit).
- * You are conducting research or require state-of-the-art results from their respective papers.
- * You have significant compute resources available for potentially complex quantization procedures.
+* Consider these if:
+ * You need extreme compression (sub-4-bit).
+ * You are conducting research or require state-of-the-art results from their respective papers.
+ * You have significant compute resources available for potentially complex quantization procedures.
We recommend consulting each methods documentation and associated papers carefully before choosing one for use in production.
## Benchmark Comparison
@@ -154,4 +154,4 @@ The key takeaways are:
| **Sub-4-bit** (VPTQ, AQLM, 2-bit GPTQ) | Extreme (>4x) | Noticeable drop, especially at 2-bit | Quantization times can be very long (AQLM, VPTQ). Performance varies. |
> [!TIP]
-> Always benchmark the performance (accuracy and speed) of the quantized model on your specific task and hardware to ensure it meets your requirements. Refer to the individual documentation pages linked above for detailed usage instructions.
\ No newline at end of file
+> Always benchmark the performance (accuracy and speed) of the quantized model on your specific task and hardware to ensure it meets your requirements. Refer to the individual documentation pages linked above for detailed usage instructions.
diff --git a/docs/source/en/quantization/torchao.md b/docs/source/en/quantization/torchao.md
index 6427866d0229..8778f9f3e5ea 100644
--- a/docs/source/en/quantization/torchao.md
+++ b/docs/source/en/quantization/torchao.md
@@ -30,7 +30,6 @@ See the table below for additional torchao features.
> [!TIP]
> Refer to the torchao [README.md](https://github.com/pytorch/ao#torchao-pytorch-architecture-optimization) for more details about the library.
-
torchao supports the [quantization techniques](https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md) below.
- A16W8 Float8 Dynamic Quantization
@@ -43,7 +42,6 @@ torchao supports the [quantization techniques](https://github.com/pytorch/ao/blo
torchao also supports module level configuration by specifying a dictionary from fully qualified name of module and its corresponding quantization config. This allows skip quantizing certain layers and using different quantization config for different modules.
-
Check the table below to see if your hardware is compatible.
| Component | Compatibility |
@@ -52,8 +50,6 @@ Check the table below to see if your hardware is compatible.
| XPU Versions | ✅ pytorch2.8 |
| CPU | ✅ change `device_map="cpu"` (see examples below) |
-
-
Install torchao from PyPi or the PyTorch index with the following commands.
@@ -64,13 +60,15 @@ Install torchao from PyPi or the PyTorch index with the following commands.
# Stable release from Pypi which will default to CUDA 12.6
pip install --upgrade torchao transformers
```
+
Stable Release from the PyTorch index
-
+
```bash
pip install torchao --index-url https://download.pytorch.org/whl/cu126 # options are cpu/cu118/cu126/cu128
```
+
@@ -118,6 +116,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -146,6 +145,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -177,13 +177,14 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
### A100 GPU
-
+
```py
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
@@ -210,6 +211,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -245,6 +247,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -276,13 +279,14 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
### Intel XPU
-
+
```py
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
@@ -309,6 +313,7 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -340,14 +345,14 @@ input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
-
### CPU
-
+
```py
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
@@ -373,6 +378,7 @@ input_ids = tokenizer(input_text, return_tensors="pt")
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
@@ -404,12 +410,14 @@ input_ids = tokenizer(input_text, return_tensors="pt")
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
### Per Module Quantization
#### 1. Skip quantization for certain layers
With `ModuleFqnToConfig` we can specify a default configuration for all layers while skipping quantization for certain layers.
+
```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
@@ -438,6 +446,7 @@ print(output_text)
```
#### 2. Quantizing different layers with different quantization configs
+
```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
@@ -485,7 +494,6 @@ Note: autoquant is for GPU only right now.
Create a [`TorchAoConfig`] and set to `"autoquant"`. Set the `cache_implementation` to `"static"` to automatically [torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) the forward method. Finally, call `finalize_autoquant` on the quantized model to finalize the quantization and log the input shapes.
-
```py
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
@@ -509,7 +517,6 @@ quantized_model.finalize_autoquant()
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
-
## Serialization
torchao implements [torch.Tensor subclasses](https://pytorch.org/docs/stable/notes/extending.html#subclassing-torch-tensor) for maximum flexibility in supporting new quantized torch.Tensor formats. [Safetensors](https://huggingface.co/docs/safetensors/en/index) serialization and deserialization does not work with torchao.
@@ -518,15 +525,16 @@ To avoid arbitrary user code execution, torchao sets `weights_only=True` in [tor
-
+
```py
# don't serialize model with Safetensors
output_dir = "llama3-8b-int4wo-128"
quantized_model.save_pretrained("llama3-8b-int4wo-128", safe_serialization=False)
```
+
-
+
```py
# don't serialize model with Safetensors
USER_ID = "your_huggingface_user_id"
@@ -534,13 +542,14 @@ REPO_ID = "llama3-8b-int4wo-128"
quantized_model.push_to_hub(f"{USER_ID}/llama3-8b-int4wo-128", safe_serialization=False)
tokenizer.push_to_hub(f"{USER_ID}/llama3-8b-int4wo-128")
```
+
-
## Loading quantized models
Loading a quantized model depends on the quantization scheme. For quantization schemes, like int8 and float8, you can quantize the model on any device and also load it on any device. The example below demonstrates quantizing a model on the CPU and then loading it on CUDA or XPU.
+
```py
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
@@ -574,6 +583,7 @@ output = reloaded_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
+
For int4, the model can only be loaded on the same device it was quantized on because the layout is specific to the device. The example below demonstrates quantizing and loading a model on the CPU.
```py
@@ -641,8 +651,6 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
>
> All configuration objects accept parameters for customization (e.g., `group_size`, `scheme`, `layout`).
-
-
## Resources
For a better sense of expected performance, view the [benchmarks](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks) for various models with CUDA and XPU backends. You can also run the code below to benchmark a model yourself.
diff --git a/docs/source/en/run_scripts.md b/docs/source/en/run_scripts.md
index c3a4787575c0..594eb84b02a1 100644
--- a/docs/source/en/run_scripts.md
+++ b/docs/source/en/run_scripts.md
@@ -52,6 +52,7 @@ Start with a smaller dataset by including the `max_train_samples`, `max_eval_sam
> [!WARNING]
> Not all example scripts support the `max_predict_samples` parameter. Run the command below to check whether a script supports it or not.
+>
> ```bash
> examples/pytorch/summarization/run_summarization.py -h
> ```
@@ -104,7 +105,7 @@ torchrun \
...
```
-PyTorch supports TPUs, hardware designed to accelerate performance, through the [PyTorch/XLA](https://github.com/pytorch/xla/blob/master/README.md) package. Launch the `xla_spawn.py` script and use `num _cores` to set the number of TPU cores to train with.
+PyTorch supports TPUs, hardware designed to accelerate performance, through the [PyTorch/XLA](https://github.com/pytorch/xla/blob/master/README.md) package. Launch the `xla_spawn.py` script and use `num_cores` to set the number of TPU cores to train with.
```bash
python xla_spawn.py --num_cores 8 pytorch/summarization/run_summarization.py \
diff --git a/docs/source/en/serialization.md b/docs/source/en/serialization.md
index 831f163bed18..1fefe08d5ca9 100644
--- a/docs/source/en/serialization.md
+++ b/docs/source/en/serialization.md
@@ -38,6 +38,7 @@ pip install optimum[exporters]
> [!TIP]
> Refer to the [Export a model to ONNX with optimum.exporters.onnx](https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli) guide for all available arguments or with the command below.
+>
> ```bash
> optimum-cli export onnx --help
> ```
@@ -50,7 +51,7 @@ optimum-cli export onnx --model distilbert/distilbert-base-uncased-distilled-squ
You should see logs indicating the progress and showing where the resulting `model.onnx` is saved.
-```bash
+```text
Validating ONNX model distilbert_base_uncased_squad_onnx/model.onnx...
-[✓] ONNX model output names match reference model (start_logits, end_logits)
- Validating ONNX Model output "start_logits":
diff --git a/docs/source/en/serving.md b/docs/source/en/serving.md
index f421a284950a..4287c5d2d5ec 100644
--- a/docs/source/en/serving.md
+++ b/docs/source/en/serving.md
@@ -16,7 +16,7 @@ rendered properly in your Markdown viewer.
# Serving
-Transformer models can be efficiently deployed using libraries such as vLLM, Text Generation Inference (TGI), and others. These libraries are designed for production-grade user-facing services, and can scale to multiple servers and millions of concurrent users. Refer to [Transformers as Backend for Inference Servers](./transformers_as_backends) for usage examples.
+Transformer models can be efficiently deployed using libraries such as vLLM, Text Generation Inference (TGI), and others. These libraries are designed for production-grade user-facing services, and can scale to multiple servers and millions of concurrent users. Refer to [Transformers as Backend for Inference Servers](./transformers_as_backend) for usage examples.
> [!TIP]
> Responses API is now supported as an experimental API! Read more about it [here](#responses-api).
@@ -24,19 +24,20 @@ Transformer models can be efficiently deployed using libraries such as vLLM, Tex
You can also serve transformer models with the `transformers serve` CLI. With Continuous Batching, `serve` now delivers solid throughput and latency well suited for evaluation, experimentation, and moderate-load local or self-hosted deployments. While vLLM, SGLang, or other inference engines remain our recommendations for large-scale production, `serve` avoids the extra runtime and operational overhead, and is on track to gain more production-oriented features.
In this document, we dive into the different supported endpoints and modalities; we also cover the setup of several user interfaces that can be used on top of `transformers serve` in the following guides:
-- [Jan (text and MCP user interface)](./jan.md)
-- [Cursor (IDE)](./cursor.md)
-- [Open WebUI (text, image, speech user interface)](./open_webui.md)
-- [Tiny-Agents (text and MCP CLI tool)](./tiny_agents.md)
+- [Jan (text and MCP user interface)](./jan)
+- [Cursor (IDE)](./cursor)
+- [Open WebUI (text, image, speech user interface)](./open_webui)
+- [Tiny-Agents (text and MCP CLI tool)](./tiny_agents)
## Serve CLI
> [!WARNING]
> This section is experimental and subject to change in future versions
-You can serve models of diverse modalities supported by `transformers` with the `transformers serve` CLI. It spawns a local server that offers compatibility with the OpenAI SDK, which is the _de facto_ standard for LLM conversations and other related tasks. This way, you can use the server from many third party applications, or test it using the `transformers chat` CLI ([docs](conversations.md#chat-cli)).
+You can serve models of diverse modalities supported by `transformers` with the `transformers serve` CLI. It spawns a local server that offers compatibility with the OpenAI SDK, which is the _de facto_ standard for LLM conversations and other related tasks. This way, you can use the server from many third party applications, or test it using the `transformers chat` CLI ([docs](conversations#chat-cli)).
The server supports the following REST APIs:
+
- `/v1/chat/completions`
- `/v1/responses`
- `/v1/audio/transcriptions`
@@ -356,7 +357,6 @@ ResponseCompletedEvent(response=Response(id='resp_req_0', created_at=1754060400.
-
## MCP integration
The `transformers serve` server is also an MCP client, so it can interact with MCP tools in agentic use cases. This, of course, requires the use of an LLM that is designed to use tools.
@@ -382,7 +382,6 @@ transformers serve \
--attn_implementation sdpa_paged
```
-
### Performance tips
- Use an efficient attention backend when available:
@@ -401,5 +400,3 @@ transformers serve \
- `--load_in_4bit`/`--load_in_8bit` can reduce memory footprint for LoRA setups
- `--force-model ` avoids per-request model hints and helps produce stable, repeatable runs
-
-
diff --git a/docs/source/en/tasks/audio_classification.md b/docs/source/en/tasks/audio_classification.md
index 973f95e1e955..844b5caec052 100644
--- a/docs/source/en/tasks/audio_classification.md
+++ b/docs/source/en/tasks/audio_classification.md
@@ -212,7 +212,6 @@ At this point, only three steps remain:
2. Pass the training arguments to [`Trainer`] along with the model, dataset, tokenizer, data collator, and `compute_metrics` function.
3. Call [`~Trainer.train`] to fine-tune your model.
-
```py
>>> training_args = TrainingArguments(
... output_dir="my_awesome_mind_model",
diff --git a/docs/source/en/tasks/document_question_answering.md b/docs/source/en/tasks/document_question_answering.md
index d83e025c4090..2c729f76adcb 100644
--- a/docs/source/en/tasks/document_question_answering.md
+++ b/docs/source/en/tasks/document_question_answering.md
@@ -104,6 +104,7 @@ yourself with the features.
```
Here's what the individual fields represent:
+
* `id`: the example's id
* `image`: a PIL.Image.Image object containing the document image
* `query`: the question string - natural language asked question, in several languages
@@ -257,6 +258,7 @@ Once examples are encoded, however, they will look like this:
```
We'll need to find the position of the answer in the encoded input.
+
* `token_type_ids` tells us which tokens are part of the question, and which ones are part of the document's words.
* `tokenizer.cls_token_id` will help find the special token at the beginning of the input.
* `word_ids` will help match the answer found in the original `words` to the same answer in the full encoded input and determine
@@ -365,6 +367,7 @@ of the Hugging Face course for inspiration.
Congratulations! You've successfully navigated the toughest part of this guide and now you are ready to train your own model.
Training involves the following steps:
+
* Load the model with [`AutoModelForDocumentQuestionAnswering`] using the same checkpoint as in the preprocessing.
* Define your training hyperparameters in [`TrainingArguments`].
* Define a function to batch examples together, here the [`DefaultDataCollator`] will do just fine
@@ -439,6 +442,7 @@ Now that you have finetuned a LayoutLMv2 model, and uploaded it to the 🤗 Hub,
way to try out your finetuned model for inference is to use it in a [`Pipeline`].
Let's take an example:
+
```py
>>> example = dataset["test"][2]
>>> question = example["query"]["en"]
@@ -464,6 +468,7 @@ document question answering with your model, and pass the image + question combi
```
You can also manually replicate the results of the pipeline if you'd like:
+
1. Take an image and a question, prepare them for the model using the processor from your model.
2. Forward the result or preprocessing through the model.
3. The model returns `start_logits` and `end_logits`, which indicate which token is at the start of the answer and
diff --git a/docs/source/en/tasks/idefics.md b/docs/source/en/tasks/idefics.md
index 3f8915f3cc99..b03c7bccd9c2 100644
--- a/docs/source/en/tasks/idefics.md
+++ b/docs/source/en/tasks/idefics.md
@@ -18,26 +18,27 @@ rendered properly in your Markdown viewer.
[[open-in-colab]]
-While individual tasks can be tackled by fine-tuning specialized models, an alternative approach
-that has recently emerged and gained popularity is to use large models for a diverse set of tasks without fine-tuning.
-For instance, large language models can handle such NLP tasks as summarization, translation, classification, and more.
-This approach is no longer limited to a single modality, such as text, and in this guide, we will illustrate how you can
-solve image-text tasks with a large multimodal model called IDEFICS.
-
-[IDEFICS](../model_doc/idefics) is an open-access vision and language model based on [Flamingo](https://huggingface.co/papers/2204.14198),
-a state-of-the-art visual language model initially developed by DeepMind. The model accepts arbitrary sequences of image
-and text inputs and generates coherent text as output. It can answer questions about images, describe visual content,
-create stories grounded in multiple images, and so on. IDEFICS comes in two variants - [80 billion parameters](https://huggingface.co/HuggingFaceM4/idefics-80b)
-and [9 billion parameters](https://huggingface.co/HuggingFaceM4/idefics-9b), both of which are available on the 🤗 Hub. For each variant, you can also find fine-tuned instructed
+While individual tasks can be tackled by fine-tuning specialized models, an alternative approach
+that has recently emerged and gained popularity is to use large models for a diverse set of tasks without fine-tuning.
+For instance, large language models can handle such NLP tasks as summarization, translation, classification, and more.
+This approach is no longer limited to a single modality, such as text, and in this guide, we will illustrate how you can
+solve image-text tasks with a large multimodal model called IDEFICS.
+
+[IDEFICS](../model_doc/idefics) is an open-access vision and language model based on [Flamingo](https://huggingface.co/papers/2204.14198),
+a state-of-the-art visual language model initially developed by DeepMind. The model accepts arbitrary sequences of image
+and text inputs and generates coherent text as output. It can answer questions about images, describe visual content,
+create stories grounded in multiple images, and so on. IDEFICS comes in two variants - [80 billion parameters](https://huggingface.co/HuggingFaceM4/idefics-80b)
+and [9 billion parameters](https://huggingface.co/HuggingFaceM4/idefics-9b), both of which are available on the 🤗 Hub. For each variant, you can also find fine-tuned instructed
versions of the model adapted for conversational use cases.
-This model is exceptionally versatile and can be used for a wide range of image and multimodal tasks. However,
-being a large model means it requires significant computational resources and infrastructure. It is up to you to decide whether
-this approach suits your use case better than fine-tuning specialized models for each individual task.
+This model is exceptionally versatile and can be used for a wide range of image and multimodal tasks. However,
+being a large model means it requires significant computational resources and infrastructure. It is up to you to decide whether
+this approach suits your use case better than fine-tuning specialized models for each individual task.
+
+In this guide, you'll learn how to:
-In this guide, you'll learn how to:
- [Load IDEFICS](#loading-the-model) and [load the quantized version of the model](#quantized-model)
-- Use IDEFICS for:
+- Use IDEFICS for:
- [Image captioning](#image-captioning)
- [Prompted image captioning](#prompted-image-captioning)
- [Few-shot prompting](#few-shot-prompting)
@@ -47,7 +48,7 @@ In this guide, you'll learn how to:
- [Run inference in batch mode](#running-inference-in-batch-mode)
- [Run IDEFICS instruct for conversational use](#idefics-instruct-for-conversational-use)
-Before you begin, make sure you have all the necessary libraries installed.
+Before you begin, make sure you have all the necessary libraries installed.
```bash
pip install -q bitsandbytes sentencepiece accelerate transformers
@@ -59,14 +60,14 @@ To run the following examples with a non-quantized version of the model checkpoi
## Loading the model
-Let's start by loading the model's 9 billion parameters checkpoint:
+Let's start by loading the model's 9 billion parameters checkpoint:
```py
>>> checkpoint = "HuggingFaceM4/idefics-9b"
```
-Just like for other Transformers models, you need to load a processor and the model itself from the checkpoint.
-The IDEFICS processor wraps a [`LlamaTokenizer`] and IDEFICS image processor into a single processor to take care of
+Just like for other Transformers models, you need to load a processor and the model itself from the checkpoint.
+The IDEFICS processor wraps a [`LlamaTokenizer`] and IDEFICS image processor into a single processor to take care of
preparing text and image inputs for the model.
```py
@@ -79,13 +80,13 @@ preparing text and image inputs for the model.
>>> model = IdeficsForVisionText2Text.from_pretrained(checkpoint, dtype=torch.bfloat16, device_map="auto")
```
-Setting `device_map` to `"auto"` will automatically determine how to load and store the model weights in the most optimized
+Setting `device_map` to `"auto"` will automatically determine how to load and store the model weights in the most optimized
manner given existing devices.
### Quantized model
-If high-memory device availability is an issue, you can load the quantized version of the model. To load the model and the
-processor in 4bit precision, pass a `BitsAndBytesConfig` to the `from_pretrained` method and the model will be compressed
+If high-memory device availability is an issue, you can load the quantized version of the model. To load the model and the
+processor in 4bit precision, pass a `BitsAndBytesConfig` to the `from_pretrained` method and the model will be compressed
on the fly while loading.
```py
@@ -109,8 +110,8 @@ on the fly while loading.
Now that you have the model loaded in one of the suggested ways, let's move on to exploring tasks that you can use IDEFICS for.
## Image captioning
-Image captioning is the task of predicting a caption for a given image. A common application is to aid visually impaired
-people navigate through different situations, for instance, explore image content online.
+Image captioning is the task of predicting a caption for a given image. A common application is to aid visually impaired
+people navigate through different situations, for instance, explore image content online.
To illustrate the task, get an image to be captioned, e.g.:
@@ -118,10 +119,10 @@ To illustrate the task, get an image to be captioned, e.g.:
-Photo by [Hendo Wang](https://unsplash.com/@hendoo).
+Photo by [Hendo Wang](https://unsplash.com/@hendoo).
-IDEFICS accepts text and image prompts. However, to caption an image, you do not have to provide a text prompt to the
-model, only the preprocessed input image. Without a text prompt, the model will start generating text from the
+IDEFICS accepts text and image prompts. However, to caption an image, you do not have to provide a text prompt to the
+model, only the preprocessed input image. Without a text prompt, the model will start generating text from the
BOS (beginning-of-sequence) token thus creating a caption.
As image input to the model, you can use either an image object (`PIL.Image`) or a url from which the image can be retrieved.
@@ -142,15 +143,15 @@ A puppy in a flower bed
-It is a good idea to include the `bad_words_ids` in the call to `generate` to avoid errors arising when increasing
-the `max_new_tokens`: the model will want to generate a new `` or `` token when there
+It is a good idea to include the `bad_words_ids` in the call to `generate` to avoid errors arising when increasing
+the `max_new_tokens`: the model will want to generate a new `` or `` token when there
is no image being generated by the model.
You can set it on-the-fly as in this guide, or store in the `GenerationConfig` as described in the [Text generation strategies](../generation_strategies) guide.
## Prompted image captioning
-You can extend image captioning by providing a text prompt, which the model will continue given the image. Let's take
+You can extend image captioning by providing a text prompt, which the model will continue given the image. Let's take
another image to illustrate:
-
diff --git a/docs/source/en/tasks/masked_language_modeling.md b/docs/source/en/tasks/masked_language_modeling.md
index 3c024739d738..619374f91dae 100644
--- a/docs/source/en/tasks/masked_language_modeling.md
+++ b/docs/source/en/tasks/masked_language_modeling.md
@@ -150,6 +150,7 @@ To apply this preprocessing function over the entire dataset, use the 🤗 Datas
This dataset contains the token sequences, but some of these are longer than the maximum input length for the model.
You can now use a second preprocessing function to
+
- concatenate all the sequences
- split the concatenated sequences into shorter chunks defined by `block_size`, which should be both shorter than the maximum input length and short enough for your GPU RAM.
diff --git a/docs/source/en/tasks/monocular_depth_estimation.md b/docs/source/en/tasks/monocular_depth_estimation.md
index c90abce1cd57..aef9bd22c4d3 100644
--- a/docs/source/en/tasks/monocular_depth_estimation.md
+++ b/docs/source/en/tasks/monocular_depth_estimation.md
@@ -23,7 +23,7 @@ a single camera viewpoint.
Monocular depth estimation has various applications, including 3D reconstruction, augmented reality, autonomous driving,
and robotics. It is a challenging task as it requires the model to understand the complex relationships between objects
in the scene and the corresponding depth information, which can be affected by factors such as lighting conditions,
-occlusion, and texture.
+occlusion, and texture.
There are two main depth estimation categories:
@@ -143,7 +143,7 @@ Let's post-process the results to remove any padding and resize the depth map to
Pass the image and the candidate object labels to look for to the pipeline.
-Here we pass the image directly; other suitable options include a local path to an image or an image url. We also pass text descriptions for all items we want to query the image for.
+Here we pass the image directly; other suitable options include a local path to an image or an image url. We also pass text descriptions for all items we want to query the image for.
```py
>>> predictions = detector(
diff --git a/docs/source/en/testing.md b/docs/source/en/testing.md
index 497c6b019311..01658aa2beb7 100644
--- a/docs/source/en/testing.md
+++ b/docs/source/en/testing.md
@@ -16,7 +16,6 @@ rendered properly in your Markdown viewer.
# Testing
-
Let's take a look at how 🤗 Transformers models are tested and how you can write new tests and improve the existing ones.
There are 2 test suites in the repository:
@@ -51,12 +50,8 @@ RUN_SLOW=1 pytest examples/
The results can be observed [here](https://github.com/huggingface/transformers/actions).
-
-
## Running tests
-
-
### Choosing which tests to run
This document goes into many details of how tests can be run. If after reading everything, you need even more details
@@ -89,8 +84,6 @@ which tells pytest to:
- do not capture output
- run in verbose mode
-
-
### Getting the list of all tests
All tests of the test suite:
@@ -187,7 +180,6 @@ Sometimes you need to run `accelerate` tests on your models. For that you can ju
RUN_SLOW=1 pytest -m accelerate_tests tests/models/opt/test_modeling_opt.py
```
-
### Run documentation tests
In order to test whether the documentation examples are correct, you should check that the `doctests` are passing.
@@ -217,9 +209,11 @@ Example:
```
Just run the following line to automatically test every docstring example in the desired file:
+
```bash
pytest --doctest-modules
```
+
If the file has a markdown extension, you should add the `--doctest-glob="*.md"` argument.
### Run only modified tests
@@ -266,12 +260,10 @@ or `pytest.ini`/``tox.ini`` files:
looponfailroots = transformers tests
```
-This would lead to only looking for file changes in the respective directories, specified relatively to the ini-file’s
-directory.
+This would lead to only looking for file changes in the respective directories, specified relatively to the ini-file's directory.
[pytest-watch](https://github.com/joeyespo/pytest-watch) is an alternative implementation of this functionality.
-
### Skip a test module
If you want to run all test modules, except a few you can exclude them by giving an explicit list of tests to run. For
@@ -307,7 +299,6 @@ It's good to repeat the tests several times, in sequence, randomly, or in sets,
inter-dependency and state-related bugs (tear down). And the straightforward multiple repetition is just good to detect
some problems that get uncovered by randomness of DL.
-
#### Repeat tests
- [pytest-flakefinder](https://github.com/dropbox/pytest-flakefinder):
@@ -403,8 +394,6 @@ pytest -p no:sugar
or uninstall it.
-
-
#### Report each sub-test name and its progress
For a single or a group of tests via `pytest` (after `pip install pytest-pspec`):
@@ -457,7 +446,6 @@ decorators are used to set the requirements of tests CPU/GPU/XPU/TPU-wise:
Let's depict the GPU requirements in the following table:
-
| n gpus | decorator |
|--------|--------------------------------|
| `>= 0` | `@require_torch` |
@@ -466,7 +454,6 @@ Let's depict the GPU requirements in the following table:
| `< 2` | `@require_torch_non_multi_gpu` |
| `< 3` | `@require_torch_up_to_2_gpus` |
-
For example, here is a test that must be run only when there are 2 or more GPUs available and pytorch is installed:
```python no-style
@@ -520,6 +507,7 @@ Certain devices will require an additional import after importing `torch` for th
```bash
TRANSFORMERS_TEST_BACKEND="torch_npu" pytest tests/utils/test_logging.py
```
+
Alternative backends may also require the replacement of device-specific functions. For example `torch.cuda.manual_seed` may need to be replaced with a device-specific seed setter like `torch.npu.manual_seed` or `torch.xpu.manual_seed` to correctly set a random seed on the device. To specify a new backend with backend-specific device functions when running the test suite, create a Python device specification file `spec.py` in the format:
```python
@@ -536,6 +524,7 @@ MANUAL_SEED_FN = torch.npu.manual_seed
EMPTY_CACHE_FN = torch.npu.empty_cache
DEVICE_COUNT_FN = torch.npu.device_count
```
+
This format also allows for specification of any additional imports required. To use this file to replace equivalent methods in the test suite, set the environment variable `TRANSFORMERS_TEST_DEVICE_SPEC` to the path of the spec file, e.g. `TRANSFORMERS_TEST_DEVICE_SPEC=spec.py`.
Currently, only `MANUAL_SEED_FN`, `EMPTY_CACHE_FN` and `DEVICE_COUNT_FN` are supported for device-specific dispatch.
@@ -610,7 +599,6 @@ You can read [here](https://docs.pytest.org/en/stable/unittest.html) which featu
thing to remember is that most `pytest` fixtures don't work. Neither parametrization, but we use the module
`parameterized` that works in a similar way.
-
### Parametrization
Often, there is a need to run the same test multiple times, but with different arguments. It could be done from within
@@ -719,8 +707,6 @@ pytest test_this2.py::test_floor[negative--1.5--2.0] test_this2.py::test_floor[i
as in the previous example.
-
-
### Files and directories
In tests often we need to know where things are relative to the current test file, and it's not trivial since the test
@@ -843,7 +829,6 @@ otherwise.
If you need to temporary override `sys.path` to import from another test for example, you can use the
`ExtendSysPath` context manager. Example:
-
```python
import os
from transformers.testing_utils import ExtendSysPath
@@ -860,13 +845,13 @@ commit it to the main repository we need make sure it's skipped during `make tes
Methods:
-- A **skip** means that you expect your test to pass only if some conditions are met, otherwise pytest should skip
+- A **skip** means that you expect your test to pass only if some conditions are met, otherwise pytest should skip
running the test altogether. Common examples are skipping windows-only tests on non-windows platforms, or skipping
tests that depend on an external resource which is not available at the moment (for example a database).
-- A **xfail** means that you expect a test to fail for some reason. A common example is a test for a feature not yet
+- A **xfail** means that you expect a test to fail for some reason. A common example is a test for a feature not yet
implemented, or a bug not yet fixed. When a test passes despite being expected to fail (marked with
- pytest.mark.xfail), it’s an xpass and will be reported in the test summary.
+ pytest.mark.xfail), it's an xpass and will be reported in the test summary.
One of the important differences between the two is that `skip` doesn't run the test, and `xfail` does. So if the
code that's buggy causes some bad state that will affect other tests, do not use `xfail`.
@@ -893,7 +878,6 @@ or the `xfail` way:
def test_feature_x():
```
-
Here's how to skip a test based on internal checks within the test:
```python
@@ -924,7 +908,7 @@ def test_feature_x():
docutils = pytest.importorskip("docutils", minversion="0.3")
```
-- Skip a test based on a condition:
+- Skip a test based on a condition:
```python no-style
@pytest.mark.skipif(sys.version_info < (3,6), reason="requires python3.6 or higher")
@@ -1018,7 +1002,6 @@ That report is also useful to find slow outliers that aren't marked as such, or
If you notice that the test suite starts getting slow on CI, the top listing of this report will show the slowest
tests.
-
### Testing the stdout/stderr output
In order to test functions that write to `stdout` and/or `stderr`, the test can access those streams using the
@@ -1141,7 +1124,6 @@ print(cs.err, cs.out)
Also, to aid debugging test issues, by default these context managers automatically replay the captured streams on exit
from the context.
-
### Capturing logger stream
If you need to validate the output of a logger, you can use `CaptureLogger`:
@@ -1193,7 +1175,6 @@ called if anything.
This helper method creates a copy of the `os.environ` object, so the original remains intact.
-
### Getting reproducible results
In some situations you may want to remove randomness for your tests. To get identical reproducible results set, you
@@ -1241,9 +1222,6 @@ To trigger a self-push workflow CI job, you must:
4. Then you can see the job appear [here](https://github.com/huggingface/transformers/actions/workflows/self-push.yml). It may not run right away if there
is a backlog.
-
-
-
## Testing Experimental CI Features
Testing CI features can be potentially problematic as it can interfere with the normal CI functioning. Therefore if a
@@ -1306,7 +1284,7 @@ You can vote for this feature and see where it is at these CI-specific threads:
## DeepSpeed integration
-For a PR that involves the DeepSpeed integration, keep in mind our CircleCI PR CI setup doesn't have GPUs. Tests requiring GPUs are run on a different CI nightly. This means if you get a passing CI report in your PR, it doesn’t mean the DeepSpeed tests pass.
+For a PR that involves the DeepSpeed integration, keep in mind our CircleCI PR CI setup doesn't have GPUs. Tests requiring GPUs are run on a different CI nightly. This means if you get a passing CI report in your PR, it doesn't mean the DeepSpeed tests pass.
To run DeepSpeed tests:
diff --git a/docs/source/en/tiny_agents.md b/docs/source/en/tiny_agents.md
index dc53d05a4bff..7266f0236a63 100644
--- a/docs/source/en/tiny_agents.md
+++ b/docs/source/en/tiny_agents.md
@@ -42,4 +42,3 @@ Image URL: https://evalstate-flux1-schnell.hf.space/gradio_api/file=/tmp/gradio/
I have generated an image of a cat on the moon using the Flux 1 Schnell Image Generator. The image is 1024x1024 pixels and was created with 4 inference steps. Let me know if you would like to make any changes or need further assistance!
```
-
diff --git a/docs/source/en/tokenizer_summary.md b/docs/source/en/tokenizer_summary.md
index 801948f35d87..34bc16628cad 100644
--- a/docs/source/en/tokenizer_summary.md
+++ b/docs/source/en/tokenizer_summary.md
@@ -42,7 +42,7 @@ For instance, let's look at the sentence `"Don't you love 🤗 Transformers? We
A simple way of tokenizing this text is to split it by spaces, which would give:
-```
+```text
["Don't", "you", "love", "🤗", "Transformers?", "We", "sure", "do."]
```
@@ -52,7 +52,7 @@ punctuation into account so that a model does not have to learn a different repr
punctuation symbol that could follow it, which would explode the number of representations the model has to learn.
Taking punctuation into account, tokenizing our exemplary text would give:
-```
+```text
["Don", "'", "t", "you", "love", "🤗", "Transformers", "?", "We", "sure", "do", "."]
```
@@ -65,7 +65,7 @@ input that was tokenized with the same rules that were used to tokenize its trai
[spaCy](https://spacy.io/) and [Moses](http://www.statmt.org/moses/?n=Development.GetStarted) are two popular
rule-based tokenizers. Applying them on our example, *spaCy* and *Moses* would output something like:
-```
+```text
["Do", "n't", "you", "love", "🤗", "Transformers", "?", "We", "sure", "do", "."]
```
@@ -154,14 +154,14 @@ define before training the tokenizer.
As an example, let's assume that after pre-tokenization, the following set of words including their frequency has been
determined:
-```
+```text
("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)
```
Consequently, the base vocabulary is `["b", "g", "h", "n", "p", "s", "u"]`. Splitting all words into symbols of the
base vocabulary, we obtain:
-```
+```text
("h" "u" "g", 10), ("p" "u" "g", 5), ("p" "u" "n", 12), ("b" "u" "n", 4), ("h" "u" "g" "s", 5)
```
@@ -172,7 +172,7 @@ the example above `"h"` followed by `"u"` is present _10 + 5 = 15_ times (10 tim
`"u"` symbols followed by a `"g"` symbol together. Next, `"ug"` is added to the vocabulary. The set of words then
becomes
-```
+```text
("h" "ug", 10), ("p" "ug", 5), ("p" "u" "n", 12), ("b" "u" "n", 4), ("h" "ug" "s", 5)
```
@@ -183,7 +183,7 @@ BPE then identifies the next most common symbol pair. It's `"u"` followed by `"n
At this stage, the vocabulary is `["b", "g", "h", "n", "p", "s", "u", "ug", "un", "hug"]` and our set of unique words
is represented as
-```
+```text
("hug", 10), ("p" "ug", 5), ("p" "un", 12), ("b" "un", 4), ("hug" "s", 5)
```
@@ -246,7 +246,7 @@ reached the desired size. The Unigram algorithm always keeps the base characters
Because Unigram is not based on merge rules (in contrast to BPE and WordPiece), the algorithm has several ways of
tokenizing new text after training. As an example, if a trained Unigram tokenizer exhibits the vocabulary:
-```
+```text
["b", "g", "h", "n", "p", "s", "u", "ug", "un", "hug"],
```
diff --git a/docs/source/en/trainer.md b/docs/source/en/trainer.md
index 48325da6893c..32f14bc41da3 100644
--- a/docs/source/en/trainer.md
+++ b/docs/source/en/trainer.md
@@ -346,7 +346,6 @@ use_cpu: false
-
Run [accelerate_launch](https://hf.co/docs/accelerate/package_reference/cli#accelerate-launch) to start training with the configurations set in `config_file.yaml`. This file is saved to the Accelerate cache folder and automatically loaded when you run `accelerate_launch`.
The example below launches the [run_glue.py](../../../examples/pytorch/text-classification/run_glue) script with the FSDP configuration shown earlier. Parameters from the `config_file.yaml` file can also be directly set in the command line.
diff --git a/docs/source/en/training.md b/docs/source/en/training.md
index ed992e8152d9..ccee25704fa3 100644
--- a/docs/source/en/training.md
+++ b/docs/source/en/training.md
@@ -52,6 +52,7 @@ dataset = dataset.map(tokenize, batched=True)
> [!TIP]
> Fine-tune on a smaller subset of the full dataset to reduce the time it takes. The results won't be as good compared to fine-tuning on the full dataset, but it is useful to make sure everything works as expected first before committing to training on the full dataset.
+>
> ```py
> small_train = dataset["train"].shuffle(seed=42).select(range(1000))
> small_eval = dataset["test"].shuffle(seed=42).select(range(1000))
diff --git a/docs/source/en/transformers_as_backend.md b/docs/source/en/transformers_as_backend.md
index 422cc4a121e9..ce5152c2a4a7 100644
--- a/docs/source/en/transformers_as_backend.md
+++ b/docs/source/en/transformers_as_backend.md
@@ -26,12 +26,13 @@ This guide shows how to use Transformers' models as a backend to some popular in
[vLLM](https://github.com/vllm-project/vllm) is a high-performance inference engine optimized for serving LLMs at scale. It supports many Transformers' models, including all decoder-only LLMs and several vision-language models (VLMs). VLMs currently support image inputs only, with video support planned.
-vLLM automatically selects the best backend, and if a model isn’t natively supported, it falls back to the Transformers model. To explicitly use a Transformers' model, set `model_impl="transformers"`.
+vLLM automatically selects the best backend, and if a model isn't natively supported, it falls back to the Transformers model. To explicitly use a Transformers' model, set `model_impl="transformers"`.
```python
from vllm import LLM
llm = LLM(model="meta-llama/Llama-3.2-1B", model_impl="transformers")
```
+
Add `--model-impl transformers` to `vllm serve` to launch a server with a Transformers' model.
```bash
@@ -42,12 +43,11 @@ vllm serve meta-llama/Llama-3.2-1B \
Refer to the [vLLM docs](https://docs.vllm.ai/en/latest/models/supported_models.html#transformers) for more usage examples and tips on using a Transformers as the backend.
-
## SGLang
[SGLang](https://github.com/InternLM/sglang) is a high-performance, OpenAI-compatible server and runtime designed for chat-based LLMs. It offers fast inference, role-based conversation handling, and support for custom pipelines, making it great for building real-world LLM apps.
-SGLang automatically falls back to the Transformers backend if a model isn’t natively supported. To explicitly use a Transformers' model, set `impl="transformers"`.
+SGLang automatically falls back to the Transformers backend if a model isn't natively supported. To explicitly use a Transformers' model, set `impl="transformers"`.
```python
import sglang as sgl
@@ -57,12 +57,6 @@ print(llm.generate(["The capital of France is"], {"max_new_tokens": 20})[0])
```
Add `impl transformers` to `sglang.launch_server` to launch a server with a Transformers' model.
-
-
-
-
-
-
```bash
python3 -m sglang.launch_server \
@@ -133,7 +127,7 @@ class MyModel(PreTrainedModel):
3. This step is optional, but if you want to support tensor parallel and/or pipeline parallel features, add the following keys to the config.
* `base_model_tp_plan` enables [tensor parallelism](./perf_infer_gpu_multi) by mapping fully qualified layer name patterns to tensor parallel styles. Only the `"colwise"` and `"rowwise"` partitioning strategies are currently supported.
* `base_model_pp_plan` enables pipeline parallelism by mapping direct child layer names to tuples of lists of strings. The list in the first element of the tuple contains the names of the input arguments. The list in the last element of the tuple contains the names of the variables the layer outputs to in the modeling code.
-
+
Expand the code below for an example.
+
+
+
+**Transformers** হলো একটা ফ্রেমওয়ার্ক যেটা দিয়ে টেক্সট, কম্পিউটার ভিশন, অডিও, ভিডিও আর মাল্টিমোডাল—সব ধরনের মডেল তৈরি আর চালানো যায়। এটা ট্রেইনিং আর ইনফারেন্স – দুই কাজেই ব্যবহার করা হয়।
+
+Transformers মডেলের ডেফিনিশন এক জায়গায় রাখে। এর মানে হলো, একবার কোনো মডেল `transformers`-এ সাপোর্ট পেলেই সেটা সহজে বিভিন্ন ট্রেইনিং ফ্রেমওয়ার্ক (Axolotl, Unsloth, DeepSpeed, FSDP, PyTorch-Lightning ইত্যাদি), ইনফারেন্স ইঞ্জিন (vLLM, SGLang, TGI ইত্যাদি) আর অন্যান্য লাইব্রেরি (llama.cpp, mlx ইত্যাদি)-তে ব্যবহার করা যায়।
+
+আমরা চাই নতুন আর আধুনিক মডেলগুলো সবাই ব্যবহার করতে পারে। তাই মডেলের ডেফিনিশন রাখা হয়েছে সহজ, কাস্টমাইজযোগ্য আর পারফরম্যান্স-ফ্রেন্ডলি।
+
+এখন পর্যন্ত [Hugging Face Hub](https://huggingface.com/models)-এ ১০ লাখেরও বেশি Transformers [মডেল চেকপয়েন্ট](https://huggingface.co/models?library=transformers&sort=trending) আছে, যেগুলো যেকোনো সময় ব্যবহার করা যায়।
+
+আজই [Hub](https://huggingface.com/) থেকে একটা মডেল বেছে নিন আর Transformers দিয়ে শুরু করুন।
+
+
+## ইনস্টলেশন
+
+Transformers Python 3.9+ সহ কাজ করে, এবং সমর্থিত ফ্রেমওয়ার্কগুলো হলো [PyTorch](https://pytorch.org/get-started/locally/) 2.1+, [TensorFlow](https://www.tensorflow.org/install/pip) 2.6+, এবং [Flax](https://flax.readthedocs.io/en/latest/) 0.4.1+।
+
+[venv](https://docs.python.org/3/library/venv.html) বা [uv](https://docs.astral.sh/uv/) ব্যবহার করে একটি ভার্চুয়াল এনভায়রনমেন্ট তৈরি এবং সক্রিয় করুন।
+
+```py
+# venv
+python -m venv .my-env
+source .my-env/bin/activate
+# uv
+uv venv .my-env
+source .my-env/bin/activate
+```
+আপনার ভার্চুয়াল পরিবেশে Transformers ইনস্টল করুন।
+
+```py
+# pip
+pip install "transformers[torch]"
+
+# uv
+uv pip install "transformers[torch]"
+```
+যদি আপনি লাইব্রেরির সর্বশেষ পরিবর্তনগুলি চান বা অবদান রাখতে আগ্রহী হন তবে উৎস থেকে Transformers ইনস্টল করুন। তবে, সর্বশেষ সংস্করণটি স্থিতিশীল নাও হতে পারে। যদি আপনি কোনো ত্রুটির সম্মুখীন হন তবে নির্দ্বিধায় একটি [issue](https://github.com/huggingface/transformers/issues) খুলুন।
+
+```Shell
+git clone [https://github.com/huggingface/transformers.git](https://github.com/huggingface/transformers.git)
+cd transformers
+
+# pip
+pip install .[torch]
+
+# uv
+uv pip install .[torch]
+```
+
+## কুইকস্টার্ট
+
+Transformers ব্যবহার শুরু করুন এখনই [Pipeline](https://huggingface.co/docs/transformers/pipeline_tutorial) API দিয়ে। `Pipeline` হলো একটি হাই-লেভেল ইনফারেন্স ক্লাস, যা টেক্সট, অডিও, ভিশন এবং মাল্টিমোডাল টাস্ক সাপোর্ট করে। এটি ইনপুট প্রিপ্রসেসিং করে এবং সঠিক আউটপুট রিটার্ন করে।
+
+একটি পাইপলাইন তৈরি করুন এবং টেক্সট জেনারেশনের জন্য কোন মডেল ব্যবহার করবেন তা নির্দিষ্ট করুন। মডেলটি ডাউনলোড হয়ে ক্যাশে রাখা হবে, ফলে পরে সহজেই আবার ব্যবহার করতে পারবেন। সবশেষে, মডেলকে প্রম্পট করার জন্য কিছু টেক্সট দিন।
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="text-generation", model="Qwen/Qwen2.5-1.5B")
+pipeline("the secret to baking a really good cake is ")
+[{'generated_text': 'the secret to baking a really good cake is 1) to use the right ingredients and 2) to follow the recipe exactly. the recipe for the cake is as follows: 1 cup of sugar, 1 cup of flour, 1 cup of milk, 1 cup of butter, 1 cup of eggs, 1 cup of chocolate chips. if you want to make 2 cakes, how much sugar do you need? To make 2 cakes, you will need 2 cups of sugar.'}]
+```
+
+মডেলের সাথে চ্যাট করতে হলেও ব্যবহার প্যাটার্ন একই। শুধু পার্থক্য হলো, আপনাকে একটি চ্যাট হিস্ট্রি তৈরি করতে হবে (যা `Pipeline`-এ ইনপুট হিসেবে যাবে) আপনার আর সিস্টেমের মধ্যে।
+
+> [!TIP]
+> আপনি সরাসরি কমান্ড লাইন থেকেও একটি মডেলের সাথে চ্যাট করতে পারেন।
+> ```Shell
+> transformers chat Qwen/Qwen2.5-0.5B-Instruct
+> ```
+
+```Python
+import torch
+from transformers import pipeline
+
+chat = [
+ {"role": "system", "content": "You are a sassy, wise-cracking robot as imagined by Hollywood circa 1986."},
+ {"role": "user", "content": "Hey, can you tell me any fun things to do in New York?"}
+]
+
+pipeline = pipeline(task="text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct", dtype=torch.bfloat16, device_map="auto")
+response = pipeline(chat, max_new_tokens=512)
+print(response[0]["generated_text"][-1]["content"])
+
+বিভিন্ন মোডালিটি এবং কাজের জন্য Pipeline কিভাবে কাজ করে তা দেখতে নিচের উদাহরণগুলো সম্প্রসারণ করুন।
+```
+
+
+
+ +## কেন Transformers ব্যবহার করবেন না? + +* এই লাইব্রেরি নিউরাল নেটওয়ার্কের জন্য ব্লক-মডিউল টুলবক্স নয়। মডেল ফাইলের কোডে অতিরিক্ত অ্যাবস্ট্র্যাকশন intentionally করা হয়নি, যাতে গবেষকরা দ্রুত প্রতিটি মডেলের উপর কাজ করতে পারে কোনো অতিরিক্ত ফাইল বা স্তরে না গিয়ে। +* ট্রেইনিং API মূলত Transformers-এর PyTorch মডেলের সাথে কাজ করার জন্য অপটিমাইজ করা হয়েছে। সাধারণ মেশিন লার্নিং লুপের জন্য, [Accelerate](https://huggingface.co/docs/accelerate) এর মতো অন্য লাইব্রেরি ব্যবহার করা উচিত। +* [উদাহরণ স্ক্রিপ্টগুলো](https://github.com/huggingface/transformers/tree/main/examples) শুধু *উদাহরণ*। এগুলো সরাসরি আপনার ব্যবহারের ক্ষেত্রে কাজ নাও করতে পারে, তাই কোড সামঞ্জস্য করতে হতে পারে। + +## Transformers দিয়ে ১০০টি প্রজেক্ট + +Transformers শুধু প্রি-ট্রেইন্ড মডেল ব্যবহার করার টুলকিট নয়, এটি একটি কমিউনিটি, যা Hugging Face Hub-এর চারপাশে তৈরি। আমরা চাই যে ডেভেলপার, গবেষক, শিক্ষার্থী, অধ্যাপক, ইঞ্জিনিয়ার বা যে কেউ তাদের স্বপ্নের প্রজেক্ট তৈরি করতে পারে। + +Transformers 100,000 স্টার উদযাপন করতে আমরা কমিউনিটিকে তুলে ধরতে [awesome-transformers](./awesome-transformers.md) পেজ তৈরি করেছি, যেখানে Transformers দিয়ে তৈরি ১০০টি অসাধারণ প্রজেক্ট তালিকাভুক্ত আছে। + +আপনার কোনো প্রজেক্ট আছে যা তালিকায় থাকা উচিত মনে করেন? তাহলে PR খুলে যুক্ত করুন। + +## উদাহরণ মডেল + +আপনি আমাদের অধিকাংশ মডেল সরাসরি তাদের [Hub মডেল পেজ](https://huggingface.co/models) থেকে পরীক্ষা করতে পারেন। + +নিচের প্রতিটি মোডালিটি এক্সপ্যান্ড করে বিভিন্ন ব্যবহার কেসের জন্য কয়েকটি উদাহরণ মডেল দেখুন। + + +
+
+
+Transformers funge da framework di definizione dei modelli per modelli di machine learning all'avanguardia nei
+modelli di testo, visione artificiale, audio, video e multimodali, sia per l'inferenza che per l'addestramento.
+
+Centralizza la definizione del modello in modo che tale definizione sia concordata all'interno dell'ecosistema.
+`transformers` è il perno tra i framework: se una definizione di modello è supportata, sarà compatibile con la
+maggior parte dei framework di addestramento (Axolotl, Unsloth, DeepSpeed, FSDP, PyTorch-Lightning, ...), motori
+di inferenza (vLLM, SGLang, TGI, ...) e librerie di modellazione adiacenti (llama.cpp, mlx, ...) che sfruttano
+la definizione del modello da `transformers`.
+
+Ci impegniamo a sostenere nuovi modelli all'avanguardia e a democratizzarne l'utilizzo rendendo la loro definizione
+semplice, personalizzabile ed efficiente.
+
+Ci sono oltre 1 milione di Transformers [model checkpoint](https://huggingface.co/models?library=transformers&sort=trending) su [Hugging Face Hub](https://huggingface.com/models) che puoi utilizzare.
+
+Esplora oggi stesso l'[Hub](https://huggingface.com/) per trovare un modello e utilizzare Transformers per aiutarti a iniziare subito.
+
+## Installazione
+
+Transformers funziona con Python 3.9+ e [PyTorch](https://pytorch.org/get-started/locally/) 2.1+.
+
+Crea e attiva un ambiente virtuale con [venv](https://docs.python.org/3/library/venv.html) o [uv](https://docs.astral.sh/uv/), un pacchetto Python veloce basato su Rust e un gestore di progetti.
+
+```py
+# venv
+python -m venv .my-env
+source .my-env/bin/activate
+# uv
+uv venv .my-env
+source .my-env/bin/activate
+```
+
+Installa Transformers nel tuo ambiente virtuale.
+
+```py
+# pip
+pip install "transformers[torch]"
+
+# uv
+uv pip install "transformers[torch]"
+```
+
+Installa Transformers dal sorgente se desideri le ultime modifiche nella libreria o sei interessato a contribuire. Tuttavia, la versione *più recente* potrebbe non essere stabile. Non esitare ad aprire una [issue](https://github.com/huggingface/transformers/issues) se riscontri un errore.
+
+```shell
+git clone https://github.com/huggingface/transformers.git
+cd transformers
+
+# pip
+pip install .[torch]
+
+# uv
+uv pip install .[torch]
+```
+
+## Quickstart
+
+Inizia subito a utilizzare Transformers con l'API [Pipeline](https://huggingface.co/docs/transformers/pipeline_tutorial). Pipeline è una classe di inferenza di alto livello che supporta attività di testo, audio, visione e multimodali. Gestisce la pre-elaborazione dell'input e restituisce l'output appropriato.
+
+Istanziare una pipeline e specificare il modello da utilizzare per la generazione di testo. Il modello viene scaricato e memorizzato nella cache in modo da poterlo riutilizzare facilmente. Infine, passare del testo per attivare il modello.
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="text-generation", model="Qwen/Qwen2.5-1.5B")
+pipeline("il segreto per preparare una torta davvero buona è ")
+[{'generated_text': 'il segreto per preparare una torta davvero buona è 1) usare gli ingredienti giusti e 2) seguire alla lettera la ricetta. la ricetta della torta è la seguente: 1 tazza di zucchero, 1 tazza di farina, 1 tazza di latte, 1 tazza di burro, 1 tazza di uova, 1 tazza di gocce di cioccolato. se vuoi preparare 2 torte, quanto zucchero ti serve? Per preparare 2 torte, avrete bisogno di 2 tazze di zucchero.'}]
+```
+
+Per chattare con un modello, lo schema di utilizzo è lo stesso. L'unica differenza è che è necessario creare una cronologia delle chat (l'input per `Pipeline`) tra l'utente e il sistema.
+
+> [!TIP]
+> È anche possibile chattare con un modello direttamente dalla riga di comando.
+> ```shell
+> transformers chat Qwen/Qwen2.5-0.5B-Instruct
+> ```
+
+```py
+import torch
+from transformers import pipeline
+
+chat = [
+ {"role": "system", "content": "Sei un robot sfacciato e spiritoso, proprio come lo immaginava Hollywood nel 1986."},
+ {"role": "user", "content": "Ehi, mi puoi suggerire qualcosa di divertente da fare a New York?"}
+]
+
+pipeline = pipeline(task="text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct", dtype=torch.bfloat16, device_map="auto")
+response = pipeline(chat, max_new_tokens=512)
+print(response[0]["generated_text"][-1]["content"])
+```
+
+Espandi gli esempi riportati di seguito per vedere come funziona `Pipeline` per diverse modalità e attività.
+
+
+
+ +## Perché non dovrei usare Transformers? + +- Questa libreria non è un toolbox modulare di blocchi costitutivi per reti neurali. Il codice nei file dei modelli non è stato rifattorizzato con ulteriori astrazioni di proposito, in modo che i ricercatori possano iterare rapidamente su ciascuno dei modelli senza dover approfondire ulteriori astrazioni/file. +- L'API di addestramento è ottimizzata per funzionare con i modelli PyTorch forniti da Transformers. Per i loop generici di machine learning, è necessario utilizzare un'altra libreria come [Accelerate](https://huggingface.co/docs/accelerate). +- Gli [script di esempio](https://github.com/huggingface/transformers/tree/main/examples) sono solo *esempi*. Potrebbero non funzionare immediatamente nel vostro caso specifico e potrebbe essere necessario adattare il codice affinché funzioni. + +## 100 progetti che usano Transformers + +Transformers è più di un semplice toolkit per l'utilizzo di modelli preaddestrati, è una comunità di progetti costruita attorno ad esso e all' +Hugging Face Hub. Vogliamo che Transformers consenta a sviluppatori, ricercatori, studenti, professori, ingegneri e chiunque altro +di realizzare i propri progetti dei sogni. + +Per celebrare le 100.000 stelle di Transformers, abbiamo voluto puntare i riflettori sulla +comunità con la pagina [awesome-transformers](./awesome-transformers.md), che elenca 100 +incredibili progetti realizzati con Transformers. + +Se possiedi o utilizzi un progetto che ritieni debba essere inserito nell'elenco, apri una PR per aggiungerlo! + +## Modelli di esempio + +È possibile testare la maggior parte dei nostri modelli direttamente sulle loro [pagine dei modelli Hub](https://huggingface.co/models). + +Espandi ciascuna modalità qui sotto per vedere alcuni modelli di esempio per vari casi d'uso. + +
-
- To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`.
-
-
+ > [!TIP]
+ > To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`.
token (`str`, `optional`):
Whether to use authentication token to load the remote folder. Useful to load private repositories
@@ -151,11 +142,11 @@ def load_adapter(
offload_index (`int`, `optional`):
`offload_index` argument to be passed to `accelerate.dispatch_model` method.
peft_config (`dict[str, Any]`, *optional*):
- The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts
- methods. This argument is used in case users directly pass PEFT state dicts
+ The configuration of the adapter to add, supported adapters are all non-prompt learning configs (LoRA,
+ IA³, etc). This argument is used in case users directly pass PEFT state dicts.
adapter_state_dict (`dict[str, torch.Tensor]`, *optional*):
The state dict of the adapter to load. This argument is used in case users directly pass PEFT state
- dicts
+ dicts.
low_cpu_mem_usage (`bool`, *optional*, defaults to `False`):
Reduce memory usage while loading the PEFT adapter. This should also speed up the loading process.
Requires PEFT version 0.13.0 or higher.
@@ -320,10 +311,12 @@ def add_adapter(self, adapter_config, adapter_name: Optional[str] = None) -> Non
name is assigned to the adapter to follow the convention of PEFT library (in PEFT we use "default" as the
default adapter name).
+ Note that the newly added adapter is not automatically activated. To activate it, use `model.set_adapter`.
+
Args:
adapter_config (`~peft.PeftConfig`):
- The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts
- methods
+ The configuration of the adapter to add, supported adapters are non-prompt learning methods (LoRA,
+ IA³, etc.).
adapter_name (`str`, *optional*, defaults to `"default"`):
The name of the adapter to add. If no name is passed, a default name is assigned to the adapter.
"""
@@ -470,13 +463,6 @@ def active_adapters(self) -> list[str]:
return active_adapters
- def active_adapter(self) -> str:
- warnings.warn(
- "The `active_adapter` method is deprecated and will be removed in a future version.", FutureWarning
- )
-
- return self.active_adapters()[0]
-
def get_adapter_state_dict(self, adapter_name: Optional[str] = None, state_dict: Optional[dict] = None) -> dict:
"""
If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
@@ -564,34 +550,47 @@ def _dispatch_accelerate_model(
def delete_adapter(self, adapter_names: Union[list[str], str]) -> None:
"""
- Delete an adapter's LoRA layers from the underlying model.
+ Delete a PEFT adapter from the underlying model.
Args:
adapter_names (`Union[list[str], str]`):
The name(s) of the adapter(s) to delete.
-
- Example:
-
- ```py
- from diffusers import AutoPipelineForText2Image
- import torch
-
- pipeline = AutoPipelineForText2Image.from_pretrained(
- "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
- ).to("cuda")
- pipeline.load_lora_weights(
- "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic"
- )
- pipeline.delete_adapters("cinematic")
- ```
"""
check_peft_version(min_version=MIN_PEFT_VERSION)
+ min_version_delete_adapter = "0.18.0"
if not self._hf_peft_config_loaded:
raise ValueError("No adapter loaded. Please load an adapter first.")
- from peft.tuners.tuners_utils import BaseTunerLayer
+ # TODO: delete old version once support for PEFT < 0.18.0 is dropped
+ def old_delete_adapter(model, adapter_name, prefix=None):
+ from peft.tuners.tuners_utils import BaseTunerLayer
+ from peft.utils import ModulesToSaveWrapper
+
+ has_modules_to_save = False
+ for module in model.modules():
+ if isinstance(module, ModulesToSaveWrapper):
+ has_modules_to_save |= True
+ continue
+ if isinstance(module, BaseTunerLayer):
+ if hasattr(module, "delete_adapter"):
+ module.delete_adapter(adapter_name)
+ else:
+ raise ValueError(
+ "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1"
+ )
+
+ if has_modules_to_save:
+ logger.warning(
+ "The deleted adapter contains modules_to_save, which could not be deleted. For this to work, PEFT version "
+ f">= {min_version_delete_adapter} is required."
+ )
+
+ if version.parse(importlib.metadata.version("peft")) >= version.parse(min_version_delete_adapter):
+ from peft.functional import delete_adapter
+ else:
+ delete_adapter = old_delete_adapter
if isinstance(adapter_names, str):
adapter_names = [adapter_names]
@@ -603,16 +602,9 @@ def delete_adapter(self, adapter_names: Union[list[str], str]) -> None:
f"The following adapter(s) are not present and cannot be deleted: {', '.join(missing_adapters)}"
)
- for adapter_name in adapter_names:
- for module in self.modules():
- if isinstance(module, BaseTunerLayer):
- if hasattr(module, "delete_adapter"):
- module.delete_adapter(adapter_name)
- else:
- raise ValueError(
- "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1"
- )
-
+ prefixes = [f"{self.peft_config[adapter_name].peft_type.value.lower()}_" for adapter_name in adapter_names]
+ for adapter_name, prefix in zip(adapter_names, prefixes):
+ delete_adapter(self, adapter_name=adapter_name, prefix=prefix)
# For transformers integration - we need to pop the adapter from the config
if getattr(self, "_hf_peft_config_loaded", False) and hasattr(self, "peft_config"):
self.peft_config.pop(adapter_name, None)
diff --git a/src/transformers/integrations/sdpa_attention.py b/src/transformers/integrations/sdpa_attention.py
index f6c6f2785c3f..e2eb69b2db8f 100644
--- a/src/transformers/integrations/sdpa_attention.py
+++ b/src/transformers/integrations/sdpa_attention.py
@@ -2,7 +2,7 @@
import torch
-from ..utils import is_torch_xpu_available, logging
+from ..utils import is_torch_npu_available, is_torch_xpu_available, logging
from ..utils.import_utils import is_torch_greater_or_equal
@@ -12,6 +12,7 @@
_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_8 = is_torch_greater_or_equal("2.8", accept_dev=True)
_is_torch_xpu_available = is_torch_xpu_available()
+_is_torch_npu_available = is_torch_npu_available()
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -35,8 +36,12 @@ def use_gqa_in_sdpa(attention_mask: Optional[torch.Tensor], key: torch.Tensor) -
# 2.xpu
# - torch version >= 2.8
# - key is not a torch.fx.Proxy (otherwise it will fail with a tracing error)
+ # 3.npu
+ # - npu is not supported gqa currently
if _is_torch_xpu_available:
return _is_torch_greater_or_equal_than_2_8 and not isinstance(key, torch.fx.Proxy)
+ if _is_torch_npu_available:
+ return False
return _is_torch_greater_or_equal_than_2_5 and attention_mask is None and not isinstance(key, torch.fx.Proxy)
@@ -80,6 +85,14 @@ def sdpa_attention_forward(
if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor):
is_causal = is_causal.item()
+ # When `is_causal = False` and the `attention_mask` is not of boolean type, the Ascend NPU's SDPA interface cannot utilize the FlashAttentionScore operator，
+ # and falls back to small-operator concatenation. To invoke the FlashAttentionScore, the attention_mask must be converted to boolean type.
+ # This adaptation ensures the `attention_mask` meets the requirement for using FlashAttentionScore.
+ if _is_torch_npu_available:
+ if attention_mask is not None and attention_mask.dtype != torch.bool:
+ # Convert to boolean type, making sdpa to force call FlashAttentionScore to improve performance.
+ attention_mask = torch.logical_not(attention_mask.bool()).to(query.device)
+
attn_output = torch.nn.functional.scaled_dot_product_attention(
query,
key,
diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py
index 3f9d40f13388..e746ed60a7e4 100644
--- a/src/transformers/integrations/tensor_parallel.py
+++ b/src/transformers/integrations/tensor_parallel.py
@@ -1009,7 +1009,7 @@ def add_tensor_parallel_hooks_to_module(
def shard_and_distribute_module(
- model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh, set_param=True
+ model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh
): # TODO: rename to shard_and_distribute_param
r"""
This function is called in `from_pretrained` when loading a model's checkpoints.
@@ -1103,8 +1103,6 @@ def distribute_model(model, distributed_config, device_mesh, tp_size):
raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}")
for name, module in model.named_modules():
if not getattr(module, "_is_hooked", False):
- from transformers.integrations.tensor_parallel import add_tensor_parallel_hooks_to_module
-
plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False)
add_tensor_parallel_hooks_to_module(
model=model,
diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py
index 1899a6de8af8..99306bd94c88 100644
--- a/src/transformers/masking_utils.py
+++ b/src/transformers/masking_utils.py
@@ -26,7 +26,7 @@
if is_torch_flex_attn_available():
- from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size # noqa: N811
+ 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
@@ -43,7 +43,7 @@
logger = logging.get_logger(__name__)
-def and_masks(*mask_functions: list[Callable]) -> Callable:
+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}")
@@ -57,7 +57,7 @@ def and_mask(batch_idx, head_idx, q_idx, kv_idx):
return and_mask
-def or_masks(*mask_functions: list[Callable]) -> Callable:
+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}")
@@ -625,6 +625,7 @@ class AttentionMaskInterface(GeneralInterface):
"sdpa": sdpa_mask,
"eager": eager_mask,
"flash_attention_2": flash_attention_mask,
+ "flash_attention_3": flash_attention_mask,
"flex_attention": flex_attention_mask,
}
diff --git a/src/transformers/model_debugging_utils.py b/src/transformers/model_debugging_utils.py
index 9f763c83c66d..2c7b47c04fd5 100644
--- a/src/transformers/model_debugging_utils.py
+++ b/src/transformers/model_debugging_utils.py
@@ -21,6 +21,7 @@
from io import StringIO
from typing import Optional
+from .utils import logging
from .utils.import_utils import is_torch_available, requires
@@ -28,6 +29,7 @@
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
@@ -35,7 +37,6 @@
_torch_distributed_available = True
else:
_torch_distributed_available = False
-from .utils import logging
logger = logging.get_logger(__name__)
@@ -224,7 +225,7 @@ def prune_intermediate_layers(node):
prune_intermediate_layers(child)
-def log_model_debug_trace(debug_path, model):
+def log_model_debug_trace(debug_path: Optional[str], model):
if debug_path:
try:
os.makedirs(debug_path, exist_ok=True)
@@ -269,8 +270,8 @@ def clean(val):
def _attach_debugger_logic(
model,
- debug_path: Optional[str] = ".",
- do_prune_layers: Optional[bool] = True,
+ debug_path: str = ".",
+ do_prune_layers: bool = True,
use_repr: bool = True,
):
"""
@@ -283,7 +284,7 @@ def _attach_debugger_logic(
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 asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors
+ `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__
@@ -399,8 +400,8 @@ def top_wrapped_forward(*inps, **kws):
def model_addition_debugger_context(
model,
debug_path: Optional[str] = None,
- do_prune_layers: Optional[bool] = True,
- use_repr: Optional[bool] = True,
+ do_prune_layers: bool = True,
+ use_repr: bool = True,
):
"""
# Model addition debugger - context manager for model adders
diff --git a/src/transformers/modelcard.py b/src/transformers/modelcard.py
index 8c68d8b8af10..dd3a0b401733 100644
--- a/src/transformers/modelcard.py
+++ b/src/transformers/modelcard.py
@@ -794,8 +794,7 @@ def parse_log_history(log_history):
if idx > 0:
eval_results = {}
for key, value in log_history[idx].items():
- if key.startswith("eval_"):
- key = key[5:]
+ key = key.removeprefix("eval_")
if key not in ["runtime", "samples_per_second", "steps_per_second", "epoch", "step"]:
camel_cased_key = " ".join([part.capitalize() for part in key.split("_")])
eval_results[camel_cased_key] = value
diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py
index 37554773a85f..5312b0dd9cd0 100644
--- a/src/transformers/modeling_flash_attention_utils.py
+++ b/src/transformers/modeling_flash_attention_utils.py
@@ -124,7 +124,7 @@ def _lazy_define_process_function(flash_function):
return partial(_process_flash_attention_kwargs, supports_mapping=supports_mapping)
-def lazy_import_flash_attention(implementation: Optional[str]):
+def lazy_import_flash_attention(implementation: Optional[str], force_import: Optional[bool] = False):
"""
Lazily import flash attention and return the respective functions + flags.
@@ -132,11 +132,11 @@ def lazy_import_flash_attention(implementation: Optional[str]):
work without preloading. See `load_and_register_kernel` in `integrations.hub_kernels`.
"""
global _flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn
- if any(k is None for k in [_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn]):
+ if force_import or any(k is None for k in [_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn]):
_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn = _lazy_imports(implementation)
global _process_flash_kwargs_fn
- if _process_flash_kwargs_fn is None:
+ if force_import or _process_flash_kwargs_fn is None:
_process_flash_kwargs_fn = _lazy_define_process_function(_flash_varlen_fn)
return (_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn), _process_flash_kwargs_fn
diff --git a/src/transformers/modeling_gguf_pytorch_utils.py b/src/transformers/modeling_gguf_pytorch_utils.py
index 9b90fb82afa2..08aaac3617ff 100644
--- a/src/transformers/modeling_gguf_pytorch_utils.py
+++ b/src/transformers/modeling_gguf_pytorch_utils.py
@@ -243,6 +243,17 @@ def process(self, weights, name, **kwargs):
return GGUFTensor(weights, name, {})
+class Lfm2TensorProcessor(TensorProcessor):
+ def __init__(self, config=None):
+ super().__init__(config=config)
+
+ def process(self, weights, name, **kwargs):
+ if "shortconv.conv.weight" in name:
+ ## GGUF shape is [hidden_dim, L_cache], HF expects [hidden_dim, 1, L_cache]
+ weights = np.expand_dims(weights, axis=1) ## equivalent to unsqueeze(1)
+ return GGUFTensor(weights, name, {})
+
+
TENSOR_PROCESSORS = {
"llama": LlamaTensorProcessor,
"qwen2moe": Qwen2MoeTensorProcessor,
@@ -255,6 +266,7 @@ def process(self, weights, name, **kwargs):
"nemotron": NemotronTensorProcessor,
"gemma2": Gemma2TensorProcessor,
"gemma3": Gemma2TensorProcessor,
+ "lfm2": Lfm2TensorProcessor,
}
@@ -459,6 +471,19 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo
if parsed_parameters["config"]["model_type"] == "gemma3":
parsed_parameters["config"]["model_type"] = "gemma3_text"
+ if parsed_parameters["config"]["model_type"] == "lfm2":
+ gguf_num_key_value_heads = parsed_parameters["config"]["num_key_value_heads"]
+ # LFM2 GGUF checkpoint defines num_key_value_heads as a list of integers .e.g [0, 0, 8, 0, 0, 8, 0, 0, 8, 0, 8, 0, 8, 0, 8, 0] but we need to set it to the max value for HF
+ parsed_parameters["config"]["num_key_value_heads"] = max(gguf_num_key_value_heads)
+ ## we already read the correct intermediate_size from the GGUF checkpoint so we need to set block_auto_adjust_ff_dim to False
+ parsed_parameters["config"]["block_auto_adjust_ff_dim"] = False
+
+ ## llama.cpp defines the layers that are full-attention by looking at num_key_value_heads
+ ## we need to set the full_attn_idxs to the layers that are full-attention
+ parsed_parameters["config"]["full_attn_idxs"] = [
+ i for i, num_kv_heads in enumerate(gguf_num_key_value_heads) if num_kv_heads > 0
+ ]
+
# retrieve config vocab_size from tokenizer
# Please refer to https://github.com/huggingface/transformers/issues/32526 for more details
if "vocab_size" not in parsed_parameters["config"]:
diff --git a/src/transformers/modeling_outputs.py b/src/transformers/modeling_outputs.py
index 597e20b28ca8..1747f6fa477b 100755
--- a/src/transformers/modeling_outputs.py
+++ b/src/transformers/modeling_outputs.py
@@ -1651,7 +1651,7 @@ class Seq2SeqTSPredictionOutput(ModelOutput):
"""
loss: Optional[torch.FloatTensor] = None
- params: Optional[tuple[torch.FloatTensor]] = None
+ params: Optional[tuple[torch.FloatTensor, ...]] = None
past_key_values: Optional[EncoderDecoderCache] = None
decoder_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
decoder_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
diff --git a/src/transformers/modeling_rope_utils.py b/src/transformers/modeling_rope_utils.py
index 34c136980234..c0070df6ee17 100644
--- a/src/transformers/modeling_rope_utils.py
+++ b/src/transformers/modeling_rope_utils.py
@@ -98,17 +98,30 @@ def _compute_default_rope_parameters(
Computes the inverse frequencies according to the original RoPE implementation
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*): If less than 1.0, inverse frequencies will be returned for
+ the first fraction of the head_dim. Defaults to 1.0.
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_theta
- partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
dim = int(head_dim * partial_rotary_factor)
@@ -128,11 +141,24 @@ def _compute_linear_scaling_rope_parameters(
Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*): If less than 1.0, inverse frequencies will be returned for
+ the first fraction of the head_dim. Defaults to 1.0.
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).
@@ -156,20 +182,43 @@ def _compute_dynamic_ntk_parameters(
) -> tuple["torch.Tensor", float]:
"""
Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla
+
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+ * max_position_embeddings (`int`): The default sequence length used to update the dynamic RoPE at
+ inference time
+ * rope_scaling (`dict[str, float]`): The standard RoPE scaling parameters, from which `factor`
+ will be accessed. The value of `factor` is used to determine the new base frequency, along with the
+ current sequence length (seq_len), the maximum positional embeddings (max_position_embeddings), and the
+ computed dimensionality (dim) of the rotary embeddings. If seq_len <= max_position_embeddings, this
+ factor has no effect. If seq_len <= max_position_embeddings, this factor effectively stretches the
+ context window using an exponent derived from `dim`.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*): If less than 1.0, inverse frequencies will be returned for
+ the first fraction of the head_dim. Defaults to 1.0.
device (`torch.device`):
The device to use for initialization of the inverse frequencies.
seq_len (`int`, *optional*):
- The current sequence length, used to update the dynamic RoPE at inference time.
+ The current sequence length, used to update the dynamic RoPE at inference time. If `None` or shorter than
+ max_position_embeddings, this value will be overridden by max_position_embeddings.
+
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).
"""
# TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
base = config.rope_theta
- partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
dim = int(head_dim * partial_rotary_factor)
max_position_embeddings = config.max_position_embeddings
@@ -200,20 +249,58 @@ def _compute_yarn_parameters(
"""
Computes the inverse frequencies with NTK scaling. Please refer to the
[original paper](https://huggingface.co/papers/2309.00071)
+
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+ * max_position_embeddings (`int`): The maximum length of the positional embeddings.
+ * rope_scaling (`dict[str, float | int]`): The standard RoPE scaling parameters, from which the following
+ keys will be accessed:
+ * `attention_factor` (`float`, *optional*): The scaling factor to be applied to the computed cos/sin.
+ If None, the value is inferred from `factor`, `mscale`, and `mscale_all_dim` as avaialble.
+ * `beta_fast` (`float`, *optional*, defaults to 32): Parameter to set the boundary for extrapolation
+ (only) in the linear ramp function.
+ * `beta_slow` (`float`, *optional*, defaults to 1): Parameter to set the boundary for interpolation
+ (only) in the linear ramp function.
+ * `factor` (`float`, *optional*): The scaling factor applied when interpolating the position IDs to
+ extend the possible context length. Additionally, if `attention_factor` is None, the log of this
+ value is used to compute a value for `attention_factor`, possibly in conjunciton with `mscale` and
+ `mscale_all_dim`, if provided.
+ * `mscale` (`float`, *optional*): If `attention_factor` is None and both `mscale` and
+ `mscale_all_dim` are provided, `mscale` acts scalar augmenting `log(factor)` when computing the
+ numerator for the inferred value of `attention_factor`. If not provided, `attention_factor` will be
+ calculated based on `factor` only.
+ * `mscale_all_dim` (`float`, *optional*): If `attention_factor` is None and both `mscale` and
+ `mscale_all_dim` are provided, `mscale_all_dim` acts scalar augmenting `log(factor)` when computing
+ the denominator for the inferred value of `attention_factor`. If not provided, `attention_factor`
+ will be calculated based on `factor` only.
+ * `original_max_position_embeddings` (`int`, *optional*): The original max position embeddings used
+ during pretraining. If not provided, the function falls back to `max_position_embeddings`.
+ * `truncate` (`bool`, *optional*): Whether to truncate the correction range.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*, defaults to 1.0): If less than 1.0, inverse frequencies
+ will be returned for the first fraction of the head_dim.
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.
"""
base = config.rope_theta
- partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
dim = int(head_dim * partial_rotary_factor)
factor = config.rope_scaling["factor"]
@@ -237,7 +324,7 @@ def get_mscale(scale, mscale=1):
attention_factor = get_mscale(factor)
# Optional config options
- # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly)
+ # beta_fast/beta_slow: as suggested in the paper, default to 32 and 1 respectively
beta_fast = config.rope_scaling.get("beta_fast") or 32
beta_slow = config.rope_scaling.get("beta_slow") or 1
@@ -287,20 +374,49 @@ def _compute_longrope_parameters(
"""
Computes the inverse frequencies with LongRoPE scaling. Please refer to the
[original implementation](https://github.com/microsoft/LongRoPE)
+
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+ * max_position_embeddings (`int`): The maximum length of the positional embeddings.
+ * original_max_position_embeddings (`int`, *optional*): The original max position embeddings used during
+ pretraining. If not provided, defaults to `max_position_embeddings`.
+ * rope_scaling (`dict[str, float]`): The standard RoPE scaling parameters, from which the following keys
+ will be accessed:
+ * `attention_factor` (`float`, *optional*): The scaling factor to be applied on the attention
+ computation. If unspecified, it defaults to value recommended by the implementation, inferred from
+ the value of `factor`.
+ * `factor` (`float`, *optional*): The scaling factor to apply to the RoPE embeddings. If both
+ `max_position_embeddings` and `original_max_position_embeddings` are provided, this value will be
+ overridden s the ratio between those values.
+ * `long_factor` (`float`, *optional*): The scale factor applied when computing the inverse
+ frequencies if `seq_len` is provided and greater than `original_max_position_embeddings`.
+ * `short_factor` (`float`, *optional*): The scale factor applied when computing the inverse
+ frequencies if `seq_len` is None or less-than-or-equal-to `original_max_position_embeddings`.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*, defaults to 1.0): If less than 1.0, inverse frequencies
+ will be returned for the first fraction of the head_dim.
device (`torch.device`):
The device to use for initialization of the inverse frequencies.
seq_len (`int`, *optional*):
The current sequence length.
+
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.
"""
# TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
base = config.rope_theta
- partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
dim = int(head_dim * partial_rotary_factor)
long_factor = config.rope_scaling["long_factor"]
@@ -311,9 +427,8 @@ def _compute_longrope_parameters(
# NOTE: Phi3 (and potentially other models) modify `max_position_embeddings` and have a
# `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two
# values to compute the default attention scaling factor, instead of using `factor`.
- if hasattr(config, "original_max_position_embeddings"):
- original_max_position_embeddings = config.original_max_position_embeddings
- factor = config.max_position_embeddings / config.original_max_position_embeddings
+ if original_max_position_embeddings := getattr(config, "original_max_position_embeddings", None):
+ factor = config.max_position_embeddings / original_max_position_embeddings
else:
original_max_position_embeddings = config.max_position_embeddings
@@ -343,7 +458,31 @@ def _compute_llama3_parameters(
Args:
config ([`~transformers.PretrainedConfig`]):
- The model configuration.
+ The model configuration. This function assumes that the config will provide at least the following
+ properties:
+
+ * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived.
+ * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly.
+ * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly.
+ * rope_scaling (`dict[str, float | int]`): The standard RoPE scaling parameters, from which the following
+ keys will be accessed:
+ * `factor` (`float`, *optional*): The scaling factor applied to the inverse frequencies when 1) the
+ wavelength is greater than `low_freq_wavelen` prior to smoothing, and 2) to all inverse frequencies
+ during smoothing.
+ * `high_freq_factor` (`float`): The scale factor used to compute `high_freq_wavelen` and
+ the value for the denominator of the smoothing factor prior to the `low_freq_factor` shift.
+ * `low_freq_factor` (`float`): The scale factor used to compute `low_freq_wavelen` and
+ the shift applied to the numerator and denominator of the smoothing factor.
+ frequencies if `seq_len` is None or less-than-or-equal-to `original_max_position_embeddings`.
+ * `original_max_position_embeddings` (`int`): The original max position embeddings used
+ during pretraining. If not provided, the function falls back to `max_position_embeddings`.
+
+ Additionally, this function will make use of the following properties if they are found in the config:
+
+ * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be
+ derived as hidden_size // num_attention_heads.
+ * partial_rotary_factor (`float`, *optional*): If less than 1.0, inverse frequencies will be returned for
+ the first fraction of the head_dim. Defaults to 1.0.
device (`torch.device`):
The device to use for initialization of the inverse frequencies.
seq_len (`int`, *optional*):
@@ -527,7 +666,7 @@ def _validate_longrope_parameters(config: PretrainedConfig, ignore_keys: Optiona
received_keys = set(rope_scaling.keys())
_check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)
- partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
dim = int(head_dim * partial_rotary_factor)
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 12c3e7cd99ef..188ac3782b1f 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -19,13 +19,10 @@
import gc
import importlib.metadata
import inspect
-import itertools
import json
import os
import re
-import shutil
import sys
-import tempfile
import warnings
from abc import abstractmethod
from collections import defaultdict
@@ -40,6 +37,9 @@
import torch
from huggingface_hub import split_torch_state_dict_into_shards
from packaging import version
+from safetensors import safe_open
+from safetensors.torch import load_file as safe_load_file
+from safetensors.torch import save_file as safe_save_file
from torch import Tensor, nn
from torch.distributions import constraints
from torch.utils.checkpoint import checkpoint
@@ -103,14 +103,12 @@
is_optimum_available,
is_peft_available,
is_remote_url,
- is_safetensors_available,
is_torch_flex_attn_available,
is_torch_greater_or_equal,
is_torch_mlu_available,
is_torch_npu_available,
is_torch_xla_available,
is_torch_xpu_available,
- is_torchao_available,
logging,
)
from .utils.generic import _CAN_RECORD_REGISTRY, GeneralInterface, OutputRecorder
@@ -125,9 +123,6 @@
from .utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
-if is_torchao_available():
- from torchao.quantization import Int4WeightOnlyConfig
-
if is_accelerate_available():
from accelerate import dispatch_model, infer_auto_device_map
from accelerate.hooks import add_hook_to_module
@@ -136,7 +131,6 @@
extract_model_from_parallel,
get_balanced_memory,
get_max_memory,
- load_offloaded_weights,
offload_weight,
save_offload_index,
)
@@ -145,11 +139,6 @@
if accelerate_version >= version.parse("0.31"):
from accelerate.utils.modeling import get_state_dict_from_offload
-if is_safetensors_available():
- from safetensors import safe_open
- from safetensors.torch import load_file as safe_load_file
- from safetensors.torch import save_file as safe_save_file
-
if is_peft_available():
from .utils import find_adapter_config_file
@@ -414,24 +403,11 @@ def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True):
index_present = os.path.isfile(index_file)
safe_index_present = os.path.isfile(safe_index_file)
- if not index_present and not (safe_index_present and is_safetensors_available()):
- filenames = (
- (WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME) if is_safetensors_available() else (WEIGHTS_INDEX_NAME,)
- )
+ if not index_present and not safe_index_present:
+ filenames = (WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME)
raise ValueError(f"Can't find a checkpoint index ({' or '.join(filenames)}) in {folder}.")
- load_safe = False
- if safe_index_present:
- if prefer_safe:
- if is_safetensors_available():
- load_safe = True # load safe due to preference
- else:
- logger.warning(
- f"Cannot load sharded checkpoint at {folder} safely since safetensors is not installed!"
- )
- elif not index_present:
- load_safe = True # load safe since we have no other choice
-
+ load_safe = safe_index_present and (prefer_safe or not index_present)
load_index = safe_index_file if load_safe else index_file
with open(load_index, "r", encoding="utf-8") as f:
@@ -504,7 +480,7 @@ def load_state_dict(
Reads a `safetensor` or a `.bin` checkpoint file. We load the checkpoint on "cpu" by default.
"""
# Use safetensors if possible
- if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
+ if checkpoint_file.endswith(".safetensors"):
with safe_open(checkpoint_file, framework="pt") as f:
metadata = f.metadata()
@@ -575,26 +551,6 @@ def load_state_dict(
)
-def set_initialized_submodules(model, state_dict_keys):
- """
- Sets the `_is_hf_initialized` flag in all submodules of a given model when all its weights are in the loaded state
- dict.
- """
- state_dict_keys = set(state_dict_keys)
- not_initialized_submodules = {}
- for module_name, module in model.named_modules():
- if module_name == "":
- # When checking if the root module is loaded there's no need to prepend module_name.
- module_keys = set(module.state_dict())
- else:
- module_keys = {f"{module_name}.{k}" for k in module.state_dict()}
- if module_keys.issubset(state_dict_keys):
- module._is_hf_initialized = True
- else:
- not_initialized_submodules[module_name] = module
- return not_initialized_submodules
-
-
def _end_ptr(tensor: torch.Tensor) -> int:
# extract the end of the pointer if the tensor is a slice of a bigger tensor
if tensor.nelement():
@@ -682,6 +638,7 @@ def _infer_parameter_dtype(
QuantizationMethod.HQQ,
QuantizationMethod.QUARK,
QuantizationMethod.MXFP4,
+ QuantizationMethod.BITS_AND_BYTES,
}:
return True, None
else:
@@ -715,17 +672,12 @@ def _load_state_dict_into_meta_model(
model: "PreTrainedModel",
state_dict: dict,
shard_file: str,
- expected_keys: list[str],
reverse_renaming_mapping: dict[str, str],
device_map: Optional[dict] = None,
disk_offload_folder: Optional[str] = None,
disk_offload_index: Optional[dict] = None,
- cpu_offload_folder: Optional[str] = None,
- cpu_offload_index: Optional[dict] = None,
hf_quantizer: Optional[HfQuantizer] = None,
- is_safetensors: bool = False,
keep_in_fp32_regex: Optional[re.Pattern] = None,
- unexpected_keys: Optional[list[str]] = None, # passing `unexpected` for cleanup from quantization items
device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None,
) -> tuple[Optional[dict], Optional[dict]]:
"""Load parameters from `meta_state_dict` into the model. The parameters of the `meta_state_dict` are on the meta
@@ -741,18 +693,13 @@ def _load_state_dict_into_meta_model(
device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)])
is_quantized = hf_quantizer is not None
- is_hqq_or_bnb = is_quantized and hf_quantizer.quantization_config.quant_method in {
- QuantizationMethod.HQQ,
- QuantizationMethod.BITS_AND_BYTES,
- }
- is_meta_state_dict = shard_file.endswith(".safetensors") and not is_hqq_or_bnb
- file_pointer = None
- if is_meta_state_dict:
- file_pointer = safe_open(shard_file, framework="pt", device=tensor_device)
+ is_safetensors = shard_file.endswith(".safetensors")
+ is_meta_state_dict = is_safetensors
+ file_pointer = safe_open(shard_file, framework="pt", device=tensor_device) if is_meta_state_dict else None
+ params_to_load = list(state_dict.keys())
- for param_name, empty_param in state_dict.items():
- if param_name not in expected_keys: # when loading from ckpt, we skip param if doesnt exist in modeling
- continue
+ for param_name in params_to_load:
+ empty_param = state_dict[param_name]
# we need to use serialized_param_name as file pointer is untouched
if is_meta_state_dict:
# This is the name of the parameter as it appears on disk file
@@ -769,19 +716,8 @@ def _load_state_dict_into_meta_model(
)
if device_mesh is not None:
- if (
- not is_quantized
- or (not hf_quantizer.requires_parameters_quantization)
- or (
- not hf_quantizer.check_quantized_param(
- model,
- param,
- param_name,
- state_dict,
- device_map=device_map,
- )
- )
- ): # In this case, the param is already on the correct device!
+ if not is_quantized or not hf_quantizer.param_needs_quantization(model, param_name):
+ # In this case, the param is already on the correct device!
shard_and_distribute_module(
model,
param,
@@ -792,7 +728,8 @@ def _load_state_dict_into_meta_model(
device_mesh.get_local_rank(),
device_mesh,
)
- else: # we have a device mesh but the param needs to be quantized, so we shard inside create_quantized_param:
+ else:
+ # we have a device mesh but the param needs to be quantized, so we shard inside create_quantized_param
sharding_kwargs = {
"empty_param": empty_param,
"casting_dtype": casting_dtype,
@@ -805,8 +742,6 @@ def _load_state_dict_into_meta_model(
param,
param_name,
device_mesh.get_local_rank(),
- state_dict,
- unexpected_keys,
**sharding_kwargs,
)
else:
@@ -828,22 +763,7 @@ def _load_state_dict_into_meta_model(
if param_device == "disk":
if not is_safetensors:
disk_offload_index = offload_weight(param, param_name, disk_offload_folder, disk_offload_index)
- elif param_device == "cpu" and cpu_offload_index is not None:
- cpu_offload_index = offload_weight(param, param_name, cpu_offload_folder, cpu_offload_index)
- elif (
- not is_quantized
- or (not hf_quantizer.requires_parameters_quantization)
- or (
- not hf_quantizer.check_quantized_param(
- model,
- param,
- param_name,
- state_dict,
- param_device=param_device,
- device_map=device_map,
- )
- )
- ):
+ elif not is_quantized or not hf_quantizer.param_needs_quantization(model, param_name):
if is_fsdp_enabled():
param_device = "cpu" if is_local_dist_rank_0() else "meta"
@@ -851,9 +771,7 @@ def _load_state_dict_into_meta_model(
else:
# TODO naming is stupid it loads it as well
- hf_quantizer.create_quantized_param(
- model, param, param_name, param_device, state_dict, unexpected_keys
- )
+ hf_quantizer.create_quantized_param(model, param, param_name, param_device)
# For quantized modules with FSDP/DeepSpeed Stage 3, we need to quantize the parameter on the GPU
# and then cast it to CPU to avoid excessive memory usage on each GPU
@@ -876,10 +794,14 @@ def _load_state_dict_into_meta_model(
value = type(value)(value.data.to(param_to), **val_kwargs, **value.__dict__)
setattr(module, param_type, value)
+ # Remove the param from the state dict if it was not loaded on the fly to avoid wasting memory
+ if not is_meta_state_dict:
+ del state_dict[param_name]
+
if file_pointer is not None:
file_pointer.__exit__(None, None, None)
- return disk_offload_index, cpu_offload_index
+ return disk_offload_index
def load_shard_file(args):
@@ -887,46 +809,26 @@ def load_shard_file(args):
shard_file,
state_dict,
disk_only_shard_files,
- is_hqq_or_bnb,
is_quantized,
device_map,
hf_quantizer,
key_renaming_mapping,
weights_only,
- model_to_load,
- expected_keys,
+ model,
reverse_key_renaming_mapping,
disk_offload_folder,
disk_offload_index,
- cpu_offload_folder,
- cpu_offload_index,
- is_offloaded_safetensors,
keep_in_fp32_regex,
- unexpected_keys,
device_mesh,
) = args
# Skip the load for shards that only contain disk-offloaded weights
if shard_file in disk_only_shard_files:
- return [], disk_offload_index, cpu_offload_index
+ return [], disk_offload_index
map_location = "cpu"
- if (
- shard_file.endswith(".safetensors")
- and not is_hqq_or_bnb
- and not (is_deepspeed_zero3_enabled() and not is_quantized)
- ):
+ if shard_file.endswith(".safetensors") and not (is_deepspeed_zero3_enabled() and not is_quantized):
map_location = "meta"
- elif (
- device_map is not None
- and hf_quantizer is not None
- and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO
- and (
- hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"]
- or isinstance(hf_quantizer.quantization_config.quant_type, Int4WeightOnlyConfig)
- )
- ):
- map_location = torch.device([d for d in device_map.values() if d not in ["disk"]][0])
# If shard_file is "", we use the existing state_dict instead of loading it
if shard_file != "":
@@ -938,30 +840,24 @@ def load_shard_file(args):
state_dict = {key_renaming_mapping[k]: v for k, v in state_dict.items() if k in key_renaming_mapping}
error_msgs = []
-
if is_deepspeed_zero3_enabled() and not is_quantized:
- error_msgs += _load_state_dict_into_zero3_model(model_to_load, state_dict)
+ error_msgs += _load_state_dict_into_zero3_model(model, state_dict)
# Skip it with fsdp on ranks other than 0
elif not (is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized):
- disk_offload_index, cpu_offload_index = _load_state_dict_into_meta_model(
- model_to_load,
+ disk_offload_index = _load_state_dict_into_meta_model(
+ model,
state_dict,
shard_file,
- expected_keys,
reverse_key_renaming_mapping,
device_map=device_map,
disk_offload_folder=disk_offload_folder,
disk_offload_index=disk_offload_index,
- cpu_offload_folder=cpu_offload_folder,
- cpu_offload_index=cpu_offload_index,
hf_quantizer=hf_quantizer,
- is_safetensors=is_offloaded_safetensors,
keep_in_fp32_regex=keep_in_fp32_regex,
- unexpected_keys=unexpected_keys,
device_mesh=device_mesh,
)
- return error_msgs, disk_offload_index, cpu_offload_index
+ return error_msgs, disk_offload_index
def load_shard_files_with_threadpool(args_list):
@@ -978,18 +874,13 @@ def load_shard_files_with_threadpool(args_list):
with logging.tqdm(total=len(args_list), desc="Loading checkpoint shards") as pbar:
futures = [executor.submit(load_shard_file, arg) for arg in args_list]
for future in as_completed(futures):
- result = future.result()
- (
- _error_msgs,
- disk_offload_index,
- cpu_offload_index,
- ) = result
+ _error_msgs, disk_offload_index = future.result()
error_msgs += _error_msgs
pbar.update(1)
- return error_msgs, disk_offload_index, cpu_offload_index
+ return error_msgs, disk_offload_index
def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:
@@ -1190,7 +1081,12 @@ def _get_resolved_checkpoint_files(
is_sharded = True
if not local_files_only and not is_offline_mode():
if resolved_archive_file is not None:
- if filename in [WEIGHTS_NAME, WEIGHTS_INDEX_NAME]:
+ # In a CI environment (CircleCI / Github Actions workflow runs) or in a pytest run,
+ # we set `DISABLE_SAFETENSORS_CONVERSION=true` to prevent the conversion.
+ if (
+ filename in [WEIGHTS_NAME, WEIGHTS_INDEX_NAME]
+ and os.getenv("DISABLE_SAFETENSORS_CONVERSION", None) != "true"
+ ):
# If the PyTorch file was found, check if there is a safetensors file on the repository
# If there is no safetensors file on the repositories, start an auto conversion
safe_weights_name = SAFE_WEIGHTS_INDEX_NAME if is_sharded else SAFE_WEIGHTS_NAME
@@ -1481,20 +1377,18 @@ def _get_device_map(
def _find_missing_and_unexpected_keys(
- cls,
model: "PreTrainedModel",
original_checkpoint_keys: list[str],
checkpoint_keys: list[str],
loading_base_model_from_task_state_dict: bool,
hf_quantizer: Optional[HfQuantizer],
- device_map: dict,
) -> tuple[list[str], list[str]]:
"""Find missing keys (keys that are part of the model parameters but were NOT found in the loaded state dict keys) and unexpected keys
(keys found in the loaded state dict keys, but that are NOT part of the model parameters)
"""
prefix = model.base_model_prefix
- # Compute expected keys, i.e. keys that the FULL model (not model_to_load) expects
+ # Compute expected keys, i.e. keys that the full model expects
expected_keys = list(model.state_dict().keys())
if hf_quantizer is not None:
expected_keys = hf_quantizer.update_expected_keys(model, expected_keys, checkpoint_keys)
@@ -1512,12 +1406,6 @@ def _find_missing_and_unexpected_keys(
model_buffers = {n for n, _ in model.named_buffers()}
unexpected_keys = sorted(unexpected_keys - model_buffers)
- # Old checkpoints may have keys for rotary_emb.inv_freq for each layer, however we moved this buffer to the main model
- # (so the buffer name has changed). Remove them in such a case
- has_inv_freq_buffers = any(buffer.endswith("rotary_emb.inv_freq") for buffer in model_buffers)
- if has_inv_freq_buffers:
- unexpected_keys = [k for k in unexpected_keys if "rotary_emb.inv_freq" not in k]
-
tied_params = find_tied_parameters(model)
for group in tied_params:
missing_in_group = [k for k in missing_keys if k in group]
@@ -1526,16 +1414,7 @@ def _find_missing_and_unexpected_keys(
if hf_quantizer is not None:
missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix)
- unexpected_keys = hf_quantizer.update_unexpected_keys(model, unexpected_keys, prefix)
-
- # Model-specific exceptions for missing and unexpected keys (e.g. if the modeling change over time, or any other reason...)
- if cls._keys_to_ignore_on_load_missing is not None:
- for pattern in cls._keys_to_ignore_on_load_missing:
- missing_keys = [k for k in missing_keys if re.search(pattern, k) is None]
-
- if cls._keys_to_ignore_on_load_unexpected is not None:
- for pattern in cls._keys_to_ignore_on_load_unexpected:
- unexpected_keys = [k for k in unexpected_keys if re.search(pattern, k) is None]
+ unexpected_keys = hf_quantizer.update_unexpected_keys(model, unexpected_keys)
return missing_keys, unexpected_keys
@@ -1721,7 +1600,7 @@ def create_extended_attention_mask_for_decoder(input_shape, attention_mask, devi
def get_extended_attention_mask(
self,
attention_mask: Tensor,
- input_shape: tuple[int],
+ input_shape: tuple[int, ...],
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
) -> Tensor:
@@ -1959,7 +1838,7 @@ def get_input_embeddings(self) -> nn.Module:
)
def set_input_embeddings(self, value: nn.Module):
- """Fallback setter that handles **~70 %** of models in the code‑base.
+ """Fallback setter that handles **~70%** of models in the code-base.
Order of attempts:
1. `self.model.embed_tokens`
@@ -2305,8 +2184,6 @@ def tp_plan(self, plan: dict[str, str]):
if hasattr(self, "named_parameters"):
model_param_names = [name for name, _ in self.named_parameters()]
if model_param_names: # Only validate if model has parameters
- import re
-
for layer_pattern in plan.keys():
# Convert pattern to regex (replace * with .*)
regex_pattern = layer_pattern.replace("*", r"\d+")
@@ -2332,8 +2209,6 @@ def tp_plan(self, plan: dict[str, str]):
flexible_matched = True
break
if not flexible_matched:
- import warnings
-
warnings.warn(
f"Layer pattern '{layer_pattern}' does not match any parameters in the model. "
f"This rule may not be applied during tensor parallelization."
@@ -2778,42 +2653,46 @@ def _check_and_adjust_attn_implementation(
None to sdpa (to potentially eager).
"""
applicable_attn_implementation = attn_implementation
+
# If FA not installed, do not fail but use kernels instead
if (
- applicable_attn_implementation == "flash_attention_2"
+ attn_implementation is not None
+ and attn_implementation.startswith("flash_attention")
and self._supports_flash_attn
- and not is_flash_attn_2_available()
+ and not (is_flash_attn_2_available() or is_flash_attn_3_available())
and is_kernels_available()
):
- applicable_attn_implementation = "kernels-community/flash-attn"
+ if attn_implementation.endswith("2"):
+ applicable_attn_implementation = "kernels-community/flash-attn"
+ else:
+ applicable_attn_implementation = "kernels-community/vllm-flash-attn3"
+
if is_kernel(applicable_attn_implementation):
try:
load_and_register_kernel(applicable_attn_implementation)
# log that we used kernel fallback if successful
- if attn_implementation == "flash_attention_2":
+ if attn_implementation.startswith("flash_attention"):
logger.warning_once(
- "You do not have `flash_attn` installed, using `kernels-community/flash-attn` from the `kernels` "
- "library instead!"
+ f"You do not have `flash_attn` installed, using `{applicable_attn_implementation}` "
+ "from the `kernels` library instead!"
)
except Exception as e:
- if attn_implementation == "flash_attention_2":
- self._flash_attn_2_can_dispatch() # will fail as fa2 is not available but raise the proper exception
- logger.warning_once(
- f"Could not find a kernel matching `{applicable_attn_implementation}` compatible with your device in the "
- f"hub:\n{e}.\nUsing default attention implementation instead (sdpa if available, eager otherwise)."
- )
- try:
- self._sdpa_can_dispatch(is_init_check)
- applicable_attn_implementation = "sdpa"
- except (ValueError, ImportError) as e:
- applicable_attn_implementation = "eager"
+ # raise the proper exception for requested flash attention
+ if attn_implementation.startswith("flash_attention"):
+ if attn_implementation.endswith("2"):
+ self._flash_attn_2_can_dispatch()
+ else:
+ self._flash_attn_3_can_dispatch()
+
+ # error properly out if a kernel was specifically requested
+ raise e
else:
applicable_attn_implementation = self.get_correct_attn_implementation(
applicable_attn_implementation, is_init_check
)
# preload flash attention here to allow compile with fullgraph
if applicable_attn_implementation.startswith("flash_attention"):
- lazy_import_flash_attention(applicable_attn_implementation)
+ lazy_import_flash_attention(applicable_attn_implementation, force_import=True)
return applicable_attn_implementation
@@ -3558,7 +3437,7 @@ def _get_resized_lm_head(
self,
old_lm_head: nn.Linear,
new_num_tokens: Optional[int] = None,
- transposed: Optional[bool] = False,
+ transposed: bool = False,
mean_resizing: bool = True,
) -> nn.Linear:
"""
@@ -3715,7 +3594,7 @@ def _init_added_lm_head_weights_with_mean(
old_lm_head_dim,
old_num_tokens,
added_num_tokens,
- transposed=False,
+ transposed: bool = False,
):
if transposed:
# Transpose to the desired shape for the function.
@@ -3993,8 +3872,6 @@ def save_pretrained(
"`save_config` is deprecated and will be removed in v5 of Transformers. Use `is_main_process` instead."
)
is_main_process = kwargs.pop("save_config")
- if safe_serialization and not is_safetensors_available():
- raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.")
# we need to check against tp_size, not tp_plan, as tp_plan is substituted to the class one
if self._tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"):
@@ -4263,7 +4140,7 @@ def save_pretrained(
if _is_dtensor_available and isinstance(state_dict[tensor], DTensor):
full_tensor = state_dict[tensor].full_tensor()
# to get the correctly ordered tensor we need to repack if packed
- if _get_parameter_tp_plan(tensor, self._tp_plan) in ("local_packed_rowwise",):
+ if _get_parameter_tp_plan(tensor, self._tp_plan) == "local_packed_rowwise":
full_tensor = repack_weights(full_tensor, -1, self._tp_size, 2)
shard[tensor] = full_tensor.contiguous() # only do contiguous after it's permuted correctly
else:
@@ -4365,9 +4242,9 @@ def get_memory_footprint(self, return_buffers=True):
are tensors that do not require gradients and not registered as parameters. E.g. mean and std in batch
norm layers. Please see: https://discuss.pytorch.org/t/what-pytorch-means-by-buffers/120266/2
"""
- mem = sum([param.nelement() * param.element_size() for param in self.parameters()])
+ mem = sum(param.nelement() * param.element_size() for param in self.parameters())
if return_buffers:
- mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()])
+ mem_bufs = sum(buf.nelement() * buf.element_size() for buf in self.buffers())
mem = mem + mem_bufs
return mem
@@ -4591,9 +4468,6 @@ def from_pretrained(
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- resume_download:
- Deprecated and ignored. All downloads are now resumed by default when possible.
- Will be removed in v5 of Transformers.
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.
@@ -4683,10 +4557,6 @@ def from_pretrained(
If provided, it has to contain dimension named `"tp"` in case it's > 1 dimensional, this dimension will be used for tensor parallelism
offload_folder (`str` or `os.PathLike`, *optional*):
If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- offload_state_dict (`bool`, *optional*):
- If `True`, will temporarily offload the CPU state dict to the hard drive to avoid getting out of CPU
- RAM if the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to
- `True` when there is some disk offload.
offload_buffers (`bool`, *optional*):
Whether or not to offload the buffers with the model parameters.
quantization_config (`Union[QuantizationConfigMixin,Dict]`, *optional*):
@@ -4764,7 +4634,6 @@ def from_pretrained(
device_map = kwargs.pop("device_map", None)
max_memory = kwargs.pop("max_memory", None)
offload_folder = kwargs.pop("offload_folder", None)
- offload_state_dict = kwargs.pop("offload_state_dict", False)
offload_buffers = kwargs.pop("offload_buffers", False)
load_in_8bit = kwargs.pop("load_in_8bit", False)
load_in_4bit = kwargs.pop("load_in_4bit", False)
@@ -4859,9 +4728,6 @@ def from_pretrained(
if token is not None and adapter_kwargs is not None and "token" not in adapter_kwargs:
adapter_kwargs["token"] = token
- if use_safetensors is None and not is_safetensors_available():
- use_safetensors = False
-
if gguf_file is not None and not is_accelerate_available():
raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -5058,12 +4924,7 @@ def from_pretrained(
is_quantized = hf_quantizer is not None
is_from_file = pretrained_model_name_or_path is not None or gguf_file is not None
- if (
- is_safetensors_available()
- and is_from_file
- and not is_sharded
- and checkpoint_files[0].endswith(".safetensors")
- ):
+ if is_from_file and not is_sharded and checkpoint_files[0].endswith(".safetensors"):
with safe_open(checkpoint_files[0], framework="pt") as f:
metadata = f.metadata()
@@ -5159,6 +5020,10 @@ def _assign_original_dtype(module):
config._pre_quantization_dtype = original_dtype
_assign_original_dtype(model)
+ # Torchao needs access to all metadata later
+ if hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO:
+ hf_quantizer.set_metadata(checkpoint_files)
+
if _torch_distributed_available and device_mesh is not None:
model = distribute_model(model, distributed_config, device_mesh, tp_size)
@@ -5192,7 +5057,6 @@ def _assign_original_dtype(module):
sharded_metadata=sharded_metadata,
device_map=device_map,
disk_offload_folder=offload_folder,
- offload_state_dict=offload_state_dict,
dtype=dtype,
hf_quantizer=hf_quantizer,
keep_in_fp32_regex=keep_in_fp32_regex,
@@ -5346,6 +5210,14 @@ def _get_key_renaming_mapping(
prefix = self.base_model_prefix
_prefix = f"{prefix}."
+ if loading_task_model_from_base_state_dict:
+ task_specific_expected_keys, base_model_keys = [], []
+ for key in self.state_dict():
+ if key.startswith(_prefix):
+ base_model_keys.append(key[len(_prefix) :])
+ else:
+ task_specific_expected_keys.append(key)
+
renamed_keys = {}
key_renaming_mapping = {}
for key in checkpoint_keys:
@@ -5363,6 +5235,13 @@ def _get_key_renaming_mapping(
# In this case, we need to add the prefix to the keys, to match them to the expected keys
if loading_task_model_from_base_state_dict:
+ # small sanity check: if we find a key that is only part of the task-specific keys, we raise
+ # (if it's also part of the base model, we do not raise and assume it comes from there)
+ if new_key in task_specific_expected_keys and new_key not in base_model_keys:
+ raise ValueError(
+ "The state dictionary of the model you are trying to load is corrupted. Are you sure it was "
+ "properly saved?"
+ )
new_key = ".".join([prefix, new_key])
# In this case we need to remove the prefix from the key to match them to the expected keys, and use
# only the keys starting with the prefix
@@ -5416,7 +5295,6 @@ def _load_pretrained_model(
sharded_metadata: Optional[dict] = None,
device_map: Optional[dict] = None,
disk_offload_folder: Optional[str] = None,
- offload_state_dict: Optional[bool] = None,
dtype: Optional[torch.dtype] = None,
hf_quantizer: Optional[HfQuantizer] = None,
keep_in_fp32_regex: Optional[re.Pattern] = None,
@@ -5430,10 +5308,6 @@ def _load_pretrained_model(
QuantizationMethod.HQQ,
QuantizationMethod.QUARK,
}
- is_hqq_or_bnb = is_quantized and hf_quantizer.quantization_config.quant_method in {
- QuantizationMethod.HQQ,
- QuantizationMethod.BITS_AND_BYTES,
- }
# Get all the keys of the state dicts that we have to initialize the model
if sharded_metadata is not None:
@@ -5447,7 +5321,6 @@ def _load_pretrained_model(
# Check if we are in a special state, i.e. loading from a state dict coming from a different architecture
prefix = model.base_model_prefix
- _prefix = f"{prefix}."
has_prefix_module = any(s.startswith(prefix) for s in original_checkpoint_keys) if len(prefix) > 0 else False
expects_prefix_module = hasattr(model, prefix) if len(prefix) > 0 else False
loading_task_model_from_base_state_dict = not has_prefix_module and expects_prefix_module
@@ -5464,13 +5337,7 @@ def _load_pretrained_model(
# Find missing and unexpected keys from the state dict
missing_keys, unexpected_keys = _find_missing_and_unexpected_keys(
- cls,
- model,
- original_checkpoint_keys,
- checkpoint_keys,
- loading_base_model_from_task_state_dict,
- hf_quantizer,
- device_map,
+ model, original_checkpoint_keys, checkpoint_keys, loading_base_model_from_task_state_dict, hf_quantizer
)
# Find all the keys with shape mismatch (if we ignore the mismatch, the weights need to be newly initialized the
# same way as missing keys)
@@ -5484,16 +5351,18 @@ def _load_pretrained_model(
weights_only,
)
- # We need to update both the mapping and the list of checkpoint keys to remove the mismatched ones
- key_renaming_mapping = {k: v for k, v in key_renaming_mapping.items() if v not in mismatched_keys}
+ # We need to update both the mapping and the list of checkpoint keys to remove the mismatched and unexpected ones
+ key_renaming_mapping = {
+ k: v for k, v in key_renaming_mapping.items() if v not in mismatched_keys and v not in unexpected_keys
+ }
checkpoint_keys = list(key_renaming_mapping.values())
# Move missing (and potentially mismatched) keys back to cpu from meta device (because they won't be moved when
# loading the weights as they are not in the loaded state dict)
- model._move_missing_keys_from_meta_to_cpu(missing_keys + mismatched_keys, unexpected_keys, dtype, hf_quantizer)
+ model._move_missing_keys_from_meta_to_cpu(missing_keys + mismatched_keys, dtype, hf_quantizer)
# correctly initialize the missing (and potentially mismatched) keys
- model._initialize_missing_keys(checkpoint_keys, ignore_mismatched_sizes, is_quantized)
+ model._initialize_missing_keys(missing_keys + mismatched_keys, is_quantized)
# Set some modules to fp32 if needed
if keep_in_fp32_regex is not None:
@@ -5502,29 +5371,6 @@ def _load_pretrained_model(
# param = param.to(torch.float32) does not work here as only in the local scope.
param.data = param.data.to(torch.float32)
- # Make sure we are able to load base models as well as derived models (specific task models, with heads)
- model_to_load = model
- # In this case, we load a ForTaskModel with keys from a BaseModel -> only load keys to the BaseModel
- if loading_task_model_from_base_state_dict:
- model_to_load = getattr(model, prefix)
- # Here we need to remove the prefix we added to correctly find missing/unexpected keys, as we will load
- # in the submodule
- key_renaming_mapping = {k: v[len(_prefix) :] for k, v in key_renaming_mapping.items()}
- checkpoint_keys = list(key_renaming_mapping.values())
- # We need to update the device map as well
- if device_map is not None:
- device_map = {k[len(_prefix) :] if k.startswith(_prefix) else k: v for k, v in device_map.items()}
- # small sanity check: the base model should not contain task-specific head keys
- task_specific_expected_keys = [s for s in model.state_dict() if not s.startswith(_prefix)]
- base_model_expected_keys = list(model_to_load.state_dict().keys())
- if any(
- key in task_specific_expected_keys and key not in base_model_expected_keys for key in checkpoint_keys
- ):
- raise ValueError(
- "The state dictionary of the model you are trying to load is corrupted. Are you sure it was "
- "properly saved?"
- )
-
# Get reverse key mapping
reverse_key_renaming_mapping = {v: k for k, v in key_renaming_mapping.items()}
@@ -5534,8 +5380,6 @@ def _load_pretrained_model(
disk_only_shard_files = []
# Prepare parameters offloading if needed
if device_map is not None and "disk" in device_map.values():
- if offload_state_dict is None:
- offload_state_dict = True
if disk_offload_folder is not None:
os.makedirs(disk_offload_folder, exist_ok=True)
is_offloaded_safetensors = checkpoint_files is not None and checkpoint_files[0].endswith(".safetensors")
@@ -5573,31 +5417,22 @@ def _load_pretrained_model(
else:
disk_offload_index = {}
- # This offload index if for params that are supposed to be on the "cpu", either with or without a device_map
- # It allows to load parameters one-by-one from the state dict, avoiding a memory peak of 2 x state_dict_size,
- # i.e. 1x to load it, and 1x to copy it to model
- cpu_offload_folder = None
- cpu_offload_index = None
- if offload_state_dict:
- cpu_offload_folder = tempfile.mkdtemp()
- cpu_offload_index = {}
-
# To be able to iterate, even if we don't use it if the state_dict is already provided
elif state_dict is not None:
checkpoint_files = [""]
# Compute expected model keys
- expected_keys = list(model_to_load.state_dict().keys())
+ expected_keys = list(model.state_dict().keys())
if hf_quantizer is not None:
- expected_keys = hf_quantizer.update_expected_keys(model_to_load, expected_keys, checkpoint_keys)
+ expected_keys = hf_quantizer.update_expected_keys(model, expected_keys, checkpoint_keys)
if logger.level >= logging.WARNING:
- verify_tp_plan(expected_keys, getattr(model_to_load, "_tp_plan", None))
+ verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None))
# Warmup cuda to load the weights much faster on devices
if device_map is not None and not is_hqq_or_quark:
expanded_device_map = expand_device_map(device_map, expected_keys)
- caching_allocator_warmup(model_to_load, expanded_device_map, hf_quantizer)
+ caching_allocator_warmup(model, expanded_device_map, hf_quantizer)
# Prepare and compatabilize arguments for serial and parallel shard loading
args_list = [
@@ -5605,22 +5440,16 @@ def _load_pretrained_model(
shard_file,
state_dict,
disk_only_shard_files,
- is_hqq_or_bnb,
is_quantized,
device_map,
hf_quantizer,
key_renaming_mapping,
weights_only,
- model_to_load,
- expected_keys,
+ model,
reverse_key_renaming_mapping,
disk_offload_folder,
disk_offload_index,
- cpu_offload_folder,
- cpu_offload_index,
- is_offloaded_safetensors,
keep_in_fp32_regex,
- unexpected_keys,
device_mesh,
)
for shard_file in checkpoint_files
@@ -5632,40 +5461,20 @@ def _load_pretrained_model(
os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES
and not is_deepspeed_zero3_enabled()
):
- _error_msgs, disk_offload_index, cpu_offload_index = load_shard_files_with_threadpool(args_list)
+ _error_msgs, disk_offload_index = load_shard_files_with_threadpool(args_list)
error_msgs += _error_msgs
else:
if len(args_list) > 1:
args_list = logging.tqdm(args_list, desc="Loading checkpoint shards")
for args in args_list:
- _error_msgs, disk_offload_index, cpu_offload_index = load_shard_file(args)
+ _error_msgs, disk_offload_index = load_shard_file(args)
error_msgs += _error_msgs
- # Adjust offloaded weights name and save if needed
- if disk_offload_index is not None and len(disk_offload_index) > 0:
- if loading_task_model_from_base_state_dict:
- # We need to add the prefix of the base model
- prefix = cls.base_model_prefix
- if not is_offloaded_safetensors:
- for weight_name in disk_offload_index:
- shutil.move(
- os.path.join(disk_offload_folder, f"{weight_name}.dat"),
- os.path.join(disk_offload_folder, f"{prefix}.{weight_name}.dat"),
- )
- disk_offload_index = {f"{prefix}.{key}": value for key, value in disk_offload_index.items()}
- if not is_offloaded_safetensors:
- save_offload_index(disk_offload_index, disk_offload_folder)
- disk_offload_index = None
-
- # one-at-a-time param loading for the cpu offloaded params
- if offload_state_dict:
- # Load back temporarily offloaded state dict
- load_offloaded_weights(model_to_load, cpu_offload_index, cpu_offload_folder)
- shutil.rmtree(cpu_offload_folder)
-
- if hf_quantizer is not None:
- missing_keys = hf_quantizer.update_missing_keys_after_loading(model_to_load, missing_keys, prefix)
+ # Save offloaded index if needed
+ if disk_offload_index is not None and len(disk_offload_index) > 0 and not is_offloaded_safetensors:
+ save_offload_index(disk_offload_index, disk_offload_folder)
+ disk_offload_index = None
# Post-processing for tensor parallelism
if device_mesh is not None:
@@ -5700,6 +5509,11 @@ def _load_pretrained_model(
device_mesh,
)
+ # Remove potential model-specific exceptions from the warnings
+ missing_keys, unexpected_keys = model._adjust_missing_and_unexpected_keys(
+ missing_keys, unexpected_keys, loading_task_model_from_base_state_dict
+ )
+
# All potential warnings/infos
if len(error_msgs) > 0:
error_msg = "\n\t".join(error_msgs)
@@ -5720,21 +5534,12 @@ def _load_pretrained_model(
f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
" (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
)
- else:
- logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
if len(missing_keys) > 0:
logger.warning(
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
" TRAIN this model on a down-stream task to be able to use it for predictions and inference."
)
- elif len(mismatched_keys) == 0:
- logger.info(
- f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
- f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
- f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
- " training."
- )
if len(mismatched_keys) > 0:
mismatched_warning = "\n".join(
[
@@ -5803,7 +5608,7 @@ def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=Fal
for name, module in self.named_modules():
if remove_prefix:
_prefix = f"{self.base_model_prefix}."
- name = name[len(_prefix) :] if name.startswith(_prefix) else name
+ name = name.removeprefix(_prefix)
elif add_prefix:
name = ".".join([self.base_model_prefix, name]) if len(name) > 0 else self.base_model_prefix
@@ -6022,12 +5827,8 @@ def is_backend_compatible(cls):
return cls._supports_attention_backend
def _move_missing_keys_from_meta_to_cpu(
- self,
- missing_keys: list[str],
- unexpected_keys: list[str],
- dtype: Optional[torch.dtype],
- hf_quantizer: Optional[HfQuantizer],
- ) -> "PreTrainedModel":
+ self, missing_keys: list[str], dtype: torch.dtype, hf_quantizer: Optional[HfQuantizer]
+ ) -> None:
"""Move the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts) back
from meta device to cpu.
"""
@@ -6047,56 +5848,90 @@ def _move_missing_keys_from_meta_to_cpu(
# Buffers are not initialized on the meta device, so we still need this check to avoid overwriting them
if param.device == torch.device("meta"):
value = torch.empty_like(param, dtype=dtype, device="cpu")
- if (
- not is_quantized
- or (getattr(hf_quantizer, "requires_parameters_quantization", False))
- or not hf_quantizer.check_quantized_param(self, param_value=value, param_name=key, state_dict={})
- ):
+ if not is_quantized or not hf_quantizer.param_needs_quantization(self, key):
_load_parameter_into_model(self, key, value)
else:
- hf_quantizer.create_quantized_param(self, value, key, "cpu", model_state_dict, unexpected_keys)
+ hf_quantizer.create_quantized_param(self, value, key, "cpu")
- def _initialize_missing_keys(
- self,
- loaded_keys: list[str],
- ignore_mismatched_sizes: bool,
- is_quantized: bool,
- ) -> "PreTrainedModel":
+ def _initialize_missing_keys(self, missing_keys: list[str], is_quantized: bool) -> None:
"""Initialize the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts), according to
`_initialize_weights`. Indeed, since the corresponding weights are missing from the state dict, they will not be replaced and need to
be initialized correctly (i.e. weight initialization distribution).
Also take care of setting the `_is_hf_initialized` flag for keys that are not missing.
"""
- if not ignore_mismatched_sizes:
- not_initialized_submodules = set_initialized_submodules(self, loaded_keys)
- # If we're about to tie the output embeds to the input embeds we don't need to init them
+ for key in self.state_dict():
+ # If it's part of the keys that will be loaded, mark it as already initialized
+ if key not in missing_keys:
+ param_or_buffer = self.get_parameter_or_buffer(key)
+ param_or_buffer._is_hf_initialized = True
+
+ def set_is_initialized_for_modules(module):
+ # A module is already initialized if and only if all its children are also already initialized, and all
+ # its immediate `nn.Parameter` and persistent buffers are also already initialized
if (
- hasattr(self.config.get_text_config(decoder=True), "tie_word_embeddings")
- and self.config.get_text_config(decoder=True).tie_word_embeddings
+ all(getattr(child, "_is_hf_initialized", False) for child in module.children())
+ and all(getattr(param, "_is_hf_initialized", False) for param in module.parameters(recurse=False))
+ and all(
+ getattr(buffer, "_is_hf_initialized", False)
+ for buffer in module.buffers(recurse=False)
+ if buffer not in module._non_persistent_buffers_set
+ )
):
- output_embeddings = self.get_output_embeddings()
- if output_embeddings is not None:
- # Still need to initialize if there is a bias term since biases are not tied.
- if not hasattr(output_embeddings, "bias") or output_embeddings.bias is None:
- output_embeddings._is_hf_initialized = True
- else:
- not_initialized_submodules = dict(self.named_modules())
+ module._is_hf_initialized = True
+
+ # Set the flag on the modules as well. We do it recursively (depth-first), as it's more efficient (we do not
+ # need to check the entire state dict of each module, only the immediate children, so we only iterate once over
+ # each param)
+ self.apply(set_is_initialized_for_modules)
+
# This will only initialize submodules that are not marked as initialized by the line above.
if is_deepspeed_zero3_enabled() and not is_quantized:
import deepspeed
not_initialized_parameters = list(
- set(
- itertools.chain.from_iterable(
- submodule.parameters(recurse=False) for submodule in not_initialized_submodules.values()
- )
- )
+ {v for v in self.state_dict().values() if not getattr(v, "_is_hf_initialized", False)}
)
with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0):
self.initialize_weights()
else:
self.initialize_weights()
+ def _adjust_missing_and_unexpected_keys(
+ self, missing_keys: list[str], unexpected_keys: list[str], loading_task_model_from_base_state_dict: bool
+ ) -> tuple[list[str], list[str]]:
+ """Adjust the `missing_keys` and `unexpected_keys` based on current model's exception rules, to avoid
+ raising unneeded warnings/errors.
+ """
+ # Old checkpoints may have keys for rotary_emb.inv_freq for each layer, however we moved this buffer to the main model
+ # (so the buffer name has changed). Remove them in such a case. This is another exception that was not added to
+ # `_keys_to_ignore_on_load_unexpected` as it touches many models -> we add it manually to the existing patterns
+ has_inv_freq_buffers = any(buffer.endswith("rotary_emb.inv_freq") for buffer, _ in self.named_buffers())
+ additional_unexpected_patterns = [r"rotary_emb\.inv_freq"] if has_inv_freq_buffers else []
+
+ missing_patterns = self._keys_to_ignore_on_load_missing or []
+ unexpected_patterns = (self._keys_to_ignore_on_load_unexpected or []) + additional_unexpected_patterns
+ ignore_missing_regex, ignore_unexpected_regex = None, None
+ if len(missing_patterns) > 0:
+ ignore_missing_regex = re.compile("|".join(rf"({pattern})" for pattern in missing_patterns))
+ if len(unexpected_patterns) > 0:
+ ignore_unexpected_regex = re.compile("|".join(rf"({pattern})" for pattern in unexpected_patterns))
+
+ # Clean-up missing keys
+ if ignore_missing_regex is not None:
+ missing_keys = [key for key in missing_keys if ignore_missing_regex.search(key) is None]
+
+ # Clean-up unexpected keys
+ if ignore_unexpected_regex is not None:
+ unexpected_keys = [key for key in unexpected_keys if ignore_unexpected_regex.search(key) is None]
+
+ # Note: only the unexpected keys should remove the added prefix here, to correctly display the original name
+ # in the warnings. For missing keys, we should show the prefix in the warning as it's part of the final model
+ if loading_task_model_from_base_state_dict:
+ _prefix = f"{self.base_model_prefix}."
+ unexpected_keys = [k.removeprefix(_prefix) for k in unexpected_keys]
+
+ return missing_keys, unexpected_keys
+
def get_parameter_or_buffer(self, target: str):
"""
Return the parameter or buffer given by `target` if it exists, otherwise throw an error. This combines
diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py
index 5c391e7162f4..c721f24a506d 100644
--- a/src/transformers/models/__init__.py
+++ b/src/transformers/models/__init__.py
@@ -48,6 +48,7 @@
from .blip import *
from .blip_2 import *
from .bloom import *
+ from .blt import *
from .bridgetower import *
from .bros import *
from .byt5 import *
@@ -107,6 +108,8 @@
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 *
@@ -183,6 +186,7 @@
from .led import *
from .levit import *
from .lfm2 import *
+ from .lfm2_vl import *
from .lightglue import *
from .lilt import *
from .llama import *
@@ -251,6 +255,7 @@
from .owlv2 import *
from .owlvit import *
from .paligemma import *
+ from .parakeet import *
from .patchtsmixer import *
from .patchtst import *
from .pegasus import *
@@ -281,6 +286,7 @@
from .qwen3 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 *
diff --git a/src/transformers/models/altclip/configuration_altclip.py b/src/transformers/models/altclip/configuration_altclip.py
index 5e8c0f2a262e..474fc48081b5 100755
--- a/src/transformers/models/altclip/configuration_altclip.py
+++ b/src/transformers/models/altclip/configuration_altclip.py
@@ -303,7 +303,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -335,7 +335,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/aria/image_processing_aria.py b/src/transformers/models/aria/image_processing_aria.py
index 4fc2fcf7ec6b..659ed5f112d8 100644
--- a/src/transformers/models/aria/image_processing_aria.py
+++ b/src/transformers/models/aria/image_processing_aria.py
@@ -43,12 +43,12 @@
logger = logging.get_logger(__name__)
-def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.array]:
+def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.ndarray]:
"""
Divides an image into patches of a specified size.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The input image.
patch_size (`int`):
The size of each patch.
@@ -56,7 +56,7 @@ def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) ->
The channel dimension format of the input image.
Returns:
- list: A list of np.array representing the patches.
+ list: A list of np.ndarray representing the patches.
"""
patches = []
height, width = get_image_size(image, channel_dim=input_data_format)
@@ -342,12 +342,12 @@ def preprocess(
def _resize_for_patching(
self, image: np.ndarray, target_resolution: tuple, resample, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Resizes an image to a target resolution while maintaining aspect ratio.
Args:
- image (np.array):
+ image (np.ndarray):
The input image.
target_resolution (tuple):
The target resolution (height, width) of the image.
@@ -357,7 +357,7 @@ def _resize_for_patching(
The channel dimension format of the input image.
Returns:
- np.array: The resized and padded image.
+ np.ndarray: The resized and padded image.
"""
new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
@@ -375,7 +375,7 @@ def _get_padding_size(self, original_resolution: tuple, target_resolution: tuple
def _pad_for_patching(
self, image: np.ndarray, target_resolution: tuple, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pad an image to a target resolution while maintaining aspect ratio.
"""
@@ -460,12 +460,12 @@ def get_image_patches(
resample: PILImageResampling,
data_format: ChannelDimension,
input_data_format: ChannelDimension,
- ) -> list[np.array]:
+ ) -> list[np.ndarray]:
"""
Process an image with variable resolutions by dividing it into patches.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The input image to be processed.
grid_pinpoints (list[tuple[int, int]]):
A list of possible resolutions as tuples.
@@ -479,7 +479,7 @@ def get_image_patches(
The channel dimension format of the input image.
Returns:
- `list[np.array]`: A list of NumPy arrays containing the processed image patches.
+ `list[np.ndarray]`: A list of NumPy arrays containing the processed image patches.
"""
if not isinstance(grid_pinpoints, list):
raise TypeError("grid_pinpoints must be a list of possible resolutions.")
diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py
index a626d2cd4b82..02f2f884dadf 100644
--- a/src/transformers/models/aria/modular_aria.py
+++ b/src/transformers/models/aria/modular_aria.py
@@ -725,12 +725,12 @@ def preprocess(
def _resize_for_patching(
self, image: np.ndarray, target_resolution: tuple, resample, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Resizes an image to a target resolution while maintaining aspect ratio.
Args:
- image (np.array):
+ image (np.ndarray):
The input image.
target_resolution (tuple):
The target resolution (height, width) of the image.
@@ -740,7 +740,7 @@ def _resize_for_patching(
The channel dimension format of the input image.
Returns:
- np.array: The resized and padded image.
+ np.ndarray: The resized and padded image.
"""
new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
@@ -758,7 +758,7 @@ def _get_padding_size(self, original_resolution: tuple, target_resolution: tuple
def _pad_for_patching(
self, image: np.ndarray, target_resolution: tuple, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pad an image to a target resolution while maintaining aspect ratio.
"""
@@ -843,12 +843,12 @@ def get_image_patches(
resample: PILImageResampling,
data_format: ChannelDimension,
input_data_format: ChannelDimension,
- ) -> list[np.array]:
+ ) -> list[np.ndarray]:
"""
Process an image with variable resolutions by dividing it into patches.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The input image to be processed.
grid_pinpoints (list[tuple[int, int]]):
A list of possible resolutions as tuples.
@@ -862,7 +862,7 @@ def get_image_patches(
The channel dimension format of the input image.
Returns:
- `list[np.array]`: A list of NumPy arrays containing the processed image patches.
+ `list[np.ndarray]`: A list of NumPy arrays containing the processed image patches.
"""
if not isinstance(grid_pinpoints, list):
raise TypeError("grid_pinpoints must be a list of possible resolutions.")
diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py
index 38f38cd31b40..f6a12e7cef98 100644
--- a/src/transformers/models/auto/configuration_auto.py
+++ b/src/transformers/models/auto/configuration_auto.py
@@ -65,6 +65,7 @@
("blip-2", "Blip2Config"),
("blip_2_qformer", "Blip2QFormerConfig"),
("bloom", "BloomConfig"),
+ ("blt", "BltConfig"),
("bridgetower", "BridgeTowerConfig"),
("bros", "BrosConfig"),
("camembert", "CamembertConfig"),
@@ -126,6 +127,9 @@
("dots1", "Dots1Config"),
("dpr", "DPRConfig"),
("dpt", "DPTConfig"),
+ ("edgetam", "EdgeTamConfig"),
+ ("edgetam_video", "EdgeTamVideoConfig"),
+ ("edgetam_vision_model", "EdgeTamVisionConfig"),
("efficientformer", "EfficientFormerConfig"),
("efficientloftr", "EfficientLoFTRConfig"),
("efficientnet", "EfficientNetConfig"),
@@ -222,6 +226,7 @@
("led", "LEDConfig"),
("levit", "LevitConfig"),
("lfm2", "Lfm2Config"),
+ ("lfm2_vl", "Lfm2VlConfig"),
("lightglue", "LightGlueConfig"),
("lilt", "LiltConfig"),
("llama", "LlamaConfig"),
@@ -294,6 +299,8 @@
("owlv2", "Owlv2Config"),
("owlvit", "OwlViTConfig"),
("paligemma", "PaliGemmaConfig"),
+ ("parakeet_ctc", "ParakeetCTCConfig"),
+ ("parakeet_encoder", "ParakeetEncoderConfig"),
("patchtsmixer", "PatchTSMixerConfig"),
("patchtst", "PatchTSTConfig"),
("pegasus", "PegasusConfig"),
@@ -328,6 +335,7 @@
("qwen3", "Qwen3Config"),
("qwen3_moe", "Qwen3MoeConfig"),
("qwen3_next", "Qwen3NextConfig"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeConfig"),
("qwen3_vl", "Qwen3VLConfig"),
("qwen3_vl_moe", "Qwen3VLMoeConfig"),
("qwen3_vl_moe_text", "Qwen3VLMoeTextConfig"),
@@ -366,6 +374,7 @@
("shieldgemma2", "ShieldGemma2Config"),
("siglip", "SiglipConfig"),
("siglip2", "Siglip2Config"),
+ ("siglip2_vision_model", "Siglip2VisionConfig"),
("siglip_vision_model", "SiglipVisionConfig"),
("smollm3", "SmolLM3Config"),
("smolvlm", "SmolVLMConfig"),
@@ -488,6 +497,7 @@
("blip-2", "BLIP-2"),
("blip_2_qformer", "BLIP-2 QFormer"),
("bloom", "BLOOM"),
+ ("blt", "Blt"),
("bort", "BORT"),
("bridgetower", "BridgeTower"),
("bros", "BROS"),
@@ -556,6 +566,9 @@
("dots1", "dots1"),
("dpr", "DPR"),
("dpt", "DPT"),
+ ("edgetam", "EdgeTAM"),
+ ("edgetam_video", "EdgeTamVideo"),
+ ("edgetam_vision_model", "EdgeTamVisionModel"),
("efficientformer", "EfficientFormer"),
("efficientloftr", "EfficientLoFTR"),
("efficientnet", "EfficientNet"),
@@ -657,6 +670,7 @@
("led", "LED"),
("levit", "LeViT"),
("lfm2", "Lfm2"),
+ ("lfm2_vl", "Lfm2Vl"),
("lightglue", "LightGlue"),
("lilt", "LiLT"),
("llama", "LLaMA"),
@@ -739,6 +753,9 @@
("owlv2", "OWLv2"),
("owlvit", "OWL-ViT"),
("paligemma", "PaliGemma"),
+ ("parakeet", "Parakeet"),
+ ("parakeet_ctc", "Parakeet"),
+ ("parakeet_encoder", "ParakeetEncoder"),
("patchtsmixer", "PatchTSMixer"),
("patchtst", "PatchTST"),
("pegasus", "Pegasus"),
@@ -774,6 +791,7 @@
("qwen3", "Qwen3"),
("qwen3_moe", "Qwen3MoE"),
("qwen3_next", "Qwen3Next"),
+ ("qwen3_omni_moe", "Qwen3OmniMoE"),
("qwen3_vl", "Qwen3VL"),
("qwen3_vl_moe", "Qwen3VLMoe"),
("qwen3_vl_moe_text", "Qwen3VLMoe"),
@@ -958,6 +976,7 @@
("glm4v_moe_text", "glm4v_moe"),
("idefics3_vision", "idefics3"),
("siglip_vision_model", "siglip"),
+ ("siglip2_vision_model", "siglip2"),
("aimv2_vision_model", "aimv2"),
("smolvlm_vision", "smolvlm"),
("chinese_clip_vision_model", "chinese_clip"),
@@ -970,12 +989,15 @@
("qwen3_vl_moe_text", "qwen3_vl_moe"),
("sam_vision_model", "sam"),
("sam2_vision_model", "sam2"),
+ ("edgetam_vision_model", "edgetam"),
("sam2_hiera_det_model", "sam2"),
("sam_hq_vision_model", "sam_hq"),
("llama4_text", "llama4"),
("blip_2_qformer", "blip_2"),
("fastspeech2_conformer_with_hifigan", "fastspeech2_conformer"),
("perception_encoder", "perception_lm"),
+ ("parakeet_encoder", "parakeet"),
+ ("parakeet_ctc", "parakeet"),
]
)
diff --git a/src/transformers/models/auto/feature_extraction_auto.py b/src/transformers/models/auto/feature_extraction_auto.py
index 0307aeba077f..6d4c4f554d9d 100644
--- a/src/transformers/models/auto/feature_extraction_auto.py
+++ b/src/transformers/models/auto/feature_extraction_auto.py
@@ -81,6 +81,8 @@
("moshi", "EncodecFeatureExtractor"),
("nat", "ViTFeatureExtractor"),
("owlvit", "OwlViTFeatureExtractor"),
+ ("parakeet_ctc", "ParakeetFeatureExtractor"),
+ ("parakeet_encoder", "ParakeetFeatureExtractor"),
("perceiver", "PerceiverFeatureExtractor"),
("phi4_multimodal", "Phi4MultimodalFeatureExtractor"),
("poolformer", "PoolFormerFeatureExtractor"),
diff --git a/src/transformers/models/auto/image_processing_auto.py b/src/transformers/models/auto/image_processing_auto.py
index ebaa4a30849d..4b71712dfc7b 100644
--- a/src/transformers/models/auto/image_processing_auto.py
+++ b/src/transformers/models/auto/image_processing_auto.py
@@ -91,6 +91,7 @@
("dinov3_vit", (None, "DINOv3ViTImageProcessorFast")),
("donut-swin", ("DonutImageProcessor", "DonutImageProcessorFast")),
("dpt", ("DPTImageProcessor", "DPTImageProcessorFast")),
+ ("edgetam", (None, "Sam2ImageProcessorFast")),
("efficientformer", ("EfficientFormerImageProcessor", None)),
("efficientloftr", ("EfficientLoFTRImageProcessor", "EfficientLoFTRImageProcessorFast")),
("efficientnet", ("EfficientNetImageProcessor", "EfficientNetImageProcessorFast")),
@@ -120,6 +121,7 @@
("layoutlmv2", ("LayoutLMv2ImageProcessor", "LayoutLMv2ImageProcessorFast")),
("layoutlmv3", ("LayoutLMv3ImageProcessor", "LayoutLMv3ImageProcessorFast")),
("levit", ("LevitImageProcessor", "LevitImageProcessorFast")),
+ ("lfm2_vl", (None, "Lfm2VlImageProcessorFast")),
("lightglue", ("LightGlueImageProcessor", None)),
("llama4", ("Llama4ImageProcessor", "Llama4ImageProcessorFast")),
("llava", ("LlavaImageProcessor", "LlavaImageProcessorFast")),
@@ -564,9 +566,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
)
image_processor_class = get_image_processor_class_from_name(image_processor_type)
else:
- image_processor_type_slow = (
- image_processor_type[:-4] if image_processor_type.endswith("Fast") else image_processor_type
- )
+ image_processor_type_slow = image_processor_type.removesuffix("Fast")
image_processor_class = get_image_processor_class_from_name(image_processor_type_slow)
if image_processor_class is None and image_processor_type.endswith("Fast"):
raise ValueError(
diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py
index 93420820fb9e..298834bebe93 100644
--- a/src/transformers/models/auto/modeling_auto.py
+++ b/src/transformers/models/auto/modeling_auto.py
@@ -72,6 +72,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("blip-2", "Blip2Model"),
("blip_2_qformer", "Blip2QFormerModel"),
("bloom", "BloomModel"),
+ ("blt", "BltModel"),
("bridgetower", "BridgeTowerModel"),
("bros", "BrosModel"),
("camembert", "CamembertModel"),
@@ -130,6 +131,9 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("dots1", "Dots1Model"),
("dpr", "DPRQuestionEncoder"),
("dpt", "DPTModel"),
+ ("edgetam", "EdgeTamModel"),
+ ("edgetam_video", "EdgeTamVideoModel"),
+ ("edgetam_vision_model", "EdgeTamVisionModel"),
("efficientformer", "EfficientFormerModel"),
("efficientloftr", "EfficientLoFTRModel"),
("efficientnet", "EfficientNetModel"),
@@ -222,6 +226,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("led", "LEDModel"),
("levit", "LevitModel"),
("lfm2", "Lfm2Model"),
+ ("lfm2_vl", "Lfm2VlModel"),
("lightglue", "LightGlueForKeypointMatching"),
("lilt", "LiltModel"),
("llama", "LlamaModel"),
@@ -293,6 +298,8 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("owlv2", "Owlv2Model"),
("owlvit", "OwlViTModel"),
("paligemma", "PaliGemmaModel"),
+ ("parakeet_ctc", "ParakeetForCTC"),
+ ("parakeet_encoder", "ParakeetEncoder"),
("patchtsmixer", "PatchTSMixerModel"),
("patchtst", "PatchTSTModel"),
("pegasus", "PegasusModel"),
@@ -356,6 +363,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("sew-d", "SEWDModel"),
("siglip", "SiglipModel"),
("siglip2", "Siglip2Model"),
+ ("siglip2_vision_model", "Siglip2VisionModel"),
("siglip_vision_model", "SiglipVisionModel"),
("smollm3", "SmolLM3Model"),
("smolvlm", "SmolVLMModel"),
@@ -631,6 +639,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("blenderbot", "BlenderbotForCausalLM"),
("blenderbot-small", "BlenderbotSmallForCausalLM"),
("bloom", "BloomForCausalLM"),
+ ("blt", "BltForCausalLM"),
("camembert", "CamembertForCausalLM"),
("code_llama", "LlamaForCausalLM"),
("codegen", "CodeGenForCausalLM"),
@@ -1026,6 +1035,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("janus", "JanusForConditionalGeneration"),
("kosmos-2", "Kosmos2ForConditionalGeneration"),
("kosmos-2.5", "Kosmos2_5ForConditionalGeneration"),
+ ("lfm2_vl", "Lfm2VlForConditionalGeneration"),
("llama4", "Llama4ForConditionalGeneration"),
("llava", "LlavaForConditionalGeneration"),
("llava_next", "LlavaNextForConditionalGeneration"),
@@ -1596,6 +1606,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("data2vec-audio", "Data2VecAudioForCTC"),
("hubert", "HubertForCTC"),
("mctct", "MCTCTForCTC"),
+ ("parakeet_ctc", "ParakeetForCTC"),
("sew", "SEWForCTC"),
("sew-d", "SEWDForCTC"),
("unispeech", "UniSpeechForCTC"),
@@ -1649,6 +1660,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
("musicgen", "MusicgenForConditionalGeneration"),
("musicgen_melody", "MusicgenMelodyForConditionalGeneration"),
("qwen2_5_omni", "Qwen2_5OmniForConditionalGeneration"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeForConditionalGeneration"),
("seamless_m4t", "SeamlessM4TForTextToSpeech"),
("seamless_m4t_v2", "SeamlessM4Tv2ForTextToSpeech"),
("vits", "VitsModel"),
@@ -1700,6 +1712,8 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
MODEL_FOR_MASK_GENERATION_MAPPING_NAMES = OrderedDict(
[
+ ("edgetam", "EdgeTamModel"),
+ ("edgetam_video", "EdgeTamModel"),
("sam", "SamModel"),
("sam2", "Sam2Model"),
("sam2_video", "Sam2Model"),
diff --git a/src/transformers/models/auto/processing_auto.py b/src/transformers/models/auto/processing_auto.py
index 13583c55002f..11862a5896b9 100644
--- a/src/transformers/models/auto/processing_auto.py
+++ b/src/transformers/models/auto/processing_auto.py
@@ -66,6 +66,7 @@
("deepseek_vl", "DeepseekVLProcessor"),
("deepseek_vl_hybrid", "DeepseekVLHybridProcessor"),
("dia", "DiaProcessor"),
+ ("edgetam", "Sam2Processor"),
("emu3", "Emu3Processor"),
("evolla", "EvollaProcessor"),
("flava", "FlavaProcessor"),
@@ -93,6 +94,7 @@
("kyutai_speech_to_text", "KyutaiSpeechToTextProcessor"),
("layoutlmv2", "LayoutLMv2Processor"),
("layoutlmv3", "LayoutLMv3Processor"),
+ ("lfm2_vl", "Lfm2VlProcessor"),
("llama4", "Llama4Processor"),
("llava", "LlavaProcessor"),
("llava_next", "LlavaNextProcessor"),
@@ -120,6 +122,7 @@
("qwen2_5_vl", "Qwen2_5_VLProcessor"),
("qwen2_audio", "Qwen2AudioProcessor"),
("qwen2_vl", "Qwen2VLProcessor"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeProcessor"),
("qwen3_vl", "Qwen3VLProcessor"),
("qwen3_vl_moe", "Qwen3VLProcessor"),
("sam", "SamProcessor"),
diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py
index 7858ae587946..d0c3af490d71 100644
--- a/src/transformers/models/auto/tokenization_auto.py
+++ b/src/transformers/models/auto/tokenization_auto.py
@@ -105,6 +105,7 @@
("blip", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("blip-2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("bloom", (None, "BloomTokenizerFast" if is_tokenizers_available() else None)),
+ ("blt", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
("bridgetower", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
("bros", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("byt5", ("ByT5Tokenizer", None)),
@@ -501,6 +502,7 @@
("owlv2", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
("owlvit", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
("paligemma", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("parakeet", ("ParakeetCTCTokenizer", None)),
(
"pegasus",
(
@@ -585,6 +587,7 @@
"Qwen2TokenizerFast" if is_tokenizers_available() else None,
),
),
+ ("qwen3_omni_moe", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
("qwen3_vl", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
("qwen3_vl_moe", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
("rag", ("RagTokenizer", None)),
@@ -1139,7 +1142,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
# Otherwise we have to be creative.
# if model is an encoder decoder, the encoder tokenizer class is used by default
if isinstance(config, EncoderDecoderConfig):
- if type(config.decoder) is not type(config.encoder): # noqa: E721
+ if type(config.decoder) is not type(config.encoder):
logger.warning(
f"The encoder model config class: {config.encoder.__class__} is different from the decoder model "
f"config class: {config.decoder.__class__}. It is not recommended to use the "
diff --git a/src/transformers/models/auto/video_processing_auto.py b/src/transformers/models/auto/video_processing_auto.py
index 551de914626e..84bbc8e6fdb1 100644
--- a/src/transformers/models/auto/video_processing_auto.py
+++ b/src/transformers/models/auto/video_processing_auto.py
@@ -56,6 +56,7 @@
("qwen2_5_omni", "Qwen2VLVideoProcessor"),
("qwen2_5_vl", "Qwen2VLVideoProcessor"),
("qwen2_vl", "Qwen2VLVideoProcessor"),
+ ("qwen3_omni_moe", "Qwen2VLVideoProcessor"),
("qwen3_vl", "Qwen3VLVideoProcessor"),
("qwen3_vl_moe", "Qwen3VLVideoProcessor"),
("sam2_video", "Sam2VideoVideoProcessor"),
diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index 09f00845524d..60bf385bf494 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -531,7 +531,7 @@ def __init__(self, config: BambaConfig, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/bamba/modular_bamba.py b/src/transformers/models/bamba/modular_bamba.py
index f2495b446aa5..5ae5313d21b8 100644
--- a/src/transformers/models/bamba/modular_bamba.py
+++ b/src/transformers/models/bamba/modular_bamba.py
@@ -288,7 +288,7 @@ def __init__(self, config: BambaConfig, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/bark/modeling_bark.py b/src/transformers/models/bark/modeling_bark.py
index 8770e3e0691b..af57f7826734 100644
--- a/src/transformers/models/bark/modeling_bark.py
+++ b/src/transformers/models/bark/modeling_bark.py
@@ -595,7 +595,7 @@ class BarkSemanticModel(BarkCausalModel):
def generate(
self,
input_ids: torch.Tensor,
- semantic_generation_config: BarkSemanticGenerationConfig = None,
+ semantic_generation_config: Optional[BarkSemanticGenerationConfig] = None,
history_prompt: Optional[dict[str, torch.Tensor]] = None,
attention_mask: Optional[torch.Tensor] = None,
**kwargs,
@@ -780,8 +780,8 @@ def preprocess_histories(
def generate(
self,
semantic_output: torch.Tensor,
- semantic_generation_config: BarkSemanticGenerationConfig = None,
- coarse_generation_config: BarkCoarseGenerationConfig = None,
+ semantic_generation_config: Optional[BarkSemanticGenerationConfig] = None,
+ coarse_generation_config: Optional[BarkCoarseGenerationConfig] = None,
codebook_size: int = 1024,
history_prompt: Optional[dict[str, torch.Tensor]] = None,
return_output_lengths: Optional[bool] = None,
@@ -1192,8 +1192,8 @@ def forward(
def generate(
self,
coarse_output: torch.Tensor,
- semantic_generation_config: BarkSemanticGenerationConfig = None,
- coarse_generation_config: BarkCoarseGenerationConfig = None,
+ semantic_generation_config: Optional[BarkSemanticGenerationConfig] = None,
+ coarse_generation_config: Optional[BarkCoarseGenerationConfig] = None,
fine_generation_config: BarkFineGenerationConfig = None,
codebook_size: int = 1024,
history_prompt: Optional[dict[str, torch.Tensor]] = None,
diff --git a/src/transformers/models/beit/image_processing_beit_fast.py b/src/transformers/models/beit/image_processing_beit_fast.py
index e10dc552cf37..4518043e6841 100644
--- a/src/transformers/models/beit/image_processing_beit_fast.py
+++ b/src/transformers/models/beit/image_processing_beit_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -38,16 +39,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class BeitFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
r"""
do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):
diff --git a/src/transformers/models/big_bird/modeling_big_bird.py b/src/transformers/models/big_bird/modeling_big_bird.py
index f42b1eeaeeb1..eb89d9872be8 100755
--- a/src/transformers/models/big_bird/modeling_big_bird.py
+++ b/src/transformers/models/big_bird/modeling_big_bird.py
@@ -1272,14 +1272,14 @@ def _get_single_block_row_attention(
if block_id == to_end_block_id - 2:
illegal_blocks.append(1)
- selected_random_blokcs = []
+ selected_random_blocks = []
for i in range(to_end_block_id - to_start_block_id):
if perm_block[i] not in illegal_blocks:
- selected_random_blokcs.append(perm_block[i])
- if len(selected_random_blokcs) == num_rand_blocks:
+ selected_random_blocks.append(perm_block[i])
+ if len(selected_random_blocks) == num_rand_blocks:
break
- return np.array(selected_random_blokcs, dtype=np.int32)
+ return np.array(selected_random_blocks, dtype=np.int32)
# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->BigBird
@@ -2877,7 +2877,6 @@ def forward(
logits_mask = self.prepare_question_mask(question_lengths, seqlen)
if token_type_ids is None:
token_type_ids = torch.ones(logits_mask.size(), dtype=int, device=logits_mask.device) - logits_mask
- logits_mask = logits_mask
logits_mask[:, 0] = False
logits_mask.unsqueeze_(2)
diff --git a/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py b/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py
index 90f3c886ad93..e419af75da38 100755
--- a/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py
+++ b/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py
@@ -1088,14 +1088,14 @@ def _get_single_block_row_attention(
if block_id == to_end_block_id - 2:
illegal_blocks.append(1)
- selected_random_blokcs = []
+ selected_random_blocks = []
for i in range(to_end_block_id - to_start_block_id):
if perm_block[i] not in illegal_blocks:
- selected_random_blokcs.append(perm_block[i])
- if len(selected_random_blokcs) == num_rand_blocks:
+ selected_random_blocks.append(perm_block[i])
+ if len(selected_random_blocks) == num_rand_blocks:
break
- return np.array(selected_random_blokcs, dtype=np.int32)
+ return np.array(selected_random_blocks, dtype=np.int32)
class BigBirdPegasusEncoderAttention(nn.Module):
diff --git a/src/transformers/models/biogpt/modeling_biogpt.py b/src/transformers/models/biogpt/modeling_biogpt.py
index 8690082625a7..7b9937420025 100755
--- a/src/transformers/models/biogpt/modeling_biogpt.py
+++ b/src/transformers/models/biogpt/modeling_biogpt.py
@@ -871,6 +871,7 @@ def forward(
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.Tensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
) -> Union[tuple, SequenceClassifierOutputWithPast]:
r"""
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -894,7 +895,8 @@ def forward(
cache_position=cache_position,
)
hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.score(hidden_states[:, slice_indices, :])
if input_ids is not None:
batch_size, sequence_length = input_ids.shape[:2]
diff --git a/src/transformers/models/biogpt/modular_biogpt.py b/src/transformers/models/biogpt/modular_biogpt.py
index 001c1de65756..8d95b2a2d051 100644
--- a/src/transformers/models/biogpt/modular_biogpt.py
+++ b/src/transformers/models/biogpt/modular_biogpt.py
@@ -693,6 +693,7 @@ def forward(
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.Tensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
) -> Union[tuple, SequenceClassifierOutputWithPast]:
r"""
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -716,7 +717,8 @@ def forward(
cache_position=cache_position,
)
hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.score(hidden_states[:, slice_indices, :])
if input_ids is not None:
batch_size, sequence_length = input_ids.shape[:2]
diff --git a/src/transformers/models/blt/__init__.py b/src/transformers/models/blt/__init__.py
new file mode 100644
index 000000000000..703b81ecdd09
--- /dev/null
+++ b/src/transformers/models/blt/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_blt import *
+ from .modeling_blt import *
+ from .tokenization_blt import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/blt/configuration_blt.py b/src/transformers/models/blt/configuration_blt.py
new file mode 100644
index 000000000000..0bc6718e5bd1
--- /dev/null
+++ b/src/transformers/models/blt/configuration_blt.py
@@ -0,0 +1,423 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Blt model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class BltLocalEncoderConfig(PretrainedConfig):
+ """
+ Configuration class for the Blt Local Encoder component.
+ """
+
+ model_type = "blt_local_encoder"
+
+ def __init__(
+ self,
+ vocab_size=260,
+ cross_attn_all_layers=False,
+ cross_attn_k=2,
+ hidden_size_global=2048,
+ hidden_size=1024,
+ num_attention_heads=16,
+ num_key_value_heads=None,
+ num_hidden_layers=1,
+ rms_norm_eps=1e-5,
+ dropout=0.0,
+ max_position_embeddings=24576,
+ rope_theta=500000.0,
+ rope_scaling=None,
+ hidden_act="silu",
+ intermediate_size=2816,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.cross_attn_all_layers = cross_attn_all_layers
+ self.cross_attn_k = cross_attn_k
+ self.hidden_size_global = hidden_size_global
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads or num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.intermediate_size = intermediate_size or int(8 * hidden_size / 3)
+ self.num_hidden_layers = num_hidden_layers
+ self.rms_norm_eps = rms_norm_eps
+ self.dropout = dropout
+ self.max_position_embeddings = max_position_embeddings
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+
+ # Remove tie_word_embeddings from kwargs to avoid duplicate parameter error
+ kwargs.pop("tie_word_embeddings", None)
+ super().__init__(**kwargs, tie_word_embeddings=False)
+
+
+class BltLocalDecoderConfig(PretrainedConfig):
+ """
+ Configuration class for the Blt Local Decoder component.
+ """
+
+ model_type = "blt_local_decoder"
+
+ def __init__(
+ self,
+ vocab_size=260,
+ cross_attn_all_layers=True,
+ cross_attn_k=2,
+ hidden_size_global=2048,
+ hidden_size=1024,
+ num_attention_heads=16,
+ num_key_value_heads=None,
+ num_hidden_layers=9,
+ rms_norm_eps=1e-5,
+ dropout=0.0,
+ max_position_embeddings=24576,
+ rope_theta=500000.0,
+ rope_scaling=None,
+ hidden_act="silu",
+ intermediate_size=2816,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.cross_attn_all_layers = cross_attn_all_layers
+ self.cross_attn_k = cross_attn_k
+ self.hidden_size_global = hidden_size_global
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads or num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.intermediate_size = intermediate_size or int(8 * hidden_size / 3)
+ self.num_hidden_layers = num_hidden_layers
+ self.rms_norm_eps = rms_norm_eps
+ self.dropout = dropout
+ self.max_position_embeddings = max_position_embeddings
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+
+ # Remove tie_word_embeddings from kwargs to avoid duplicate parameter error
+ kwargs.pop("tie_word_embeddings", None)
+ super().__init__(**kwargs, tie_word_embeddings=False)
+
+
+class BltGlobalTransformerConfig(PretrainedConfig):
+ """
+ Configuration class for the Blt Global Transformer component.
+ """
+
+ model_type = "blt_global_transformer"
+
+ def __init__(
+ self,
+ hidden_size=2048,
+ num_attention_heads=16,
+ num_key_value_heads=None,
+ num_hidden_layers=25,
+ rms_norm_eps=1e-5,
+ dropout=0.0,
+ max_position_embeddings=4096,
+ rope_theta=500000.0,
+ rope_scaling=None,
+ hidden_act="silu",
+ intermediate_size=5632,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads or num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.intermediate_size = intermediate_size or int(8 * hidden_size / 3)
+ self.num_hidden_layers = num_hidden_layers
+ self.rms_norm_eps = rms_norm_eps
+ self.dropout = dropout
+ self.max_position_embeddings = max_position_embeddings
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+
+ # Remove tie_word_embeddings from kwargs to avoid duplicate parameter error
+ kwargs.pop("tie_word_embeddings", None)
+ super().__init__(**kwargs, tie_word_embeddings=False)
+
+
+class BltPatcherConfig(PretrainedConfig):
+ r"""
+ Configuration class for the Blt Patcher/Entropy model component.
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 260):
+ Vocabulary size of the Blt patcher model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling the patcher model.
+ hidden_size (`int`, *optional*, defaults to 768):
+ Dimension of the hidden representations.
+ num_hidden_layers (`int`, *optional*, defaults to 14):
+ Number of hidden layers in the Transformer decoder.
+ num_attention_heads (`int`, *optional*, defaults to 12):
+ Number of attention heads for each attention layer in the Transformer decoder.
+ num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+ `num_attention_heads`.
+ max_position_embeddings (`int`, *optional*, defaults to 8192):
+ The maximum sequence length that this model might ever be used with.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+ The epsilon used by the rms normalization layers.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ intermediate_size (`int`, *optional*, defaults to 2048):
+ Dimension of the MLP representations.
+ rope_scaling (`dict`, *optional*):
+ Dictionary containing the RoPE scaling configuration.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ """
+
+ model_type = "blt_patcher"
+
+ def __init__(
+ self,
+ vocab_size=260,
+ hidden_size=768,
+ num_hidden_layers=14,
+ num_attention_heads=12,
+ num_key_value_heads=None,
+ max_position_embeddings=8192,
+ rms_norm_eps=1e-5,
+ dropout=0.0,
+ rope_theta=10000.0,
+ intermediate_size=2048,
+ rope_scaling=None,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+ self.max_position_embeddings = max_position_embeddings
+ self.rms_norm_eps = rms_norm_eps
+ self.dropout = dropout
+ self.rope_theta = rope_theta
+ self.hidden_act = "silu" # Blt uses silu activation
+ self.intermediate_size = intermediate_size or int(8 * self.hidden_size / 3)
+ self.rope_scaling = rope_scaling
+ self.initializer_range = initializer_range
+
+ # Remove tie_word_embeddings from kwargs to avoid duplicate parameter error
+ kwargs.pop("tie_word_embeddings", None)
+ super().__init__(**kwargs, tie_word_embeddings=False)
+
+
+class BltConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`BltModel`]. It is used to instantiate a
+ Blt model according to the specified arguments, defining the model architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 260):
+ Vocabulary size of the Blt model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`BltModel`].
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
+ The maximum sequence length that this model might ever be used with.
+ patch_in_forward (`bool`, *optional*, defaults to `True`):
+ Whether to perform patching during the forward pass.
+ patch_size (`int`, *optional*, defaults to 4):
+ Size of the patches used in the patching mechanism.
+ patching_mode (`str`, *optional*, defaults to `"entropy"`):
+ The mode used for patching, such as entropy-based patching.
+ patching_threshold (`float`, *optional*, defaults to 1.34):
+ Threshold value used for determining when to apply patches.
+ patching_batch_size (`int`, *optional*, defaults to 1):
+ Batch size used during the patching process.
+ max_patch_length (`int`, *optional*):
+ Maximum length of patches that can be generated.
+ cross_attn_k (`int`, *optional*, defaults to 2):
+ Number of cross-attention heads used in the model.
+ encoder_hash_byte_group_size (`list`, *optional*):
+ List of byte group sizes used in the encoder hash function.
+ encoder_hash_byte_group_vocab (`int`, *optional*, defaults to 500002):
+ Vocabulary size for the encoder hash byte groups.
+ encoder_hash_byte_group_nb_functions (`int`, *optional*, defaults to 1):
+ Number of hash functions used in the encoder byte grouping.
+ patcher_config (`BltPatcherConfig`, *optional*):
+ Configuration for the patcher component of the model.
+ encoder_config (`BltLocalEncoderConfig`, *optional*):
+ Configuration for the local encoder component of the model.
+ decoder_config (`BltLocalDecoderConfig`, *optional*):
+ Configuration for the local decoder component of the model.
+ global_config (`BltGlobalTransformerConfig`, *optional*):
+ Configuration for the global transformer component of the model.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether to tie weight embeddings.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rope_theta (`float`, *optional*, defaults to 500000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`dict`, *optional*):
+ Dictionary containing the RoPE scaling configuration.
+
+ ```python
+ >>> from transformers import BltModel, BltConfig
+
+ >>> # Initializing a Blt configuration
+ >>> configuration = BltConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = BltModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+
+ Checkpoint: [facebook/blt](https://huggingface.co/facebook/blt)
+ """
+
+ model_type = "blt"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ sub_configs = {
+ "patcher_config": BltPatcherConfig,
+ "encoder_config": BltLocalEncoderConfig,
+ "decoder_config": BltLocalDecoderConfig,
+ "global_config": BltGlobalTransformerConfig,
+ }
+
+ def __init__(
+ self,
+ vocab_size=260,
+ max_position_embeddings=4096,
+ patch_in_forward=True,
+ patch_size=4,
+ patching_mode="entropy",
+ patching_threshold=1.335442066192627,
+ patching_batch_size=1,
+ max_patch_length=None,
+ cross_attn_k=2,
+ encoder_hash_byte_group_size=None,
+ encoder_hash_byte_group_vocab=500002,
+ encoder_hash_byte_group_nb_functions=1,
+ patcher_config=None,
+ encoder_config=None,
+ decoder_config=None,
+ global_config=None,
+ tie_word_embeddings=False,
+ initializer_range=0.02,
+ rope_theta=500000.0,
+ rope_scaling=None,
+ **kwargs,
+ ):
+ # Basic model configuration
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+
+ # Patching configuration
+ self.patch_in_forward = patch_in_forward
+ self.patch_size = patch_size
+ self.patching_mode = patching_mode
+ self.patching_threshold = patching_threshold
+ self.patching_batch_size = patching_batch_size
+ self.max_patch_length = max_patch_length
+ self.patching_device = kwargs.get("patching_device", "cuda")
+ self.realtime_patching = kwargs.get("realtime_patching", True)
+ self.patching_threshold_add = kwargs.get("patching_threshold_add")
+ self.monotonicity = kwargs.get("monotonicity", False)
+
+ # Cross attention configurations
+ self.cross_attn_k = cross_attn_k
+
+ # Encoder configurations
+ self.encoder_hash_byte_group_size = encoder_hash_byte_group_size or [3, 4, 5, 6, 7, 8]
+ self.encoder_hash_byte_group_vocab = encoder_hash_byte_group_vocab
+ self.encoder_hash_byte_group_nb_functions = encoder_hash_byte_group_nb_functions
+
+ # Initialize component configurations
+ if patcher_config is None:
+ self.patcher_config = BltPatcherConfig(initializer_range=initializer_range)
+ logger.info("patcher_config is None, using default Blt patcher config")
+ elif isinstance(patcher_config, dict):
+ patcher_config.setdefault("initializer_range", initializer_range)
+ self.patcher_config = BltPatcherConfig(**patcher_config)
+ elif isinstance(patcher_config, BltPatcherConfig):
+ self.patcher_config = patcher_config
+
+ if encoder_config is None:
+ self.encoder_config = BltLocalEncoderConfig(initializer_range=initializer_range)
+ logger.info("encoder_config is None, using default Blt encoder config")
+ elif isinstance(encoder_config, dict):
+ encoder_config.setdefault("initializer_range", initializer_range)
+ self.encoder_config = BltLocalEncoderConfig(**encoder_config)
+ elif isinstance(encoder_config, BltLocalEncoderConfig):
+ self.encoder_config = encoder_config
+
+ if decoder_config is None:
+ self.decoder_config = BltLocalDecoderConfig(initializer_range=initializer_range)
+ logger.info("decoder_config is None, using default Blt decoder config")
+ elif isinstance(decoder_config, dict):
+ decoder_config.setdefault("initializer_range", initializer_range)
+ self.decoder_config = BltLocalDecoderConfig(**decoder_config)
+ elif isinstance(decoder_config, BltLocalDecoderConfig):
+ self.decoder_config = decoder_config
+
+ if global_config is None:
+ self.global_config = BltGlobalTransformerConfig(initializer_range=initializer_range)
+ logger.info("global_config is None, using default Blt global config")
+ elif isinstance(global_config, dict):
+ global_config.setdefault("initializer_range", initializer_range)
+ self.global_config = BltGlobalTransformerConfig(**global_config)
+ elif isinstance(global_config, BltGlobalTransformerConfig):
+ self.global_config = global_config
+
+ # Determine if token embedding projection is needed based on dimension mismatch (7b)
+ encoder_cross_output_size = self.encoder_config.hidden_size * self.cross_attn_k
+ self.global_config.encoder_cross_output_size = (
+ encoder_cross_output_size if encoder_cross_output_size != self.global_config.hidden_size else None
+ )
+
+ # Remove tie_word_embeddings from kwargs to avoid duplicate parameter error
+ kwargs.pop("tie_word_embeddings", None)
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = [
+ "BltConfig",
+ "BltPatcherConfig",
+ "BltLocalEncoderConfig",
+ "BltLocalDecoderConfig",
+ "BltGlobalTransformerConfig",
+]
diff --git a/src/transformers/models/blt/convert_blt_weights_to_hf.py b/src/transformers/models/blt/convert_blt_weights_to_hf.py
new file mode 100644
index 000000000000..f9decff3a1f8
--- /dev/null
+++ b/src/transformers/models/blt/convert_blt_weights_to_hf.py
@@ -0,0 +1,487 @@
+import argparse
+import json
+import logging
+import os
+from typing import Any, Optional
+
+import torch
+from huggingface_hub import hf_hub_download, upload_folder
+from safetensors.torch import load_file, save_file
+from tokenizers import Tokenizer, decoders, pre_tokenizers, processors
+from tokenizers.models import BPE
+
+from transformers import PreTrainedTokenizerFast
+from transformers.convert_slow_tokenizer import bytes_to_unicode
+from transformers.utils import logging as transformers_logging
+
+
+logger = transformers_logging.get_logger(__name__)
+transformers_logging.set_verbosity_info()
+
+
+def merge_configurations(config_path: str, entropy_params_path: str) -> dict[str, Any]:
+ logger.info("Merging configurations")
+
+ with open(config_path, "r") as f:
+ main_config = json.load(f)
+
+ with open(entropy_params_path, "r") as f:
+ entropy_data = json.load(f)
+
+ entropy_model_params = entropy_data.get("entropy_model", {})
+ patcher_args = entropy_data.get("data", {}).get("patcher_args", {})
+
+ unified_config = main_config.copy()["args"]
+
+ for key in ["vocab_size", "dim", "n_layers", "n_heads", "max_seqlen"]:
+ if key in unified_config and not isinstance(unified_config[key], int):
+ unified_config[key] = int(unified_config[key])
+
+ patch_size = patcher_args.get("patch_size", 8)
+ if isinstance(patch_size, float):
+ patch_size = int(patch_size)
+
+ # Create patcher config
+ patcher_hidden_size = int(entropy_model_params.get("dim", 512))
+ patcher_multiple_of = int(entropy_model_params.get("multiple_of", 256))
+ patcher_intermediate_size = patcher_multiple_of * (
+ (int(8 * patcher_hidden_size / 3) + patcher_multiple_of - 1) // patcher_multiple_of
+ )
+
+ patcher_config = {
+ "vocab_size": int(entropy_model_params.get("vocab_size", 256)),
+ "hidden_size": patcher_hidden_size,
+ "num_hidden_layers": int(entropy_model_params.get("n_layers", 8)),
+ "num_attention_heads": int(entropy_model_params.get("n_heads", 8)),
+ "num_key_value_heads": int(entropy_model_params.get("n_kv_heads"))
+ if entropy_model_params.get("n_kv_heads") is not None
+ else None,
+ "max_position_embeddings": int(entropy_model_params.get("max_seqlen", 1024)),
+ "norm_eps": entropy_model_params.get("norm_eps", 1e-5),
+ "dropout": entropy_model_params.get("dropout", 0.0),
+ "rope_theta": entropy_model_params.get("rope_theta", 10000.0),
+ "attn_impl": entropy_model_params.get("attn_impl", "sdpa"),
+ "attn_bias_type": entropy_model_params.get("attn_bias_type", "causal"),
+ "intermediate_size": patcher_intermediate_size,
+ }
+
+ # Create encoder config
+ encoder_hidden_size = unified_config.get("dim_local_encoder", 1024)
+ encoder_multiple_of = unified_config.get("multiple_of", 256)
+ encoder_intermediate_size = encoder_multiple_of * (
+ (int(8 * encoder_hidden_size / 3) + encoder_multiple_of - 1) // encoder_multiple_of
+ )
+
+ encoder_config = {
+ "vocab_size": unified_config.get("vocab_size", 256),
+ "cross_attn_all_layers": unified_config.get("cross_attn_all_layers_encoder", False),
+ "cross_attn_k": unified_config.get("cross_attn_k", 2),
+ "hidden_size_global": unified_config.get("dim_global", 2048),
+ "pm_size": unified_config.get("pm_size", 0),
+ "hidden_size": encoder_hidden_size,
+ "num_attention_heads": unified_config.get("n_heads_local_encoder", 16),
+ "num_key_value_heads": unified_config.get("n_kv_heads"),
+ "num_hidden_layers": unified_config.get("n_layers_local_encoder", 1),
+ "norm_eps": unified_config.get("norm_eps", 1e-5),
+ "dropout": unified_config.get("dropout", 0.0),
+ "max_position_embeddings": unified_config.get("max_encoder_seq_length")
+ or unified_config.get("max_seqlen", 1024),
+ "rope_theta": unified_config.get("rope_theta", 10000.0),
+ "rope_scaling": {"rope_type": "default"},
+ "hidden_act": unified_config.get("hidden_act", "silu"),
+ "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"),
+ "intermediate_size": encoder_intermediate_size,
+ }
+
+ # Create decoder config
+ decoder_hidden_size = unified_config.get("dim_local_decoder", 1024)
+ decoder_multiple_of = unified_config.get("multiple_of", 256)
+ decoder_intermediate_size = decoder_multiple_of * (
+ (int(8 * decoder_hidden_size / 3) + decoder_multiple_of - 1) // decoder_multiple_of
+ )
+
+ decoder_config = {
+ "vocab_size": unified_config.get("vocab_size", 256),
+ "cross_attn_all_layers": unified_config.get("cross_attn_all_layers_decoder", False),
+ "cross_attn_k": unified_config.get("cross_attn_k", 2),
+ "hidden_size_global": unified_config.get("dim_global", 2048),
+ "hidden_size": decoder_hidden_size,
+ "num_attention_heads": unified_config.get("n_heads_local_decoder", 16),
+ "num_key_value_heads": unified_config.get("n_kv_heads"),
+ "num_hidden_layers": unified_config.get("n_layers_local_decoder", 9),
+ "norm_eps": unified_config.get("norm_eps", 1e-5),
+ "dropout": unified_config.get("dropout", 0.0),
+ "max_position_embeddings": unified_config.get("max_encoder_seq_length")
+ or unified_config.get("max_seqlen", 1024),
+ "rope_theta": unified_config.get("rope_theta", 10000.0),
+ "rope_scaling": {"rope_type": "default"},
+ "hidden_act": unified_config.get("hidden_act", "silu"),
+ "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"),
+ "intermediate_size": decoder_intermediate_size,
+ }
+
+ # Create global transformer config
+ global_hidden_size = unified_config.get("dim_global", 2048)
+ global_multiple_of = unified_config.get("multiple_of", 256)
+ global_intermediate_size = global_multiple_of * (
+ (int(8 * global_hidden_size / 3) + global_multiple_of - 1) // global_multiple_of
+ )
+
+ global_config = {
+ "hidden_size": global_hidden_size,
+ "num_attention_heads": unified_config.get("n_heads_global", 16),
+ "num_key_value_heads": unified_config.get("n_kv_heads_global"),
+ "num_hidden_layers": unified_config.get("n_layers_global", 25),
+ "norm_eps": unified_config.get("norm_eps", 1e-5),
+ "dropout": unified_config.get("dropout", 0.0),
+ "max_position_embeddings": unified_config.get("max_seqlen", 1024),
+ "rope_theta": unified_config.get("rope_theta", 10000.0),
+ "rope_scaling": {"rope_type": "default"},
+ "hidden_act": unified_config.get("hidden_act", "silu"),
+ "_attn_implementation": unified_config.get("_attn_implementation", "sdpa"),
+ "intermediate_size": global_intermediate_size,
+ }
+
+ # Create main config with sub-configs
+ main_config_dict = {
+ "model_type": "blt",
+ "vocab_size": unified_config.get("vocab_size", 256),
+ "max_position_embeddings": unified_config.get("max_seqlen", 1024),
+ "patch_in_forward": True,
+ "realtime_patching": True,
+ "patching_mode": "entropy",
+ "patch_size": patch_size,
+ "patching_threshold": patcher_args.get("threshold", 0.5),
+ "patching_threshold_add": patcher_args.get("threshold_add", 0.0),
+ "max_patch_length": patcher_args.get("max_patch_length"),
+ "patching_batch_size": patcher_args.get("patching_batch_size", 1),
+ "patching_device": patcher_args.get("patching_device", "cuda"),
+ "monotonicity": patcher_args.get("monotonicity", False),
+ "cross_attn_k": unified_config.get("cross_attn_k", 2),
+ "encoder_hash_byte_group_size": unified_config.get("encoder_hash_byte_group_size"),
+ "encoder_hash_byte_group_vocab": unified_config.get("encoder_hash_byte_group_vocab", 30000),
+ "encoder_hash_byte_group_nb_functions": unified_config.get("encoder_hash_byte_group_nb_functions", 3),
+ "pm_size": unified_config.get("pm_size", 0),
+ "patcher_config": patcher_config,
+ "encoder_config": encoder_config,
+ "decoder_config": decoder_config,
+ "global_config": global_config,
+ }
+
+ main_config_dict["tie_word_embeddings"] = False
+
+ logger.info(f"Merged configuration with {len(main_config_dict)} parameters")
+ return main_config_dict
+
+
+def apply_weight_mapping(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ component_mappings = {
+ ".attention.": ".self_attn.",
+ ".feed_forward.": ".mlp.",
+ ".attention_norm.": ".input_layernorm.",
+ ".ffn_norm.": ".post_attention_layernorm.",
+ ".tok_embeddings.": ".embed_tokens.",
+ ".cross_attn_norm_q.": ".q_norm.",
+ ".cross_attn_norm_kv.": ".k_norm.",
+ ".w1.": ".gate_proj.",
+ ".w2.": ".down_proj.",
+ ".w3.": ".up_proj.",
+ ".wq.": ".q_proj.",
+ ".wk.": ".k_proj.",
+ ".wv.": ".v_proj.",
+ ".wo.": ".o_proj.",
+ ".output.": ".lm_head.",
+ }
+
+ new_state_dict = {}
+
+ for old_key, tensor in state_dict.items():
+ new_key = old_key
+
+ for old_pattern, new_pattern in component_mappings.items():
+ if old_pattern in new_key:
+ new_key = new_key.replace(old_pattern, new_pattern)
+
+ new_state_dict[new_key] = tensor
+
+ return new_state_dict
+
+
+def convert_hash_embeddings_to_fused(
+ unified_weights: dict[str, torch.Tensor], config: dict[str, Any]
+) -> dict[str, torch.Tensor]:
+ """Convert ModuleList hash embeddings to nn.embedding format"""
+ original_keys_format = [
+ key
+ for key in unified_weights.keys()
+ if "encoder_hash_tok_embedding." in key and ".weight" in key and key.split(".")[-2].isdigit()
+ ]
+
+ num_embeddings = config.get("encoder_hash_byte_group_nb_functions", 1) * len(
+ config.get("encoder_hash_byte_group_size", [3, 4, 5, 6, 7, 8])
+ )
+ vocab_size = config.get("encoder_hash_byte_group_vocab", 500002)
+ hidden_size = config.get("encoder_config", {}).get("hidden_size", 1024)
+
+ fused_weight = torch.zeros(vocab_size * num_embeddings, hidden_size)
+
+ sorted_keys = sorted(original_keys_format, key=lambda k: int(k.split(".")[-2]))
+
+ for i, old_key in enumerate(sorted_keys):
+ start_idx = i * vocab_size
+ end_idx = (i + 1) * vocab_size
+ fused_weight[start_idx:end_idx] = unified_weights[old_key]
+ logger.info(f"Copied {old_key} to indices {start_idx}:{end_idx}")
+ del unified_weights[old_key]
+
+ fused_key = "model.encoder_hash_tok_embedding.weight"
+ unified_weights[fused_key] = fused_weight
+
+ return unified_weights
+
+
+def merge_weights(weights_path: str, entropy_weights_path: str) -> dict[str, torch.Tensor]:
+ main_weights = load_file(weights_path)
+
+ entropy_weights = torch.load(entropy_weights_path, map_location="cpu", weights_only=True)
+
+ if "model" in entropy_weights:
+ entropy_weights = entropy_weights["model"]
+ elif "state_dict" in entropy_weights:
+ entropy_weights = entropy_weights["state_dict"]
+
+ unified_weights = main_weights.copy()
+
+ for key, tensor in entropy_weights.items():
+ patcher_key = f"patcher.{key}"
+ unified_weights[patcher_key] = tensor
+
+ unified_weights = apply_weight_mapping(unified_weights)
+
+ decoder_lm_head_key = "local_decoder.lm_head.weight"
+ top_lm_head_key = "lm_head.weight"
+ unified_weights[top_lm_head_key] = unified_weights[decoder_lm_head_key]
+ del unified_weights[decoder_lm_head_key]
+
+ prefixed_weights = {}
+ for key, tensor in unified_weights.items():
+ if key == top_lm_head_key:
+ prefixed_weights[key] = tensor
+ elif not key.startswith("model."):
+ prefixed_weights[f"model.{key}"] = tensor
+ else:
+ prefixed_weights[key] = tensor
+
+ unified_weights = prefixed_weights
+
+ return unified_weights
+
+
+def create_tokenizer_config(output_dir: str, config: dict[str, Any]):
+ tokenizer_config = {
+ "tokenizer_class": "PreTrainedTokenizerFast",
+ "vocab_size": config.get("vocab_size", 256),
+ "model_max_length": config.get("max_seqlen", 1024),
+ "model_input_names": ["input_ids", "attention_mask"],
+ "add_bos_token": True,
+ "add_eos_token": True,
+ "bos_token": "",
+ "eos_token": "",
+ "pad_token": "",
+ "unk_token": "",
+ }
+
+ tokenizer_path = os.path.join(output_dir, "tokenizer_config.json")
+ with open(tokenizer_path, "w") as f:
+ json.dump(tokenizer_config, f, indent=2)
+
+
+def create_tokenizer_json(output_dir: str, config: dict[str, Any]):
+ byte_encoder = bytes_to_unicode()
+
+ vocab: dict[str, int] = {}
+ vocab[""] = 0
+ vocab[""] = 1
+ vocab[""] = 2
+ vocab[""] = 3
+
+ offset = 4
+ for byte_val, unicode_char in byte_encoder.items():
+ vocab[unicode_char] = byte_val + offset
+
+ backend = Tokenizer(
+ BPE(vocab=vocab, merges=[], continuing_subword_prefix="", end_of_word_suffix="", fuse_unk=False)
+ )
+ backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+ backend.decoder = decoders.ByteLevel()
+
+ bos = config.get("bos_token", "")
+ backend.post_processor = processors.TemplateProcessing(
+ single=f"{bos}:0 $A:0",
+ pair=f"{bos}:0 $A:0 $B:1",
+ special_tokens=[(bos, 1)],
+ )
+
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=backend,
+ bos_token=config.get("bos_token", ""),
+ eos_token=config.get("eos_token", ""),
+ pad_token=config.get("pad_token", ""),
+ unk_token=config.get("unk_token", ""),
+ )
+
+ tokenizer.add_bos_token = bool(config.get("add_bos_token", True))
+ tokenizer.add_eos_token = bool(config.get("add_eos_token", True))
+
+ tokenizer.save_pretrained(output_dir)
+ logger.info(f"Saved tokenizer.json to {os.path.join(output_dir, 'tokenizer.json')}")
+
+
+def push_to_hub(
+ local_dir: str,
+ repo_id: str,
+ commit_message: str = "Upload converted Blt model",
+ private: bool = False,
+ token: Optional[str] = None,
+) -> None:
+ try:
+ upload_folder(
+ folder_path=local_dir,
+ repo_id=repo_id,
+ commit_message=commit_message,
+ repo_type="model",
+ token=token,
+ )
+ logger.info(f"Successfully pushed model to {repo_id}")
+
+ except Exception as e:
+ logger.error(f"Failed to push model to Hub: {e}")
+ raise
+
+
+def convert_hf_blt_to_unified(
+ model_id: str,
+ output_dir: str,
+ config_name: str = "config.json",
+ weights_name: str = "model.bin",
+ cache_dir: Optional[str] = None,
+ push_to_hub_repo: Optional[str] = None,
+ hub_private: bool = False,
+ hub_token: Optional[str] = None,
+) -> None:
+ # Download model files
+ config_path = hf_hub_download(repo_id=model_id, filename="config.json", cache_dir=cache_dir)
+ weights_path = hf_hub_download(repo_id=model_id, filename="model.safetensors", cache_dir=cache_dir)
+ entropy_params_path = hf_hub_download(repo_id=model_id, filename="entropy_model/params.json", cache_dir=cache_dir)
+ entropy_weights_path = hf_hub_download(
+ repo_id=model_id, filename="entropy_model/consolidated.pth", cache_dir=cache_dir
+ )
+
+ unified_config = merge_configurations(config_path, entropy_params_path)
+ unified_weights = merge_weights(weights_path, entropy_weights_path)
+
+ unified_weights = convert_hash_embeddings_to_fused(unified_weights, unified_config)
+
+ os.makedirs(output_dir, exist_ok=True)
+
+ config_path = os.path.join(output_dir, config_name)
+ with open(config_path, "w") as f:
+ json.dump(unified_config, f, indent=2)
+
+ if weights_name.endswith(".bin"):
+ weights_name = weights_name.replace(".bin", ".safetensors")
+
+ weights_path = os.path.join(output_dir, weights_name)
+ save_file(unified_weights, weights_path)
+
+ create_tokenizer_json(output_dir=output_dir, config=unified_config)
+
+ create_tokenizer_config(output_dir, unified_config)
+
+ logger.info(f"Conversion completed, model saved to: {output_dir}")
+
+ if push_to_hub_repo:
+ push_to_hub(
+ local_dir=output_dir,
+ repo_id=push_to_hub_repo,
+ commit_message="Upload Blt model converted",
+ private=hub_private,
+ token=hub_token,
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Convert Blt models from HuggingFace Hub format to unified format",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--model_id",
+ type=str,
+ default="facebook/blt-7b",
+ )
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ default="./blt_converted",
+ )
+ parser.add_argument(
+ "--config_name",
+ type=str,
+ default="config.json",
+ )
+ parser.add_argument(
+ "--weights_name",
+ type=str,
+ default="model.bin",
+ )
+ parser.add_argument(
+ "--cache_dir",
+ type=str,
+ default=None,
+ )
+ parser.add_argument(
+ "--debug",
+ action="store_true",
+ default=True,
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ type=str,
+ default=None,
+ )
+ parser.add_argument(
+ "--hub_private",
+ action="store_true",
+ default=False,
+ )
+ parser.add_argument(
+ "--hub_token",
+ type=str,
+ default="hf_token",
+ )
+
+ args = parser.parse_args()
+
+ transformers_logging.set_verbosity_debug()
+ logging.basicConfig(level=logging.DEBUG)
+
+ try:
+ convert_hf_blt_to_unified(
+ model_id=args.model_id,
+ output_dir=args.output_dir,
+ config_name=args.config_name,
+ weights_name=args.weights_name,
+ cache_dir=args.cache_dir,
+ push_to_hub_repo=False, # args.push_to_hub,
+ hub_private=args.hub_private,
+ hub_token=args.hub_token,
+ )
+ except Exception as e:
+ logger.error(f"Conversion failed: {e}")
+ raise
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
new file mode 100644
index 000000000000..1e677dda4a98
--- /dev/null
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -0,0 +1,1311 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/blt/modular_blt.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_blt.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Optional, Union
+
+import torch
+import torch.distributions
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import OutputRecorder, check_model_inputs
+from .configuration_blt import (
+ BltConfig,
+ BltGlobalTransformerConfig,
+ BltLocalDecoderConfig,
+ BltLocalEncoderConfig,
+ BltPatcherConfig,
+)
+
+
+class BltMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ # Ignore copy
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class BltRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ BltRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class BltRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: BltConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.repeat_interleave(freqs, 2, dim=-1) # diff from Llama: we interleave() instead of cat()
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+# Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer
+class BltTransformerLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = BltSelfAttention(config=config, layer_idx=layer_idx)
+ self.mlp = BltMLP(config)
+ self.input_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ self.layer_idx = layer_idx
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cross_attention_states: Optional[torch.Tensor] = None,
+ cross_attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*):
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+ query_sequence_length, key_sequence_length)` if default attention is used.
+
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ 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:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half(x):
+ # Split and rotate. Note that this function is different from e.g. Llama.
+ x1 = x[..., ::2]
+ x2 = x[..., 1::2]
+ rot_x = torch.stack([-x2, x1], dim=-1).flatten(-2)
+ return rot_x
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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.
+ position_ids (`torch.Tensor`, *optional*):
+ Deprecated and unused.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class BltSelfAttention(nn.Module):
+ def __init__(self, config: BltConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.num_heads = config.num_attention_heads
+ self.dropout = config.dropout
+ self.hidden_size = config.hidden_size
+ self.num_key_value_heads = config.num_key_value_heads
+ self.head_dim = config.hidden_size // self.num_heads
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.rope_theta = config.rope_theta
+ self.layer_idx = layer_idx
+
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+ self.is_causal = True
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ position_embeddings: torch.Tensor,
+ use_cache: bool = False,
+ past_key_values=None,
+ cache_position=None,
+ **kwargs,
+ ):
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class BltCrossAttention(nn.Module):
+ """Cross-attention module for Blt, following transformers style"""
+
+ def __init__(self, config: BltConfig, layer_idx: int, hidden_size: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.num_heads = self.config.num_attention_heads
+ self.num_key_value_heads = self.config.num_key_value_heads
+ self.dropout = config.dropout
+ self.hidden_size = config.hidden_size
+ self.head_dim = config.hidden_size // self.num_heads
+ self.layer_idx = layer_idx
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+ self.q_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+ self.k_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+ self.is_causal = False
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cross_attention_states: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Cache] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+ """Input shape: Batch x Time x Channel"""
+ bsz, q_len, _ = hidden_states.size()
+ query_states = self.q_norm(hidden_states)
+ query_states = self.q_proj(query_states)
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+ if cross_attention_states is not None:
+ cross_attention_states = self.k_norm(cross_attention_states)
+ key_states = self.k_proj(cross_attention_states)
+ value_states = self.v_proj(cross_attention_states)
+ key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(
+ key_states, value_states, self.layer_idx, {"cache_position": cache_position}
+ )
+ elif cache_position[0] != 0:
+ key_states, value_states = (
+ past_key_values.layers[self.layer_idx].keys,
+ past_key_values.layers[self.layer_idx].values,
+ )
+ else:
+ raise ValueError(
+ "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!"
+ )
+ attention_interface: Callable = eager_attention_forward
+
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ attn_output = attn_output + hidden_states
+ return attn_output, attn_weights
+
+
+@auto_docstring
+class BltPreTrainedModel(PreTrainedModel):
+ config: BltConfig
+ base_model_prefix = ""
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["BltTransformerLayer"]
+ _can_compile_fullgraph = False # static cache cannot have different shapes for each layer
+ _supports_sdpa = True
+ _supports_flash_attn = False
+ _supports_flex_attn = False
+ _supports_attention_backend = False
+ _can_record_outputs = {
+ "hidden_states": OutputRecorder(BltTransformerLayer, index=0, layer_name="local_decoder"),
+ "attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_decoder"),
+ }
+
+
+class BltLocalEncoder(BltPreTrainedModel):
+ config: BltLocalEncoderConfig
+ _can_record_outputs = {
+ "encoder_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_encoder"),
+ }
+
+ def __init__(self, config: BltLocalEncoderConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+ self.config = config
+ self.layers = nn.ModuleList(
+ [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+ self.patch_embedding_projection = nn.Linear(
+ in_features=config.hidden_size,
+ out_features=config.hidden_size * config.cross_attn_k,
+ bias=False,
+ )
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.cross_attn_layers = nn.ModuleList()
+ layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1
+ for layer_idx in range(layers_to_add):
+ self.cross_attn_layers.append(
+ BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size)
+ )
+
+ self.post_init()
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ patch_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ num_patches: Optional[int] = None,
+ patch_ids: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ batch_size = inputs_embeds.shape[0]
+ hidden_states = F.dropout(inputs_embeds, p=self.config.dropout, training=self.training)
+
+ if position_ids is None:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+
+ for idx, layer in enumerate(self.layers):
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ if idx == len(self.layers) - 1 or self.config.cross_attn_all_layers:
+ patch_embeds = self.patch_reduce(hidden_states, num_patches, patch_ids)
+ patch_embeds = self.patch_embedding_projection(patch_embeds)
+ patch_embeds = patch_embeds.reshape(
+ batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size
+ )
+ layer_idx = idx if self.config.cross_attn_all_layers else 0
+ cross_attention_output, _ = self.cross_attn_layers[layer_idx](
+ hidden_states=patch_embeds,
+ cross_attention_states=hidden_states,
+ attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+ patch_embeds = patch_embeds + cross_attention_output
+ encoder_cross_states = patch_embeds
+ return hidden_states, encoder_cross_states
+
+ def patch_reduce(self, hidden_states, max_num_patches, patch_ids):
+ """
+ Reduce variable length patches to single embedding per patch
+ Note: this works with variable number of patches for different sequences in the batch
+ It handles variable length patches by assuming that patch_lengths will be 0 for any
+ extra patches on the *right*. Since there can be a variable number of patches
+ this function also return the number of patches for each sequence in the batch.
+ Any embeddings on the right that are not allocated to a patch
+ (i.e. if the sum(patch_lengths[i]) < seq_len for any i)
+ will be sent to a dummy patch, which is trimmed before returning.
+ """
+ batch_size = hidden_states.shape[0]
+ embedding_dim = hidden_states.shape[-1]
+
+ patch_ids = patch_ids.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])
+
+ reduced_embeddings = torch.zeros(
+ (batch_size, max_num_patches, embedding_dim), dtype=hidden_states.dtype, device=hidden_states.device
+ )
+
+ reduced_embeddings = reduced_embeddings.scatter_reduce(
+ src=hidden_states,
+ dim=1,
+ index=patch_ids,
+ reduce="amax",
+ include_self=False,
+ )
+ reduced_embeddings = reduced_embeddings[:, :max_num_patches, :]
+
+ return reduced_embeddings
+
+
+class BltLocalDecoder(BltPreTrainedModel):
+ config: BltLocalDecoderConfig
+
+ def __init__(self, config: BltLocalDecoderConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+ self.config = config
+ self.cross_attn_decoder = True
+ self.layers = nn.ModuleList(
+ [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+ self.patch_embedding_projection = nn.Linear(
+ in_features=config.hidden_size_global,
+ out_features=config.hidden_size * config.cross_attn_k,
+ bias=False,
+ )
+ self.norm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.cross_attn_layers = nn.ModuleList()
+ layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1
+ for layer_idx in range(layers_to_add):
+ self.cross_attn_layers.append(
+ BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size)
+ )
+
+ self.post_init()
+
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ patch_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ batch_size = inputs_embeds.shape[0]
+ hidden_states = inputs_embeds
+ patch_embeds = self.patch_embedding_projection(patch_embeds)
+ patch_embeds = patch_embeds.reshape(
+ batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size
+ )
+
+ if patch_embeds is not None and not self.cross_attn_decoder:
+ hidden_states = hidden_states + patch_embeds
+
+ if position_ids is None:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+
+ for i, layer in enumerate(self.layers):
+ if i == 0 or self.config.cross_attn_all_layers:
+ cross_attention_output, _ = self.cross_attn_layers[i](
+ hidden_states=hidden_states,
+ cross_attention_states=patch_embeds,
+ attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+ hidden_states = hidden_states + cross_attention_output
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ logits = self.norm(hidden_states)
+ return logits
+
+
+class BltGlobalTransformer(BltPreTrainedModel):
+ config: BltGlobalTransformerConfig
+ _can_record_outputs = {
+ "global_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="global_transformer"),
+ }
+
+ def __init__(self, config: BltGlobalTransformerConfig):
+ super().__init__(config)
+ self.config = config
+ self.layers = nn.ModuleList()
+ for layer_idx in range(config.num_hidden_layers):
+ self.layers.append(BltTransformerLayer(config, layer_idx))
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+
+ # Create token embedding projection (use nn.Identity() when no projection needed)
+ if getattr(config, "encoder_cross_output_size", None) is not None:
+ self.token_embedding_projection = nn.Linear(
+ config.encoder_cross_output_size, config.hidden_size, bias=False
+ )
+ else:
+ self.token_embedding_projection = nn.Identity()
+
+ self.post_init()
+
+ def forward(
+ self,
+ input_embeds: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ batch_size, seq_len, _ = input_embeds.shape
+ hidden_states = self.token_embedding_projection(input_embeds)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+ if position_ids is None:
+ position_ids = (
+ torch.arange(input_embeds.shape[1], device=input_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ for i, layer in enumerate(self.layers):
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ return hidden_states
+
+
+def process_patch_lengths(patch_lengths: torch.Tensor, max_patch_length: Optional[int]) -> torch.Tensor:
+ """
+ Splits patch lengths into smaller segments if they exceed `max_patch_length`.
+ Pads the result to uniform length across the batch.
+
+ Args:
+ patch_lengths (torch.Tensor): [batch_size, num_patches] tensor of patch lengths.
+ max_patch_length (int, optional): Maximum allowed length per patch.
+
+ Returns:
+ torch.Tensor: [batch_size, max_len] tensor of split and padded patch lengths.
+ """
+ if max_patch_length is None:
+ return patch_lengths
+
+ batch_size = patch_lengths.size(0)
+ processed = []
+
+ for seq in patch_lengths:
+ splits = []
+ for length in seq[seq > 0]:
+ length = length.item()
+ full_chunks, remainder = divmod(length, max_patch_length)
+ splits.extend([max_patch_length] * full_chunks)
+ if remainder:
+ splits.append(remainder)
+ processed.append(splits)
+
+ # Find max length to pad to
+ max_len = max(len(splits) for splits in processed)
+ padded = torch.zeros((batch_size, max_len), dtype=patch_lengths.dtype, device=patch_lengths.device)
+
+ for i, splits in enumerate(processed):
+ if splits:
+ padded[i, : len(splits)] = torch.tensor(splits, dtype=patch_lengths.dtype, device=patch_lengths.device)
+
+ # Trim zero columns
+ if (padded != 0).any(dim=0).sum() < padded.shape[1]:
+ last_nonzero = (padded != 0).any(dim=0).nonzero().max().item() + 1
+ padded = padded[:, :last_nonzero]
+
+ return padded
+
+
+class BltPatcher(BltPreTrainedModel):
+ config: BltPatcherConfig
+
+ def __init__(self, config: BltPatcherConfig):
+ super().__init__(config)
+ self.rotary_emb = BltRotaryEmbedding(config=self.config)
+ self.layers = nn.ModuleList()
+ for layer_idx in range(self.config.num_hidden_layers):
+ self.layers.append(BltTransformerLayer(self.config, layer_idx))
+ self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size)
+ self.norm = BltRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+ self.lm_head = nn.Linear(
+ self.config.hidden_size,
+ self.config.vocab_size,
+ bias=False,
+ )
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ patch_size: Optional[int] = None,
+ threshold: Optional[float] = None,
+ max_patch_length: Optional[int] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache()
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for layer in self.layers:
+ hidden_states = layer(hidden_states, position_embeddings=position_embeddings, attention_mask=causal_mask)
+
+ logits = self.lm_head(self.norm(hidden_states))
+ prediction_entropies = torch.distributions.Categorical(logits=logits).entropy()
+
+ batch_size, sequence_length = inputs_embeds.shape[:2]
+ if patch_size is not None:
+ patch_lengths = self.patch_lengths_from_entropies(
+ entropies=prediction_entropies,
+ sequence_length=sequence_length,
+ patch_size=patch_size,
+ threshold=threshold,
+ )
+ else:
+ patch_lengths = torch.ones(
+ (batch_size, sequence_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device
+ )
+ patch_lengths = process_patch_lengths(patch_lengths, max_patch_length)
+ return prediction_entropies, patch_lengths, logits
+
+ @staticmethod
+ def patch_lengths_from_entropies(
+ entropies,
+ sequence_length,
+ patch_size=None,
+ threshold=None,
+ ):
+ """
+ Computes patch lengths from token entropies.
+
+ Depending on whether a threshold is provided, the function uses either:
+ - Thresholding the entropy values (when `threshold` is set).
+ """
+
+ batch_size = entropies.shape[0]
+
+ # Always include token 0 and 1 as starting tokens
+ init_tokens = (
+ torch.tensor([0, 1], dtype=torch.long, device=entropies.device).unsqueeze(0).repeat(batch_size, 1)
+ )
+ offset = init_tokens.shape[1]
+
+ # Ignore first token entropy (BOS)
+ entropies = entropies[:, 1:]
+
+ # Threshold the entropy values to define patch start points
+ patch_mask = entropies > threshold
+
+ seq_len = patch_mask.shape[1]
+
+ # Create patch IDs (token indices), and add a sentinel to ensure alignment
+ token_indices = torch.arange(seq_len, device=entropies.device).unsqueeze(0).expand(batch_size, -1)
+ sentinel = torch.full_like(token_indices, seq_len)
+ padded_indices = torch.cat([token_indices, sentinel], dim=1)
+
+ # Pad mask with inverse to align sentinel correctly
+ padded_mask = torch.cat([patch_mask, ~patch_mask], dim=1)
+
+ # Select indices where mask is True
+ patch_starts = padded_indices[padded_mask].reshape(batch_size, seq_len)
+ max_valid_patches = patch_mask.sum(dim=1).max()
+ patch_starts = patch_starts[:, :max_valid_patches]
+
+ # Offset patch starts to account for the two initial tokens
+ patch_start_ids = torch.cat((init_tokens, patch_starts + offset), dim=1)
+
+ # Compute patch end positions by shifting start positions
+ last_token = torch.full_like(patch_start_ids[:, :1], sequence_length - 1)
+ patch_ends = torch.cat((patch_start_ids[:, 1:] - 1, last_token), dim=1)
+
+ patch_lengths = patch_ends - patch_start_ids + 1
+
+ return patch_lengths
+
+
+def rolling_polynomial_hash(token_tensor, prime: int = 1000000007):
+ """
+ A polynomial rolling hash algorithm that converts sequences
+ of tokens into hash values. The hash is computed as:
+ hash = (token_0 * prime^0 + token_1 * prime^1 + ... + token_n * prime^n)
+
+ The rolling hash allows the model to efficiently
+ identify and encode recurring byte-level patterns in the input text.
+
+ Args:
+ token_tensor (torch.Tensor): [batch_size, seq_len, group_size] containing token IDs to hash
+ prime (int): Prime number used as the base for the polynomial hash.
+
+ Returns:
+ torch.Tensor: Hash values of shape [batch_size, seq_len] where each value
+ represents the hash of the corresponding token group
+
+ Example:
+ >>> tokens = torch.tensor([[1, 2, 3], [4, 5, 6]])
+ >>> hashes = rolling_polynomial_hash(tokens, prime=31)
+ >>> # hash[0] = 1*31^0 + 2*31^1 + 3*31^2
+ >>> # hash[1] = 4*31^0 + 5*31^1 + 6*31^2
+ """
+ prime_tensor = torch.tensor(prime, dtype=torch.int64, device=token_tensor.device)
+ powers = torch.arange(token_tensor.shape[-1], device=token_tensor.device)
+ prime_powers = prime_tensor**powers
+ return torch.sum(token_tensor * prime_powers, dim=-1)
+
+
+def byte_group_hash_function(
+ token_ids: torch.Tensor, group_size: int = 2, prime: int = 1000000007, max_hash: int = 30000
+):
+ """Hash token groups and map to range [0, max_hash]."""
+ with torch.no_grad():
+ batch_size, seq_len = token_ids.shape
+ # Add padding for sliding window
+ padding = torch.zeros(batch_size, group_size - 1, dtype=torch.int64, device=token_ids.device)
+ padded_tokens = torch.cat([padding, token_ids], dim=1)
+
+ # Create sliding windows and compute hashes
+ windows = padded_tokens.unfold(1, group_size, 1)
+ hashes = rolling_polynomial_hash(windows, prime)
+ hash_values = hashes % max_hash
+
+ return hash_values
+
+
+def compute_hash_embeddings(
+ local_encoder_tokens: torch.Tensor,
+ local_encoder,
+ encoder_hash_tok_embedding: nn.Embedding,
+ encoder_hash_byte_group_nb_functions: int,
+ encoder_hash_byte_group_size: list,
+ encoder_hash_byte_group_vocab: int,
+) -> torch.Tensor:
+ """Compute token embeddings enhanced with hash-based embeddings."""
+ # Available primes for hash functions
+ primes = [
+ 1000000007,
+ 5915587277,
+ 1500450271,
+ 3267000013,
+ 5754853343,
+ 4093082899,
+ 9576890767,
+ 3628273133,
+ 2860486313,
+ 5463458053,
+ 3367900313,
+ ]
+
+ embeddings = local_encoder.embed_tokens(local_encoder_tokens)
+ embedding_idx = 0
+ for func_nb in range(encoder_hash_byte_group_nb_functions):
+ prime = primes[func_nb % len(primes)] # Cycle through primes if more functions than primes
+ for group_size in encoder_hash_byte_group_size:
+ hash_ids = byte_group_hash_function(local_encoder_tokens, group_size, prime, encoder_hash_byte_group_vocab)
+ # Apply offset to get the correct slice of the fused embedding
+ offset_hash_ids = hash_ids + embedding_idx * encoder_hash_byte_group_vocab
+ embeddings += encoder_hash_tok_embedding(offset_hash_ids)
+ embedding_idx += 1
+
+ return embeddings
+
+
+def _prepare_patch_cross_attention_mask(
+ patch_ids: torch.Tensor,
+ num_patches: int,
+ sequence_length: int,
+ patches_as_queries: bool = False,
+ cross_attn_k: int = 1,
+ dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Prepare cross-attention mask for patch-based attention, following mllama's robust approach.
+
+ This function creates masks that control which patches can attend to which other patches,
+ with support for query/key role swapping and cross-attention multipliers.
+
+ Args:
+ patch_ids (torch.Tensor): Tensor of shape [batch_size, seq_len] containing patch ids.
+ num_patches (int): Total number of patches.
+ sequence_length (int): Length of the sequence.
+ patches_as_queries (bool): If True, patches are used as queries, otherwise as keys.
+ cross_attn_k (int): Cross-attention multiplier for repeating patches.
+ dtype (torch.dtype): Data type for the output mask.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ - cross_attention_mask: 4D tensor [batch_size, 1, q_len, kv_len]
+ """
+ batch_size, seq_len = patch_ids.shape
+ device = patch_ids.device
+
+ # Determine query and key lengths based on configuration
+ if patches_as_queries:
+ q_len = num_patches * cross_attn_k
+ kv_len = sequence_length
+ # Create patch-to-sequence mapping
+ q_patch_ids = (
+ torch.arange(num_patches, device=device)
+ .unsqueeze(0)
+ .unsqueeze(-1)
+ .expand(batch_size, num_patches, seq_len)
+ )
+ kv_patch_ids = patch_ids.unsqueeze(1).expand(batch_size, num_patches, seq_len)
+ else:
+ q_len = sequence_length
+ kv_len = num_patches * cross_attn_k
+ # Create sequence-to-patch mapping
+ q_patch_ids = patch_ids.unsqueeze(-1).expand(batch_size, seq_len, num_patches)
+ kv_patch_ids = (
+ torch.arange(num_patches, device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, seq_len, num_patches)
+ )
+
+ # Create base attention mask - boolean mask where True means "should attend"
+ # Exact patch matching
+ cross_attention_mask = q_patch_ids == kv_patch_ids
+
+ # Handle cross_attn_k multiplier by repeating along appropriate dimension
+ repeat_dim = 1 if patches_as_queries else -1
+ cross_attention_mask = cross_attention_mask.repeat_interleave(cross_attn_k, dim=repeat_dim)
+
+ # Validate dimensions
+ expected_shape = (batch_size, q_len, kv_len)
+ if cross_attention_mask.shape != expected_shape:
+ raise ValueError(
+ f"Cross attention mask shape {cross_attention_mask.shape} doesn't match expected {expected_shape}"
+ )
+
+ # Reshape so it can be used by attn module - add head dimension
+ cross_attention_mask = cross_attention_mask.unsqueeze(1) # [batch_size, 1, q_len, kv_len]
+
+ # Invert the mask (following mllama pattern exactly)
+ # True -> 0.0 (attend), False -> 1.0 (will become -inf)
+ inverted_cross_attn_mask = 1.0 - cross_attention_mask.to(dtype)
+ cross_attention_mask = inverted_cross_attn_mask.masked_fill(
+ inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min
+ )
+
+ return cross_attention_mask
+
+
+class BltModel(BltPreTrainedModel):
+ def __init__(self, config: BltConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+
+ self.config = config
+ self.local_encoder = BltLocalEncoder(config.encoder_config)
+ self.global_transformer = BltGlobalTransformer(config.global_config)
+ self.local_decoder = BltLocalDecoder(config.decoder_config)
+ num_embeddings = config.encoder_hash_byte_group_nb_functions * len(config.encoder_hash_byte_group_size)
+ total_vocab_size = config.encoder_hash_byte_group_vocab * num_embeddings
+ self.encoder_hash_tok_embedding = nn.Embedding(total_vocab_size, config.encoder_config.hidden_size)
+ if self.config.patch_in_forward:
+ self.patcher = BltPatcher(config.patcher_config)
+ self.patcher.eval()
+ for param in self.patcher.parameters():
+ param.requires_grad = False
+ else:
+ self.patcher = None
+ self.post_init()
+
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ patch_lengths: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = 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")
+
+ # Extract input embeddings as early as possible
+ if inputs_embeds is not None:
+ encoder_embeds = inputs_embeds
+ batch_size, sequence_length, _ = inputs_embeds.shape
+ else:
+ batch_size, sequence_length = input_ids.shape
+ encoder_embeds = compute_hash_embeddings(
+ input_ids,
+ self.local_encoder,
+ self.encoder_hash_tok_embedding,
+ self.config.encoder_hash_byte_group_nb_functions,
+ self.config.encoder_hash_byte_group_size,
+ self.config.encoder_hash_byte_group_vocab,
+ )
+
+ if patch_lengths is None:
+ if self.config.patching_mode == "entropy" and self.patcher is not None:
+ if input_ids is None:
+ raise ValueError("input_ids is required for entropy-based patching")
+ _, patch_lengths, _ = self.patcher(
+ input_ids,
+ patch_size=self.config.patch_size,
+ threshold=self.config.patching_threshold,
+ max_patch_length=self.config.max_patch_length,
+ patching_batch_size=self.config.patching_batch_size,
+ device=input_ids.device,
+ )
+ else:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ dtype = input_ids.dtype if input_ids is not None else inputs_embeds.dtype
+ patch_lengths = process_patch_lengths(
+ torch.ones((batch_size, sequence_length + 1), dtype=dtype, device=device),
+ self.config.max_patch_length,
+ )
+ patch_ids = self._patch_ids_from_lengths(patch_lengths, sequence_length)
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + encoder_embeds.shape[1], device=encoder_embeds.device
+ )
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=encoder_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ cross_attn_mask_enc = _prepare_patch_cross_attention_mask(
+ patch_ids=patch_ids,
+ num_patches=patch_lengths.shape[1],
+ sequence_length=sequence_length,
+ patches_as_queries=True,
+ cross_attn_k=self.config.cross_attn_k,
+ dtype=encoder_embeds.dtype,
+ )
+ encoder_hidden_states, encoder_cross_states = self.local_encoder(
+ input_ids=input_ids,
+ inputs_embeds=encoder_embeds,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ encoder_attention_mask=cross_attn_mask_enc,
+ num_patches=patch_lengths.shape[1],
+ patch_ids=patch_ids,
+ **kwargs,
+ )
+ encoder_cross_states = encoder_cross_states.view(batch_size, patch_lengths.shape[1], -1)
+ global_cache_position = torch.arange(0, encoder_cross_states.shape[1], device=encoder_cross_states.device)
+ global_position_ids = global_cache_position.unsqueeze(0)
+ global_causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=encoder_cross_states,
+ attention_mask=None,
+ cache_position=global_cache_position,
+ past_key_values=None,
+ position_ids=None,
+ )
+
+ global_hidden_states = self.global_transformer(
+ input_embeds=encoder_cross_states,
+ attention_mask=global_causal_mask,
+ position_ids=global_position_ids,
+ **kwargs,
+ )
+ decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length)
+ cross_attn_mask_dec = _prepare_patch_cross_attention_mask(
+ patch_ids=decoder_patch_ids,
+ num_patches=patch_lengths.shape[1],
+ sequence_length=sequence_length,
+ patches_as_queries=False,
+ cross_attn_k=self.config.cross_attn_k,
+ dtype=encoder_embeds.dtype,
+ )
+ output = self.local_decoder(
+ input_ids=input_ids,
+ inputs_embeds=encoder_hidden_states,
+ patch_embeds=global_hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ encoder_attention_mask=cross_attn_mask_dec,
+ **kwargs,
+ )
+ return BaseModelOutputWithPast(
+ last_hidden_state=output,
+ past_key_values=past_key_values,
+ )
+
+ def get_input_embeddings(self):
+ return self.local_encoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.local_encoder.embed_tokens = value
+
+ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) -> torch.Tensor:
+ batch_size = patch_lengths.shape[0]
+ patch_starts = torch.cat(
+ [
+ torch.zeros(batch_size, 1, dtype=patch_lengths.dtype, device=patch_lengths.device),
+ patch_lengths.cumsum(dim=-1)[:, :-1],
+ ],
+ dim=-1,
+ )
+ token_positions = torch.arange(seq_len, device=patch_lengths.device)
+ return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
+
+
+@auto_docstring(
+ custom_intro="""
+ The Blt Text Model with a language modeling head on top.
+ """
+)
+class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
+ config: BltConfig
+ _can_compile_fullgraph = False
+ base_model_prefix = "model"
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config: BltConfig):
+ super().__init__(config.get_text_config())
+ self.text_config = config.get_text_config()
+ self.vocab_size = config.vocab_size
+ self.model = BltModel(config)
+ self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False)
+
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ cross_attention_states: Optional[torch.LongTensor] = None, # Keep for compatibility
+ cross_attention_mask: Optional[torch.LongTensor] = None,
+ full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, CausalLMOutputWithPast]:
+ r"""
+ cross_attention_states (`torch.FloatTensor`, *optional*):
+ Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+ the language model will attend to.
+ cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+ Cross-attention mask to control the interaction between text tokens and image tiles.
+ This 4D tensor defines which image tiles each text token should attend to.
+
+ For each text token (in seq_length):
+ - 1 indicates the token **should attend** to the corresponding image tile
+ - 0 indicates the token **should not attend** to the corresponding image tile
+ full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*):
+ A tuple containing two tensors that mask out rows in the cross-attention mechanism:
+ - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1.
+ A value of 0 indicates that the corresponding text token's entire row in the cross-attention
+ matrix should be masked out (all image tokens ignored).
+ - The second tensor has the same shape and is used internally to apply the masking during
+ the forward pass of cross-attention layers.
+ This mask is derived from the cross_attention_mask and is used to handle cases where a text token
+ should not attend to any image token.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BltForCausalLM
+
+ >>> model = BltForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
+ >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
+
+ >>> prompt = "If I had to write a haiku, it would be:"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
+ >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ >>> print(result)
+ If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
+ I love the idea of snowflakes gently falling, each one
+ ```
+ """
+ # Call parent forward but exclude cross_attention_states from model call
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ cross_attention_mask=cross_attention_mask,
+ full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["BltPreTrainedModel", "BltModel", "BltPatcher", "BltForCausalLM"]
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
new file mode 100644
index 000000000000..00b1211fdb08
--- /dev/null
+++ b/src/transformers/models/blt/modular_blt.py
@@ -0,0 +1,1015 @@
+# coding=utf-8
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Blt modular model, inheriting from Mllama where appropriate."""
+
+from typing import Callable, Optional, Union
+
+import torch
+import torch.distributions
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...cache_utils import Cache, DynamicCache
+from ...masking_utils import create_causal_mask
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import OutputRecorder, check_model_inputs
+from ..cohere2.modeling_cohere2 import (
+ Cohere2RotaryEmbedding,
+ rotate_half, # noqa: F401
+)
+from ..mllama.modeling_mllama import (
+ MllamaForCausalLM,
+ MllamaPreTrainedModel,
+ MllamaSelfAttentionDecoderLayer,
+ MllamaTextCrossAttention,
+ MllamaTextMLP,
+ MllamaTextRMSNorm,
+ MllamaTextSelfAttention,
+ eager_attention_forward,
+)
+from .configuration_blt import (
+ BltConfig,
+ BltGlobalTransformerConfig,
+ BltLocalDecoderConfig,
+ BltLocalEncoderConfig,
+ BltPatcherConfig,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def rolling_polynomial_hash(token_tensor, prime: int = 1000000007):
+ """
+ A polynomial rolling hash algorithm that converts sequences
+ of tokens into hash values. The hash is computed as:
+ hash = (token_0 * prime^0 + token_1 * prime^1 + ... + token_n * prime^n)
+
+ The rolling hash allows the model to efficiently
+ identify and encode recurring byte-level patterns in the input text.
+
+ Args:
+ token_tensor (torch.Tensor): [batch_size, seq_len, group_size] containing token IDs to hash
+ prime (int): Prime number used as the base for the polynomial hash.
+
+ Returns:
+ torch.Tensor: Hash values of shape [batch_size, seq_len] where each value
+ represents the hash of the corresponding token group
+
+ Example:
+ >>> tokens = torch.tensor([[1, 2, 3], [4, 5, 6]])
+ >>> hashes = rolling_polynomial_hash(tokens, prime=31)
+ >>> # hash[0] = 1*31^0 + 2*31^1 + 3*31^2
+ >>> # hash[1] = 4*31^0 + 5*31^1 + 6*31^2
+ """
+ prime_tensor = torch.tensor(prime, dtype=torch.int64, device=token_tensor.device)
+ powers = torch.arange(token_tensor.shape[-1], device=token_tensor.device)
+ prime_powers = prime_tensor**powers
+ return torch.sum(token_tensor * prime_powers, dim=-1)
+
+
+def byte_group_hash_function(
+ token_ids: torch.Tensor, group_size: int = 2, prime: int = 1000000007, max_hash: int = 30000
+):
+ """Hash token groups and map to range [0, max_hash]."""
+ with torch.no_grad():
+ batch_size, seq_len = token_ids.shape
+ # Add padding for sliding window
+ padding = torch.zeros(batch_size, group_size - 1, dtype=torch.int64, device=token_ids.device)
+ padded_tokens = torch.cat([padding, token_ids], dim=1)
+
+ # Create sliding windows and compute hashes
+ windows = padded_tokens.unfold(1, group_size, 1)
+ hashes = rolling_polynomial_hash(windows, prime)
+ hash_values = hashes % max_hash
+
+ return hash_values
+
+
+def compute_hash_embeddings(
+ local_encoder_tokens: torch.Tensor,
+ local_encoder,
+ encoder_hash_tok_embedding: nn.Embedding,
+ encoder_hash_byte_group_nb_functions: int,
+ encoder_hash_byte_group_size: list,
+ encoder_hash_byte_group_vocab: int,
+) -> torch.Tensor:
+ """Compute token embeddings enhanced with hash-based embeddings."""
+ # Available primes for hash functions
+ primes = [
+ 1000000007,
+ 5915587277,
+ 1500450271,
+ 3267000013,
+ 5754853343,
+ 4093082899,
+ 9576890767,
+ 3628273133,
+ 2860486313,
+ 5463458053,
+ 3367900313,
+ ]
+
+ embeddings = local_encoder.embed_tokens(local_encoder_tokens)
+ embedding_idx = 0
+ for func_nb in range(encoder_hash_byte_group_nb_functions):
+ prime = primes[func_nb % len(primes)] # Cycle through primes if more functions than primes
+ for group_size in encoder_hash_byte_group_size:
+ hash_ids = byte_group_hash_function(local_encoder_tokens, group_size, prime, encoder_hash_byte_group_vocab)
+ # Apply offset to get the correct slice of the fused embedding
+ offset_hash_ids = hash_ids + embedding_idx * encoder_hash_byte_group_vocab
+ embeddings += encoder_hash_tok_embedding(offset_hash_ids)
+ embedding_idx += 1
+
+ return embeddings
+
+
+def _prepare_patch_cross_attention_mask(
+ patch_ids: torch.Tensor,
+ num_patches: int,
+ sequence_length: int,
+ patches_as_queries: bool = False,
+ cross_attn_k: int = 1,
+ dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Prepare cross-attention mask for patch-based attention, following mllama's robust approach.
+
+ This function creates masks that control which patches can attend to which other patches,
+ with support for query/key role swapping and cross-attention multipliers.
+
+ Args:
+ patch_ids (torch.Tensor): Tensor of shape [batch_size, seq_len] containing patch ids.
+ num_patches (int): Total number of patches.
+ sequence_length (int): Length of the sequence.
+ patches_as_queries (bool): If True, patches are used as queries, otherwise as keys.
+ cross_attn_k (int): Cross-attention multiplier for repeating patches.
+ dtype (torch.dtype): Data type for the output mask.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ - cross_attention_mask: 4D tensor [batch_size, 1, q_len, kv_len]
+ """
+ batch_size, seq_len = patch_ids.shape
+ device = patch_ids.device
+
+ # Determine query and key lengths based on configuration
+ if patches_as_queries:
+ q_len = num_patches * cross_attn_k
+ kv_len = sequence_length
+ # Create patch-to-sequence mapping
+ q_patch_ids = (
+ torch.arange(num_patches, device=device)
+ .unsqueeze(0)
+ .unsqueeze(-1)
+ .expand(batch_size, num_patches, seq_len)
+ )
+ kv_patch_ids = patch_ids.unsqueeze(1).expand(batch_size, num_patches, seq_len)
+ else:
+ q_len = sequence_length
+ kv_len = num_patches * cross_attn_k
+ # Create sequence-to-patch mapping
+ q_patch_ids = patch_ids.unsqueeze(-1).expand(batch_size, seq_len, num_patches)
+ kv_patch_ids = (
+ torch.arange(num_patches, device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, seq_len, num_patches)
+ )
+
+ # Create base attention mask - boolean mask where True means "should attend"
+ # Exact patch matching
+ cross_attention_mask = q_patch_ids == kv_patch_ids
+
+ # Handle cross_attn_k multiplier by repeating along appropriate dimension
+ repeat_dim = 1 if patches_as_queries else -1
+ cross_attention_mask = cross_attention_mask.repeat_interleave(cross_attn_k, dim=repeat_dim)
+
+ # Validate dimensions
+ expected_shape = (batch_size, q_len, kv_len)
+ if cross_attention_mask.shape != expected_shape:
+ raise ValueError(
+ f"Cross attention mask shape {cross_attention_mask.shape} doesn't match expected {expected_shape}"
+ )
+
+ # Reshape so it can be used by attn module - add head dimension
+ cross_attention_mask = cross_attention_mask.unsqueeze(1) # [batch_size, 1, q_len, kv_len]
+
+ # Invert the mask (following mllama pattern exactly)
+ # True -> 0.0 (attend), False -> 1.0 (will become -inf)
+ inverted_cross_attn_mask = 1.0 - cross_attention_mask.to(dtype)
+ cross_attention_mask = inverted_cross_attn_mask.masked_fill(
+ inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min
+ )
+
+ return cross_attention_mask
+
+
+def process_patch_lengths(patch_lengths: torch.Tensor, max_patch_length: Optional[int]) -> torch.Tensor:
+ """
+ Splits patch lengths into smaller segments if they exceed `max_patch_length`.
+ Pads the result to uniform length across the batch.
+
+ Args:
+ patch_lengths (torch.Tensor): [batch_size, num_patches] tensor of patch lengths.
+ max_patch_length (int, optional): Maximum allowed length per patch.
+
+ Returns:
+ torch.Tensor: [batch_size, max_len] tensor of split and padded patch lengths.
+ """
+ if max_patch_length is None:
+ return patch_lengths
+
+ batch_size = patch_lengths.size(0)
+ processed = []
+
+ for seq in patch_lengths:
+ splits = []
+ for length in seq[seq > 0]:
+ length = length.item()
+ full_chunks, remainder = divmod(length, max_patch_length)
+ splits.extend([max_patch_length] * full_chunks)
+ if remainder:
+ splits.append(remainder)
+ processed.append(splits)
+
+ # Find max length to pad to
+ max_len = max(len(splits) for splits in processed)
+ padded = torch.zeros((batch_size, max_len), dtype=patch_lengths.dtype, device=patch_lengths.device)
+
+ for i, splits in enumerate(processed):
+ if splits:
+ padded[i, : len(splits)] = torch.tensor(splits, dtype=patch_lengths.dtype, device=patch_lengths.device)
+
+ # Trim zero columns
+ if (padded != 0).any(dim=0).sum() < padded.shape[1]:
+ last_nonzero = (padded != 0).any(dim=0).nonzero().max().item() + 1
+ padded = padded[:, :last_nonzero]
+
+ return padded
+
+
+class BltMLP(MllamaTextMLP):
+ pass
+
+
+class BltRMSNorm(MllamaTextRMSNorm):
+ pass
+
+
+class BltRotaryEmbedding(Cohere2RotaryEmbedding):
+ pass
+
+
+class BltTransformerLayer(MllamaSelfAttentionDecoderLayer):
+ def __init__(self, config, layer_idx: int):
+ super().__init__()
+
+ self.self_attn = BltSelfAttention(config=config, layer_idx=layer_idx)
+ self.mlp = BltMLP(config)
+ self.input_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+class BltSelfAttention(MllamaTextSelfAttention):
+ def __init__(self, config: BltConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.is_causal = True
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ position_embeddings: torch.Tensor,
+ use_cache: bool = False,
+ past_key_values=None,
+ cache_position=None,
+ **kwargs,
+ ):
+ return super().forward(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=position_embeddings,
+ use_cache=use_cache,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+
+class BltCrossAttention(MllamaTextCrossAttention):
+ """Cross-attention module for Blt, following transformers style"""
+
+ def __init__(self, config: BltConfig, layer_idx: int, hidden_size: Optional[int] = None):
+ super().__init__()
+ self.is_causal = False
+ self.q_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+ self.k_norm = BltRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cross_attention_states: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Cache] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ bsz, q_len, _ = hidden_states.size()
+ query_states = self.q_norm(hidden_states)
+ query_states = self.q_proj(query_states)
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+ if cross_attention_states is not None:
+ cross_attention_states = self.k_norm(cross_attention_states)
+ key_states = self.k_proj(cross_attention_states)
+ value_states = self.v_proj(cross_attention_states)
+ key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(
+ key_states, value_states, self.layer_idx, {"cache_position": cache_position}
+ )
+ elif cache_position[0] != 0:
+ key_states, value_states = (
+ past_key_values.layers[self.layer_idx].keys,
+ past_key_values.layers[self.layer_idx].values,
+ )
+ else:
+ raise ValueError(
+ "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!"
+ )
+ attention_interface: Callable = eager_attention_forward
+
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ attn_output = attn_output + hidden_states
+ return attn_output, attn_weights
+
+
+@auto_docstring
+class BltPreTrainedModel(MllamaPreTrainedModel):
+ config: BltConfig
+ _supports_attention_backend = False
+ _supports_flash_attn = False
+ _supports_flex_attn = False
+ _no_split_modules = ["BltTransformerLayer"]
+ _can_record_outputs = {
+ "hidden_states": OutputRecorder(BltTransformerLayer, index=0, layer_name="local_decoder"),
+ "attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_decoder"),
+ }
+
+ def _init_weights(self, module):
+ raise AttributeError("No need to inherit it!")
+
+ def _update_causal_mask(self, module):
+ raise AttributeError("No need to inherit it!")
+
+ def _prepare_4d_causal_attention_mask_with_cache_position(self, module):
+ raise AttributeError("No need to inherit it!")
+
+
+class BltLocalEncoder(BltPreTrainedModel):
+ config: BltLocalEncoderConfig
+ _can_record_outputs = {
+ "encoder_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="local_encoder"),
+ }
+
+ def __init__(self, config: BltLocalEncoderConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+ self.config = config
+ self.layers = nn.ModuleList(
+ [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+ self.patch_embedding_projection = nn.Linear(
+ in_features=config.hidden_size,
+ out_features=config.hidden_size * config.cross_attn_k,
+ bias=False,
+ )
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.cross_attn_layers = nn.ModuleList()
+ layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1
+ for layer_idx in range(layers_to_add):
+ self.cross_attn_layers.append(
+ BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size)
+ )
+
+ self.post_init()
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ patch_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ num_patches: Optional[int] = None,
+ patch_ids: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ batch_size = inputs_embeds.shape[0]
+ hidden_states = F.dropout(inputs_embeds, p=self.config.dropout, training=self.training)
+
+ if position_ids is None:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+
+ for idx, layer in enumerate(self.layers):
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ if idx == len(self.layers) - 1 or self.config.cross_attn_all_layers:
+ patch_embeds = self.patch_reduce(hidden_states, num_patches, patch_ids)
+ patch_embeds = self.patch_embedding_projection(patch_embeds)
+ patch_embeds = patch_embeds.reshape(
+ batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size
+ )
+ layer_idx = idx if self.config.cross_attn_all_layers else 0
+ cross_attention_output, _ = self.cross_attn_layers[layer_idx](
+ hidden_states=patch_embeds,
+ cross_attention_states=hidden_states,
+ attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+ patch_embeds = patch_embeds + cross_attention_output
+ encoder_cross_states = patch_embeds
+ return hidden_states, encoder_cross_states
+
+ def patch_reduce(self, hidden_states, max_num_patches, patch_ids):
+ """
+ Reduce variable length patches to single embedding per patch
+ Note: this works with variable number of patches for different sequences in the batch
+ It handles variable length patches by assuming that patch_lengths will be 0 for any
+ extra patches on the *right*. Since there can be a variable number of patches
+ this function also return the number of patches for each sequence in the batch.
+ Any embeddings on the right that are not allocated to a patch
+ (i.e. if the sum(patch_lengths[i]) < seq_len for any i)
+ will be sent to a dummy patch, which is trimmed before returning.
+ """
+ batch_size = hidden_states.shape[0]
+ embedding_dim = hidden_states.shape[-1]
+
+ patch_ids = patch_ids.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])
+
+ reduced_embeddings = torch.zeros(
+ (batch_size, max_num_patches, embedding_dim), dtype=hidden_states.dtype, device=hidden_states.device
+ )
+
+ reduced_embeddings = reduced_embeddings.scatter_reduce(
+ src=hidden_states,
+ dim=1,
+ index=patch_ids,
+ reduce="amax",
+ include_self=False,
+ )
+ reduced_embeddings = reduced_embeddings[:, :max_num_patches, :]
+
+ return reduced_embeddings
+
+
+class BltLocalDecoder(BltPreTrainedModel):
+ config: BltLocalDecoderConfig
+
+ def __init__(self, config: BltLocalDecoderConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+ self.config = config
+ self.cross_attn_decoder = True
+ self.layers = nn.ModuleList(
+ [BltTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+ self.patch_embedding_projection = nn.Linear(
+ in_features=config.hidden_size_global,
+ out_features=config.hidden_size * config.cross_attn_k,
+ bias=False,
+ )
+ self.norm = BltRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.cross_attn_layers = nn.ModuleList()
+ layers_to_add = config.num_hidden_layers if config.cross_attn_all_layers else 1
+ for layer_idx in range(layers_to_add):
+ self.cross_attn_layers.append(
+ BltCrossAttention(config=config, layer_idx=layer_idx, hidden_size=config.hidden_size)
+ )
+
+ self.post_init()
+
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ patch_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ batch_size = inputs_embeds.shape[0]
+ hidden_states = inputs_embeds
+ patch_embeds = self.patch_embedding_projection(patch_embeds)
+ patch_embeds = patch_embeds.reshape(
+ batch_size, patch_embeds.shape[1] * self.config.cross_attn_k, self.config.hidden_size
+ )
+
+ if patch_embeds is not None and not self.cross_attn_decoder:
+ hidden_states = hidden_states + patch_embeds
+
+ if position_ids is None:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+
+ for i, layer in enumerate(self.layers):
+ if i == 0 or self.config.cross_attn_all_layers:
+ cross_attention_output, _ = self.cross_attn_layers[i](
+ hidden_states=hidden_states,
+ cross_attention_states=patch_embeds,
+ attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+ hidden_states = hidden_states + cross_attention_output
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ logits = self.norm(hidden_states)
+ return logits
+
+
+class BltGlobalTransformer(BltPreTrainedModel):
+ config: BltGlobalTransformerConfig
+ _can_record_outputs = {
+ "global_attentions": OutputRecorder(BltSelfAttention, index=1, layer_name="global_transformer"),
+ }
+
+ def __init__(self, config: BltGlobalTransformerConfig):
+ super().__init__(config)
+ self.config = config
+ self.layers = nn.ModuleList()
+ for layer_idx in range(config.num_hidden_layers):
+ self.layers.append(BltTransformerLayer(config, layer_idx))
+ self.rotary_emb = BltRotaryEmbedding(config=config)
+
+ # Create token embedding projection (use nn.Identity() when no projection needed)
+ if getattr(config, "encoder_cross_output_size", None) is not None:
+ self.token_embedding_projection = nn.Linear(
+ config.encoder_cross_output_size, config.hidden_size, bias=False
+ )
+ else:
+ self.token_embedding_projection = nn.Identity()
+
+ self.post_init()
+
+ def forward(
+ self,
+ input_embeds: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ batch_size, seq_len, _ = input_embeds.shape
+ hidden_states = self.token_embedding_projection(input_embeds)
+ hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
+ if position_ids is None:
+ position_ids = (
+ torch.arange(input_embeds.shape[1], device=input_embeds.device).unsqueeze(0).expand(batch_size, -1)
+ )
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+ for i, layer in enumerate(self.layers):
+ hidden_states = layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ return hidden_states
+
+
+class BltPatcher(BltPreTrainedModel):
+ config: BltPatcherConfig
+
+ def __init__(self, config: BltPatcherConfig):
+ super().__init__(config)
+ self.rotary_emb = BltRotaryEmbedding(config=self.config)
+ self.layers = nn.ModuleList()
+ for layer_idx in range(self.config.num_hidden_layers):
+ self.layers.append(BltTransformerLayer(self.config, layer_idx))
+ self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size)
+ self.norm = BltRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+ self.lm_head = nn.Linear(
+ self.config.hidden_size,
+ self.config.vocab_size,
+ bias=False,
+ )
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ patch_size: Optional[int] = None,
+ threshold: Optional[float] = None,
+ max_patch_length: Optional[int] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache()
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for layer in self.layers:
+ hidden_states = layer(hidden_states, position_embeddings=position_embeddings, attention_mask=causal_mask)
+
+ logits = self.lm_head(self.norm(hidden_states))
+ prediction_entropies = torch.distributions.Categorical(logits=logits).entropy()
+
+ batch_size, sequence_length = inputs_embeds.shape[:2]
+ if patch_size is not None:
+ patch_lengths = self.patch_lengths_from_entropies(
+ entropies=prediction_entropies,
+ sequence_length=sequence_length,
+ patch_size=patch_size,
+ threshold=threshold,
+ )
+ else:
+ patch_lengths = torch.ones(
+ (batch_size, sequence_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device
+ )
+ patch_lengths = process_patch_lengths(patch_lengths, max_patch_length)
+ return prediction_entropies, patch_lengths, logits
+
+ @staticmethod
+ def patch_lengths_from_entropies(
+ entropies,
+ sequence_length,
+ patch_size=None,
+ threshold=None,
+ ):
+ """
+ Computes patch lengths from token entropies.
+
+ Depending on whether a threshold is provided, the function uses either:
+ - Thresholding the entropy values (when `threshold` is set).
+ """
+
+ batch_size = entropies.shape[0]
+
+ # Always include token 0 and 1 as starting tokens
+ init_tokens = (
+ torch.tensor([0, 1], dtype=torch.long, device=entropies.device).unsqueeze(0).repeat(batch_size, 1)
+ )
+ offset = init_tokens.shape[1]
+
+ # Ignore first token entropy (BOS)
+ entropies = entropies[:, 1:]
+
+ # Threshold the entropy values to define patch start points
+ patch_mask = entropies > threshold
+
+ seq_len = patch_mask.shape[1]
+
+ # Create patch IDs (token indices), and add a sentinel to ensure alignment
+ token_indices = torch.arange(seq_len, device=entropies.device).unsqueeze(0).expand(batch_size, -1)
+ sentinel = torch.full_like(token_indices, seq_len)
+ padded_indices = torch.cat([token_indices, sentinel], dim=1)
+
+ # Pad mask with inverse to align sentinel correctly
+ padded_mask = torch.cat([patch_mask, ~patch_mask], dim=1)
+
+ # Select indices where mask is True
+ patch_starts = padded_indices[padded_mask].reshape(batch_size, seq_len)
+ max_valid_patches = patch_mask.sum(dim=1).max()
+ patch_starts = patch_starts[:, :max_valid_patches]
+
+ # Offset patch starts to account for the two initial tokens
+ patch_start_ids = torch.cat((init_tokens, patch_starts + offset), dim=1)
+
+ # Compute patch end positions by shifting start positions
+ last_token = torch.full_like(patch_start_ids[:, :1], sequence_length - 1)
+ patch_ends = torch.cat((patch_start_ids[:, 1:] - 1, last_token), dim=1)
+
+ patch_lengths = patch_ends - patch_start_ids + 1
+
+ return patch_lengths
+
+
+class BltModel(BltPreTrainedModel):
+ def __init__(self, config: BltConfig):
+ super().__init__(config)
+ self.gradient_checkpointing = False
+
+ self.config = config
+ self.local_encoder = BltLocalEncoder(config.encoder_config)
+ self.global_transformer = BltGlobalTransformer(config.global_config)
+ self.local_decoder = BltLocalDecoder(config.decoder_config)
+ num_embeddings = config.encoder_hash_byte_group_nb_functions * len(config.encoder_hash_byte_group_size)
+ total_vocab_size = config.encoder_hash_byte_group_vocab * num_embeddings
+ self.encoder_hash_tok_embedding = nn.Embedding(total_vocab_size, config.encoder_config.hidden_size)
+ if self.config.patch_in_forward:
+ self.patcher = BltPatcher(config.patcher_config)
+ self.patcher.eval()
+ for param in self.patcher.parameters():
+ param.requires_grad = False
+ else:
+ self.patcher = None
+ self.post_init()
+
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ patch_lengths: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = 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")
+
+ # Extract input embeddings as early as possible
+ if inputs_embeds is not None:
+ encoder_embeds = inputs_embeds
+ batch_size, sequence_length, _ = inputs_embeds.shape
+ else:
+ batch_size, sequence_length = input_ids.shape
+ encoder_embeds = compute_hash_embeddings(
+ input_ids,
+ self.local_encoder,
+ self.encoder_hash_tok_embedding,
+ self.config.encoder_hash_byte_group_nb_functions,
+ self.config.encoder_hash_byte_group_size,
+ self.config.encoder_hash_byte_group_vocab,
+ )
+
+ if patch_lengths is None:
+ if self.config.patching_mode == "entropy" and self.patcher is not None:
+ if input_ids is None:
+ raise ValueError("input_ids is required for entropy-based patching")
+ _, patch_lengths, _ = self.patcher(
+ input_ids,
+ patch_size=self.config.patch_size,
+ threshold=self.config.patching_threshold,
+ max_patch_length=self.config.max_patch_length,
+ patching_batch_size=self.config.patching_batch_size,
+ device=input_ids.device,
+ )
+ else:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ dtype = input_ids.dtype if input_ids is not None else inputs_embeds.dtype
+ patch_lengths = process_patch_lengths(
+ torch.ones((batch_size, sequence_length + 1), dtype=dtype, device=device),
+ self.config.max_patch_length,
+ )
+ patch_ids = self._patch_ids_from_lengths(patch_lengths, sequence_length)
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + encoder_embeds.shape[1], device=encoder_embeds.device
+ )
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=encoder_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ cross_attn_mask_enc = _prepare_patch_cross_attention_mask(
+ patch_ids=patch_ids,
+ num_patches=patch_lengths.shape[1],
+ sequence_length=sequence_length,
+ patches_as_queries=True,
+ cross_attn_k=self.config.cross_attn_k,
+ dtype=encoder_embeds.dtype,
+ )
+ encoder_hidden_states, encoder_cross_states = self.local_encoder(
+ input_ids=input_ids,
+ inputs_embeds=encoder_embeds,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ encoder_attention_mask=cross_attn_mask_enc,
+ num_patches=patch_lengths.shape[1],
+ patch_ids=patch_ids,
+ **kwargs,
+ )
+ encoder_cross_states = encoder_cross_states.view(batch_size, patch_lengths.shape[1], -1)
+ global_cache_position = torch.arange(0, encoder_cross_states.shape[1], device=encoder_cross_states.device)
+ global_position_ids = global_cache_position.unsqueeze(0)
+ global_causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=encoder_cross_states,
+ attention_mask=None,
+ cache_position=global_cache_position,
+ past_key_values=None,
+ position_ids=None,
+ )
+
+ global_hidden_states = self.global_transformer(
+ input_embeds=encoder_cross_states,
+ attention_mask=global_causal_mask,
+ position_ids=global_position_ids,
+ **kwargs,
+ )
+ decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length)
+ cross_attn_mask_dec = _prepare_patch_cross_attention_mask(
+ patch_ids=decoder_patch_ids,
+ num_patches=patch_lengths.shape[1],
+ sequence_length=sequence_length,
+ patches_as_queries=False,
+ cross_attn_k=self.config.cross_attn_k,
+ dtype=encoder_embeds.dtype,
+ )
+ output = self.local_decoder(
+ input_ids=input_ids,
+ inputs_embeds=encoder_hidden_states,
+ patch_embeds=global_hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ encoder_attention_mask=cross_attn_mask_dec,
+ **kwargs,
+ )
+ return BaseModelOutputWithPast(
+ last_hidden_state=output,
+ past_key_values=past_key_values,
+ )
+
+ def get_input_embeddings(self):
+ return self.local_encoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.local_encoder.embed_tokens = value
+
+ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) -> torch.Tensor:
+ batch_size = patch_lengths.shape[0]
+ patch_starts = torch.cat(
+ [
+ torch.zeros(batch_size, 1, dtype=patch_lengths.dtype, device=patch_lengths.device),
+ patch_lengths.cumsum(dim=-1)[:, :-1],
+ ],
+ dim=-1,
+ )
+ token_positions = torch.arange(seq_len, device=patch_lengths.device)
+ return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
+
+
+class BltForCausalLM(MllamaForCausalLM):
+ config: BltConfig
+ _can_compile_fullgraph = False
+ base_model_prefix = "model"
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config: BltConfig):
+ super().__init__(config)
+ self.vocab_size = config.vocab_size
+ self.model = BltModel(config)
+ self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False)
+ self.post_init()
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ cross_attention_states: Optional[torch.LongTensor] = None, # Keep for compatibility
+ cross_attention_mask: Optional[torch.LongTensor] = None,
+ full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, CausalLMOutputWithPast]:
+ # Call parent forward but exclude cross_attention_states from model call
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ cross_attention_mask=cross_attention_mask,
+ full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "BltPreTrainedModel",
+ "BltModel",
+ "BltPatcher",
+ "BltForCausalLM",
+]
diff --git a/src/transformers/models/bridgetower/image_processing_bridgetower_fast.py b/src/transformers/models/bridgetower/image_processing_bridgetower_fast.py
index 44da5d4486e7..5be6f9f6c54b 100644
--- a/src/transformers/models/bridgetower/image_processing_bridgetower_fast.py
+++ b/src/transformers/models/bridgetower/image_processing_bridgetower_fast.py
@@ -18,6 +18,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -31,13 +32,7 @@
reorder_images,
)
from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling
-from ...utils import auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import auto_docstring
def make_pixel_mask(
diff --git a/src/transformers/models/chameleon/image_processing_chameleon_fast.py b/src/transformers/models/chameleon/image_processing_chameleon_fast.py
index 39aa4ec87b00..1d102614f7df 100644
--- a/src/transformers/models/chameleon/image_processing_chameleon_fast.py
+++ b/src/transformers/models/chameleon/image_processing_chameleon_fast.py
@@ -19,17 +19,13 @@
import numpy as np
import PIL
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast
from ...image_utils import ImageInput, PILImageResampling, SizeDict
-from ...utils import auto_docstring, is_torchvision_v2_available, logging
+from ...utils import auto_docstring, logging
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/chinese_clip/configuration_chinese_clip.py b/src/transformers/models/chinese_clip/configuration_chinese_clip.py
index e7c98d0d2d9f..c628107048b9 100644
--- a/src/transformers/models/chinese_clip/configuration_chinese_clip.py
+++ b/src/transformers/models/chinese_clip/configuration_chinese_clip.py
@@ -307,7 +307,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -339,7 +339,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/chinese_clip/convert_chinese_clip_original_pytorch_to_hf.py b/src/transformers/models/chinese_clip/convert_chinese_clip_original_pytorch_to_hf.py
index adc9300ef512..2ec838a7da63 100644
--- a/src/transformers/models/chinese_clip/convert_chinese_clip_original_pytorch_to_hf.py
+++ b/src/transformers/models/chinese_clip/convert_chinese_clip_original_pytorch_to_hf.py
@@ -105,7 +105,7 @@ def convert_chinese_clip_checkpoint(checkpoint_path, pytorch_dump_folder_path, c
hf_model = ChineseCLIPModel(config).eval()
pt_weights = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["state_dict"]
- pt_weights = {(name[7:] if name.startswith("module.") else name): value for name, value in pt_weights.items()}
+ pt_weights = {(name.removeprefix("module.")): value for name, value in pt_weights.items()}
copy_text_model_and_projection(hf_model, pt_weights)
copy_vision_model_and_projection(hf_model, pt_weights)
diff --git a/src/transformers/models/clap/feature_extraction_clap.py b/src/transformers/models/clap/feature_extraction_clap.py
index e333248c18ed..33daac615c07 100644
--- a/src/transformers/models/clap/feature_extraction_clap.py
+++ b/src/transformers/models/clap/feature_extraction_clap.py
@@ -152,7 +152,7 @@ def to_dict(self) -> dict[str, Any]:
del output["mel_filters_slaney"]
return output
- def _np_extract_fbank_features(self, waveform: np.ndarray, mel_filters: Optional[np.array] = None) -> np.ndarray:
+ def _np_extract_fbank_features(self, waveform: np.ndarray, mel_filters: Optional[np.ndarray] = None) -> np.ndarray:
"""
Compute the log-mel spectrogram of the provided `waveform` using the Hann window. In CLAP, two different filter
banks are used depending on the truncation pattern:
@@ -199,7 +199,7 @@ def _random_mel_fusion(self, mel, total_frames, chunk_frames):
mel_fusion = np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back], axis=0)
return mel_fusion
- def _get_input_mel(self, waveform: np.ndarray, max_length, truncation, padding) -> np.array:
+ def _get_input_mel(self, waveform: np.ndarray, max_length, truncation, padding) -> np.ndarray:
"""
Extracts the mel spectrogram and prepares it for the mode based on the `truncation` and `padding` arguments.
Four different path are possible:
diff --git a/src/transformers/models/clip/configuration_clip.py b/src/transformers/models/clip/configuration_clip.py
index 0b4fe6ba37f6..e343715e29ee 100644
--- a/src/transformers/models/clip/configuration_clip.py
+++ b/src/transformers/models/clip/configuration_clip.py
@@ -296,7 +296,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -328,7 +328,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/clipseg/configuration_clipseg.py b/src/transformers/models/clipseg/configuration_clipseg.py
index 60b14eb7efbb..e338d278577a 100644
--- a/src/transformers/models/clipseg/configuration_clipseg.py
+++ b/src/transformers/models/clipseg/configuration_clipseg.py
@@ -307,7 +307,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -339,7 +339,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/cohere2_vision/image_processing_cohere2_vision_fast.py b/src/transformers/models/cohere2_vision/image_processing_cohere2_vision_fast.py
index afe76134bc8d..322e98dbd0f5 100644
--- a/src/transformers/models/cohere2_vision/image_processing_cohere2_vision_fast.py
+++ b/src/transformers/models/cohere2_vision/image_processing_cohere2_vision_fast.py
@@ -24,6 +24,7 @@
import numpy as np
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -34,13 +35,7 @@
)
from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, ImageInput, PILImageResampling, SizeDict
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import TensorType, auto_docstring
class Cohere2VisionFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
diff --git a/src/transformers/models/colpali/configuration_colpali.py b/src/transformers/models/colpali/configuration_colpali.py
index 84be59aef09b..be7eaf47b428 100644
--- a/src/transformers/models/colpali/configuration_colpali.py
+++ b/src/transformers/models/colpali/configuration_colpali.py
@@ -83,9 +83,7 @@ def __init__(
f"The model type `{vlm_config['model_type']}` is not supported. Please provide a valid model type."
)
vlm_config = CONFIG_MAPPING[vlm_config["model_type"]](**vlm_config)
- elif isinstance(vlm_config, PretrainedConfig):
- vlm_config = vlm_config
- else:
+ elif not isinstance(vlm_config, PretrainedConfig):
raise TypeError(
f"Invalid type for `vlm_config`. Expected `PretrainedConfig`, `dict`, or `None`, but got {type(vlm_config)}."
)
diff --git a/src/transformers/models/colqwen2/configuration_colqwen2.py b/src/transformers/models/colqwen2/configuration_colqwen2.py
index d9a42df4c97e..21f6e46f1f00 100644
--- a/src/transformers/models/colqwen2/configuration_colqwen2.py
+++ b/src/transformers/models/colqwen2/configuration_colqwen2.py
@@ -75,9 +75,7 @@ def __init__(
"The `model_type` key is missing in the `vlm_config` dictionary. Please provide the model type."
)
vlm_config = CONFIG_MAPPING[vlm_config["model_type"]](**vlm_config)
- elif isinstance(vlm_config, PretrainedConfig):
- vlm_config = vlm_config
- else:
+ elif not isinstance(vlm_config, PretrainedConfig):
raise TypeError(
f"Invalid type for `vlm_config`. Expected `PretrainedConfig`, `dict`, or `None`, but got {type(vlm_config)}."
)
diff --git a/src/transformers/models/conditional_detr/image_processing_conditional_detr_fast.py b/src/transformers/models/conditional_detr/image_processing_conditional_detr_fast.py
index 5b9fe6325517..351d4fa1470f 100644
--- a/src/transformers/models/conditional_detr/image_processing_conditional_detr_fast.py
+++ b/src/transformers/models/conditional_detr/image_processing_conditional_detr_fast.py
@@ -10,6 +10,7 @@
import torch
from torch import nn
from torchvision.io import read_image
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -33,7 +34,7 @@
validate_annotations,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, logging
+from ...utils import TensorType, auto_docstring, logging
from ...utils.import_utils import requires
from .image_processing_conditional_detr import (
compute_segments,
@@ -43,12 +44,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
@@ -433,13 +428,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/convnext/image_processing_convnext_fast.py b/src/transformers/models/convnext/image_processing_convnext_fast.py
index a1002d950399..3ab00c0fd091 100644
--- a/src/transformers/models/convnext/image_processing_convnext_fast.py
+++ b/src/transformers/models/convnext/image_processing_convnext_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -37,16 +38,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class ConvNextFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
crop_pct (`float`, *optional*):
diff --git a/src/transformers/models/cpmant/modeling_cpmant.py b/src/transformers/models/cpmant/modeling_cpmant.py
index 1930cc0e8793..15881a64eb37 100755
--- a/src/transformers/models/cpmant/modeling_cpmant.py
+++ b/src/transformers/models/cpmant/modeling_cpmant.py
@@ -351,7 +351,7 @@ def forward(
output_hidden_states: Optional[bool] = None,
past_key_values: Optional[Cache] = None,
use_cache: Optional[bool] = None,
- cache_postion: Optional[torch.Tensor] = None,
+ cache_position: Optional[torch.Tensor] = None,
):
"""
Args:
@@ -492,16 +492,16 @@ def _position_bucket(self, relative_position, num_buckets=32, max_distance=128):
relative_position = torch.abs(relative_position)
max_exact = num_buckets // 2
is_small = relative_position < max_exact
- relative_postion_if_large = max_exact + (
+ relative_position_if_large = max_exact + (
torch.log(relative_position.float() / max_exact)
/ math.log(max_distance / max_exact)
* (num_buckets - max_exact)
).to(torch.int32)
- relative_postion_if_large = torch.min(
- relative_postion_if_large,
- torch.full_like(relative_postion_if_large, num_buckets - 1),
+ relative_position_if_large = torch.min(
+ relative_position_if_large,
+ torch.full_like(relative_position_if_large, num_buckets - 1),
)
- relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_postion_if_large)
+ relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_position_if_large)
return relative_buckets
diff --git a/src/transformers/models/csm/generation_csm.py b/src/transformers/models/csm/generation_csm.py
index 400c023e0284..cf8bc141f5d1 100644
--- a/src/transformers/models/csm/generation_csm.py
+++ b/src/transformers/models/csm/generation_csm.py
@@ -15,7 +15,7 @@
import os
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
import torch
import torch.nn as nn
@@ -90,7 +90,7 @@ def _get_stopping_criteria(
return kept_criteria
def _prepare_generation_config(
- self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: dict
+ self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: Any
) -> tuple[GenerationConfig, dict]:
"""
This method overrides [~generation.utils.GenerationMixin._prepare_generation_config].
diff --git a/src/transformers/models/csm/processing_csm.py b/src/transformers/models/csm/processing_csm.py
index 0f929f6a2a0c..95596f4a3a9e 100644
--- a/src/transformers/models/csm/processing_csm.py
+++ b/src/transformers/models/csm/processing_csm.py
@@ -152,7 +152,6 @@ def _get_encoded_length(audio_length, kernel_sizes=None, strides=None, dilations
padding_left = padding_total
padding_right = extra_padding
else:
- padding_left = padding_left
padding_right = padding_right + extra_padding
cur_length = cur_length + padding_left + padding_right
diff --git a/src/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py b/src/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py
index f65389d1d18a..5752c1fb7aa9 100644
--- a/src/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py
+++ b/src/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py
@@ -283,11 +283,9 @@ def convert_cvt_checkpoint(cvt_model, image_size, cvt_file_name, pytorch_dump_fo
num_labels = 1000
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.loads(Path(hf_hub_download(repo_id, img_labels_file, repo_type="dataset")).read_text())
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
config = CvtConfig(num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/d_fine/modeling_d_fine.py b/src/transformers/models/d_fine/modeling_d_fine.py
index 5cc2f5e221d1..cdc008e3c7bb 100644
--- a/src/transformers/models/d_fine/modeling_d_fine.py
+++ b/src/transformers/models/d_fine/modeling_d_fine.py
@@ -459,6 +459,12 @@ def _init_weights(self, module):
nn.init.constant_(layer.layers[-1].weight, 0)
nn.init.constant_(layer.layers[-1].bias, 0)
+ if hasattr(module, "reg_scale"):
+ module.reg_scale.fill_(self.config.reg_scale)
+
+ if hasattr(module, "up"):
+ module.up.fill_(self.config.up)
+
if isinstance(module, DFineMultiscaleDeformableAttention):
nn.init.constant_(module.sampling_offsets.weight.data, 0.0)
default_dtype = torch.get_default_dtype()
@@ -496,6 +502,10 @@ def _init_weights(self, module):
init.constant_(module.reg_conf.layers[-1].bias, 0)
init.constant_(module.reg_conf.layers[-1].weight, 0)
+ if isinstance(module, nn.LayerNorm):
+ module.weight.data.fill_(1.0)
+ module.bias.data.zero_()
+
if hasattr(module, "weight_embedding") and self.config.learn_initial_query:
nn.init.xavier_uniform_(module.weight_embedding.weight)
if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0:
@@ -1833,8 +1843,6 @@ def __init__(
self, config: DFineConfig, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 1.0
):
super().__init__()
- in_channels = in_channels
- out_channels = out_channels
activation = config.activation_function
hidden_channels = int(out_channels * expansion)
diff --git a/src/transformers/models/d_fine/modular_d_fine.py b/src/transformers/models/d_fine/modular_d_fine.py
index 52ac7fef7b0d..9a41fb23308e 100644
--- a/src/transformers/models/d_fine/modular_d_fine.py
+++ b/src/transformers/models/d_fine/modular_d_fine.py
@@ -635,6 +635,12 @@ def _init_weights(self, module):
nn.init.constant_(layer.layers[-1].weight, 0)
nn.init.constant_(layer.layers[-1].bias, 0)
+ if hasattr(module, "reg_scale"):
+ module.reg_scale.fill_(self.config.reg_scale)
+
+ if hasattr(module, "up"):
+ module.up.fill_(self.config.up)
+
if isinstance(module, DFineMultiscaleDeformableAttention):
nn.init.constant_(module.sampling_offsets.weight.data, 0.0)
default_dtype = torch.get_default_dtype()
@@ -672,6 +678,10 @@ def _init_weights(self, module):
init.constant_(module.reg_conf.layers[-1].bias, 0)
init.constant_(module.reg_conf.layers[-1].weight, 0)
+ if isinstance(module, nn.LayerNorm):
+ module.weight.data.fill_(1.0)
+ module.bias.data.zero_()
+
if hasattr(module, "weight_embedding") and self.config.learn_initial_query:
nn.init.xavier_uniform_(module.weight_embedding.weight)
if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0:
@@ -1100,8 +1110,6 @@ def __init__(
self, config: DFineConfig, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 1.0
):
super().__init__()
- in_channels = in_channels
- out_channels = out_channels
activation = config.activation_function
hidden_channels = int(out_channels * expansion)
diff --git a/src/transformers/models/deberta_v2/modeling_deberta_v2.py b/src/transformers/models/deberta_v2/modeling_deberta_v2.py
index 9d06f00c0ce6..dd04dd947738 100644
--- a/src/transformers/models/deberta_v2/modeling_deberta_v2.py
+++ b/src/transformers/models/deberta_v2/modeling_deberta_v2.py
@@ -253,7 +253,6 @@ def forward(
if rel_att is not None:
attention_scores = attention_scores + rel_att
- attention_scores = attention_scores
attention_scores = attention_scores.view(
-1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
)
@@ -914,7 +913,7 @@ def forward(self, sequence_output, word_embeddings):
@auto_docstring
class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
_tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]
- _keys_to_ignore_on_load_unexpected = r"mask_predictions.*"
+ _keys_to_ignore_on_load_unexpected = [r"mask_predictions.*"]
def __init__(self, config):
super().__init__(config)
diff --git a/src/transformers/models/deepseek_vl/image_processing_deepseek_vl.py b/src/transformers/models/deepseek_vl/image_processing_deepseek_vl.py
index 7ab4e98012ac..12aa7caf892e 100644
--- a/src/transformers/models/deepseek_vl/image_processing_deepseek_vl.py
+++ b/src/transformers/models/deepseek_vl/image_processing_deepseek_vl.py
@@ -38,12 +38,7 @@
valid_images,
validate_preprocess_arguments,
)
-from ...utils import (
- TensorType,
- filter_out_non_signature_kwargs,
- is_vision_available,
- logging,
-)
+from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
if is_vision_available():
@@ -358,7 +353,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
diff --git a/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py b/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py
index 22d8e0928a6e..ce884da8d08b 100644
--- a/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py
+++ b/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py
@@ -29,11 +29,7 @@
from ...modeling_outputs import ModelOutput
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
-from ...utils import (
- TransformersKwargs,
- auto_docstring,
- can_return_tuple,
-)
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
from ..auto import AutoModel
from .configuration_deepseek_vl import DeepseekVLConfig
diff --git a/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py b/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py
index 7c7d6df82424..865e13fa964f 100644
--- a/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py
+++ b/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid.py
@@ -39,12 +39,7 @@
valid_images,
validate_preprocess_arguments,
)
-from ...utils import (
- TensorType,
- filter_out_non_signature_kwargs,
- is_vision_available,
- logging,
-)
+from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
if is_vision_available():
@@ -431,7 +426,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
diff --git a/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid_fast.py b/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid_fast.py
index db9c9ad987c1..c04e006e358d 100644
--- a/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid_fast.py
+++ b/src/transformers/models/deepseek_vl_hybrid/image_processing_deepseek_vl_hybrid_fast.py
@@ -21,6 +21,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -39,13 +40,7 @@
pil_torch_interpolation_mapping,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import TensorType, auto_docstring
class DeepseekVLHybridFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
diff --git a/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py b/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py
index cae509e14d64..d9a85654e901 100644
--- a/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py
+++ b/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py
@@ -29,11 +29,7 @@
from ...modeling_outputs import ModelOutput
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
-from ...utils import (
- TransformersKwargs,
- auto_docstring,
- can_return_tuple,
-)
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
from ..auto import AutoModel
from .configuration_deepseek_vl_hybrid import DeepseekVLHybridConfig
diff --git a/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py b/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py
index d97b00f7fbd2..0da40603c2e9 100644
--- a/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py
+++ b/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py
@@ -16,6 +16,7 @@
import torch
import torch.nn as nn
+from torchvision.transforms.v2 import functional as F
from ...cache_utils import Cache
from ...image_processing_utils_fast import (
@@ -53,7 +54,6 @@
auto_docstring,
can_return_tuple,
filter_out_non_signature_kwargs,
- is_torchvision_v2_available,
logging,
)
from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel
@@ -70,12 +70,6 @@
from ..sam.modeling_sam import SamLayerNorm, SamVisionNeck
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/deformable_detr/image_processing_deformable_detr_fast.py b/src/transformers/models/deformable_detr/image_processing_deformable_detr_fast.py
index cd07f8db350b..8458d02d58a5 100644
--- a/src/transformers/models/deformable_detr/image_processing_deformable_detr_fast.py
+++ b/src/transformers/models/deformable_detr/image_processing_deformable_detr_fast.py
@@ -9,6 +9,7 @@
import torch
from torchvision.io import read_image
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -32,17 +33,11 @@
validate_annotations,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, logging
+from ...utils import TensorType, auto_docstring, logging
from ...utils.import_utils import requires
from .image_processing_deformable_detr import get_size_with_aspect_ratio
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
@@ -427,13 +422,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/deprecated/gptsan_japanese/tokenization_gptsan_japanese.py b/src/transformers/models/deprecated/gptsan_japanese/tokenization_gptsan_japanese.py
index c67b27f64fa1..1025fdf75fb4 100644
--- a/src/transformers/models/deprecated/gptsan_japanese/tokenization_gptsan_japanese.py
+++ b/src/transformers/models/deprecated/gptsan_japanese/tokenization_gptsan_japanese.py
@@ -495,7 +495,7 @@ def checku2e(x):
candidates.append((self.vocab[wd], wd, e))
if len(candidates) > 0:
# the smallest token_id is adopted
- _, wd, e = sorted(candidates, key=lambda x: x[0])[0]
+ _, wd, e = min(candidates, key=lambda x: x[0])
result.append(wd)
pos = e
else:
diff --git a/src/transformers/models/deprecated/graphormer/collating_graphormer.py b/src/transformers/models/deprecated/graphormer/collating_graphormer.py
index 19bcaac3f572..88657bab435d 100644
--- a/src/transformers/models/deprecated/graphormer/collating_graphormer.py
+++ b/src/transformers/models/deprecated/graphormer/collating_graphormer.py
@@ -14,7 +14,7 @@
import pyximport
pyximport.install(setup_args={"include_dirs": np.get_include()})
- from . import algos_graphormer # noqa E402
+ from . import algos_graphormer
def convert_to_single_emb(x, offset: int = 512):
diff --git a/src/transformers/models/deprecated/mctct/modeling_mctct.py b/src/transformers/models/deprecated/mctct/modeling_mctct.py
index 253b09c1c43c..16f59d3d1dfa 100755
--- a/src/transformers/models/deprecated/mctct/modeling_mctct.py
+++ b/src/transformers/models/deprecated/mctct/modeling_mctct.py
@@ -96,7 +96,7 @@ def __init__(self, config):
def forward(self, input_features):
# NOTE: in reference to the NOTE in __init__, right now it just calculates padding as if
# there will be just one conv layer.
- padding = sum([size // 2 for size in self.kernel_size]) # (7, 7) -> (3, 3)
+ padding = sum(size // 2 for size in self.kernel_size) # (7, 7) -> (3, 3)
input_features = torch.nn.functional.pad(input_features, (0, 0, padding, padding), "constant", 0)
hidden_states = input_features.transpose(1, 2).contiguous() # -> Batch x Frame x Time
diff --git a/src/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py b/src/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py
index 19c3fb0bd485..49d07391320d 100644
--- a/src/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py
+++ b/src/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py
@@ -558,8 +558,8 @@ def _get_new_num_tokens_layer(self, new_num_tokens, layer):
new_num_tokens_layer = (
new_num_tokens
- - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[:layer]])
- - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[layer + 1 :]])
+ - sum(emb.weight.shape[0] for emb in embeddings.emb_layers[:layer])
+ - sum(emb.weight.shape[0] for emb in embeddings.emb_layers[layer + 1 :])
)
return new_num_tokens_layer, layer
diff --git a/src/transformers/models/deprecated/tvlt/feature_extraction_tvlt.py b/src/transformers/models/deprecated/tvlt/feature_extraction_tvlt.py
index 3c65f4314616..b9350d31a019 100644
--- a/src/transformers/models/deprecated/tvlt/feature_extraction_tvlt.py
+++ b/src/transformers/models/deprecated/tvlt/feature_extraction_tvlt.py
@@ -202,7 +202,7 @@ def __call__(
# Create audio attention mask
max_patch_len = max(
- [ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len for feature in audio_features]
+ ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len for feature in audio_features
) # The maximum number of audio patches in a batch
if return_attention_mask:
audio_mask = [
diff --git a/src/transformers/models/deprecated/tvlt/image_processing_tvlt.py b/src/transformers/models/deprecated/tvlt/image_processing_tvlt.py
index c0e1a33f091b..01fb42429a96 100644
--- a/src/transformers/models/deprecated/tvlt/image_processing_tvlt.py
+++ b/src/transformers/models/deprecated/tvlt/image_processing_tvlt.py
@@ -395,7 +395,7 @@ def preprocess(
f"number of frames must not be greater than the maximum frames of the model {self.num_frames}."
)
- max_num_frames = max([len(video) for video in videos])
+ max_num_frames = max(len(video) for video in videos)
num_patches_per_image = (size["shortest_edge"] // patch_size[0]) ** 2
video_masks = np.array(
[
diff --git a/src/transformers/models/deprecated/van/convert_van_to_pytorch.py b/src/transformers/models/deprecated/van/convert_van_to_pytorch.py
index ec43af68d76c..9f97d1c0c296 100644
--- a/src/transformers/models/deprecated/van/convert_van_to_pytorch.py
+++ b/src/transformers/models/deprecated/van/convert_van_to_pytorch.py
@@ -168,11 +168,9 @@ def convert_weights_and_push(save_directory: Path, model_name: Optional[str] = N
num_labels = 1000
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
ImageNetPreTrainedConfig = partial(VanConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/deprecated/xlm_prophetnet/modeling_xlm_prophetnet.py b/src/transformers/models/deprecated/xlm_prophetnet/modeling_xlm_prophetnet.py
index 3c4dc3de8393..36f6e6097bc3 100644
--- a/src/transformers/models/deprecated/xlm_prophetnet/modeling_xlm_prophetnet.py
+++ b/src/transformers/models/deprecated/xlm_prophetnet/modeling_xlm_prophetnet.py
@@ -1233,7 +1233,7 @@ class XLMProphetNetEncoder(XLMProphetNetPreTrainedModel):
embeddings instead of randomly initialized word embeddings.
"""
- def __init__(self, config: XLMProphetNetConfig, word_embeddings: nn.Embedding = None):
+ def __init__(self, config: XLMProphetNetConfig, word_embeddings: Optional[nn.Embedding] = None):
super().__init__(config)
self.word_embeddings = (
diff --git a/src/transformers/models/depth_pro/configuration_depth_pro.py b/src/transformers/models/depth_pro/configuration_depth_pro.py
index 6bc14a0e154f..69bfffeb93f1 100644
--- a/src/transformers/models/depth_pro/configuration_depth_pro.py
+++ b/src/transformers/models/depth_pro/configuration_depth_pro.py
@@ -188,7 +188,6 @@ def __init__(
sub_config.update({"image_size": patch_size})
sub_config = CONFIG_MAPPING[sub_config["model_type"]](**sub_config)
elif isinstance(sub_config, PretrainedConfig):
- sub_config = sub_config
image_size = getattr(sub_config, "image_size", None)
if image_size != patch_size:
raise ValueError(
diff --git a/src/transformers/models/depth_pro/image_processing_depth_pro_fast.py b/src/transformers/models/depth_pro/image_processing_depth_pro_fast.py
index 76c1a53e0073..bc621e0ffc26 100644
--- a/src/transformers/models/depth_pro/image_processing_depth_pro_fast.py
+++ b/src/transformers/models/depth_pro/image_processing_depth_pro_fast.py
@@ -30,7 +30,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
requires_backends,
)
@@ -41,10 +40,7 @@
from .modeling_depth_pro import DepthProDepthEstimatorOutput
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from torchvision.transforms.v2 import functional as F
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/depth_pro/modeling_depth_pro.py b/src/transformers/models/depth_pro/modeling_depth_pro.py
index 52de04d42df7..7c32703b7c25 100644
--- a/src/transformers/models/depth_pro/modeling_depth_pro.py
+++ b/src/transformers/models/depth_pro/modeling_depth_pro.py
@@ -299,7 +299,6 @@ def forward(
scaled_images_features = []
for i in range(self.n_scaled_images):
hidden_state = scaled_images_last_hidden_state[i]
- batch_size = batch_size
padding = torch_int(self.merge_padding_value * (1 / self.scaled_images_ratios[i]))
output_height = base_height * 2**i
output_width = base_width * 2**i
diff --git a/src/transformers/models/detr/image_processing_detr_fast.py b/src/transformers/models/detr/image_processing_detr_fast.py
index 96a89a98074c..ffe040898497 100644
--- a/src/transformers/models/detr/image_processing_detr_fast.py
+++ b/src/transformers/models/detr/image_processing_detr_fast.py
@@ -23,6 +23,7 @@
import torch
from torch import nn
from torchvision.io import read_image
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -49,7 +50,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from ...utils.import_utils import requires
@@ -61,12 +61,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC)
@@ -450,13 +444,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/dia/generation_dia.py b/src/transformers/models/dia/generation_dia.py
index bf18c775eed6..c297de7203d4 100644
--- a/src/transformers/models/dia/generation_dia.py
+++ b/src/transformers/models/dia/generation_dia.py
@@ -109,7 +109,7 @@ def _get_logits_processor(
return merged_processors
def _prepare_generation_config(
- self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: dict
+ self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: Any
) -> tuple[GenerationConfig, dict]:
generation_config, model_kwargs = super()._prepare_generation_config(
generation_config, use_model_defaults, **kwargs
diff --git a/src/transformers/models/diffllama/modular_diffllama.py b/src/transformers/models/diffllama/modular_diffllama.py
index fc0b7a9172d3..253b99edff0d 100644
--- a/src/transformers/models/diffllama/modular_diffllama.py
+++ b/src/transformers/models/diffllama/modular_diffllama.py
@@ -439,7 +439,7 @@ class DiffLlamaForTokenClassification(LlamaForTokenClassification):
__all__ = [
"DiffLlamaPreTrainedModel",
- "DiffLlamaModel", # noqa: F822
+ "DiffLlamaModel",
"DiffLlamaForCausalLM",
"DiffLlamaForSequenceClassification",
"DiffLlamaForQuestionAnswering",
diff --git a/src/transformers/models/dinov3_vit/image_processing_dinov3_vit_fast.py b/src/transformers/models/dinov3_vit/image_processing_dinov3_vit_fast.py
index cdb68044bfc4..7c080485ed00 100644
--- a/src/transformers/models/dinov3_vit/image_processing_dinov3_vit_fast.py
+++ b/src/transformers/models/dinov3_vit/image_processing_dinov3_vit_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from transformers.image_processing_base import BatchFeature
from transformers.image_processing_utils_fast import BaseImageProcessorFast, group_images_by_shape, reorder_images
@@ -24,17 +25,11 @@
from transformers.utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from transformers.utils.import_utils import requires
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/donut/image_processing_donut_fast.py b/src/transformers/models/donut/image_processing_donut_fast.py
index 7c808ab60cd4..29e06831b1b4 100644
--- a/src/transformers/models/donut/image_processing_donut_fast.py
+++ b/src/transformers/models/donut/image_processing_donut_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import group_images_by_shape, reorder_images
@@ -25,16 +26,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/dpt/configuration_dpt.py b/src/transformers/models/dpt/configuration_dpt.py
index 70e46f232022..311425fcda1c 100644
--- a/src/transformers/models/dpt/configuration_dpt.py
+++ b/src/transformers/models/dpt/configuration_dpt.py
@@ -202,9 +202,7 @@ def __init__(
if isinstance(backbone_config, dict):
logger.info("Initializing the config with a `BiT` backbone.")
backbone_config = BitConfig(**backbone_config)
- elif isinstance(backbone_config, PretrainedConfig):
- backbone_config = backbone_config
- else:
+ elif not isinstance(backbone_config, PretrainedConfig):
raise ValueError(
f"backbone_config must be a dictionary or a `PretrainedConfig`, got {backbone_config.__class__}."
)
diff --git a/src/transformers/models/dpt/image_processing_dpt_fast.py b/src/transformers/models/dpt/image_processing_dpt_fast.py
index d4848c50653c..faaddb8023c0 100644
--- a/src/transformers/models/dpt/image_processing_dpt_fast.py
+++ b/src/transformers/models/dpt/image_processing_dpt_fast.py
@@ -25,6 +25,7 @@
from typing import TYPE_CHECKING, Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_base import BatchFeature
from ...image_processing_utils_fast import BaseImageProcessorFast, DefaultFastImageProcessorKwargs
@@ -39,17 +40,12 @@
is_torch_tensor,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, requires_backends
+from ...utils import TensorType, auto_docstring, requires_backends
if TYPE_CHECKING:
from ...modeling_outputs import DepthEstimatorOutput
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
class DPTFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
diff --git a/src/transformers/models/dpt/modeling_dpt.py b/src/transformers/models/dpt/modeling_dpt.py
index 363fce92f897..cef10dd76eda 100755
--- a/src/transformers/models/dpt/modeling_dpt.py
+++ b/src/transformers/models/dpt/modeling_dpt.py
@@ -879,7 +879,7 @@ def __init__(self, config: DPTConfig):
self.config = config
# postprocessing: only required in case of a non-hierarchical backbone (e.g. ViT, BEiT)
- if config.backbone_config is not None and config.backbone_config.model_type in ["swinv2"]:
+ if config.backbone_config is not None and config.backbone_config.model_type == "swinv2":
self.reassemble_stage = None
else:
self.reassemble_stage = DPTReassembleStage(config)
diff --git a/src/transformers/models/dpt/modular_dpt.py b/src/transformers/models/dpt/modular_dpt.py
index 32ca94a2d43f..34eb08f39b68 100644
--- a/src/transformers/models/dpt/modular_dpt.py
+++ b/src/transformers/models/dpt/modular_dpt.py
@@ -32,7 +32,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
requires_backends,
)
from ..beit.image_processing_beit_fast import BeitImageProcessorFast
@@ -41,10 +40,7 @@
if TYPE_CHECKING:
from ...modeling_outputs import DepthEstimatorOutput
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from torchvision.transforms.v2 import functional as F
def get_resize_output_image_size(
diff --git a/src/transformers/models/edgetam/__init__.py b/src/transformers/models/edgetam/__init__.py
new file mode 100644
index 000000000000..d9c1a55fc5bc
--- /dev/null
+++ b/src/transformers/models/edgetam/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_edgetam import *
+ from .modeling_edgetam import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/edgetam/configuration_edgetam.py b/src/transformers/models/edgetam/configuration_edgetam.py
new file mode 100644
index 000000000000..07ccee36e932
--- /dev/null
+++ b/src/transformers/models/edgetam/configuration_edgetam.py
@@ -0,0 +1,332 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/edgetam/modular_edgetam.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_edgetam.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Meta AI Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_utils import PretrainedConfig
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+class EdgeTamVisionConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamVisionModel`]. It is used to instantiate a SAM
+ vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
+ defaults will yield a similar configuration to that of SAM 2.1 Hiera-tiny
+ [facebook/EdgeTAM](https://huggingface.co/facebook/EdgeTAM) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ backbone_config (`Union[dict, "PretrainedConfig"]`, *optional*):
+ Configuration for the vision backbone. This is used to instantiate the backbone using
+ `AutoModel.from_config`.
+ backbone_channel_list (`List[int]`, *optional*, defaults to `[384, 192, 96, 48]`):
+ The list of channel dimensions for the backbone.
+ backbone_feature_sizes (`List[List[int]]`, *optional*, defaults to `[[256, 256], [128, 128], [64, 64]]`):
+ The spatial sizes of the feature maps from the backbone.
+ fpn_hidden_size (`int`, *optional*, defaults to 256):
+ The hidden dimension of the FPN.
+ fpn_kernel_size (`int`, *optional*, defaults to 1):
+ The kernel size for the convolutions in the neck.
+ fpn_stride (`int`, *optional*, defaults to 1):
+ The stride for the convolutions in the neck.
+ fpn_padding (`int`, *optional*, defaults to 0):
+ The padding for the convolutions in the neck.
+ fpn_top_down_levels (`List[int]`, *optional*, defaults to `[2, 3]`):
+ The levels for the top-down FPN connections.
+ num_feature_levels (`int`, *optional*, defaults to 3):
+ The number of feature levels from the FPN to use.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the neck.
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon for the layer normalization.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+ """
+
+ base_config_key = "vision_config"
+ model_type = "edgetam_vision_model"
+ sub_configs = {
+ "backbone_config": AutoConfig,
+ }
+
+ def __init__(
+ self,
+ backbone_config=None,
+ backbone_channel_list=None,
+ backbone_feature_sizes=None,
+ fpn_hidden_size=256,
+ fpn_kernel_size=1,
+ fpn_stride=1,
+ fpn_padding=0,
+ fpn_top_down_levels=None,
+ num_feature_levels=3,
+ hidden_act="gelu",
+ layer_norm_eps=1e-6,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ backbone_channel_list = [384, 192, 96, 48] if backbone_channel_list is None else backbone_channel_list
+ backbone_feature_sizes = (
+ [[256, 256], [128, 128], [64, 64]] if backbone_feature_sizes is None else backbone_feature_sizes
+ )
+ fpn_top_down_levels = [2, 3] if fpn_top_down_levels is None else fpn_top_down_levels
+
+ if isinstance(backbone_config, dict):
+ backbone_config["model_type"] = backbone_config.get("model_type", "timm_wrapper")
+ backbone_config = CONFIG_MAPPING[backbone_config["model_type"]](**backbone_config)
+ elif isinstance(backbone_config, AutoConfig):
+ backbone_config = backbone_config
+ elif backbone_config is None:
+ backbone_config = AutoConfig.from_pretrained(
+ "timm/repvit_m1.dist_in1k",
+ model_args={"in_chans": 3, "features_only": True, "out_indices": [0, 1, 2, 3]},
+ )
+
+ self.backbone_config = backbone_config
+
+ # Neck
+ self.backbone_channel_list = backbone_channel_list
+ self.backbone_feature_sizes = backbone_feature_sizes
+ self.fpn_hidden_size = fpn_hidden_size
+ self.fpn_kernel_size = fpn_kernel_size
+ self.fpn_stride = fpn_stride
+ self.fpn_padding = fpn_padding
+ self.fpn_top_down_levels = fpn_top_down_levels
+ self.num_feature_levels = num_feature_levels
+
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.initializer_range = initializer_range
+
+
+class EdgeTamPromptEncoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamPromptEncoder`]. The [`EdgeTamPromptEncoder`]
+ module is used to encode the input 2D points and bounding boxes.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the hidden states.
+ image_size (`int`, *optional*, defaults to 1024):
+ The expected output resolution of the image.
+ patch_size (`int`, *optional*, defaults to 16):
+ The size (resolution) of each patch.
+ mask_input_channels (`int`, *optional*, defaults to 16):
+ The number of channels to be fed to the `MaskDecoder` module.
+ num_point_embeddings (`int`, *optional*, defaults to 4):
+ The number of point embeddings to be used.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the encoder and pooler.
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the layer normalization layers.
+ scale (`float`, *optional*, defaults to 1):
+ The scale factor for the prompt encoder.
+ """
+
+ base_config_key = "prompt_encoder_config"
+
+ def __init__(
+ self,
+ hidden_size=256,
+ image_size=1024,
+ patch_size=16,
+ mask_input_channels=16,
+ num_point_embeddings=4,
+ hidden_act="gelu",
+ layer_norm_eps=1e-6,
+ scale=1,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.hidden_size = hidden_size
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.mask_input_channels = mask_input_channels
+ self.num_point_embeddings = num_point_embeddings
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.scale = scale
+
+
+class EdgeTamMaskDecoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamMaskDecoder`]. It is used to instantiate a EDGETAM
+ memory encoder according to the specified arguments, defining the model architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the hidden states.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the EDGETAM mask decoder.
+ mlp_dim (`int`, *optional*, defaults to 2048):
+ The dimension of the MLP in the two-way transformer.
+ num_hidden_layers (`int`, *optional*, defaults to 2):
+ The number of hidden layers in the two-way transformer.
+ num_attention_heads (`int`, *optional*, defaults to 8):
+ The number of attention heads in the two-way transformer.
+ attention_downsample_rate (`int`, *optional*, defaults to 2):
+ The downsample rate for the attention layers.
+ num_multimask_outputs (`int`, *optional*, defaults to 3):
+ The number of multimask outputs.
+ iou_head_depth (`int`, *optional*, defaults to 3):
+ The depth of the IoU head.
+ iou_head_hidden_dim (`int`, *optional*, defaults to 256):
+ The hidden dimension of the IoU head.
+ dynamic_multimask_via_stability (`bool`, *optional*, defaults to `True`):
+ Whether to use dynamic multimask via stability.
+ dynamic_multimask_stability_delta (`float`, *optional*, defaults to 0.05):
+ The stability delta for the dynamic multimask.
+ dynamic_multimask_stability_thresh (`float`, *optional*, defaults to 0.98):
+ The stability threshold for the dynamic multimask.
+
+ """
+
+ base_config_key = "mask_decoder_config"
+
+ def __init__(
+ self,
+ hidden_size=256,
+ hidden_act="gelu",
+ mlp_dim=2048,
+ num_hidden_layers=2,
+ num_attention_heads=8,
+ attention_downsample_rate=2,
+ num_multimask_outputs=3,
+ iou_head_depth=3,
+ iou_head_hidden_dim=256,
+ dynamic_multimask_via_stability=True,
+ dynamic_multimask_stability_delta=0.05,
+ dynamic_multimask_stability_thresh=0.98,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.hidden_size = hidden_size
+ self.num_multimask_outputs = num_multimask_outputs
+ self.hidden_act = hidden_act
+ self.iou_head_depth = iou_head_depth
+ self.iou_head_hidden_dim = iou_head_hidden_dim
+ self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
+ self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
+ self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
+
+ # TwoWayTransformer configuration
+ self.num_hidden_layers = num_hidden_layers
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.mlp_dim = mlp_dim
+ self.attention_downsample_rate = attention_downsample_rate
+
+
+class EdgeTamConfig(PretrainedConfig):
+ r"""
+ [`EdgeTamConfig`] is the configuration class to store the configuration of a [`EdgeTamModel`]. It is used to instantiate a
+ EDGETAM model according to the specified arguments, defining the memory attention, memory encoder, and image encoder
+ configs. Instantiating a configuration defaults will yield a similar configuration to that of the SAM 2.1 Hiera-tiny
+ [facebook/edgetam.1-hiera-tiny](https://huggingface.co/facebook/edgetam.1-hiera-tiny) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vision_config (Union[`dict`, `EdgeTamVisionConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamVisionConfig`].
+ prompt_encoder_config (Union[`dict`, `EdgeTamPromptEncoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamPromptEncoderConfig`].
+ mask_decoder_config (Union[`dict`, `EdgeTamMaskDecoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamMaskDecoderConfig`].
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ Standard deviation for parameter initialization.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... EdgeTamVisionConfig,
+ ... EdgeTamPromptEncoderConfig,
+ ... EdgeTamMaskDecoderConfig,
+ ... EdgeTamModel,
+ ... )
+
+ >>> # Initializing a EdgeTamConfig with `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> configuration = EdgeTamconfig()
+
+ >>> # Initializing a EdgeTamModel (with random weights) from the `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> model = EdgeTamModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a EdgeTamConfig from a EdgeTamVisionConfig, EdgeTamPromptEncoderConfig, and EdgeTamMaskDecoderConfig
+
+ >>> # Initializing EDGETAM vision encoder, memory attention, and memory encoder configurations
+ >>> vision_config = EdgeTamVisionConfig()
+ >>> prompt_encoder_config = EdgeTamPromptEncoderConfig()
+ >>> mask_decoder_config = EdgeTamMaskDecoderConfig()
+
+ >>> config = EdgeTamConfig(vision_config, prompt_encoder_config, mask_decoder_config)
+ ```"""
+
+ model_type = "edgetam"
+ sub_configs = {
+ "vision_config": AutoConfig,
+ "prompt_encoder_config": EdgeTamPromptEncoderConfig,
+ "mask_decoder_config": EdgeTamMaskDecoderConfig,
+ }
+
+ def __init__(
+ self,
+ vision_config=None,
+ prompt_encoder_config=None,
+ mask_decoder_config=None,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ vision_config = vision_config if vision_config is not None else {}
+ prompt_encoder_config = prompt_encoder_config if prompt_encoder_config is not None else {}
+ mask_decoder_config = mask_decoder_config if mask_decoder_config is not None else {}
+
+ if isinstance(vision_config, dict):
+ vision_config["model_type"] = vision_config.get("model_type", "edgetam_vision_model")
+ vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
+ if isinstance(prompt_encoder_config, EdgeTamPromptEncoderConfig):
+ prompt_encoder_config = prompt_encoder_config.to_dict()
+ if isinstance(mask_decoder_config, EdgeTamMaskDecoderConfig):
+ mask_decoder_config = mask_decoder_config.to_dict()
+
+ self.vision_config = vision_config
+ self.prompt_encoder_config = EdgeTamPromptEncoderConfig(**prompt_encoder_config)
+ self.mask_decoder_config = EdgeTamMaskDecoderConfig(**mask_decoder_config)
+
+ self.initializer_range = initializer_range
+
+
+__all__ = ["EdgeTamConfig", "EdgeTamVisionConfig", "EdgeTamPromptEncoderConfig", "EdgeTamMaskDecoderConfig"]
diff --git a/src/transformers/models/edgetam/convert_edgetam_to_hf.py b/src/transformers/models/edgetam/convert_edgetam_to_hf.py
new file mode 100644
index 000000000000..382bc1559ec4
--- /dev/null
+++ b/src/transformers/models/edgetam/convert_edgetam_to_hf.py
@@ -0,0 +1,280 @@
+# coding=utf-8
+# 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.
+"""
+Convert SAM checkpoints from the original repository.
+
+URL: https://github.com/facebookresearch/segment-anything-2.
+"""
+
+import argparse
+import re
+
+import numpy as np
+import requests
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from transformers import (
+ EdgeTamConfig,
+ EdgeTamMaskDecoderConfig,
+ EdgeTamModel,
+ EdgeTamPromptEncoderConfig,
+ EdgeTamVisionConfig,
+ Sam2ImageProcessorFast,
+ Sam2Processor,
+ TimmWrapperConfig,
+)
+
+
+def get_config(model_name):
+ backbone_config = TimmWrapperConfig.from_pretrained(
+ "timm/repvit_m1.dist_in1k",
+ model_args={"in_chans": 3, "features_only": True, "out_indices": (0, 1, 2, 3)},
+ )
+ vision_config = EdgeTamVisionConfig(backbone_config=backbone_config)
+
+ prompt_encoder_config = EdgeTamPromptEncoderConfig()
+ mask_decoder_config = EdgeTamMaskDecoderConfig()
+ enable_temporal_pos_encoding_for_object_pointers = False
+ project_temporal_pos_encoding_in_object_pointers = False
+ enable_occlusion_spatial_embedding = False
+
+ config = EdgeTamConfig(
+ vision_config=vision_config,
+ prompt_encoder_config=prompt_encoder_config,
+ mask_decoder_config=mask_decoder_config,
+ enable_temporal_pos_encoding_for_object_pointers=enable_temporal_pos_encoding_for_object_pointers,
+ project_temporal_pos_encoding_in_object_pointers=project_temporal_pos_encoding_in_object_pointers,
+ enable_occlusion_spatial_embedding=enable_occlusion_spatial_embedding,
+ )
+
+ return config
+
+
+KEYS_TO_MODIFY_MAPPING = {
+ "iou_prediction_head.layers.0": "iou_prediction_head.proj_in",
+ "iou_prediction_head.layers.1": "iou_prediction_head.layers.0",
+ "iou_prediction_head.layers.2": "iou_prediction_head.proj_out",
+ "mask_decoder.output_upscaling.0": "mask_decoder.upscale_conv1",
+ "mask_decoder.output_upscaling.1": "mask_decoder.upscale_layer_norm",
+ "mask_decoder.output_upscaling.3": "mask_decoder.upscale_conv2",
+ "mask_downscaling.0": "mask_embed.conv1",
+ "mask_downscaling.1": "mask_embed.layer_norm1",
+ "mask_downscaling.3": "mask_embed.conv2",
+ "mask_downscaling.4": "mask_embed.layer_norm2",
+ "mask_downscaling.6": "mask_embed.conv3",
+ "dwconv": "depthwise_conv",
+ "pwconv": "pointwise_conv",
+ "fuser": "memory_fuser",
+ "point_embeddings": "point_embed",
+ "pe_layer.positional_encoding_gaussian_matrix": "shared_embedding.positional_embedding",
+ "obj_ptr_tpos_proj": "temporal_positional_encoding_projection_layer",
+ "no_obj_embed_spatial": "occlusion_spatial_embedding_parameter",
+ "sam_prompt_encoder": "prompt_encoder",
+ "sam_mask_decoder": "mask_decoder",
+ "maskmem_tpos_enc": "memory_temporal_positional_encoding",
+ "gamma": "scale",
+ "image_encoder.neck": "vision_encoder.neck",
+ "image_encoder": "vision_encoder.backbone",
+ "neck.0": "neck.conv1",
+ "neck.1": "neck.layer_norm1",
+ "neck.2": "neck.conv2",
+ "neck.3": "neck.layer_norm2",
+ "pix_feat_proj": "feature_projection",
+ "patch_embed.proj": "patch_embed.projection",
+ "no_mem_embed": "no_memory_embedding",
+ "no_mem_pos_enc": "no_memory_positional_encoding",
+ "obj_ptr": "object_pointer",
+ ".norm": ".layer_norm",
+ "trunk.": "",
+ "out_proj": "o_proj",
+ "body.": "timm_model.",
+ "ff.0": "feed_forward.layer_norm",
+ "ff.1": "feed_forward.linear1",
+ "ff.3": "feed_forward.linear2",
+}
+
+
+def replace_keys(state_dict):
+ model_state_dict = {}
+ output_hypernetworks_mlps_pattern = r".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"
+ output_mask_decoder_mlps_pattern = r"mask_decoder.transformer.layers.(\d+).mlp.layers.(\d+).*"
+ output_mask_decoder_score_head_pattern = r"mask_decoder.pred_obj_score_head.layers.(\d+).*"
+ output_vision_encoder_mlps_pattern = r"vision_encoder.backbone.blocks.(\d+).mlp.layers.(\d+).*"
+ output_vision_encoder_neck_pattern = r"vision_encoder.neck.convs.(\d+).conv"
+ output_memory_encoder_projection_pattern = r"memory_encoder.o_proj.*"
+ output_object_pointer_proj_pattern = r"object_pointer_proj.layers.(\d+).*"
+ for key, value in state_dict.items():
+ 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)
+
+ # vision_encoder.blocks.0.mlp.layers.1.weight -> vision_encoder.blocks.0.mlp.proj_out.weight
+ if re.match(output_vision_encoder_mlps_pattern, key):
+ layer_nb = int(re.match(output_vision_encoder_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "proj_out")
+
+ # mask_decoder.transformer.layers.0.mlp.layers.1.weight -> mask_decoder.transformer.layers.1.mlp.proj_out.weight
+ if re.match(output_mask_decoder_mlps_pattern, key):
+ layer_nb = int(re.match(output_mask_decoder_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("mlp.layers.0", "mlp.proj_in")
+ elif layer_nb == 1:
+ key = key.replace("mlp.layers.1", "mlp.proj_out")
+
+ # mask_decoder.pred_obj_score_head.layers.1.weight -> mask_decoder.pred_obj_score_head.proj_in.weight
+ if re.match(output_mask_decoder_score_head_pattern, key):
+ layer_nb = int(re.match(output_mask_decoder_score_head_pattern, key).group(1))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ if re.match(output_hypernetworks_mlps_pattern, key):
+ layer_nb = int(re.match(output_hypernetworks_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ # vision_encoder.neck.convs.1.conv.bias -> vision_encoder.neck.convs.1.bias
+ if re.match(output_vision_encoder_neck_pattern, key):
+ key = key.replace(".conv.", ".")
+
+ # memory_encoder.o_proj.weight -> memory_encoder.projection.weight
+ if re.match(output_memory_encoder_projection_pattern, key):
+ key = key.replace(".o_proj.", ".projection.")
+
+ if re.match(output_object_pointer_proj_pattern, key):
+ layer_nb = int(re.match(output_object_pointer_proj_pattern, key).group(1))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ key = key.replace("layers.2", "proj_out")
+
+ model_state_dict[key] = value
+
+ model_state_dict["shared_image_embedding.positional_embedding"] = model_state_dict[
+ "prompt_encoder.shared_embedding.positional_embedding"
+ ]
+ model_state_dict["prompt_encoder.point_embed.weight"] = torch.cat(
+ [model_state_dict.pop(f"prompt_encoder.point_embed.{i}.weight") for i in range(4)],
+ dim=0,
+ )
+
+ return model_state_dict
+
+
+def convert_edgetam_checkpoint(model_name, checkpoint_path, pytorch_dump_folder, push_to_hub, run_sanity_check):
+ config = get_config(model_name)
+
+ state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
+ state_dict = replace_keys(state_dict)
+
+ image_processor = Sam2ImageProcessorFast()
+ processor = Sam2Processor(image_processor=image_processor)
+ hf_model = EdgeTamModel(config)
+ hf_model.eval()
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ missing_keys, unexpected_keys = hf_model.load_state_dict(state_dict, strict=False)
+ hf_model = hf_model.to(device)
+ for pattern in EdgeTamModel._keys_to_ignore_on_load_unexpected:
+ unexpected_keys = [k for k in unexpected_keys if re.search(pattern, k) is None]
+ if missing_keys or unexpected_keys:
+ print("Missing keys:", missing_keys)
+ print("Unexpected keys:", unexpected_keys)
+ raise ValueError("Missing or unexpected keys in the state dict")
+
+ if run_sanity_check:
+ img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+
+ input_points = [[[[1000, 600]]]]
+ input_labels = [[[1]]]
+
+ inputs = processor(
+ images=np.array(raw_image), input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(device)
+
+ with torch.no_grad():
+ output = hf_model(**inputs)
+ scores = output.iou_scores.squeeze()
+
+ assert torch.allclose(scores, torch.tensor([0.0356, 0.2141, 0.9707]).cuda(), atol=1e-3)
+
+ if pytorch_dump_folder is not None:
+ processor.save_pretrained(pytorch_dump_folder)
+ hf_model.save_pretrained(pytorch_dump_folder)
+
+ if push_to_hub:
+ repo_id = f"yonigozlan/{pytorch_dump_folder.split('/')[-1]}"
+ processor.push_to_hub(repo_id)
+ hf_model.push_to_hub(repo_id)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ choices = ["EdgeTAM"]
+ parser.add_argument(
+ "--model_name",
+ default="EdgeTAM",
+ choices=choices,
+ type=str,
+ help="Name of the original model to convert",
+ )
+ parser.add_argument(
+ "--checkpoint_path",
+ type=str,
+ required=False,
+ help="Path to the original checkpoint",
+ )
+ parser.add_argument("--pytorch_dump_folder_path", default="", type=str, help="Path to the output PyTorch model.")
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether to push the model and processor to the hub after converting",
+ )
+ parser.add_argument(
+ "--run_sanity_check",
+ action="store_true",
+ help="Whether to run the sanity check after converting",
+ )
+
+ args = parser.parse_args()
+
+ hf_model_name = args.model_name.replace("_", "-")
+ checkpoint_path = (
+ hf_hub_download(f"facebook/{hf_model_name}", f"{args.model_name.lower()}.pt")
+ if args.checkpoint_path is None
+ else args.checkpoint_path
+ )
+
+ convert_edgetam_checkpoint(
+ args.model_name, checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub, args.run_sanity_check
+ )
diff --git a/src/transformers/models/edgetam/modeling_edgetam.py b/src/transformers/models/edgetam/modeling_edgetam.py
new file mode 100644
index 000000000000..d7e3ee6009cf
--- /dev/null
+++ b/src/transformers/models/edgetam/modeling_edgetam.py
@@ -0,0 +1,1252 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/edgetam/modular_edgetam.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_edgetam.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Meta AI Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import Callable, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from transformers.utils.generic import OutputRecorder, TransformersKwargs, check_model_inputs
+
+from ...activations import ACT2FN
+from ...modeling_outputs import BaseModelOutput
+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, auto_docstring
+from ..auto import AutoModel
+from .configuration_edgetam import (
+ EdgeTamConfig,
+ EdgeTamMaskDecoderConfig,
+ EdgeTamPromptEncoderConfig,
+ EdgeTamVisionConfig,
+)
+
+
+# fix this in modular
+if True:
+ from transformers.models.timm_wrapper.modeling_timm_wrapper import TimmWrapperModel
+
+
+class EdgeTamLayerNorm(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
+
+
+@dataclass
+@auto_docstring(custom_intro="Base class for the vision encoder's outputs.")
+class EdgeTamVisionEncoderOutput(ModelOutput):
+ r"""
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the model.
+ fpn_hidden_states (`tuple(torch.FloatTensor)`):
+ Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+ `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck.
+ fpn_position_encoding (`tuple(torch.FloatTensor)`):
+ Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+ `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+ one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the
+ model at the output of each stage.
+ attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
+ the self-attention heads.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ fpn_hidden_states: Optional[torch.FloatTensor] = None
+ fpn_position_encoding: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class EdgeTamAttention(nn.Module):
+ """
+ EDGETAM's attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and
+ values.
+ """
+
+ def __init__(self, config, downsample_rate=None):
+ super().__init__()
+ downsample_rate = config.attention_downsample_rate if downsample_rate is None else downsample_rate
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.internal_dim = config.hidden_size // downsample_rate
+ self.num_attention_heads = config.num_attention_heads
+ self.head_dim = self.internal_dim // config.num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.v_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_similarity: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=attention_similarity,
+ dropout=0.0,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class EdgeTamTwoWayAttentionBlock(nn.Module):
+ def __init__(self, config: EdgeTamMaskDecoderConfig, skip_first_layer_pe: bool = False):
+ """
+ A transformer block with four layers:
+ (1) self-attention of sparse inputs (2) cross attention of sparse inputs -> dense inputs (3) mlp block on
+ sparse inputs (4) cross attention of dense inputs -> sparse inputs
+
+ Arguments:
+ config (`EdgeTamMaskDecoderConfig`):
+ The configuration file used to instantiate the block
+ attention_downsample_rate (*optionalk*, int, defaults to 2):
+ The downsample ratio of the block used to reduce the inner dim of the attention.
+ skip_first_layer_pe (*optional*, bool, defaults to `False`):
+ Whether or not to skip the addition of the query_point_embedding on the first layer.
+ """
+ super().__init__()
+ self.self_attn = EdgeTamAttention(config, downsample_rate=1)
+ self.layer_norm1 = nn.LayerNorm(config.hidden_size)
+
+ self.cross_attn_token_to_image = EdgeTamAttention(config)
+ self.layer_norm2 = nn.LayerNorm(config.hidden_size)
+
+ self.mlp = EdgeTamFeedForward(
+ config.hidden_size, config.mlp_dim, config.hidden_size, num_layers=config.num_hidden_layers
+ )
+ self.layer_norm3 = nn.LayerNorm(config.hidden_size)
+
+ self.layer_norm4 = nn.LayerNorm(config.hidden_size)
+ self.cross_attn_image_to_token = EdgeTamAttention(config)
+
+ self.skip_first_layer_pe = skip_first_layer_pe
+
+ def forward(
+ self,
+ queries: Tensor,
+ keys: Tensor,
+ query_point_embedding: Tensor,
+ key_point_embedding: Tensor,
+ attention_similarity: Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ # Self attention block
+ if self.skip_first_layer_pe:
+ queries, _ = self.self_attn(query=queries, key=queries, value=queries)
+ else:
+ query = queries + query_point_embedding
+ attn_out, _ = self.self_attn(query=query, key=query, value=queries)
+ queries = queries + attn_out
+ queries = self.layer_norm1(queries)
+
+ # Cross attention block, tokens attending to image embedding
+ query = queries + query_point_embedding
+ key = keys + key_point_embedding
+
+ attn_out, _ = self.cross_attn_token_to_image(
+ query=query, key=key, value=keys, attention_similarity=attention_similarity
+ )
+ queries = queries + attn_out
+
+ queries = self.layer_norm2(queries)
+
+ # MLP block
+ mlp_out = self.mlp(queries)
+ queries = queries + mlp_out
+ queries = self.layer_norm3(queries)
+
+ # Cross attention block, image embedding attending to tokens
+ query = queries + query_point_embedding
+ key = keys + key_point_embedding
+
+ attn_out, _ = self.cross_attn_image_to_token(query=key, key=query, value=queries)
+ keys = keys + attn_out
+
+ keys = self.layer_norm4(keys)
+ return queries, keys, attn_out
+
+
+class EdgeTamFeedForward(nn.Module):
+ def __init__(
+ self,
+ input_dim: int,
+ hidden_dim: int,
+ output_dim: int,
+ num_layers: int,
+ activation: str = "relu",
+ sigmoid_output: bool = False,
+ ):
+ super().__init__()
+ self.num_layers = num_layers
+ self.activation = ACT2FN[activation]
+ self.proj_in = nn.Linear(input_dim, hidden_dim)
+ self.proj_out = nn.Linear(hidden_dim, output_dim)
+ self.layers = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers - 2)])
+ self.sigmoid_output = sigmoid_output
+
+ def forward(self, hidden_states):
+ hidden_states = self.proj_in(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ for layer in self.layers:
+ hidden_states = self.activation(layer(hidden_states))
+
+ hidden_states = self.proj_out(hidden_states)
+ if self.sigmoid_output:
+ hidden_states = F.sigmoid(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class EdgeTamPreTrainedModel(PreTrainedModel):
+ config_class = EdgeTamConfig
+ base_model_prefix = "edgetam"
+ main_input_name = "pixel_values"
+ _supports_sdpa = True
+ _supports_flash_attn_2 = True
+ _supports_attention_backend = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, (nn.LayerNorm, EdgeTamLayerNorm)):
+ module.weight.data.fill_(1.0)
+ module.bias.data.zero_()
+ if isinstance(module, EdgeTamModel):
+ if module.no_memory_embedding is not None:
+ module.no_memory_embedding.data.zero_()
+
+
+# copied and adapted from original implementation, also practically equal to DetrSinePositionEmbedding
+class EdgeTamSinePositionEmbedding(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: Optional[float] = 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: Union[torch.device, str],
+ dtype: torch.dtype,
+ mask: Optional[Tensor] = 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 EdgeTamVisionNeck(nn.Module):
+ def __init__(self, config: EdgeTamVisionConfig):
+ super().__init__()
+ self.config = config
+
+ self.position_encoding = EdgeTamSinePositionEmbedding(
+ num_pos_feats=config.fpn_hidden_size // 2, normalize=True
+ )
+ self.convs = nn.ModuleList()
+ for in_channels in config.backbone_channel_list:
+ self.convs.append(
+ nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=config.fpn_hidden_size,
+ kernel_size=config.fpn_kernel_size,
+ stride=config.fpn_stride,
+ padding=config.fpn_padding,
+ ),
+ )
+ self.fpn_top_down_levels = config.fpn_top_down_levels
+
+ def forward(self, hidden_states: torch.Tensor) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]:
+ fpn_hidden_states = ()
+ fpn_position_encoding = ()
+
+ # forward in top-down order (from low to high resolution)
+ n = len(self.convs) - 1
+ for i in range(n, -1, -1):
+ lateral_features = hidden_states[i].permute(0, 3, 1, 2)
+ lateral_features = self.convs[n - i](lateral_features)
+ if i not in self.fpn_top_down_levels or i == n:
+ prev_features = lateral_features
+ else:
+ top_down_features = F.interpolate(
+ prev_features.to(dtype=torch.float32),
+ scale_factor=2.0,
+ mode="nearest",
+ align_corners=None,
+ antialias=False,
+ ).to(lateral_features.dtype)
+ prev_features = lateral_features + top_down_features
+
+ prev_position_encoding = self.position_encoding(
+ prev_features.shape, prev_features.device, prev_features.dtype
+ ).to(prev_features.dtype)
+
+ fpn_hidden_states += (prev_features,)
+ fpn_position_encoding += (prev_position_encoding,)
+
+ return fpn_hidden_states, fpn_position_encoding
+
+
+@auto_docstring(
+ custom_intro="""
+ The vision model from EdgeTAM without any head or projection on top.
+ """
+)
+class EdgeTamVisionModel(EdgeTamPreTrainedModel):
+ config_class = EdgeTamVisionConfig
+ main_input_name = "pixel_values"
+ _can_record_outputs = {"hidden_states": TimmWrapperModel, "attentions": TimmWrapperModel}
+
+ def __init__(self, config: EdgeTamVisionConfig):
+ super().__init__(config)
+ self.config = config
+
+ self.backbone = AutoModel.from_config(config.backbone_config)
+
+ self.neck = EdgeTamVisionNeck(config)
+ self.num_feature_levels = config.num_feature_levels
+
+ self.post_init()
+
+ @check_model_inputs
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, EdgeTamVisionEncoderOutput]:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ # Forward through backbone
+ backbone_output = self.backbone(pixel_values)
+ intermediate_hidden_states = backbone_output.last_hidden_state
+ intermediate_hidden_states = [hidden_state.permute(0, 2, 3, 1) for hidden_state in intermediate_hidden_states]
+
+ fpn_hidden_states, fpn_position_encoding = self.neck(intermediate_hidden_states)
+ # Select last `num_feature_levels` feature levels from FPN and reverse order to get features from high to low resolution
+ fpn_hidden_states = fpn_hidden_states[-self.num_feature_levels :][::-1]
+ fpn_position_encoding = fpn_position_encoding[-self.num_feature_levels :][::-1]
+
+ return EdgeTamVisionEncoderOutput(
+ last_hidden_state=intermediate_hidden_states[-1],
+ fpn_hidden_states=fpn_hidden_states,
+ fpn_position_encoding=fpn_position_encoding,
+ )
+
+
+@dataclass
+@auto_docstring(custom_intro="Base class for the EdgeTam model's output.")
+class EdgeTamImageSegmentationOutput(ModelOutput):
+ r"""
+ iou_scores (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks)`):
+ The Intersection over Union (IoU) scores of the predicted masks.
+ pred_masks (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks, height, width)`):
+ The predicted low-resolution masks. This is an alias for `low_res_masks`. These masks need to be post-processed
+ by the processor to be brought to the original image size.
+ object_score_logits (`torch.FloatTensor` of shape `(batch_size, point_batch_size, 1)`):
+ Logits for the object score, indicating if an object is present.
+ image_embeddings (`tuple(torch.FloatTensor)`):
+ The features from the FPN, which are used by the mask decoder. This is a tuple of `torch.FloatTensor` where each
+ tensor has shape `(batch_size, channels, height, width)`.
+ vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`.
+ Hidden-states of the vision model at the output of each stage.
+ vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
+ Attentions weights of the vision model.
+ mask_decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
+ Attentions weights of the mask decoder.
+ """
+
+ iou_scores: Optional[torch.FloatTensor] = None
+ pred_masks: Optional[torch.FloatTensor] = None
+ object_score_logits: Optional[torch.FloatTensor] = None
+ image_embeddings: tuple[torch.FloatTensor, ...] = None
+ vision_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ vision_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+ mask_decoder_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+class EdgeTamPositionalEmbedding(nn.Module):
+ def __init__(self, config: EdgeTamPromptEncoderConfig):
+ super().__init__()
+ self.scale = config.scale
+ positional_embedding = self.scale * torch.randn((2, config.hidden_size // 2))
+ self.register_buffer("positional_embedding", positional_embedding)
+
+ def forward(self, input_coords, input_shape=None):
+ """Positionally encode points that are normalized to [0,1]."""
+ coordinates = input_coords.clone()
+
+ if input_shape is not None:
+ coordinates[:, :, :, 0] = coordinates[:, :, :, 0] / input_shape[1]
+ coordinates[:, :, :, 1] = coordinates[:, :, :, 1] / input_shape[0]
+ coordinates.to(torch.float32)
+
+ # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+ coordinates = 2 * coordinates - 1
+ coordinates = coordinates.to(self.positional_embedding.dtype)
+ coordinates = coordinates @ self.positional_embedding
+ coordinates = 2 * np.pi * coordinates
+ # outputs d_1 x ... x d_n x channel shape
+ return torch.cat([torch.sin(coordinates), torch.cos(coordinates)], dim=-1)
+
+
+class EdgeTamMaskEmbedding(nn.Module):
+ def __init__(self, config: EdgeTamPromptEncoderConfig):
+ super().__init__()
+ self.mask_input_channels = config.mask_input_channels // 4
+ self.activation = ACT2FN[config.hidden_act]
+ self.conv1 = nn.Conv2d(1, self.mask_input_channels, kernel_size=2, stride=2)
+ self.conv2 = nn.Conv2d(self.mask_input_channels, config.mask_input_channels, kernel_size=2, stride=2)
+ self.conv3 = nn.Conv2d(config.mask_input_channels, config.hidden_size, kernel_size=1)
+ self.layer_norm1 = EdgeTamLayerNorm(
+ self.mask_input_channels, eps=config.layer_norm_eps, data_format="channels_first"
+ )
+ self.layer_norm2 = EdgeTamLayerNorm(
+ self.mask_input_channels * 4, eps=config.layer_norm_eps, data_format="channels_first"
+ )
+
+ def forward(self, masks):
+ hidden_states = self.conv1(masks)
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ dense_embeddings = self.conv3(hidden_states)
+ return dense_embeddings
+
+
+class EdgeTamPromptEncoder(nn.Module):
+ def __init__(self, config: EdgeTamPromptEncoderConfig):
+ super().__init__()
+ self.shared_embedding = EdgeTamPositionalEmbedding(config)
+ self.mask_embed = EdgeTamMaskEmbedding(config)
+ self.no_mask_embed = nn.Embedding(1, config.hidden_size)
+
+ self.image_embedding_size = (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ self.mask_input_size = (4 * config.image_size // config.patch_size, 4 * config.image_size // config.patch_size)
+ self.input_image_size = config.image_size
+
+ self.point_embed = nn.Embedding(config.num_point_embeddings, config.hidden_size)
+ self.hidden_size = config.hidden_size
+ self.not_a_point_embed = nn.Embedding(1, config.hidden_size)
+
+ def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -> torch.Tensor:
+ """Embeds point prompts."""
+ points = points + 0.5 # Shift to center of pixel
+ if pad:
+ points = torch.nn.functional.pad(points, (0, 0, 0, 1), mode="constant", value=0)
+ labels = torch.nn.functional.pad(labels, (0, 1), mode="constant", value=-1)
+ input_shape = (self.input_image_size, self.input_image_size)
+ point_embedding = self.shared_embedding(points, input_shape)
+
+ # torch.where and expanding the labels tensor is required by the ONNX export
+ point_embedding = torch.where(labels[..., None] == -1, self.not_a_point_embed.weight, point_embedding)
+
+ # This is required for the ONNX export. The dtype, device need to be explicitly
+ # specified as otherwise torch.onnx.export interprets as double
+ point_embedding = torch.where(
+ labels[..., None] != -10,
+ point_embedding,
+ torch.zeros_like(point_embedding),
+ )
+
+ # Add point embeddings for labels >= 0
+ point_embedding = point_embedding + self.point_embed(labels.clamp(min=0)) * (labels >= 0).unsqueeze(-1)
+
+ return point_embedding
+
+ def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
+ """Embeds box prompts."""
+ boxes += 0.5 # Shift to center of pixel
+ coords = boxes.view(*boxes.shape[:2], 2, 2)
+ # add padding point for consistency with the original implementation
+ coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0)
+ corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size))
+ corner_embedding[:, :, 0, :] += self.point_embed.weight[2]
+ corner_embedding[:, :, 1, :] += self.point_embed.weight[3]
+ corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :])
+ return corner_embedding
+
+ def forward(
+ self,
+ input_points: Optional[tuple[torch.Tensor, torch.Tensor]],
+ input_labels: Optional[torch.Tensor],
+ input_boxes: Optional[torch.Tensor],
+ input_masks: Optional[torch.Tensor],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Embeds different types of prompts, returning both sparse and dense embeddings.
+
+ Args:
+ points (`torch.Tensor`, *optional*):
+ point coordinates and labels to embed.
+ boxes (`torch.Tensor`, *optional*):
+ boxes to embed
+ masks (`torch.Tensor`, *optional*):
+ masks to embed
+ """
+ sparse_embeddings = None
+ batch_size = 1
+ if input_points is not None:
+ batch_size = input_points.shape[0]
+ if input_labels is None:
+ raise ValueError("If points are provided, labels must also be provided.")
+ point_embeddings = self._embed_points(input_points, input_labels, pad=(input_boxes is None))
+ sparse_embeddings = point_embeddings
+ if input_boxes is not None:
+ batch_size = input_boxes.shape[0]
+ box_embeddings = self._embed_boxes(input_boxes)
+ if sparse_embeddings is None:
+ sparse_embeddings = box_embeddings
+ else:
+ sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=2)
+ if input_masks is not None:
+ dense_embeddings = self.mask_embed(input_masks)
+ else:
+ dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
+ batch_size, -1, self.image_embedding_size[0], self.image_embedding_size[1]
+ )
+
+ return sparse_embeddings, dense_embeddings
+
+
+class EdgeTamTwoWayTransformer(nn.Module):
+ def __init__(self, config: EdgeTamMaskDecoderConfig):
+ super().__init__()
+ self.config = config
+
+ self.num_hidden_layers = config.num_hidden_layers
+ self.layers = nn.ModuleList()
+
+ for i in range(self.num_hidden_layers):
+ self.layers.append(EdgeTamTwoWayAttentionBlock(config, skip_first_layer_pe=(i == 0)))
+
+ self.final_attn_token_to_image = EdgeTamAttention(config)
+ self.layer_norm_final_attn = nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ point_embeddings: Tensor,
+ image_embeddings: Tensor,
+ image_positional_embeddings: Tensor,
+ attention_similarity: Tensor,
+ target_embedding=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, BaseModelOutput]:
+ if image_embeddings is None:
+ raise ValueError("You have to specify an image_embedding")
+
+ image_embeddings = image_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1)
+ image_positional_embeddings = image_positional_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1)
+
+ # Prepare queries
+ queries = point_embeddings
+ keys = image_embeddings
+
+ # Apply transformer blocks and final layernorm
+ for layer in self.layers:
+ if target_embedding is not None:
+ queries += target_embedding
+
+ queries, keys, _ = layer(
+ queries=queries,
+ keys=keys,
+ query_point_embedding=point_embeddings,
+ key_point_embedding=image_positional_embeddings,
+ attention_similarity=attention_similarity,
+ **kwargs,
+ )
+ # Apply the final attention layer from the points to the image
+ query = queries + point_embeddings
+ key = keys + image_positional_embeddings
+
+ attn_out, _ = self.final_attn_token_to_image(query=query, key=key, value=keys)
+
+ queries = queries + attn_out
+ queries = self.layer_norm_final_attn(queries)
+ return queries, keys
+
+
+class EdgeTamMaskDecoder(nn.Module):
+ def __init__(self, config: EdgeTamMaskDecoderConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+
+ self.num_multimask_outputs = config.num_multimask_outputs
+ self.num_mask_tokens = config.num_multimask_outputs + 1
+
+ self.iou_token = nn.Embedding(1, self.hidden_size)
+ self.mask_tokens = nn.Embedding(self.num_mask_tokens, self.hidden_size)
+
+ self.transformer = EdgeTamTwoWayTransformer(config)
+
+ # should we create a new class for this?
+ self.upscale_conv1 = nn.ConvTranspose2d(self.hidden_size, self.hidden_size // 4, kernel_size=2, stride=2)
+ self.upscale_conv2 = nn.ConvTranspose2d(self.hidden_size // 4, self.hidden_size // 8, kernel_size=2, stride=2)
+ self.upscale_layer_norm = EdgeTamLayerNorm(self.hidden_size // 4, data_format="channels_first")
+ self.activation = nn.GELU()
+
+ mlps_list = []
+ for _ in range(self.num_mask_tokens):
+ mlps_list += [EdgeTamFeedForward(self.hidden_size, self.hidden_size, self.hidden_size // 8, 3)]
+ self.output_hypernetworks_mlps = nn.ModuleList(mlps_list)
+ self.iou_prediction_head = EdgeTamFeedForward(
+ self.hidden_size,
+ config.iou_head_hidden_dim,
+ self.num_mask_tokens,
+ config.iou_head_depth,
+ sigmoid_output=True,
+ )
+
+ self.conv_s0 = nn.Conv2d(config.hidden_size, config.hidden_size // 8, kernel_size=1, stride=1)
+ self.conv_s1 = nn.Conv2d(config.hidden_size, config.hidden_size // 4, kernel_size=1, stride=1)
+
+ self.obj_score_token = nn.Embedding(1, self.hidden_size)
+ self.pred_obj_score_head = EdgeTamFeedForward(self.hidden_size, self.hidden_size, 1, 3)
+
+ self.dynamic_multimask_via_stability = config.dynamic_multimask_via_stability
+ self.dynamic_multimask_stability_delta = config.dynamic_multimask_stability_delta
+ self.dynamic_multimask_stability_thresh = config.dynamic_multimask_stability_thresh
+
+ def forward(
+ self,
+ image_embeddings: torch.Tensor,
+ image_positional_embeddings: torch.Tensor,
+ sparse_prompt_embeddings: torch.Tensor,
+ dense_prompt_embeddings: torch.Tensor,
+ multimask_output: bool,
+ high_resolution_features: list[torch.Tensor],
+ attention_similarity: Optional[torch.Tensor] = None,
+ target_embedding: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Predict masks given image and prompt embeddings.
+
+ Args:
+ image_embeddings (`torch.Tensor`):
+ The embeddings from the image encoder.
+ image_positional_embeddings (`torch.Tensor`):
+ Positional encoding with the shape of image_embeddings.
+ sparse_prompt_embeddings (`torch.Tensor`):
+ The embeddings of the points and boxes.
+ dense_prompt_embeddings (`torch.Tensor`):
+ The embeddings of the mask inputs.
+ multimask_output (`bool`):
+ Whether to return multiple masks or a single mask.
+ high_resolution_features (`list[torch.Tensor]`, *optional*):
+ The high-resolution features from the vision encoder.
+ attention_similarity (`torch.Tensor`, *optional*):
+ The attention similarity tensor.
+ target_embedding (`torch.Tensor`, *optional*):
+ The target embedding.
+ """
+ batch_size, num_channels, height, width = image_embeddings.shape
+ point_batch_size = sparse_prompt_embeddings.shape[1]
+ # Concatenate output tokens
+ output_tokens = torch.cat(
+ [
+ self.obj_score_token.weight,
+ self.iou_token.weight,
+ self.mask_tokens.weight,
+ ],
+ dim=0,
+ )
+ output_tokens = output_tokens.repeat(batch_size, point_batch_size, 1, 1)
+
+ if sparse_prompt_embeddings.shape[0] != 0:
+ tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=2)
+ else:
+ tokens = output_tokens
+ point_embeddings = tokens.to(self.iou_token.weight.dtype)
+
+ # Expand per-image data in batch direction to be per-mask
+ image_embeddings = image_embeddings + dense_prompt_embeddings
+ image_embeddings = image_embeddings.repeat_interleave(point_batch_size, dim=0)
+ image_positional_embeddings = image_positional_embeddings.repeat_interleave(point_batch_size, 0)
+ # Run the transformer
+ point_embeddings, image_embeddings = self.transformer(
+ point_embeddings=point_embeddings,
+ image_embeddings=image_embeddings,
+ image_positional_embeddings=image_positional_embeddings,
+ attention_similarity=attention_similarity,
+ target_embedding=target_embedding,
+ **kwargs,
+ )
+ iou_token_out = point_embeddings[:, :, 1, :]
+ mask_tokens_out = point_embeddings[:, :, 2 : (2 + self.num_mask_tokens), :]
+
+ # Upscale mask embeddings and predict masks using the mask tokens
+ image_embeddings = image_embeddings.transpose(2, 3).view(
+ batch_size * point_batch_size, num_channels, height, width
+ )
+
+ feat_s0, feat_s1 = high_resolution_features
+ feat_s0 = feat_s0.repeat_interleave(point_batch_size, dim=0)
+ feat_s1 = feat_s1.repeat_interleave(point_batch_size, dim=0)
+ upscaled_embedding = self.upscale_conv1(image_embeddings) + feat_s1
+ upscaled_embedding = self.activation(self.upscale_layer_norm(upscaled_embedding))
+ upscaled_embedding = self.activation(self.upscale_conv2(upscaled_embedding) + feat_s0)
+
+ hyper_in_list: list[torch.Tensor] = []
+ for i in range(self.num_mask_tokens):
+ current_mlp = self.output_hypernetworks_mlps[i]
+ hyper_in_list += [current_mlp(mask_tokens_out[:, :, i, :])]
+ hyper_in = torch.stack(hyper_in_list, dim=2)
+
+ _, num_channels, height, width = upscaled_embedding.shape
+ upscaled_embedding = upscaled_embedding.view(batch_size, point_batch_size, num_channels, height * width)
+ masks = (hyper_in @ upscaled_embedding).view(batch_size, point_batch_size, -1, height, width)
+
+ # Generate mask quality predictions
+ iou_pred = self.iou_prediction_head(iou_token_out)
+ object_score_logits = self.pred_obj_score_head(point_embeddings[:, :, 0, :])
+
+ # Select the correct mask or masks for output
+ if multimask_output:
+ mask_slice = slice(1, None)
+ masks = masks[:, :, mask_slice, :, :]
+ iou_pred = iou_pred[:, :, mask_slice]
+ elif self.dynamic_multimask_via_stability and not self.training:
+ mask_slice = slice(0, 1)
+ masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)
+ else:
+ mask_slice = slice(0, 1)
+ masks = masks[:, :, mask_slice, :, :]
+ iou_pred = iou_pred[:, :, mask_slice]
+
+ sam_tokens_out = mask_tokens_out[:, :, mask_slice] # [b, 3, c] shape
+
+ return masks, iou_pred, sam_tokens_out, object_score_logits
+
+ def _get_stability_scores(self, mask_logits):
+ """
+ Compute stability scores of the mask logits based on the IoU between upper and
+ lower thresholds.
+ """
+ mask_logits = mask_logits.flatten(-2)
+ stability_delta = self.dynamic_multimask_stability_delta
+ area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
+ area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
+ stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)
+ return stability_scores
+
+ def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
+ """
+ When outputting a single mask, if the stability score from the current single-mask
+ output (based on output token 0) falls below a threshold, we instead select from
+ multi-mask outputs (based on output token 1~3) the mask with the highest predicted
+ IoU score. This is intended to ensure a valid mask for both clicking and tracking.
+ """
+ # The best mask from multimask output tokens (1~3)
+ multimask_logits = all_mask_logits[:, :, 1:, :, :]
+ multimask_iou_scores = all_iou_scores[:, :, 1:]
+ best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) # [B, P]
+ best_scores_inds_expanded = best_scores_inds.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+ best_scores_inds_expanded = best_scores_inds_expanded.expand(
+ -1, -1, 1, multimask_logits.size(-2), multimask_logits.size(-1)
+ )
+ best_multimask_logits = torch.gather(multimask_logits, 2, best_scores_inds_expanded) # [B, P, 1, H, W]
+ best_multimask_iou_scores = torch.gather(multimask_iou_scores, 2, best_scores_inds.unsqueeze(-1)) # [B, P, 1]
+
+ # The mask from singlemask output token 0 and its stability score
+ singlemask_logits = all_mask_logits[:, :, 0:1, :, :]
+ singlemask_iou_scores = all_iou_scores[:, :, 0:1]
+ stability_scores = self._get_stability_scores(singlemask_logits)
+ is_stable = stability_scores >= self.dynamic_multimask_stability_thresh
+
+ # Dynamically fall back to best multimask output upon low stability scores.
+ mask_logits_out = torch.where(
+ is_stable[..., None, None].expand_as(singlemask_logits),
+ singlemask_logits,
+ best_multimask_logits,
+ )
+ iou_scores_out = torch.where(
+ is_stable.expand_as(singlemask_iou_scores),
+ singlemask_iou_scores,
+ best_multimask_iou_scores,
+ )
+ return mask_logits_out, iou_scores_out
+
+
+@auto_docstring(
+ custom_intro="""
+ Segment Anything Model 2 (SAM 2) for generating segmentation masks, given an input image and
+ input points and labels, boxes, or masks.
+ """
+)
+class EdgeTamModel(EdgeTamPreTrainedModel):
+ _tied_weights_keys = ["prompt_encoder.shared_embedding.positional_embedding"]
+ # need to be ignored, as it's a buffer and will not be correctly detected as tied weight
+ _keys_to_ignore_on_load_missing = ["prompt_encoder.shared_embedding.positional_embedding"]
+ _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(EdgeTamTwoWayAttentionBlock, index=2)}
+ _keys_to_ignore_on_load_unexpected = [
+ r"^memory_.*",
+ r"^mask_downsample.*",
+ r"spatial_perceiver.*",
+ r"^object_pointer_proj.*",
+ r"^temporal_positional_encoding_projection_layer.*",
+ "no_memory_positional_encoding",
+ "no_object_pointer",
+ "occlusion_spatial_embedding_parameter",
+ ]
+
+ def __init__(self, config: EdgeTamConfig):
+ super().__init__(config)
+ self.shared_image_embedding = EdgeTamPositionalEmbedding(config.prompt_encoder_config)
+ self.vision_encoder = AutoModel.from_config(config.vision_config)
+ self.prompt_encoder = EdgeTamPromptEncoder(config.prompt_encoder_config)
+ # The module using it is not a PreTrainedModel subclass so we need this
+ config.mask_decoder_config._attn_implementation = config._attn_implementation
+ self.mask_decoder = EdgeTamMaskDecoder(config.mask_decoder_config)
+
+ self.num_feature_levels = config.vision_config.num_feature_levels
+ self.backbone_feature_sizes = config.vision_config.backbone_feature_sizes
+ # a single token to indicate no memory embedding from previous frames
+ self.hidden_dim = config.vision_config.fpn_hidden_size
+ self.no_memory_embedding = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+
+ self.post_init()
+
+ def _tie_weights(self):
+ self.prompt_encoder.shared_embedding.positional_embedding.data = (
+ self.shared_image_embedding.positional_embedding.data
+ )
+
+ def get_image_wide_positional_embeddings(self) -> torch.Tensor:
+ size = self.prompt_encoder.image_embedding_size
+ target_device = self.shared_image_embedding.positional_embedding.device
+ target_dtype = self.shared_image_embedding.positional_embedding.dtype
+ grid = torch.ones(size, device=target_device, dtype=target_dtype)
+ y_embed = grid.cumsum(dim=0) - 0.5
+ x_embed = grid.cumsum(dim=1) - 0.5
+ y_embed = y_embed / size[0]
+ x_embed = x_embed / size[1]
+
+ positional_embedding = self.shared_image_embedding(torch.stack([x_embed, y_embed], dim=-1))
+ return positional_embedding.permute(2, 0, 1).unsqueeze(0) # channel x height x width
+
+ @torch.no_grad()
+ def get_image_embeddings(
+ self,
+ pixel_values: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> list[torch.Tensor]:
+ r"""
+ Returns the image embeddings by passing the pixel values through the vision encoder.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Input pixel values
+ """
+ batch_size = pixel_values.shape[0]
+ feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs)
+
+ # add no memory embedding to the last feature map
+ feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
+
+ # reshape feature maps to the same shape as the backbone feature sizes
+ image_embeddings = [
+ feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+ for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes)
+ ]
+
+ return image_embeddings
+
+ @torch.no_grad()
+ def get_prompt_embeddings(
+ self,
+ input_points: Optional[torch.FloatTensor] = None,
+ input_labels: Optional[torch.LongTensor] = None,
+ input_boxes: Optional[torch.FloatTensor] = None,
+ input_masks: Optional[torch.LongTensor] = None,
+ ):
+ r"""
+ Returns the prompt embeddings by passing the input points, labels, boxes and masks through the prompt encoder.
+
+ Args:
+ input_points (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_points_per_image, 2)`):
+ Optional input points for the prompt encoder. The padding of the point is automatically done by the
+ processor. `point_batch_size` refers to the number of masks that we want the model to predict per
+ point. The model will output `point_batch_size` times 3 masks in total.
+ input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points_per_image)`):
+ Optional input labels for the prompt encoder. The padding of the labels is automatically done by the
+ processor, or can be fed by the user.
+ input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes_per_image, 4)`):
+ Optional input boxes for the prompt encoder. The padding of the boxes is automatically done by the
+ processor. users can also pass manually the input boxes.
+ input_masks (`torch.LongTensor` of shape `(batch_size, image_size, image_size)`):
+ Optional input masks for the prompt encoder.
+ """
+ prompt_output = self.prompt_encoder(
+ input_points=input_points,
+ input_labels=input_labels,
+ input_boxes=input_boxes,
+ input_masks=input_masks,
+ )
+ return prompt_output
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ input_points: Optional[torch.FloatTensor] = None,
+ input_labels: Optional[torch.LongTensor] = None,
+ input_boxes: Optional[torch.FloatTensor] = None,
+ input_masks: Optional[torch.LongTensor] = None,
+ image_embeddings: Optional[torch.FloatTensor] = None,
+ multimask_output: bool = True,
+ attention_similarity: Optional[torch.FloatTensor] = None,
+ target_embedding: Optional[torch.FloatTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> EdgeTamImageSegmentationOutput:
+ r"""
+ input_points (`torch.FloatTensor` of shape `(batch_size, num_points, 2)`):
+ Input 2D spatial points, this is used by the prompt encoder to encode the prompt. Generally yields to much
+ better results. The points can be obtained by passing a list of list of list to the processor that will
+ create corresponding `torch` tensors of dimension 4. The first dimension is the image batch size, the
+ second dimension is the point batch size (i.e. how many segmentation masks do we want the model to predict
+ per input point), the third dimension is the number of points per segmentation mask (it is possible to pass
+ multiple points for a single mask), and the last dimension is the x (vertical) and y (horizontal)
+ coordinates of the point. If a different number of points is passed either for each image, or for each
+ mask, the processor will create "PAD" points that will correspond to the (0, 0) coordinate, and the
+ computation of the embedding will be skipped for these points using the labels.
+ input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points)`):
+ Input labels for the points, this is used by the prompt encoder to encode the prompt. According to the
+ official implementation, there are 3 types of labels
+
+ - `1`: the point is a point that contains the object of interest
+ - `0`: the point is a point that does not contain the object of interest
+ - `-1`: the point corresponds to the background
+
+ We added the label:
+
+ - `-10`: the point is a padding point, thus should be ignored by the prompt encoder
+
+ The padding labels should be automatically done by the processor.
+ input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes, 4)`):
+ Input boxes for the points, this is used by the prompt encoder to encode the prompt. Generally yields to
+ much better generated masks. The boxes can be obtained by passing a list of list of list to the processor,
+ that will generate a `torch` tensor, with each dimension corresponding respectively to the image batch
+ size, the number of boxes per image and the coordinates of the top left and bottom right point of the box.
+ In the order (`x1`, `y1`, `x2`, `y2`):
+
+ - `x1`: the x coordinate of the top left point of the input box
+ - `y1`: the y coordinate of the top left point of the input box
+ - `x2`: the x coordinate of the bottom right point of the input box
+ - `y2`: the y coordinate of the bottom right point of the input box
+ input_masks (`torch.FloatTensor` of shape `(batch_size, image_size, image_size)`):
+ SAM model also accepts segmentation masks as input. The mask will be embedded by the prompt encoder to
+ generate a corresponding embedding, that will be fed later on to the mask decoder. These masks needs to be
+ manually fed by the user, and they need to be of shape (`batch_size`, `image_size`, `image_size`).
+ image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_channels, window_size, window_size)`):
+ Image embeddings, this is used by the mask decoder to generate masks and iou scores. For more memory
+ efficient computation, users can first retrieve the image embeddings using the `get_image_embeddings`
+ method, and then feed them to the `forward` method instead of feeding the `pixel_values`.
+ multimask_output (`bool`, *optional*):
+ In the original implementation and paper, the model always outputs 3 masks per image (or per point / per
+ bounding box if relevant). However, it is possible to just output a single mask, that corresponds to the
+ "best" mask, by specifying `multimask_output=False`.
+ attention_similarity (`torch.FloatTensor`, *optional*):
+ Attention similarity tensor, to be provided to the mask decoder for target-guided attention in case the
+ model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048).
+ target_embedding (`torch.FloatTensor`, *optional*):
+ Embedding of the target concept, to be provided to the mask decoder for target-semantic prompting in case
+ the model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048).
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoModel, AutoProcessor
+
+ >>> model = AutoModel.from_pretrained("danelcsb/edgetam.1_hiera_tiny")
+ >>> processor = AutoProcessor.from_pretrained("danelcsb/edgetam.1_hiera_tiny")
+
+ >>> img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png"
+ >>> raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ >>> input_points = [[[400, 650]]] # 2D location of a window on the car
+ >>> inputs = processor(images=raw_image, input_points=input_points, return_tensors="pt")
+
+ >>> # Get segmentation mask
+ >>> outputs = model(**inputs)
+
+ >>> # Postprocess masks
+ >>> masks = processor.post_process_masks(
+ ... outputs.pred_masks, inputs["original_sizes"], inputs["reshaped_input_sizes"]
+ ... )
+ ```
+ """
+ if not ((pixel_values is None) ^ (image_embeddings is None)):
+ raise ValueError("Exactly one of pixel_values or image_embeddings must be provided.")
+ if input_points is not None and input_boxes is not None:
+ if input_points.shape[1] != input_boxes.shape[1]:
+ raise ValueError(
+ f"You should provide as many bounding boxes as input points per box. Got {input_points.shape[1]} and {input_boxes.shape[1]}."
+ )
+
+ image_positional_embeddings = self.get_image_wide_positional_embeddings()
+ # repeat with batch size
+ batch_size = pixel_values.shape[0] if pixel_values is not None else image_embeddings[-1].shape[0]
+ image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1)
+
+ vision_attentions = None
+ vision_hidden_states = None
+
+ if pixel_values is not None:
+ feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features(
+ pixel_values,
+ **kwargs,
+ )
+
+ # add no memory embedding to the last feature map
+ feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
+
+ # reshape feature maps to the same shape as the backbone feature sizes
+ image_embeddings = [
+ feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+ for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes)
+ ]
+
+ if input_points is not None and input_labels is None:
+ input_labels = torch.ones_like(input_points[:, :, :, 0], dtype=torch.int, device=input_points.device)
+
+ if input_points is None and input_boxes is None:
+ # If no points are provide, pad with an empty point (with label -1)
+ input_points = torch.zeros(
+ batch_size, 1, 1, 2, dtype=image_embeddings[-1].dtype, device=image_embeddings[-1].device
+ )
+ input_labels = -torch.ones(batch_size, 1, 1, dtype=torch.int32, device=image_embeddings[-1].device)
+
+ if input_masks is not None:
+ # If mask_inputs is provided, downsize it into low-res mask input if needed
+ # and feed it as a dense mask prompt into the SAM mask encoder
+ if input_masks.shape[-2:] != self.prompt_encoder.mask_input_size:
+ input_masks = F.interpolate(
+ input_masks.float(),
+ size=self.prompt_encoder.mask_input_size,
+ align_corners=False,
+ mode="bilinear",
+ antialias=True, # use antialias for downsampling
+ ).to(input_masks.dtype)
+
+ sparse_embeddings, dense_embeddings = self.prompt_encoder(
+ input_points=input_points,
+ input_labels=input_labels,
+ input_boxes=input_boxes,
+ input_masks=input_masks,
+ )
+ low_res_multimasks, iou_scores, _, object_score_logits = self.mask_decoder(
+ image_embeddings=image_embeddings[-1],
+ image_positional_embeddings=image_positional_embeddings,
+ sparse_prompt_embeddings=sparse_embeddings,
+ dense_prompt_embeddings=dense_embeddings,
+ multimask_output=multimask_output,
+ high_resolution_features=image_embeddings[:-1],
+ attention_similarity=attention_similarity,
+ target_embedding=target_embedding,
+ **kwargs,
+ )
+
+ return EdgeTamImageSegmentationOutput(
+ iou_scores=iou_scores,
+ pred_masks=low_res_multimasks,
+ object_score_logits=object_score_logits,
+ image_embeddings=image_embeddings,
+ vision_hidden_states=vision_hidden_states,
+ vision_attentions=vision_attentions,
+ )
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[
+ list[torch.Tensor],
+ list[torch.Tensor],
+ Optional[tuple[torch.FloatTensor, ...]],
+ Optional[tuple[torch.FloatTensor, ...]],
+ ]:
+ r"""
+ Extract and preprocess image features using the vision encoder.
+
+ Args:
+ pixel_values (`torch.FloatTensor`):
+ Input pixel values of shape `(batch_size, num_channels, height, width)`.
+
+ Returns:
+ `tuple`: A tuple containing:
+ - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels.
+ - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level.
+ - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder.
+ - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder.
+ """
+ vision_outputs: EdgeTamVisionEncoderOutput = self.vision_encoder(
+ pixel_values,
+ **kwargs,
+ )
+
+ feature_maps = vision_outputs.fpn_hidden_states
+ feature_maps_position_embeddings = vision_outputs.fpn_position_encoding
+
+ # precompute projected level 0 and level 1 features in SAM decoder
+ # to avoid running it again on every SAM click
+ feature_maps = list(feature_maps)
+ feature_maps[0] = self.mask_decoder.conv_s0(feature_maps[0])
+ feature_maps[1] = self.mask_decoder.conv_s1(feature_maps[1])
+
+ # flatten NxCxHxW to HWxNxC
+ feature_maps = [feature_map.flatten(2).permute(2, 0, 1) for feature_map in feature_maps]
+ feature_maps_position_embeddings = [
+ feature_map_position_embedding.flatten(2).permute(2, 0, 1)
+ for feature_map_position_embedding in feature_maps_position_embeddings
+ ]
+
+ return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions
+
+
+__all__ = ["EdgeTamModel", "EdgeTamVisionModel", "EdgeTamPreTrainedModel"]
diff --git a/src/transformers/models/edgetam/modular_edgetam.py b/src/transformers/models/edgetam/modular_edgetam.py
new file mode 100644
index 000000000000..e26d58d96b81
--- /dev/null
+++ b/src/transformers/models/edgetam/modular_edgetam.py
@@ -0,0 +1,261 @@
+# coding=utf-8
+# Copyright 2025 The Meta AI Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch SAM 2 model."""
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from transformers.models.sam2.configuration_sam2 import Sam2Config, Sam2MaskDecoderConfig, Sam2PromptEncoderConfig
+from transformers.models.sam2.modeling_sam2 import (
+ Sam2Attention,
+ Sam2FeedForward,
+ Sam2LayerNorm,
+ Sam2Model,
+ Sam2PreTrainedModel,
+ Sam2TwoWayAttentionBlock,
+ Sam2VisionEncoderOutput,
+ Sam2VisionModel,
+)
+from transformers.utils.generic import TransformersKwargs, check_model_inputs
+
+from ...configuration_utils import PretrainedConfig
+from ...processing_utils import Unpack
+from ...utils import (
+ auto_docstring,
+)
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+# fix this in modular
+if True:
+ from transformers.models.timm_wrapper.modeling_timm_wrapper import TimmWrapperModel
+
+
+class EdgeTamVisionConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamVisionModel`]. It is used to instantiate a SAM
+ vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
+ defaults will yield a similar configuration to that of SAM 2.1 Hiera-tiny
+ [facebook/EdgeTAM](https://huggingface.co/facebook/EdgeTAM) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ backbone_config (`Union[dict, "PretrainedConfig"]`, *optional*):
+ Configuration for the vision backbone. This is used to instantiate the backbone using
+ `AutoModel.from_config`.
+ backbone_channel_list (`List[int]`, *optional*, defaults to `[384, 192, 96, 48]`):
+ The list of channel dimensions for the backbone.
+ backbone_feature_sizes (`List[List[int]]`, *optional*, defaults to `[[256, 256], [128, 128], [64, 64]]`):
+ The spatial sizes of the feature maps from the backbone.
+ fpn_hidden_size (`int`, *optional*, defaults to 256):
+ The hidden dimension of the FPN.
+ fpn_kernel_size (`int`, *optional*, defaults to 1):
+ The kernel size for the convolutions in the neck.
+ fpn_stride (`int`, *optional*, defaults to 1):
+ The stride for the convolutions in the neck.
+ fpn_padding (`int`, *optional*, defaults to 0):
+ The padding for the convolutions in the neck.
+ fpn_top_down_levels (`List[int]`, *optional*, defaults to `[2, 3]`):
+ The levels for the top-down FPN connections.
+ num_feature_levels (`int`, *optional*, defaults to 3):
+ The number of feature levels from the FPN to use.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the neck.
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon for the layer normalization.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+ """
+
+ base_config_key = "vision_config"
+ model_type = "edgetam_vision_model"
+ sub_configs = {
+ "backbone_config": AutoConfig,
+ }
+
+ def __init__(
+ self,
+ backbone_config=None,
+ backbone_channel_list=None,
+ backbone_feature_sizes=None,
+ fpn_hidden_size=256,
+ fpn_kernel_size=1,
+ fpn_stride=1,
+ fpn_padding=0,
+ fpn_top_down_levels=None,
+ num_feature_levels=3,
+ hidden_act="gelu",
+ layer_norm_eps=1e-6,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ backbone_channel_list = [384, 192, 96, 48] if backbone_channel_list is None else backbone_channel_list
+ backbone_feature_sizes = (
+ [[256, 256], [128, 128], [64, 64]] if backbone_feature_sizes is None else backbone_feature_sizes
+ )
+ fpn_top_down_levels = [2, 3] if fpn_top_down_levels is None else fpn_top_down_levels
+
+ if isinstance(backbone_config, dict):
+ backbone_config["model_type"] = backbone_config.get("model_type", "timm_wrapper")
+ backbone_config = CONFIG_MAPPING[backbone_config["model_type"]](**backbone_config)
+ elif isinstance(backbone_config, AutoConfig):
+ backbone_config = backbone_config
+ elif backbone_config is None:
+ backbone_config = AutoConfig.from_pretrained(
+ "timm/repvit_m1.dist_in1k",
+ model_args={"in_chans": 3, "features_only": True, "out_indices": [0, 1, 2, 3]},
+ )
+
+ self.backbone_config = backbone_config
+
+ # Neck
+ self.backbone_channel_list = backbone_channel_list
+ self.backbone_feature_sizes = backbone_feature_sizes
+ self.fpn_hidden_size = fpn_hidden_size
+ self.fpn_kernel_size = fpn_kernel_size
+ self.fpn_stride = fpn_stride
+ self.fpn_padding = fpn_padding
+ self.fpn_top_down_levels = fpn_top_down_levels
+ self.num_feature_levels = num_feature_levels
+
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.initializer_range = initializer_range
+
+
+class EdgeTamPromptEncoderConfig(Sam2PromptEncoderConfig):
+ pass
+
+
+class EdgeTamMaskDecoderConfig(Sam2MaskDecoderConfig):
+ pass
+
+
+class EdgeTamConfig(Sam2Config):
+ pass
+
+
+class EdgeTamLayerNorm(Sam2LayerNorm):
+ pass
+
+
+class EdgeTamVisionEncoderOutput(Sam2VisionEncoderOutput):
+ pass
+
+
+class EdgeTamAttention(Sam2Attention):
+ pass
+
+
+class EdgeTamTwoWayAttentionBlock(Sam2TwoWayAttentionBlock):
+ pass
+
+
+class EdgeTamFeedForward(Sam2FeedForward):
+ pass
+
+
+@auto_docstring
+class EdgeTamPreTrainedModel(Sam2PreTrainedModel):
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, (nn.LayerNorm, EdgeTamLayerNorm)):
+ module.weight.data.fill_(1.0)
+ module.bias.data.zero_()
+ if isinstance(module, EdgeTamModel):
+ if module.no_memory_embedding is not None:
+ module.no_memory_embedding.data.zero_()
+
+
+@auto_docstring(
+ custom_intro="""
+ The vision model from EdgeTAM without any head or projection on top.
+ """
+)
+class EdgeTamVisionModel(Sam2VisionModel):
+ config_class = EdgeTamVisionConfig
+ main_input_name = "pixel_values"
+ _can_record_outputs = {"hidden_states": TimmWrapperModel, "attentions": TimmWrapperModel}
+
+ def get_input_embeddings(self):
+ raise NotImplementedError("Can't get input embeddings from timm wrapper model")
+
+ @check_model_inputs
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, EdgeTamVisionEncoderOutput]:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ # Forward through backbone
+ backbone_output = self.backbone(pixel_values)
+ intermediate_hidden_states = backbone_output.last_hidden_state
+ intermediate_hidden_states = [hidden_state.permute(0, 2, 3, 1) for hidden_state in intermediate_hidden_states]
+
+ fpn_hidden_states, fpn_position_encoding = self.neck(intermediate_hidden_states)
+ # Select last `num_feature_levels` feature levels from FPN and reverse order to get features from high to low resolution
+ fpn_hidden_states = fpn_hidden_states[-self.num_feature_levels :][::-1]
+ fpn_position_encoding = fpn_position_encoding[-self.num_feature_levels :][::-1]
+
+ return EdgeTamVisionEncoderOutput(
+ last_hidden_state=intermediate_hidden_states[-1],
+ fpn_hidden_states=fpn_hidden_states,
+ fpn_position_encoding=fpn_position_encoding,
+ )
+
+
+class EdgeTamModel(Sam2Model):
+ _keys_to_ignore_on_load_unexpected = [
+ r"^memory_.*",
+ r"^mask_downsample.*",
+ r"spatial_perceiver.*",
+ r"^object_pointer_proj.*",
+ r"^temporal_positional_encoding_projection_layer.*",
+ "no_memory_positional_encoding",
+ "no_object_pointer",
+ "occlusion_spatial_embedding_parameter",
+ ]
+
+ def get_input_embeddings(self):
+ raise NotImplementedError("Can't get input embeddings from timm wrapper model")
+
+
+__all__ = [
+ "EdgeTamModel",
+ "EdgeTamVisionModel",
+ "EdgeTamPreTrainedModel",
+ "EdgeTamConfig",
+ "EdgeTamVisionConfig",
+ "EdgeTamPromptEncoderConfig",
+ "EdgeTamMaskDecoderConfig",
+]
diff --git a/src/transformers/models/edgetam_video/__init__.py b/src/transformers/models/edgetam_video/__init__.py
new file mode 100644
index 000000000000..669dd64ec304
--- /dev/null
+++ b/src/transformers/models/edgetam_video/__init__.py
@@ -0,0 +1,29 @@
+# coding=utf-8
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_edgetam_video import *
+ from .modeling_edgetam_video import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/edgetam_video/configuration_edgetam_video.py b/src/transformers/models/edgetam_video/configuration_edgetam_video.py
new file mode 100644
index 000000000000..954864397dcb
--- /dev/null
+++ b/src/transformers/models/edgetam_video/configuration_edgetam_video.py
@@ -0,0 +1,435 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/edgetam_video/modular_edgetam_video.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_edgetam_video.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# 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 ...configuration_utils import PretrainedConfig
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+class EdgeTamVideoPromptEncoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamVideoPromptEncoder`]. The [`EdgeTamVideoPromptEncoder`]
+ module is used to encode the input 2D points and bounding boxes.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the hidden states.
+ image_size (`int`, *optional*, defaults to 1024):
+ The expected output resolution of the image.
+ patch_size (`int`, *optional*, defaults to 16):
+ The size (resolution) of each patch.
+ mask_input_channels (`int`, *optional*, defaults to 16):
+ The number of channels to be fed to the `MaskDecoder` module.
+ num_point_embeddings (`int`, *optional*, defaults to 4):
+ The number of point embeddings to be used.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the encoder and pooler.
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the layer normalization layers.
+ scale (`float`, *optional*, defaults to 1):
+ The scale factor for the prompt encoder.
+ """
+
+ base_config_key = "prompt_encoder_config"
+
+ def __init__(
+ self,
+ hidden_size=256,
+ image_size=1024,
+ patch_size=16,
+ mask_input_channels=16,
+ num_point_embeddings=4,
+ hidden_act="gelu",
+ layer_norm_eps=1e-6,
+ scale=1,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.hidden_size = hidden_size
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.mask_input_channels = mask_input_channels
+ self.num_point_embeddings = num_point_embeddings
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.scale = scale
+
+
+class EdgeTamVideoMaskDecoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`EdgeTamVideoMaskDecoder`]. It is used to instantiate a EDGETAM_VIDEO
+ memory encoder according to the specified arguments, defining the model architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the hidden states.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the EDGETAM_VIDEO mask decoder.
+ mlp_dim (`int`, *optional*, defaults to 2048):
+ The dimension of the MLP in the two-way transformer.
+ num_hidden_layers (`int`, *optional*, defaults to 2):
+ The number of hidden layers in the two-way transformer.
+ num_attention_heads (`int`, *optional*, defaults to 8):
+ The number of attention heads in the two-way transformer.
+ attention_downsample_rate (`int`, *optional*, defaults to 2):
+ The downsample rate for the attention layers.
+ num_multimask_outputs (`int`, *optional*, defaults to 3):
+ The number of multimask outputs.
+ iou_head_depth (`int`, *optional*, defaults to 3):
+ The depth of the IoU head.
+ iou_head_hidden_dim (`int`, *optional*, defaults to 256):
+ The hidden dimension of the IoU head.
+ dynamic_multimask_via_stability (`bool`, *optional*, defaults to `True`):
+ Whether to use dynamic multimask via stability.
+ dynamic_multimask_stability_delta (`float`, *optional*, defaults to 0.05):
+ The stability delta for the dynamic multimask.
+ dynamic_multimask_stability_thresh (`float`, *optional*, defaults to 0.98):
+ The stability threshold for the dynamic multimask.
+
+ """
+
+ base_config_key = "mask_decoder_config"
+
+ def __init__(
+ self,
+ hidden_size=256,
+ hidden_act="gelu",
+ mlp_dim=2048,
+ num_hidden_layers=2,
+ num_attention_heads=8,
+ attention_downsample_rate=2,
+ num_multimask_outputs=3,
+ iou_head_depth=3,
+ iou_head_hidden_dim=256,
+ dynamic_multimask_via_stability=True,
+ dynamic_multimask_stability_delta=0.05,
+ dynamic_multimask_stability_thresh=0.98,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.hidden_size = hidden_size
+ self.num_multimask_outputs = num_multimask_outputs
+ self.hidden_act = hidden_act
+ self.iou_head_depth = iou_head_depth
+ self.iou_head_hidden_dim = iou_head_hidden_dim
+ self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
+ self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
+ self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
+
+ # TwoWayTransformer configuration
+ self.num_hidden_layers = num_hidden_layers
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.mlp_dim = mlp_dim
+ self.attention_downsample_rate = attention_downsample_rate
+
+
+class EdgeTamVideoConfig(PretrainedConfig):
+ r"""
+ [`EdgeTamVideoConfig`] is the configuration class to store the configuration of a [`EdgeTamVideoModel`]. It is used to instantiate a
+ EDGETAM model according to the specified arguments, defining the memory attention, memory encoder, and image encoder
+ configs. Instantiating a configuration defaults will yield a similar configuration to that of the SAM 2.1 Hiera-tiny
+ [facebook/EdgeTAM](https://huggingface.co/facebook/EdgeTAM) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vision_config (Union[`dict`, `EdgeTamVideoVisionConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamVideoVisionConfig`].
+ prompt_encoder_config (Union[`dict`, `EdgeTamVideoPromptEncoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamVideoPromptEncoderConfig`].
+ mask_decoder_config (Union[`dict`, `EdgeTamVideoMaskDecoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamMaskDecoderConfig`].
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ Standard deviation for parameter initialization.
+ num_maskmem (`int`, *optional*, defaults to 7):
+ The number of memory slots for the mask memory.
+ image_size (`int`, *optional*, defaults to 1024):
+ The size of the input images.
+ sigmoid_scale_for_mem_enc (`float`, *optional*, defaults to 20.0):
+ Scale factor for the sigmoid function in the memory encoder.
+ sigmoid_bias_for_mem_enc (`float`, *optional*, defaults to -10.0):
+ Bias for the sigmoid function in the memory encoder.
+ enable_occlusion_spatial_embedding (`bool`, *optional*, defaults to `True`):
+ Whether to enable spatial embedding for occlusions.
+ multimask_output_in_sam (`bool`, *optional*, defaults to `True`):
+ Whether to output multiple masks from the SAM head.
+ multimask_min_pt_num (`int`, *optional*, defaults to 0):
+ The minimum number of points to trigger multimask output.
+ multimask_max_pt_num (`int`, *optional*, defaults to 1):
+ The maximum number of points to trigger multimask output.
+ multimask_output_for_tracking (`bool`, *optional*, defaults to `True`):
+ Whether to use multimask output for tracking.
+ max_object_pointers_in_encoder (`int`, *optional*, defaults to 16):
+ The maximum number of object pointers in the encoder.
+ enable_temporal_pos_encoding_for_object_pointers (`bool`, *optional*, defaults to `True`):
+ Whether to enable temporal positional encoding for object pointers.
+ memory_attention_hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the memory attention hidden states.
+ memory_attention_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the memory attention module.
+ memory_attention_num_attention_heads (`int`, *optional*, defaults to 1):
+ Number of attention heads for each attention layer in the memory attention.
+ memory_attention_downsample_rate (`int`, *optional*, defaults to 1):
+ The downsample rate for the attention layers.
+ memory_attention_mlp_hidden_size (`int`, *optional*, defaults to 2048):
+ The dimension of the feedforward network in the memory attention module.
+ memory_attention_mlp_hidden_act (`str`, *optional*, defaults to `"relu"`):
+ The non-linear activation function in the feedforward network in the memory attention module.
+ memory_attention_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout rate for the memory attention module.
+ memory_attention_rope_theta (`float`, *optional*, defaults to 10000):
+ The Rope theta parameter.
+ memory_attention_rope_feat_sizes (`Tuple[int, int]`, *optional*, defaults to `[64, 64]`):
+ The feature sizes for the Rope positional encoding.
+ memory_attention_rope_k_sizes (`List[int]`, *optional*, defaults to `[16, 16]`):
+ The key feature sizes for the RoPE positional encoding in memory attention.
+ memory_attention_rope_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout rate for the Rope positional encoding.
+ perceiver_resampler_num_latents (`int`, *optional*, defaults to 256):
+ The number of 1D latent tokens in the perceiver resampler.
+ perceiver_resampler_num_latents_2d (`int`, *optional*, defaults to 256):
+ The number of 2D latent tokens in the perceiver resampler.
+ perceiver_resampler_hidden_size (`int`, *optional*, defaults to 64):
+ The hidden size of the perceiver resampler.
+ perceiver_resampler_mlp_intermediate_size (`int`, *optional*, defaults to 256):
+ The intermediate size of the feedforward network in the perceiver resampler.
+ perceiver_resampler_num_attention_heads (`int`, *optional*, defaults to 1):
+ The number of attention heads in the perceiver resampler.
+ perceiver_resampler_attention_head_dim (`int`, *optional*, defaults to 64):
+ The dimension of each attention head in the perceiver resampler.
+ perceiver_resampler_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the perceiver resampler.
+ perceiver_resampler_hidden_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout rate for the hidden layers in the perceiver resampler.
+ perceiver_resampler_attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout rate for the attention layers in the perceiver resampler.
+ memory_encoder_hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the memory encoder hidden states.
+ memory_encoder_output_channels (`int`, *optional*, defaults to 64):
+ The number of output channels for the memory encoder.
+ mask_downsampler_embed_dim (`int`, *optional*, defaults to 256):
+ The dimension of the mask downsampler embedding.
+ memory_fuser_intermediate_dim (`int`, *optional*, defaults to 1024):
+ The intermediate dimension of the memory fuser feedforward network.
+ mask_downsampler_kernel_size (`int`, *optional*, defaults to 3):
+ The kernel size for the mask downsampler.
+ mask_downsampler_stride (`int`, *optional*, defaults to 2):
+ The stride for the mask downsampler.
+ mask_downsampler_padding (`int`, *optional*, defaults to 1):
+ The padding for the mask downsampler.
+ mask_downsampler_total_stride (`int`, *optional*, defaults to 16):
+ The total stride for the mask downsampler.
+ mask_downsampler_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the mask downsampler.
+ memory_fuser_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the memory fuser.
+ memory_fuser_embed_dim (`int`, *optional*, defaults to 256):
+ The dimension of the memory fuser embedding.
+ memory_fuser_kernel_size (`int`, *optional*, defaults to 7):
+ The kernel size for the memory fuser.
+ memory_fuser_padding (`int`, *optional*, defaults to 3):
+ The padding for the memory fuser.
+ memory_fuser_layer_scale_init_value (`float`, *optional*, defaults to 1e-06):
+ The initial value for the layer scale in the memory fuser.
+ memory_fuser_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the memory fuser.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... EdgeTamVisionConfig,
+ ... EdgeTamVideoPromptEncoderConfig,
+ ... EdgeTamVideoMaskDecoderConfig,
+ ... EdgeTamVideoModel,
+ ... EdgeTamVideoConfig,
+ ... )
+
+ >>> # Initializing a EdgeTamVideoConfig with `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> configuration = EdgeTamVideoConfig()
+
+ >>> # Initializing a EdgeTamVideoModel (with random weights) from the `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> model = EdgeTamVideoModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a EdgeTamConfig from a EdgeTamVisionConfig, EdgeTamPromptEncoderConfig, and EdgeTamMaskDecoderConfig
+
+ >>> # Initializing EDGETAM vision encoder, memory attention, and memory encoder configurations
+ >>> vision_config = EdgeTamVisionConfig()
+ >>> prompt_encoder_config = EdgeTamVideoPromptEncoderConfig()
+ >>> mask_decoder_config = EdgeTamVideoMaskDecoderConfig()
+
+ >>> config = EdgeTamVideoConfig(vision_config, prompt_encoder_config, mask_decoder_config)
+ ```"""
+
+ model_type = "edgetam_video"
+ sub_configs = {
+ "vision_config": AutoConfig,
+ "prompt_encoder_config": EdgeTamVideoPromptEncoderConfig,
+ "mask_decoder_config": EdgeTamVideoMaskDecoderConfig,
+ }
+
+ def __init__(
+ self,
+ vision_config=None,
+ prompt_encoder_config=None,
+ mask_decoder_config=None,
+ initializer_range=0.02,
+ num_maskmem=7,
+ image_size=1024,
+ sigmoid_scale_for_mem_enc=20.0,
+ sigmoid_bias_for_mem_enc=-10.0,
+ enable_occlusion_spatial_embedding=True,
+ multimask_output_in_sam=True,
+ multimask_min_pt_num=0,
+ multimask_max_pt_num=1,
+ multimask_output_for_tracking=True,
+ max_object_pointers_in_encoder=16,
+ enable_temporal_pos_encoding_for_object_pointers=True,
+ # memory attention
+ memory_attention_hidden_size=256,
+ memory_attention_num_layers=2,
+ memory_attention_num_attention_heads=1,
+ memory_attention_downsample_rate=1,
+ memory_attention_mlp_hidden_size=2048,
+ memory_attention_mlp_hidden_act="relu",
+ memory_attention_dropout=0.1,
+ memory_attention_rope_theta=10000,
+ memory_attention_rope_feat_sizes=None,
+ memory_attention_rope_k_sizes=None,
+ memory_attention_rope_dropout=0.1,
+ # spatial perceiver resampler
+ perceiver_resampler_num_latents=256,
+ perceiver_resampler_num_latents_2d=256,
+ perceiver_resampler_hidden_size=64,
+ perceiver_resampler_mlp_intermediate_size=256,
+ perceiver_resampler_num_attention_heads=1,
+ perceiver_resampler_attention_head_dim=64,
+ perceiver_resampler_num_layers=2,
+ perceiver_resampler_hidden_dropout=0.0,
+ perceiver_resampler_attention_dropout=0.0,
+ # memory encoder
+ memory_encoder_hidden_size=256,
+ memory_encoder_output_channels=64,
+ mask_downsampler_embed_dim=256,
+ memory_fuser_intermediate_dim=1024,
+ mask_downsampler_kernel_size=3,
+ mask_downsampler_stride=2,
+ mask_downsampler_padding=1,
+ mask_downsampler_total_stride=16,
+ mask_downsampler_hidden_act="gelu",
+ memory_fuser_num_layers=2,
+ memory_fuser_embed_dim=256,
+ memory_fuser_kernel_size=7,
+ memory_fuser_padding=3,
+ memory_fuser_layer_scale_init_value=1e-6,
+ memory_fuser_hidden_act="gelu",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ vision_config = vision_config if vision_config is not None else {}
+ prompt_encoder_config = prompt_encoder_config if prompt_encoder_config is not None else {}
+ mask_decoder_config = mask_decoder_config if mask_decoder_config is not None else {}
+ memory_attention_rope_feat_sizes = (
+ [64, 64] if memory_attention_rope_feat_sizes is None else memory_attention_rope_feat_sizes
+ )
+ memory_attention_rope_k_sizes = (
+ [16, 16] if memory_attention_rope_k_sizes is None else memory_attention_rope_k_sizes
+ )
+
+ if isinstance(vision_config, dict):
+ vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
+ vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
+ if isinstance(prompt_encoder_config, EdgeTamVideoPromptEncoderConfig):
+ prompt_encoder_config = prompt_encoder_config.to_dict()
+ if isinstance(mask_decoder_config, EdgeTamVideoMaskDecoderConfig):
+ mask_decoder_config = mask_decoder_config.to_dict()
+
+ self.vision_config = vision_config
+ self.prompt_encoder_config = EdgeTamVideoPromptEncoderConfig(**prompt_encoder_config)
+ self.mask_decoder_config = EdgeTamVideoMaskDecoderConfig(**mask_decoder_config)
+
+ self.initializer_range = initializer_range
+ self.num_maskmem = num_maskmem # default 1 input frame + 6 previous frames
+ self.image_size = image_size
+ self.sigmoid_scale_for_mem_enc = sigmoid_scale_for_mem_enc # scale factor for mask sigmoid prob
+ self.sigmoid_bias_for_mem_enc = sigmoid_bias_for_mem_enc # bias factor for mask sigmoid prob
+ self.enable_occlusion_spatial_embedding = enable_occlusion_spatial_embedding
+ self.multimask_output_in_sam = multimask_output_in_sam
+ self.multimask_min_pt_num = multimask_min_pt_num
+ self.multimask_max_pt_num = multimask_max_pt_num
+ self.multimask_output_for_tracking = multimask_output_for_tracking
+ self.max_object_pointers_in_encoder = max_object_pointers_in_encoder
+ self.enable_temporal_pos_encoding_for_object_pointers = enable_temporal_pos_encoding_for_object_pointers
+
+ # memory attention
+ self.memory_attention_hidden_size = memory_attention_hidden_size
+ self.memory_attention_num_layers = memory_attention_num_layers
+ self.memory_attention_num_attention_heads = memory_attention_num_attention_heads
+ self.memory_attention_downsample_rate = memory_attention_downsample_rate
+ self.memory_attention_mlp_hidden_size = memory_attention_mlp_hidden_size
+ self.memory_attention_mlp_hidden_act = memory_attention_mlp_hidden_act
+ self.memory_attention_dropout = memory_attention_dropout
+ self.memory_attention_rope_theta = memory_attention_rope_theta
+ self.memory_attention_rope_feat_sizes = memory_attention_rope_feat_sizes
+ self.memory_attention_rope_k_sizes = memory_attention_rope_k_sizes
+ self.memory_attention_rope_dropout = memory_attention_rope_dropout
+
+ # spatial perceiver resampler
+ self.perceiver_resampler_num_latents = perceiver_resampler_num_latents
+ self.perceiver_resampler_num_latents_2d = perceiver_resampler_num_latents_2d
+ self.perceiver_resampler_hidden_size = perceiver_resampler_hidden_size
+ self.perceiver_resampler_mlp_intermediate_size = perceiver_resampler_mlp_intermediate_size
+ self.perceiver_resampler_attention_head_dim = perceiver_resampler_attention_head_dim
+ self.perceiver_resampler_num_attention_heads = perceiver_resampler_num_attention_heads
+ self.perceiver_resampler_num_layers = perceiver_resampler_num_layers
+ self.perceiver_resampler_hidden_dropout = perceiver_resampler_hidden_dropout
+ self.perceiver_resampler_attention_dropout = perceiver_resampler_attention_dropout
+
+ # memory encoder
+ self.memory_encoder_hidden_size = memory_encoder_hidden_size
+ self.memory_encoder_output_channels = memory_encoder_output_channels
+ self.mask_downsampler_embed_dim = mask_downsampler_embed_dim
+ self.mask_downsampler_kernel_size = mask_downsampler_kernel_size
+ self.mask_downsampler_stride = mask_downsampler_stride
+ self.mask_downsampler_padding = mask_downsampler_padding
+ self.mask_downsampler_total_stride = mask_downsampler_total_stride
+ self.mask_downsampler_hidden_act = mask_downsampler_hidden_act
+ self.memory_fuser_num_layers = memory_fuser_num_layers
+ self.memory_fuser_embed_dim = memory_fuser_embed_dim
+ self.memory_fuser_intermediate_dim = memory_fuser_intermediate_dim
+ self.memory_fuser_kernel_size = memory_fuser_kernel_size
+ self.memory_fuser_padding = memory_fuser_padding
+ self.memory_fuser_layer_scale_init_value = memory_fuser_layer_scale_init_value
+ self.memory_fuser_hidden_act = memory_fuser_hidden_act
+
+
+__all__ = ["EdgeTamVideoMaskDecoderConfig", "EdgeTamVideoPromptEncoderConfig", "EdgeTamVideoConfig"]
diff --git a/src/transformers/models/edgetam_video/convert_edgetam_video_to_hf.py b/src/transformers/models/edgetam_video/convert_edgetam_video_to_hf.py
new file mode 100644
index 000000000000..6290bef5e1c8
--- /dev/null
+++ b/src/transformers/models/edgetam_video/convert_edgetam_video_to_hf.py
@@ -0,0 +1,320 @@
+# coding=utf-8
+# 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.
+"""
+Convert SAM checkpoints from the original repository.
+
+URL: https://github.com/facebookresearch/segment-anything-2.
+"""
+
+import argparse
+import re
+
+import numpy as np
+import requests
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from transformers import (
+ EdgeTamVideoConfig,
+ EdgeTamVideoMaskDecoderConfig,
+ EdgeTamVideoModel,
+ EdgeTamVideoPromptEncoderConfig,
+ EdgeTamVisionConfig,
+ Sam2ImageProcessorFast,
+ Sam2VideoProcessor,
+ Sam2VideoVideoProcessor,
+ TimmWrapperConfig,
+)
+
+
+def get_config(model_name):
+ backbone_config = TimmWrapperConfig.from_pretrained(
+ "timm/repvit_m1.dist_in1k",
+ model_args={"in_chans": 3, "features_only": True, "out_indices": (0, 1, 2, 3)},
+ )
+ vision_config = EdgeTamVisionConfig(backbone_config=backbone_config)
+
+ prompt_encoder_config = EdgeTamVideoPromptEncoderConfig()
+ mask_decoder_config = EdgeTamVideoMaskDecoderConfig()
+ enable_temporal_pos_encoding_for_object_pointers = False
+ enable_occlusion_spatial_embedding = False
+
+ config = EdgeTamVideoConfig(
+ vision_config=vision_config,
+ prompt_encoder_config=prompt_encoder_config,
+ mask_decoder_config=mask_decoder_config,
+ enable_temporal_pos_encoding_for_object_pointers=enable_temporal_pos_encoding_for_object_pointers,
+ enable_occlusion_spatial_embedding=enable_occlusion_spatial_embedding,
+ )
+
+ return config
+
+
+KEYS_TO_MODIFY_MAPPING = {
+ "iou_prediction_head.layers.0": "iou_prediction_head.proj_in",
+ "iou_prediction_head.layers.1": "iou_prediction_head.layers.0",
+ "iou_prediction_head.layers.2": "iou_prediction_head.proj_out",
+ "mask_decoder.output_upscaling.0": "mask_decoder.upscale_conv1",
+ "mask_decoder.output_upscaling.1": "mask_decoder.upscale_layer_norm",
+ "mask_decoder.output_upscaling.3": "mask_decoder.upscale_conv2",
+ "mask_downscaling.0": "mask_embed.conv1",
+ "mask_downscaling.1": "mask_embed.layer_norm1",
+ "mask_downscaling.3": "mask_embed.conv2",
+ "mask_downscaling.4": "mask_embed.layer_norm2",
+ "mask_downscaling.6": "mask_embed.conv3",
+ "dwconv": "depthwise_conv",
+ "pwconv": "pointwise_conv",
+ "fuser": "memory_fuser",
+ "point_embeddings": "point_embed",
+ "pe_layer.positional_encoding_gaussian_matrix": "shared_embedding.positional_embedding",
+ "obj_ptr_tpos_proj": "temporal_positional_encoding_projection_layer",
+ "no_obj_embed_spatial": "occlusion_spatial_embedding_parameter",
+ "sam_prompt_encoder": "prompt_encoder",
+ "sam_mask_decoder": "mask_decoder",
+ "maskmem_tpos_enc": "memory_temporal_positional_encoding",
+ "gamma": "scale",
+ "image_encoder.neck": "vision_encoder.neck",
+ "image_encoder": "vision_encoder.backbone",
+ "neck.0": "neck.conv1",
+ "neck.1": "neck.layer_norm1",
+ "neck.2": "neck.conv2",
+ "neck.3": "neck.layer_norm2",
+ "pix_feat_proj": "feature_projection",
+ "patch_embed.proj": "patch_embed.projection",
+ "no_mem_embed": "no_memory_embedding",
+ "no_mem_pos_enc": "no_memory_positional_encoding",
+ "obj_ptr": "object_pointer",
+ ".norm": ".layer_norm",
+ "trunk.": "",
+ "out_proj": "o_proj",
+ "body.": "timm_model.",
+ "ff.0": "mlp.layer_norm",
+ "ff.1": "mlp.up_proj",
+ "ff.3": "mlp.down_proj",
+}
+
+
+def replace_keys(state_dict):
+ model_state_dict = {}
+ output_hypernetworks_mlps_pattern = r".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"
+ output_mask_decoder_mlps_pattern = r"mask_decoder.transformer.layers.(\d+).mlp.layers.(\d+).*"
+ output_mask_decoder_score_head_pattern = r"mask_decoder.pred_obj_score_head.layers.(\d+).*"
+ output_vision_encoder_mlps_pattern = r"vision_encoder.backbone.blocks.(\d+).mlp.layers.(\d+).*"
+ output_vision_encoder_neck_pattern = r"vision_encoder.neck.convs.(\d+).conv"
+ output_memory_encoder_projection_pattern = r"memory_encoder.o_proj.*"
+ memory_attention_pattern = r"memory_attention.*"
+ output_object_pointer_proj_pattern = r"object_pointer_proj.layers.(\d+).*"
+ output_memory_encoder_mask_downsampler_pattern = r"memory_encoder.mask_downsampler.encoder.(\d+).*"
+ perceiver_resampler_patterns = {
+ r"spatial_perceiver.latents": r"spatial_perceiver.latents_1d",
+ r"spatial_perceiver.latents_1d_2d": r"spatial_perceiver.latents_2d",
+ r"spatial_perceiver.layers.(\d+).attn.layer_norm_x": r"spatial_perceiver.layers.\1.layer_norm_input",
+ r"spatial_perceiver.layers.(\d+).attn.layer_norm_latents": r"spatial_perceiver.layers.\1.layer_norm_latents",
+ r"spatial_perceiver.layers.(\d+).self_attn.layer_norm": r"spatial_perceiver.layers.\1.layer_norm_self",
+ r"spatial_perceiver.layers.(\d+).attn.to_q": r"spatial_perceiver.layers.\1.cross_attention.q_proj",
+ r"spatial_perceiver.layers.(\d+).attn.to_kv": r"spatial_perceiver.layers.\1.cross_attention.kv_proj_combined",
+ r"spatial_perceiver.layers.(\d+).attn.to_out": r"spatial_perceiver.layers.\1.cross_attention.o_proj",
+ r"spatial_perceiver.layers.(\d+).self_attn.to_q": r"spatial_perceiver.layers.\1.self_attention.q_proj",
+ r"spatial_perceiver.layers.(\d+).self_attn.to_kv": r"spatial_perceiver.layers.\1.self_attention.kv_proj_combined",
+ r"spatial_perceiver.layers.(\d+).self_attn.to_out": r"spatial_perceiver.layers.\1.self_attention.o_proj",
+ r"spatial_perceiver.layers.(\d+).attn": r"spatial_perceiver.layers.\1.cross_attention",
+ r"spatial_perceiver.layers.(\d+).self_attn": r"spatial_perceiver.layers.\1.self_attention",
+ }
+
+ for key, value in state_dict.items():
+ 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)
+
+ for pattern, replacement in perceiver_resampler_patterns.items():
+ if re.match(pattern, key):
+ key = re.sub(pattern, replacement, key)
+
+ # vision_encoder.blocks.0.mlp.layers.1.weight -> vision_encoder.blocks.0.mlp.proj_out.weight
+ if re.match(output_vision_encoder_mlps_pattern, key):
+ layer_nb = int(re.match(output_vision_encoder_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "proj_out")
+
+ if re.match(memory_attention_pattern, key):
+ key = key.replace("linear1", "mlp.up_proj")
+ key = key.replace("linear2", "mlp.down_proj")
+
+ # mask_decoder.transformer.layers.0.mlp.layers.1.weight -> mask_decoder.transformer.layers.1.mlp.proj_out.weight
+ if re.match(output_mask_decoder_mlps_pattern, key):
+ layer_nb = int(re.match(output_mask_decoder_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("mlp.layers.0", "mlp.proj_in")
+ elif layer_nb == 1:
+ key = key.replace("mlp.layers.1", "mlp.proj_out")
+
+ # mask_decoder.pred_obj_score_head.layers.1.weight -> mask_decoder.pred_obj_score_head.proj_in.weight
+ if re.match(output_mask_decoder_score_head_pattern, key):
+ layer_nb = int(re.match(output_mask_decoder_score_head_pattern, key).group(1))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ if re.match(output_hypernetworks_mlps_pattern, key):
+ layer_nb = int(re.match(output_hypernetworks_mlps_pattern, key).group(2))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ # vision_encoder.neck.convs.1.conv.bias -> vision_encoder.neck.convs.1.bias
+ if re.match(output_vision_encoder_neck_pattern, key):
+ key = key.replace(".conv.", ".")
+
+ # memory_encoder.o_proj.weight -> memory_encoder.projection.weight
+ if re.match(output_memory_encoder_projection_pattern, key):
+ key = key.replace(".o_proj.", ".projection.")
+
+ if re.match(output_object_pointer_proj_pattern, key):
+ layer_nb = int(re.match(output_object_pointer_proj_pattern, key).group(1))
+ if layer_nb == 0:
+ key = key.replace("layers.0", "proj_in")
+ elif layer_nb == 1:
+ key = key.replace("layers.1", "layers.0")
+ elif layer_nb == 2:
+ key = key.replace("layers.2", "proj_out")
+
+ key = key.replace("layers.2", "proj_out")
+
+ if re.match(output_memory_encoder_mask_downsampler_pattern, key):
+ layer_nb = int(re.match(output_memory_encoder_mask_downsampler_pattern, key).group(1))
+ if layer_nb == 12:
+ key = key.replace(f"encoder.{layer_nb}", "final_conv")
+ elif layer_nb % 3 == 0:
+ key = key.replace(f"encoder.{layer_nb}", f"layers.{layer_nb // 3}.conv")
+ elif layer_nb % 3 == 1:
+ key = key.replace(f"encoder.{layer_nb}", f"layers.{layer_nb // 3}.layer_norm")
+ if "kv_proj_combined" in key:
+ # Split the weight tensor in half along dimension 0 (output dimension)
+ k_weight, v_weight = torch.chunk(value, 2, dim=0)
+ # Create the k_proj and v_proj keys
+ k_key = key.replace("kv_proj_combined", "k_proj")
+ v_key = key.replace("kv_proj_combined", "v_proj")
+ model_state_dict[k_key] = k_weight
+ model_state_dict[v_key] = v_weight
+ continue
+
+ model_state_dict[key] = value
+
+ model_state_dict["shared_image_embedding.positional_embedding"] = model_state_dict[
+ "prompt_encoder.shared_embedding.positional_embedding"
+ ]
+ model_state_dict["prompt_encoder.point_embed.weight"] = torch.cat(
+ [model_state_dict.pop(f"prompt_encoder.point_embed.{i}.weight") for i in range(4)],
+ dim=0,
+ )
+
+ return model_state_dict
+
+
+def convert_edgetam_checkpoint(model_name, checkpoint_path, pytorch_dump_folder, push_to_hub, run_sanity_check):
+ config = get_config(model_name)
+
+ state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
+ state_dict = replace_keys(state_dict)
+
+ image_processor = Sam2ImageProcessorFast()
+ video_processor = Sam2VideoVideoProcessor()
+ processor = Sam2VideoProcessor(image_processor=image_processor, video_processor=video_processor)
+ hf_model = EdgeTamVideoModel(config)
+ hf_model.eval()
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ missing_keys, unexpected_keys = hf_model.load_state_dict(state_dict, strict=True)
+ hf_model = hf_model.to(device)
+ print("Missing keys:", missing_keys)
+ print("Unexpected keys:", unexpected_keys)
+
+ if run_sanity_check:
+ img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+
+ input_points = [[[[1000, 600]]]]
+ input_labels = [[[1]]]
+
+ inputs = processor(
+ images=np.array(raw_image), input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(device)
+
+ with torch.no_grad():
+ output = hf_model._single_frame_forward(**inputs)
+ scores = output.iou_scores.squeeze()
+
+ assert torch.allclose(scores, torch.tensor([0.0356, 0.2141, 0.9707]).cuda(), atol=1e-3)
+
+ if pytorch_dump_folder is not None:
+ processor.save_pretrained(pytorch_dump_folder)
+ hf_model.save_pretrained(pytorch_dump_folder)
+
+ if push_to_hub:
+ repo_id = f"yonigozlan/{pytorch_dump_folder.split('/')[-1]}"
+ processor.push_to_hub(repo_id)
+ hf_model.push_to_hub(repo_id)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ choices = ["EdgeTAM"]
+ parser.add_argument(
+ "--model_name",
+ default="EdgeTAM",
+ choices=choices,
+ type=str,
+ help="Name of the original model to convert",
+ )
+ parser.add_argument(
+ "--checkpoint_path",
+ type=str,
+ required=False,
+ help="Path to the original checkpoint",
+ )
+ parser.add_argument("--pytorch_dump_folder_path", default="", type=str, help="Path to the output PyTorch model.")
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether to push the model and processor to the hub after converting",
+ )
+ parser.add_argument(
+ "--run_sanity_check",
+ action="store_true",
+ help="Whether to run the sanity check after converting",
+ )
+
+ args = parser.parse_args()
+
+ hf_model_name = args.model_name.replace("_", "-")
+ checkpoint_path = (
+ hf_hub_download(f"facebook/{hf_model_name}", f"{args.model_name.lower()}.pt")
+ if args.checkpoint_path is None
+ else args.checkpoint_path
+ )
+
+ convert_edgetam_checkpoint(
+ args.model_name, checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub, args.run_sanity_check
+ )
diff --git a/src/transformers/models/edgetam_video/modeling_edgetam_video.py b/src/transformers/models/edgetam_video/modeling_edgetam_video.py
new file mode 100644
index 000000000000..3ba7ab4ebf2f
--- /dev/null
+++ b/src/transformers/models/edgetam_video/modeling_edgetam_video.py
@@ -0,0 +1,3062 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/edgetam_video/modular_edgetam_video.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_edgetam_video.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# 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 math
+from collections import OrderedDict
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from tqdm import tqdm
+
+from transformers.utils.generic import OutputRecorder
+
+from ...activations import ACT2FN
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import compile_compatible_method_lru_cache
+from ...utils import ModelOutput, auto_docstring
+from ...utils.generic import TransformersKwargs
+from ..auto import AutoModel
+from .configuration_edgetam_video import (
+ EdgeTamVideoConfig,
+ EdgeTamVideoMaskDecoderConfig,
+ EdgeTamVideoPromptEncoderConfig,
+)
+
+
+class EdgeTamVideoLayerNorm(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
+
+
+# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt)
+class EdgeTamVideoMemoryFuserCXBlock(GradientCheckpointingLayer):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.depthwise_conv = nn.Conv2d(
+ config.memory_fuser_embed_dim,
+ config.memory_fuser_embed_dim,
+ kernel_size=config.memory_fuser_kernel_size,
+ padding=config.memory_fuser_padding,
+ groups=config.memory_fuser_embed_dim,
+ ) # depthwise conv
+ self.layer_norm = EdgeTamVideoLayerNorm(config.memory_fuser_embed_dim, eps=1e-6, data_format="channels_first")
+ self.activation = ACT2FN[config.memory_fuser_hidden_act]
+ self.pointwise_conv1 = nn.Linear(
+ config.memory_fuser_embed_dim, config.memory_fuser_intermediate_dim
+ ) # pointwise/1x1 convs, implemented with linear layers
+ self.pointwise_conv2 = nn.Linear(config.memory_fuser_intermediate_dim, config.memory_fuser_embed_dim)
+ self.scale = nn.Parameter(
+ config.memory_fuser_layer_scale_init_value * torch.ones(config.memory_fuser_embed_dim),
+ requires_grad=True,
+ )
+
+ def forward(self, hidden_states):
+ input = hidden_states
+ hidden_states = self.depthwise_conv(hidden_states)
+ hidden_states = self.layer_norm(hidden_states)
+ hidden_states = hidden_states.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+ hidden_states = self.pointwise_conv1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.pointwise_conv2(hidden_states)
+ hidden_states = self.scale * hidden_states
+ hidden_states = hidden_states.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+
+ hidden_states = input + hidden_states
+ return hidden_states
+
+
+@dataclass
+@auto_docstring(custom_intro="Base class for the vision encoder's outputs.")
+class EdgeTamVideoVisionEncoderOutput(ModelOutput):
+ r"""
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the model.
+ fpn_hidden_states (`tuple(torch.FloatTensor)`):
+ Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+ `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck.
+ fpn_position_encoding (`tuple(torch.FloatTensor)`):
+ Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+ `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+ one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the
+ model at the output of each stage.
+ attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
+ the self-attention heads.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ fpn_hidden_states: Optional[torch.FloatTensor] = None
+ fpn_position_encoding: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+class EdgeTamVideoVisionRotaryEmbedding(nn.Module):
+ """
+ Vision Rotary Position Embedding for SAM2, following transformers library standards.
+ Supports 2D (axial) rotary embeddings for spatial dimensions.
+ """
+
+ def __init__(self, config: EdgeTamVideoConfig, end_x: Optional[int] = None, end_y: Optional[int] = None):
+ super().__init__()
+ dim = config.memory_attention_hidden_size // (
+ config.memory_attention_downsample_rate * config.memory_attention_num_attention_heads
+ )
+ # Ensure even dimension for proper axial splitting
+ if dim % 4 != 0:
+ raise ValueError("Dimension must be divisible by 4 for axial RoPE")
+ end_x, end_y = config.memory_attention_rope_feat_sizes if end_x is None else (end_x, end_y)
+ freqs = 1.0 / (config.memory_attention_rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+
+ # Generate 2D position indices for axial rotary embedding
+ flattened_indices = torch.arange(end_x * end_y, dtype=torch.long)
+ x_positions = flattened_indices % end_x
+ y_positions = torch.div(flattened_indices, end_x, rounding_mode="floor")
+ freqs_x = torch.outer(x_positions, freqs).float()
+ freqs_y = torch.outer(y_positions, freqs).float()
+ inv_freq = torch.cat([freqs_x, freqs_y], dim=-1)
+ inv_freq = inv_freq.repeat_interleave(2, dim=-1)
+ # directly register the cos and sin embeddings as we have a fixed feature shape
+ self.register_buffer("rope_embeddings_cos", inv_freq.cos(), persistent=False)
+ self.register_buffer("rope_embeddings_sin", inv_freq.sin(), persistent=False)
+
+ @torch.no_grad()
+ def forward(self) -> tuple[torch.Tensor, torch.Tensor]:
+ # As the feature map size is fixed, we can just return the pre-computed embeddings.
+ return self.rope_embeddings_cos, self.rope_embeddings_sin
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class EdgeTamVideoAttention(nn.Module):
+ """
+ EDGETAM_VIDEO's attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and
+ values.
+ """
+
+ def __init__(self, config, downsample_rate=None):
+ super().__init__()
+ downsample_rate = config.attention_downsample_rate if downsample_rate is None else downsample_rate
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.internal_dim = config.hidden_size // downsample_rate
+ self.num_attention_heads = config.num_attention_heads
+ self.head_dim = self.internal_dim // config.num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.v_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_similarity: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=attention_similarity,
+ dropout=0.0,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+def rotate_pairwise(x):
+ """
+ pairwise rotation of the hidden dims of the input. Differerent from Llama Half-Tensor Rotation.
+
+ This is an optimized version of the following more explicit implementation:
+ ```python
+ x_rotated = torch.zeros_like(x, dtype=x.dtype, device=x.device)
+ x_rotated[..., ::2] = -x[..., 1::2]
+ x_rotated[..., 1::2] = x[..., ::2]
+ return x_rotated
+ ```
+ """
+ x = x.view(*x.shape[:-1], -1, 2)
+ x1, x2 = x.unbind(dim=-1)
+ x = torch.stack((-x2, x1), dim=-1)
+ return x.flatten(start_dim=-2)
+
+
+def apply_rotary_pos_emb_2d_self_attn(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary position embedding to query and key tensors for self-attention.
+
+ Args:
+ q: Query tensor of shape (..., seq_len, head_dim)
+ k: Key tensor of shape (..., seq_len, head_dim)
+ cos: Cosine position embedding of shape (seq_len, head_dim)
+ sin: Sine position embedding of shape (seq_len, head_dim)
+
+ Returns:
+ Rotated (q, k) tensors
+ """
+ # Apply RoPE to queries
+ q_embed = q.float() # force upscale to float32 as in the original implementation
+ q_embed = (q_embed * cos) + (rotate_pairwise(q_embed) * sin)
+
+ # Apply RoPE to keys (same embeddings as queries for self-attention)
+ k_embed = k.float() # force upscale to float32 as in the original implementation
+ k_embed = (k_embed * cos) + (rotate_pairwise(k_embed) * sin)
+
+ return q_embed.type_as(q), k_embed.type_as(k)
+
+
+class EdgeTamVideoRoPESelfAttention(nn.Module):
+ """Self-attention with rotary position encoding."""
+
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.internal_dim = self.hidden_size // config.memory_attention_downsample_rate
+ self.num_attention_heads = config.memory_attention_num_attention_heads
+ self.head_dim = self.internal_dim // config.memory_attention_num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.v_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+ self.dropout_p = config.memory_attention_rope_dropout
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tensor:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ # Apply rotary position encoding for self-attention
+ query, key = apply_rotary_pos_emb_2d_self_attn(query, key, cos=cos, sin=sin)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.dropout_p,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+def apply_rotary_pos_emb_2d_cross_attn(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ cos_k: torch.Tensor,
+ sin_k: torch.Tensor,
+ num_k_exclude_rope: int = 0,
+ repeat_freqs_k: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary position embedding to query and key tensors for cross-attention.
+
+ Args:
+ q: Query tensor of shape (..., seq_len, head_dim)
+ k: Key tensor of shape (..., seq_len, head_dim)
+ cos: Cosine position embedding of shape (seq_len, head_dim)
+ sin: Sine position embedding of shape (seq_len, head_dim)
+ cos_k: Cosine position embedding for keys of shape (seq_len, head_dim)
+ sin_k: Sine position embedding for keys of shape (seq_len, head_dim)
+ num_k_exclude_rope: Number of tokens at end of k to exclude from RoPE (e.g., object pointer tokens)
+ repeat_freqs_k: Frequency repetition for keys in cross-attention (e.g., for spatial memory tokens)
+
+ Returns:
+ Rotated (q, k) tensors
+ """
+ # Apply RoPE to queries (always straightforward)
+ q_embed = q.float()
+ q_embed = (q_embed * cos) + (rotate_pairwise(q_embed) * sin)
+
+ # Split keys: RoPE tokens and excluded tokens (e.g., object pointers)
+ num_total_k_tokens = k.shape[-2]
+ k_for_rope = k[..., : num_total_k_tokens - num_k_exclude_rope, :]
+ k_excluded = k[..., num_total_k_tokens - num_k_exclude_rope :, :]
+
+ # Early return if no keys need RoPE
+ if k_for_rope.shape[-2] == 0:
+ return q_embed.type_as(q), k_excluded
+
+ batch_size, num_heads, k_seq_len, channels_per_head = k_for_rope.shape
+
+ # Handle temporal/spatial token structure for memory
+ # Keys have temporal + spatial structure, only spatial tokens get RoPE
+ tokens_per_group = k_seq_len // repeat_freqs_k
+ spatial_tokens = cos_k.shape[-2]
+ temporal_tokens = tokens_per_group - spatial_tokens
+
+ # Reshape and separate temporal/spatial tokens
+ k_grouped = k_for_rope.view(batch_size, num_heads, repeat_freqs_k, tokens_per_group, channels_per_head)
+ k_temporal = k_grouped[..., :temporal_tokens, :].reshape(batch_size, num_heads, -1, channels_per_head)
+ k_spatial = k_grouped[..., temporal_tokens:, :].reshape(batch_size, num_heads, -1, channels_per_head)
+
+ # Only apply RoPE to spatial tokens
+ k_rope_input = k_spatial
+
+ # Prepare position embeddings for repeated groups
+ if repeat_freqs_k > 1:
+ cos_k = cos_k.repeat(1, 1, repeat_freqs_k, 1)
+ sin_k = sin_k.repeat(1, 1, repeat_freqs_k, 1)
+
+ # Apply RoPE to spatial tokens
+ k_spatial_embed = k_rope_input.float()
+ k_spatial_embed = (k_spatial_embed * cos_k) + (rotate_pairwise(k_spatial_embed) * sin_k)
+
+ # Reconstruct: temporal + spatial tokens back to original structure
+ k_spatial_reshaped = k_spatial_embed.view(batch_size, num_heads, repeat_freqs_k, -1, channels_per_head)
+ k_temporal_reshaped = k_temporal.view(batch_size, num_heads, repeat_freqs_k, -1, channels_per_head)
+ k_final = torch.cat([k_temporal_reshaped, k_spatial_reshaped], dim=3)
+ k_final = k_final.view(batch_size, num_heads, k_seq_len, channels_per_head)
+
+ # Combine RoPE-processed keys with excluded tokens
+ k_embed = torch.cat([k_final.type_as(k), k_excluded], dim=-2)
+ return q_embed.type_as(q), k_embed
+
+
+class EdgeTamVideoRoPECrossAttention(nn.Module):
+ """Cross-attention with rotary position encoding."""
+
+ def __init__(self, config: EdgeTamVideoConfig, kv_in_dim: int):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.internal_dim = self.hidden_size // config.memory_attention_downsample_rate
+ self.num_attention_heads = config.memory_attention_num_attention_heads
+ self.head_dim = self.internal_dim // config.memory_attention_num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.kv_in_dim = kv_in_dim
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+ self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+ self.dropout_p = config.memory_attention_rope_dropout
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ position_embeddings_k: tuple[torch.Tensor, torch.Tensor],
+ num_k_exclude_rope: int = 0,
+ rope_k_repeat: int = 0,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tensor:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ cos_k, sin_k = position_embeddings_k
+ # Apply rotary position encoding for cross-attention
+ query, key = apply_rotary_pos_emb_2d_cross_attn(
+ query,
+ key,
+ cos=cos,
+ sin=sin,
+ cos_k=cos_k,
+ sin_k=sin_k,
+ repeat_freqs_k=rope_k_repeat,
+ num_k_exclude_rope=num_k_exclude_rope,
+ )
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.dropout_p,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class EdgeTamVideoTwoWayAttentionBlock(nn.Module):
+ def __init__(self, config: EdgeTamVideoMaskDecoderConfig, skip_first_layer_pe: bool = False):
+ """
+ A transformer block with four layers:
+ (1) self-attention of sparse inputs (2) cross attention of sparse inputs -> dense inputs (3) mlp block on
+ sparse inputs (4) cross attention of dense inputs -> sparse inputs
+
+ Arguments:
+ config (`EdgeTamVideoMaskDecoderConfig`):
+ The configuration file used to instantiate the block
+ attention_downsample_rate (*optionalk*, int, defaults to 2):
+ The downsample ratio of the block used to reduce the inner dim of the attention.
+ skip_first_layer_pe (*optional*, bool, defaults to `False`):
+ Whether or not to skip the addition of the query_point_embedding on the first layer.
+ """
+ super().__init__()
+ self.self_attn = EdgeTamVideoAttention(config, downsample_rate=1)
+ self.layer_norm1 = nn.LayerNorm(config.hidden_size)
+
+ self.cross_attn_token_to_image = EdgeTamVideoAttention(config)
+ self.layer_norm2 = nn.LayerNorm(config.hidden_size)
+
+ self.mlp = EdgeTamVideoFeedForward(
+ config.hidden_size, config.mlp_dim, config.hidden_size, num_layers=config.num_hidden_layers
+ )
+ self.layer_norm3 = nn.LayerNorm(config.hidden_size)
+
+ self.layer_norm4 = nn.LayerNorm(config.hidden_size)
+ self.cross_attn_image_to_token = EdgeTamVideoAttention(config)
+
+ self.skip_first_layer_pe = skip_first_layer_pe
+
+ def forward(
+ self,
+ queries: Tensor,
+ keys: Tensor,
+ query_point_embedding: Tensor,
+ key_point_embedding: Tensor,
+ attention_similarity: Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ # Self attention block
+ if self.skip_first_layer_pe:
+ queries, _ = self.self_attn(query=queries, key=queries, value=queries)
+ else:
+ query = queries + query_point_embedding
+ attn_out, _ = self.self_attn(query=query, key=query, value=queries)
+ queries = queries + attn_out
+ queries = self.layer_norm1(queries)
+
+ # Cross attention block, tokens attending to image embedding
+ query = queries + query_point_embedding
+ key = keys + key_point_embedding
+
+ attn_out, _ = self.cross_attn_token_to_image(
+ query=query, key=key, value=keys, attention_similarity=attention_similarity
+ )
+ queries = queries + attn_out
+
+ queries = self.layer_norm2(queries)
+
+ # MLP block
+ mlp_out = self.mlp(queries)
+ queries = queries + mlp_out
+ queries = self.layer_norm3(queries)
+
+ # Cross attention block, image embedding attending to tokens
+ query = queries + query_point_embedding
+ key = keys + key_point_embedding
+
+ attn_out, _ = self.cross_attn_image_to_token(query=key, key=query, value=queries)
+ keys = keys + attn_out
+
+ keys = self.layer_norm4(keys)
+ return queries, keys, attn_out
+
+
+# copied and adapted from original implementation, also practically equal to DetrSinePositionEmbedding
+class EdgeTamVideoPositionEmbeddingSine(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: Optional[float] = 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=2)
+ def forward(
+ self,
+ shape: torch.Size,
+ device: Union[torch.device, str],
+ dtype: torch.dtype,
+ mask: Optional[Tensor] = 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 EdgeTamVideoMemoryFuser(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [EdgeTamVideoMemoryFuserCXBlock(config) for _ in range(config.memory_fuser_num_layers)]
+ )
+
+ def forward(self, hidden_states):
+ # normally hidden_states: (N, C, H, W)
+ for layer in self.layers:
+ hidden_states = layer(hidden_states)
+ return hidden_states
+
+
+class EdgeTamVideoMaskDownSamplerLayer(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig, in_channels: int, out_channels: int):
+ super().__init__()
+ self.conv = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=config.mask_downsampler_kernel_size,
+ stride=config.mask_downsampler_stride,
+ padding=config.mask_downsampler_padding,
+ )
+ self.layer_norm = EdgeTamVideoLayerNorm(out_channels, eps=1e-6, data_format="channels_first")
+ self.activation = ACT2FN[config.mask_downsampler_hidden_act]
+
+ def forward(self, x):
+ return self.activation(self.layer_norm(self.conv(x)))
+
+
+class EdgeTamVideoMaskDownSampler(nn.Module):
+ """
+ Progressively downsample a mask by total_stride, each time by stride.
+ Note that LayerNorm is applied per *token*, like in ViT.
+
+ With each downsample (by a factor stride**2), channel capacity increases by the same factor.
+ In the end, we linearly project to embed_dim channels.
+ """
+
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+
+ num_layers = int(math.log2(config.mask_downsampler_total_stride) // math.log2(config.mask_downsampler_stride))
+
+ self.layers = nn.ModuleList()
+ self.activation = ACT2FN[config.mask_downsampler_hidden_act]
+ mask_in_chans, mask_out_chans = 1, 1
+ for _ in range(num_layers):
+ mask_out_chans = mask_in_chans * (config.mask_downsampler_stride**2)
+ self.layers.append(EdgeTamVideoMaskDownSamplerLayer(config, mask_in_chans, mask_out_chans))
+ mask_in_chans = mask_out_chans
+
+ self.final_conv = nn.Conv2d(mask_out_chans, config.mask_downsampler_embed_dim, kernel_size=1)
+
+ def forward(self, x):
+ for layer in self.layers:
+ x = layer(x)
+ x = self.final_conv(x)
+ return x
+
+
+class EdgeTamVideoMemoryEncoder(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+
+ hidden_size = config.memory_encoder_hidden_size
+ output_channels = config.memory_encoder_output_channels
+ self.mask_downsampler = EdgeTamVideoMaskDownSampler(config)
+ self.feature_projection = nn.Conv2d(hidden_size, hidden_size, kernel_size=1)
+ self.memory_fuser = EdgeTamVideoMemoryFuser(config)
+ self.position_encoding = EdgeTamVideoPositionEmbeddingSine(num_pos_feats=output_channels // 2, normalize=True)
+ self.projection = nn.Conv2d(hidden_size, output_channels, kernel_size=1)
+
+ def forward(
+ self,
+ vision_features: torch.Tensor,
+ masks: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ ## Process masks
+ masks = self.mask_downsampler(masks)
+ ## Fuse pixel_features and downsampled masks
+
+ vision_features = self.feature_projection(vision_features)
+ vision_features = vision_features + masks
+ vision_features = self.memory_fuser(vision_features)
+ vision_features = self.projection(vision_features)
+
+ vision_pos_enc = self.position_encoding(vision_features.shape, vision_features.device, vision_features.dtype)
+
+ return vision_features, vision_pos_enc
+
+
+class EdgeTamVideoFeedForward(nn.Module):
+ def __init__(
+ self,
+ input_dim: int,
+ hidden_dim: int,
+ output_dim: int,
+ num_layers: int,
+ activation: str = "relu",
+ sigmoid_output: bool = False,
+ ):
+ super().__init__()
+ self.num_layers = num_layers
+ self.activation = ACT2FN[activation]
+ self.proj_in = nn.Linear(input_dim, hidden_dim)
+ self.proj_out = nn.Linear(hidden_dim, output_dim)
+ self.layers = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers - 2)])
+ self.sigmoid_output = sigmoid_output
+
+ def forward(self, hidden_states):
+ hidden_states = self.proj_in(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ for layer in self.layers:
+ hidden_states = self.activation(layer(hidden_states))
+
+ hidden_states = self.proj_out(hidden_states)
+ if self.sigmoid_output:
+ hidden_states = F.sigmoid(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class EdgeTamVideoPreTrainedModel(PreTrainedModel):
+ config_class = EdgeTamVideoConfig
+ base_model_prefix = "edgetam_video"
+ main_input_name = "pixel_values"
+ _supports_sdpa = True
+ _supports_flash_attn_2 = True
+ _supports_attention_backend = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, (nn.LayerNorm, EdgeTamVideoLayerNorm)):
+ module.weight.data.fill_(1.0)
+ module.bias.data.zero_()
+ elif isinstance(module, EdgeTamVideoModel):
+ if module.no_memory_positional_encoding is not None:
+ module.no_memory_positional_encoding.data.zero_()
+ if module.memory_temporal_positional_encoding is not None:
+ module.memory_temporal_positional_encoding.data.zero_()
+ if module.no_object_pointer is not None:
+ module.no_object_pointer.data.zero_()
+ if module.occlusion_spatial_embedding_parameter is not None:
+ module.occlusion_spatial_embedding_parameter.data.zero_()
+ if isinstance(module, EdgeTamVideoMemoryFuserCXBlock):
+ if module.scale is not None:
+ module.scale.data.zero_()
+
+
+class EdgeTamVideoInferenceCache:
+ """Cache for vision features and model constants."""
+
+ def __init__(
+ self,
+ inference_device: Union[torch.device, str] = "cpu",
+ inference_state_device: Union[torch.device, str] = "cpu",
+ max_vision_features_cache_size: int = 1,
+ ):
+ self.inference_device = inference_device
+ self.inference_state_device = inference_state_device
+ self.max_vision_features_cache_size = max_vision_features_cache_size
+
+ self._vision_features = {}
+
+ def cache_vision_features(self, frame_idx: int, features: dict):
+ """Cache vision features with automatic device management."""
+ cached = {}
+ if len(self._vision_features) >= self.max_vision_features_cache_size:
+ # remove the oldest frame
+ self._vision_features.pop(min(self._vision_features.keys()))
+
+ for key, value in features.items():
+ if isinstance(value, torch.Tensor):
+ cached[key] = value.to(self.inference_state_device, non_blocking=True)
+ elif isinstance(value, (list, tuple)) and value and isinstance(value[0], torch.Tensor):
+ cached[key] = [v.to(self.inference_state_device, non_blocking=True) for v in value]
+ else:
+ cached[key] = value
+ self._vision_features[frame_idx] = cached
+
+ def get_vision_features(self, frame_idx: int) -> Optional[dict]:
+ """Get cached vision features, automatically moved to inference device."""
+ if frame_idx not in self._vision_features:
+ return None
+
+ cached = self._vision_features[frame_idx]
+ moved = {}
+ for key, value in cached.items():
+ if isinstance(value, torch.Tensor):
+ moved[key] = value.to(self.inference_device, non_blocking=True)
+ elif isinstance(value, (list, tuple)) and value and isinstance(value[0], torch.Tensor):
+ moved[key] = [v.to(self.inference_device, non_blocking=True) for v in value]
+ else:
+ moved[key] = value
+ return moved
+
+ def clear_all(self):
+ """Clear all cached data."""
+ self._vision_features.clear()
+
+
+class EdgeTamVideoInferenceSession:
+ r"""
+ Manages video inference session parameters, state and cache.
+
+ Args:
+ video (`torch.FloatTensor`, *optional*):
+ The video to process. No need to provide when streaming.
+ video_height (`int`, *optional*):
+ The height of the video.
+ video_width (`int`, *optional*):
+ The width of the video.
+ inference_device (`torch.device`, *optional*, defaults to `"cpu"`):
+ The device to use for inference.
+ inference_state_device (`torch.device`, *optional*, defaults to `"cpu"`):
+ The device to store the inference state on.
+ video_storage_device (`torch.device`, *optional*, defaults to `"cpu"`):
+ The device to store the video on.
+ dtype (`torch.dtype`, *optional*, defaults to `"float32"`):
+ The dtype to use for the video.
+ max_vision_features_cache_size (`int`, *optional*, defaults to 1):
+ The maximum number of vision features to cache.
+ """
+
+ def __init__(
+ self,
+ video: Optional[torch.FloatTensor] = None,
+ video_height: Optional[int] = None,
+ video_width: Optional[int] = None,
+ inference_device: Union[torch.device, str] = "cpu",
+ inference_state_device: Union[torch.device, str] = "cpu",
+ video_storage_device: Union[torch.device, str] = "cpu",
+ dtype: Union[torch.dtype, str] = "float32",
+ max_vision_features_cache_size: int = 1,
+ ):
+ # store as a dictionary to avoid double memory allocation with torch.cat when adding new frames
+ self.processed_frames = (
+ dict(enumerate(video.to(video_storage_device, dtype=dtype))) if video is not None else None
+ )
+ self.video_height = video_height
+ self.video_width = video_width
+
+ self.inference_device = inference_device
+ self.inference_state_device = inference_state_device
+ self.video_storage_device = video_storage_device
+ self.dtype = dtype
+ self.max_vision_features_cache_size = max_vision_features_cache_size
+
+ # Cache for computed features
+ self.cache = EdgeTamVideoInferenceCache(
+ inference_device=self.inference_device,
+ inference_state_device=self.inference_state_device,
+ max_vision_features_cache_size=self.max_vision_features_cache_size,
+ )
+
+ # Persistent object tracking state
+ self._obj_id_to_idx = OrderedDict()
+ self._obj_idx_to_id = OrderedDict()
+ self.obj_ids = []
+
+ # Persistent user inputs
+ self.point_inputs_per_obj = {}
+ self.mask_inputs_per_obj = {}
+
+ # Persistent model outputs/history
+ self.output_dict_per_obj = {}
+ self.frames_tracked_per_obj = {}
+
+ # Session state flags
+ self.obj_with_new_inputs = []
+
+ @property
+ def num_frames(self) -> Optional[int]:
+ return len(self.processed_frames) if self.processed_frames is not None else None
+
+ # Object management
+ def obj_id_to_idx(self, obj_id: int) -> int:
+ """Map object ID to index, creating new entry if needed."""
+ obj_idx = self._obj_id_to_idx.get(obj_id, None)
+ if obj_idx is not None:
+ return obj_idx
+
+ obj_idx = len(self._obj_id_to_idx)
+ self._obj_id_to_idx[obj_id] = obj_idx
+ self._obj_idx_to_id[obj_idx] = obj_id
+ self.obj_ids = list(self._obj_id_to_idx)
+
+ self.point_inputs_per_obj[obj_idx] = {}
+ self.mask_inputs_per_obj[obj_idx] = {}
+ self.output_dict_per_obj[obj_idx] = {
+ "cond_frame_outputs": {},
+ "non_cond_frame_outputs": {},
+ }
+ self.frames_tracked_per_obj[obj_idx] = {}
+
+ return obj_idx
+
+ # Video Inference specific functions
+ def obj_idx_to_id(self, obj_idx: int) -> int:
+ """Map model-side object index to client-side object id."""
+ return self._obj_idx_to_id[obj_idx]
+
+ def get_obj_num(self) -> int:
+ """Get the total number of unique object ids received so far in this session."""
+ return len(self._obj_idx_to_id)
+
+ # Input management with device handling
+ def add_point_inputs(self, obj_idx: int, frame_idx: int, inputs: dict):
+ """Add point inputs with automatic device placement."""
+ device_inputs = {}
+ for key, value in inputs.items():
+ if isinstance(value, torch.Tensor):
+ device_inputs[key] = value.to(self.inference_device, non_blocking=True)
+ else:
+ device_inputs[key] = value
+ self.point_inputs_per_obj[obj_idx][frame_idx] = device_inputs
+
+ def remove_point_inputs(self, obj_idx: int, frame_idx: int):
+ """Remove point inputs."""
+ self.point_inputs_per_obj[obj_idx].pop(frame_idx, None)
+
+ def add_mask_inputs(self, obj_idx: int, frame_idx: int, inputs: torch.Tensor):
+ """Add mask inputs with automatic device placement."""
+ self.mask_inputs_per_obj[obj_idx][frame_idx] = inputs.to(
+ self.inference_device, dtype=self.dtype, non_blocking=True
+ )
+
+ def remove_mask_inputs(self, obj_idx: int, frame_idx: int):
+ """Remove mask inputs."""
+ self.mask_inputs_per_obj[obj_idx].pop(frame_idx, None)
+
+ # Output management with smart device placement
+ def store_output(
+ self,
+ obj_idx: int,
+ frame_idx: int,
+ output_key: Optional[str] = None,
+ output_value: Optional[Union[torch.Tensor, dict]] = None,
+ is_conditioning_frame: bool = True,
+ ):
+ """
+ Store output with smart device management.
+ If output_key is None, the output is stored as a dictionary.
+
+ Args:
+ obj_idx (int): The index of the object.
+ frame_idx (int): The index of the frame.
+ output_key (Optional[str]): The key of the output. If None, the output is stored as a dictionary.
+ output_value (Optional[Union[torch.Tensor, dict]]): The value of the output.
+ is_conditioning_frame (bool): Whether the output is for a conditioning frame.
+ """
+ storage_key = "cond_frame_outputs" if is_conditioning_frame else "non_cond_frame_outputs"
+
+ if output_key is None and isinstance(output_value, dict):
+ self.output_dict_per_obj[obj_idx][storage_key][frame_idx] = {}
+ for key, value in output_value.items():
+ self.store_output(obj_idx, frame_idx, key, value, is_conditioning_frame)
+ return
+
+ # Device placement: small tensors stay on inference device, large ones go to inference state device
+ if output_key in ["object_pointer", "object_score_logits"]: # Small tensors
+ self.output_dict_per_obj[obj_idx][storage_key][frame_idx][output_key] = output_value
+ elif isinstance(output_value, torch.Tensor): # Large tensors like masks, features
+ self.output_dict_per_obj[obj_idx][storage_key][frame_idx][output_key] = output_value.to(
+ self.inference_state_device, non_blocking=True
+ )
+ else:
+ self.output_dict_per_obj[obj_idx][storage_key][frame_idx][output_key] = output_value
+
+ def get_output(
+ self,
+ obj_idx: int,
+ frame_idx: int,
+ output_key: str,
+ is_conditioning_frame: bool = True,
+ ):
+ """
+ Get output with smart device management.
+
+ Args:
+ obj_idx (int): The index of the object.
+ frame_idx (int): The index of the frame.
+ output_key (str): The key of the output.
+ is_conditioning_frame (bool): Whether the output is for a conditioning frame.
+ """
+ storage_key = "cond_frame_outputs" if is_conditioning_frame else "non_cond_frame_outputs"
+ out = self.output_dict_per_obj[obj_idx][storage_key].get(frame_idx, None)
+ # move to inference device if needed
+ if out is None:
+ return None
+ value = out[output_key]
+ if isinstance(value, torch.Tensor):
+ value = value.to(self.inference_device, non_blocking=True)
+ return value
+
+ # Video frame management
+ def add_new_frame(self, pixel_values: torch.Tensor, frame_idx: Optional[int] = None) -> int:
+ """Add new frame with automatic device placement."""
+ pixel_values = pixel_values.to(self.video_storage_device, dtype=self.dtype, non_blocking=True)
+ if pixel_values.dim() == 4:
+ pixel_values = pixel_values.squeeze(0)
+
+ if frame_idx is None:
+ frame_idx = len(self.processed_frames) if self.processed_frames is not None else 0
+
+ if self.processed_frames is None:
+ self.processed_frames = {frame_idx: pixel_values}
+ else:
+ self.processed_frames[frame_idx] = pixel_values
+
+ return frame_idx
+
+ def get_frame(self, frame_idx: int) -> torch.Tensor:
+ """Get frame from video."""
+ return self.processed_frames[frame_idx].to(self.inference_device, non_blocking=True)
+
+ def reset_tracking_data(self):
+ """Reset tracking data but keep cache."""
+ self._obj_id_to_idx.clear()
+ self._obj_idx_to_id.clear()
+ self.obj_ids.clear()
+ self.point_inputs_per_obj.clear()
+ self.mask_inputs_per_obj.clear()
+ self.output_dict_per_obj.clear()
+ self.frames_tracked_per_obj.clear()
+ self.obj_with_new_inputs = []
+ # Note: cache and video data are preserved
+
+ def reset_inference_session(self):
+ """Reset tracking data and cache."""
+ self._obj_id_to_idx.clear()
+ self._obj_idx_to_id.clear()
+ self.obj_ids.clear()
+ self.point_inputs_per_obj.clear()
+ self.mask_inputs_per_obj.clear()
+ self.output_dict_per_obj.clear()
+ self.frames_tracked_per_obj.clear()
+ self.obj_with_new_inputs = []
+ self.cache.clear_all()
+
+
+class EdgeTamVideoMemoryAttentionMLP(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.intermediate_size = config.memory_attention_mlp_hidden_size
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size)
+ self.dropout = nn.Dropout(config.memory_attention_dropout)
+ self.act_fn = ACT2FN[config.memory_attention_mlp_hidden_act]
+
+ def forward(self, x):
+ return self.down_proj(self.dropout(self.act_fn(self.up_proj(x))))
+
+
+class EdgeTamVideoMemoryAttentionLayer(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ hidden_size = config.memory_attention_hidden_size
+ self.self_attn = EdgeTamVideoRoPESelfAttention(config)
+ self.cross_attn_image = EdgeTamVideoRoPECrossAttention(config, kv_in_dim=64)
+
+ # MLP module
+ self.mlp = EdgeTamVideoMemoryAttentionMLP(config)
+
+ self.layer_norm1 = nn.LayerNorm(hidden_size)
+ self.layer_norm2 = nn.LayerNorm(hidden_size)
+ self.layer_norm3 = nn.LayerNorm(hidden_size)
+ self.dropout1 = nn.Dropout(config.memory_attention_dropout)
+ self.dropout2 = nn.Dropout(config.memory_attention_dropout)
+ self.dropout3 = nn.Dropout(config.memory_attention_dropout)
+
+ def forward(
+ self,
+ queries: Tensor,
+ keys: Tensor,
+ key_point_embedding: Tensor,
+ rope_position_embeddings: tuple[Tensor, Tensor],
+ rope_position_embeddings_k: Optional[tuple[Tensor, Tensor]] = None,
+ num_k_exclude_rope: int = 0,
+ rope_k_repeat: int = 0,
+ ) -> torch.Tensor:
+ # Self-Attention
+ query = self.layer_norm1(queries)
+ query, _ = self.self_attn(query=query, key=query, value=query, position_embeddings=rope_position_embeddings)
+ queries = queries + self.dropout1(query)
+
+ # Cross-Attention
+ query = self.layer_norm2(queries)
+ query, _ = self.cross_attn_image(
+ query=query,
+ key=keys + key_point_embedding,
+ value=keys,
+ position_embeddings=rope_position_embeddings,
+ position_embeddings_k=rope_position_embeddings_k,
+ num_k_exclude_rope=num_k_exclude_rope,
+ rope_k_repeat=rope_k_repeat,
+ )
+ queries = queries + self.dropout2(query)
+ # MLP
+ query = self.layer_norm3(queries)
+ query = self.mlp(query)
+ queries = queries + self.dropout3(query)
+ return queries
+
+
+class EdgeTamVideoMemoryAttention(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [EdgeTamVideoMemoryAttentionLayer(config) for _ in range(config.memory_attention_num_layers)]
+ )
+ self.layer_norm = nn.LayerNorm(config.memory_attention_hidden_size)
+ self.rotary_emb = EdgeTamVideoVisionRotaryEmbedding(config=config)
+ self.rotary_emb_k = EdgeTamVideoVisionRotaryEmbedding(
+ config, end_x=config.memory_attention_rope_k_sizes[0], end_y=config.memory_attention_rope_k_sizes[1]
+ )
+
+ def forward(
+ self,
+ current_vision_features: torch.Tensor,
+ memory: torch.Tensor,
+ current_vision_position_embeddings: Optional[Tensor] = None,
+ memory_posision_embeddings: Optional[Tensor] = None,
+ num_object_pointer_tokens: int = 0,
+ num_spatial_memory_tokens: int = -1,
+ ):
+ """
+ Args:
+ current_vision_features (`torch.FloatTensor`):
+ The current vision features used for self-attention.
+ memory (`torch.FloatTensor`):
+ The memory features used for cross-attention.
+ current_vision_position_embeddings (`torch.FloatTensor`, *optional*):
+ The position embeddings for the current vision features.
+ memory_posision_embeddings (`torch.FloatTensor`, *optional*):
+ The position embeddings for the memory features.
+ num_object_pointer_tokens (`int`, *optional*, defaults to 0):
+ The number of object pointer tokens.
+ """
+ output = current_vision_features
+ if current_vision_position_embeddings is not None:
+ output = output + 0.1 * current_vision_position_embeddings
+
+ # Convert to batch first
+ output = output.transpose(0, 1)
+ memory = memory.transpose(0, 1).unsqueeze(1)
+ memory_posision_embeddings = memory_posision_embeddings.transpose(0, 1).unsqueeze(1)
+ rope_position_embeddings = self.rotary_emb()
+ rope_position_embeddings_k = self.rotary_emb_k()
+ for layer in self.layers:
+ output = layer(
+ queries=output.unsqueeze(1) if output.ndim == 3 else output,
+ keys=memory,
+ key_point_embedding=memory_posision_embeddings,
+ rope_position_embeddings=rope_position_embeddings,
+ rope_position_embeddings_k=rope_position_embeddings_k,
+ num_k_exclude_rope=num_object_pointer_tokens,
+ rope_k_repeat=num_spatial_memory_tokens,
+ )
+
+ normed_output = self.layer_norm(output)
+
+ # Convert back to seq first
+ normed_output = normed_output.transpose(0, 1)
+
+ return normed_output
+
+
+class EdgeTamVideoPerceiverMLP(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.intermediate_size = config.perceiver_resampler_mlp_intermediate_size
+
+ self.layer_norm = nn.LayerNorm(self.hidden_size)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = nn.GELU()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.layer_norm(hidden_states)
+ hidden_states = self.down_proj(self.act_fn(self.up_proj(hidden_states)))
+ return hidden_states
+
+
+class EdgeTamVideoPerceiverAttention(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.num_attention_heads = config.perceiver_resampler_num_attention_heads
+ self.head_dim = config.perceiver_resampler_attention_head_dim
+ self.attention_dropout = config.perceiver_resampler_attention_dropout
+
+ self.inner_dim = self.head_dim * self.num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.k_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.v_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.o_proj = nn.Linear(self.inner_dim, self.hidden_size, bias=False)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ # Project queries, keys, and values
+ query = self.q_proj(query)
+ key = self.k_proj(key)
+ value = self.v_proj(value)
+
+ # Reshape for multi-head attention
+ batch_size, seq_len_q = query.shape[:2]
+ query = query.view(batch_size, seq_len_q, self.num_attention_heads, self.head_dim).transpose(1, 2)
+ seq_len_kv = key.shape[1]
+ key = key.view(batch_size, seq_len_kv, self.num_attention_heads, self.head_dim).transpose(1, 2)
+ value = value.view(batch_size, seq_len_kv, self.num_attention_heads, self.head_dim).transpose(1, 2)
+
+ # Add positional encoding if provided
+ if positional_encoding is not None:
+ pos_encoding = positional_encoding.view(
+ batch_size, seq_len_kv, self.num_attention_heads, self.head_dim
+ ).transpose(1, 2)
+ key = key + pos_encoding
+ value = value + pos_encoding
+
+ # Apply attention
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, _ = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+
+ # Reshape output
+ attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, self.inner_dim)
+ return self.o_proj(attn_output)
+
+
+class EdgeTamVideoPerceiverEncoderLayer(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+
+ self.cross_attention = EdgeTamVideoPerceiverAttention(config)
+ self.mlp = EdgeTamVideoPerceiverMLP(config)
+ self.dropout = nn.Dropout(config.perceiver_resampler_hidden_dropout)
+
+ self.self_attention = EdgeTamVideoPerceiverAttention(config)
+ self.self_mlp = EdgeTamVideoPerceiverMLP(config)
+
+ # Layer norms moved from attention classes to here
+ self.layer_norm_input = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+ self.layer_norm_latents = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+ self.layer_norm_self = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+
+ def forward(
+ self,
+ latents: torch.Tensor,
+ input_features: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ # Cross attention with layer norms
+ normalized_latents = self.layer_norm_latents(latents)
+ normalized_input = self.layer_norm_input(input_features)
+ cross_attention_output = self.cross_attention(
+ query=normalized_latents,
+ key=normalized_input,
+ value=normalized_input,
+ positional_encoding=positional_encoding,
+ )
+ latents = latents + self.dropout(cross_attention_output)
+
+ mlp_output = self.mlp(latents)
+ latents = latents + mlp_output
+
+ # Self attention with layer norm
+ normalized_latents_self = self.layer_norm_self(latents)
+ self_attention_output = self.self_attention(
+ query=normalized_latents_self, key=normalized_latents_self, value=normalized_latents_self
+ )
+ latents = latents + self_attention_output
+
+ self_mlp_output = self.self_mlp(latents)
+ latents = latents + self_mlp_output
+
+ return latents
+
+
+def window_partition(hidden_state, window_size):
+ """
+ Partition into non-overlapping windows with padding if needed.
+
+ Args:
+ hidden_state (`torch.Tensor`):
+ Input tokens with [batch_size, height, width, num_channels].
+ window_size (`int`):
+ Window size.
+
+ Returns:
+ `tuple(torch.FloatTensor)` comprising various elements:
+ - windows: windows after partition with [batch_size * num_windows, window_size, window_size, num_channels].
+ - (padded_height, padded_width): padded height and width before partition
+ """
+ batch_size, height, width, num_channels = hidden_state.shape
+
+ pad_height = (window_size - height % window_size) % window_size
+ pad_width = (window_size - width % window_size) % window_size
+
+ # Noop in case pad_width == 0 and pad_height == 0.
+ hidden_state = nn.functional.pad(hidden_state, (0, 0, 0, pad_width, 0, pad_height))
+
+ padded_height, padded_width = height + pad_height, width + pad_width
+
+ hidden_state = hidden_state.view(
+ batch_size, padded_height // window_size, window_size, padded_width // window_size, window_size, num_channels
+ )
+ windows = hidden_state.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
+ return windows, (padded_height, padded_width)
+
+
+class EdgeTamVideoPerceiverResampler(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.num_latents_1d = config.perceiver_resampler_num_latents
+ self.num_latents_2d = config.perceiver_resampler_num_latents_2d
+ self.num_layers = config.perceiver_resampler_num_layers
+
+ if self.num_latents_1d > 0:
+ self.latents_1d = nn.Parameter(torch.randn(self.num_latents_1d, self.hidden_size))
+ if self.num_latents_2d > 0:
+ self.latents_2d = nn.Parameter(torch.randn(self.num_latents_2d, self.hidden_size))
+
+ self.positional_encoding = EdgeTamVideoPositionEmbeddingSine(
+ num_pos_feats=self.hidden_size // 2, normalize=True
+ )
+
+ self.layers = nn.ModuleList([EdgeTamVideoPerceiverEncoderLayer(config) for _ in range(self.num_layers)])
+
+ self.layer_norm = nn.LayerNorm(self.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ output_latents = []
+ output_positional_encodings = []
+
+ if self.num_latents_1d > 0:
+ latents_1d, pos_1d = self._forward_1d(hidden_states, positional_encoding)
+ output_latents.append(latents_1d)
+ output_positional_encodings.append(pos_1d)
+
+ if self.num_latents_2d > 0:
+ latents_2d, pos_2d = self._forward_2d(hidden_states)
+ output_latents.append(latents_2d)
+ output_positional_encodings.append(pos_2d)
+
+ combined_latents = torch.cat(output_latents, dim=1)
+
+ combined_positional_encoding = None
+ if positional_encoding is not None and output_positional_encodings:
+ combined_positional_encoding = torch.cat(output_positional_encodings, dim=1)
+
+ return combined_latents, combined_positional_encoding
+
+ def _forward_1d(
+ self,
+ hidden_states: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ batch_size = hidden_states.shape[0]
+
+ latents = self.latents_1d.unsqueeze(0).expand(batch_size, -1, -1)
+ flattened_features = hidden_states.permute(0, 2, 3, 1).flatten(1, 2)
+
+ positional_features = None
+ if positional_encoding is not None:
+ positional_features = positional_encoding.permute(0, 2, 3, 1).flatten(1, 2)
+
+ for layer in self.layers:
+ latents = layer(latents, flattened_features, positional_features)
+
+ latents = self.layer_norm(latents)
+
+ output_positional_encoding = None
+ if positional_encoding is not None:
+ output_positional_encoding = torch.zeros_like(latents)
+
+ return latents, output_positional_encoding
+
+ def _forward_2d(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, channels, height, width = hidden_states.shape
+
+ latents_2d = self.latents_2d.unsqueeze(0).expand(batch_size, -1, -1).view(-1, 1, channels)
+
+ num_windows_per_dim = int(math.sqrt(self.num_latents_2d))
+ window_size = height // num_windows_per_dim
+
+ windowed_input = hidden_states.permute(0, 2, 3, 1)
+ windowed_features, _ = window_partition(windowed_input, window_size)
+ windowed_features = windowed_features.flatten(1, 2)
+
+ for layer in self.layers:
+ latents_2d = layer(latents_2d, windowed_features, positional_encoding=None)
+
+ latents_2d = latents_2d.view(batch_size, num_windows_per_dim, num_windows_per_dim, channels).permute(
+ 0, 3, 1, 2
+ )
+
+ positional_encoding_2d = self.positional_encoding(latents_2d.shape, latents_2d.device, latents_2d.dtype).to(
+ dtype=hidden_states.dtype
+ )
+ positional_encoding_2d = positional_encoding_2d.permute(0, 2, 3, 1).flatten(1, 2)
+
+ latents_2d = latents_2d.permute(0, 2, 3, 1).flatten(1, 2)
+ latents_2d = self.layer_norm(latents_2d)
+
+ return latents_2d, positional_encoding_2d
+
+
+@dataclass
+@auto_docstring(custom_intro="Base class for the EdgeTamVideo model's output.")
+class EdgeTamVideoImageSegmentationOutput(ModelOutput):
+ r"""
+ iou_scores (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks)`):
+ The Intersection over Union (IoU) scores of the predicted masks.
+ pred_masks (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks, height, width)`):
+ The predicted low-resolution masks. This is an alias for `low_res_masks`. These masks need to be post-processed
+ by the processor to be brought to the original image size.
+ object_score_logits (`torch.FloatTensor` of shape `(batch_size, point_batch_size, 1)`):
+ Logits for the object score, indicating if an object is present.
+ image_embeddings (`tuple(torch.FloatTensor)`):
+ The features from the FPN, which are used by the mask decoder. This is a tuple of `torch.FloatTensor` where each
+ tensor has shape `(batch_size, channels, height, width)`.
+ vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`):
+ Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`.
+ Hidden-states of the vision model at the output of each stage.
+ vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
+ Attentions weights of the vision model.
+ mask_decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
+ Attentions weights of the mask decoder.
+ high_res_masks (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_masks, image_size, image_size)`, *optional*):
+ The predicted masks, upscaled to the original image size. Only used for EdgeTamVideoModel.
+ object_pointer (`torch.FloatTensor` of shape `(batch_size, point_batch_size, hidden_size)`, *optional*):
+ A tensor representing the object pointer, used for tracking in videos. Only used for EdgeTamVideoModel.
+ """
+
+ iou_scores: Optional[torch.FloatTensor] = None
+ pred_masks: Optional[torch.FloatTensor] = None
+ object_score_logits: Optional[torch.FloatTensor] = None
+ image_embeddings: tuple[torch.FloatTensor, ...] = None
+ vision_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ vision_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+ mask_decoder_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+ high_res_masks: Optional[torch.FloatTensor] = None
+ object_pointer: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+@auto_docstring(custom_intro="Base class for the Sam2 model's output.")
+class EdgeTamVideoSegmentationOutput(ModelOutput):
+ r"""
+ pred_masks (`torch.FloatTensor` of shape `(batch_size, num_masks, height, width)`):
+ The predicted masks stored at the model's resolution.
+ frame_idx (`int`):
+ The frame index of the video.
+ """
+
+ pred_masks: Optional[torch.FloatTensor] = None
+ frame_idx: Optional[int] = None
+
+
+class EdgeTamVideoPositionalEmbedding(nn.Module):
+ def __init__(self, config: EdgeTamVideoPromptEncoderConfig):
+ super().__init__()
+ self.scale = config.scale
+ positional_embedding = self.scale * torch.randn((2, config.hidden_size // 2))
+ self.register_buffer("positional_embedding", positional_embedding)
+
+ def forward(self, input_coords, input_shape=None):
+ """Positionally encode points that are normalized to [0,1]."""
+ coordinates = input_coords.clone()
+
+ if input_shape is not None:
+ coordinates[:, :, :, 0] = coordinates[:, :, :, 0] / input_shape[1]
+ coordinates[:, :, :, 1] = coordinates[:, :, :, 1] / input_shape[0]
+ coordinates.to(torch.float32)
+
+ # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+ coordinates = 2 * coordinates - 1
+ coordinates = coordinates.to(self.positional_embedding.dtype)
+ coordinates = coordinates @ self.positional_embedding
+ coordinates = 2 * np.pi * coordinates
+ # outputs d_1 x ... x d_n x channel shape
+ return torch.cat([torch.sin(coordinates), torch.cos(coordinates)], dim=-1)
+
+
+class EdgeTamVideoMaskEmbedding(nn.Module):
+ def __init__(self, config: EdgeTamVideoPromptEncoderConfig):
+ super().__init__()
+ self.mask_input_channels = config.mask_input_channels // 4
+ self.activation = ACT2FN[config.hidden_act]
+ self.conv1 = nn.Conv2d(1, self.mask_input_channels, kernel_size=2, stride=2)
+ self.conv2 = nn.Conv2d(self.mask_input_channels, config.mask_input_channels, kernel_size=2, stride=2)
+ self.conv3 = nn.Conv2d(config.mask_input_channels, config.hidden_size, kernel_size=1)
+ self.layer_norm1 = EdgeTamVideoLayerNorm(
+ self.mask_input_channels, eps=config.layer_norm_eps, data_format="channels_first"
+ )
+ self.layer_norm2 = EdgeTamVideoLayerNorm(
+ self.mask_input_channels * 4, eps=config.layer_norm_eps, data_format="channels_first"
+ )
+
+ def forward(self, masks):
+ hidden_states = self.conv1(masks)
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ dense_embeddings = self.conv3(hidden_states)
+ return dense_embeddings
+
+
+class EdgeTamVideoPromptEncoder(nn.Module):
+ def __init__(self, config: EdgeTamVideoPromptEncoderConfig):
+ super().__init__()
+ self.shared_embedding = EdgeTamVideoPositionalEmbedding(config)
+ self.mask_embed = EdgeTamVideoMaskEmbedding(config)
+ self.no_mask_embed = nn.Embedding(1, config.hidden_size)
+
+ self.image_embedding_size = (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ self.mask_input_size = (4 * config.image_size // config.patch_size, 4 * config.image_size // config.patch_size)
+ self.input_image_size = config.image_size
+
+ self.point_embed = nn.Embedding(config.num_point_embeddings, config.hidden_size)
+ self.hidden_size = config.hidden_size
+ self.not_a_point_embed = nn.Embedding(1, config.hidden_size)
+
+ def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -> torch.Tensor:
+ """Embeds point prompts."""
+ points = points + 0.5 # Shift to center of pixel
+ if pad:
+ points = torch.nn.functional.pad(points, (0, 0, 0, 1), mode="constant", value=0)
+ labels = torch.nn.functional.pad(labels, (0, 1), mode="constant", value=-1)
+ input_shape = (self.input_image_size, self.input_image_size)
+ point_embedding = self.shared_embedding(points, input_shape)
+
+ # torch.where and expanding the labels tensor is required by the ONNX export
+ point_embedding = torch.where(labels[..., None] == -1, self.not_a_point_embed.weight, point_embedding)
+
+ # This is required for the ONNX export. The dtype, device need to be explicitly
+ # specified as otherwise torch.onnx.export interprets as double
+ point_embedding = torch.where(
+ labels[..., None] != -10,
+ point_embedding,
+ torch.zeros_like(point_embedding),
+ )
+
+ # Add point embeddings for labels >= 0
+ point_embedding = point_embedding + self.point_embed(labels.clamp(min=0)) * (labels >= 0).unsqueeze(-1)
+
+ return point_embedding
+
+ def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
+ """Embeds box prompts."""
+ boxes += 0.5 # Shift to center of pixel
+ coords = boxes.view(*boxes.shape[:2], 2, 2)
+ # add padding point for consistency with the original implementation
+ coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0)
+ corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size))
+ corner_embedding[:, :, 0, :] += self.point_embed.weight[2]
+ corner_embedding[:, :, 1, :] += self.point_embed.weight[3]
+ corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :])
+ return corner_embedding
+
+ def forward(
+ self,
+ input_points: Optional[tuple[torch.Tensor, torch.Tensor]],
+ input_labels: Optional[torch.Tensor],
+ input_boxes: Optional[torch.Tensor],
+ input_masks: Optional[torch.Tensor],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Embeds different types of prompts, returning both sparse and dense embeddings.
+
+ Args:
+ points (`torch.Tensor`, *optional*):
+ point coordinates and labels to embed.
+ boxes (`torch.Tensor`, *optional*):
+ boxes to embed
+ masks (`torch.Tensor`, *optional*):
+ masks to embed
+ """
+ sparse_embeddings = None
+ batch_size = 1
+ if input_points is not None:
+ batch_size = input_points.shape[0]
+ if input_labels is None:
+ raise ValueError("If points are provided, labels must also be provided.")
+ point_embeddings = self._embed_points(input_points, input_labels, pad=(input_boxes is None))
+ sparse_embeddings = point_embeddings
+ if input_boxes is not None:
+ batch_size = input_boxes.shape[0]
+ box_embeddings = self._embed_boxes(input_boxes)
+ if sparse_embeddings is None:
+ sparse_embeddings = box_embeddings
+ else:
+ sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=2)
+ if input_masks is not None:
+ dense_embeddings = self.mask_embed(input_masks)
+ else:
+ dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
+ batch_size, -1, self.image_embedding_size[0], self.image_embedding_size[1]
+ )
+
+ return sparse_embeddings, dense_embeddings
+
+
+class EdgeTamVideoTwoWayTransformer(nn.Module):
+ def __init__(self, config: EdgeTamVideoMaskDecoderConfig):
+ super().__init__()
+ self.config = config
+
+ self.num_hidden_layers = config.num_hidden_layers
+ self.layers = nn.ModuleList()
+
+ for i in range(self.num_hidden_layers):
+ self.layers.append(EdgeTamVideoTwoWayAttentionBlock(config, skip_first_layer_pe=(i == 0)))
+
+ self.final_attn_token_to_image = EdgeTamVideoAttention(config)
+ self.layer_norm_final_attn = nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ point_embeddings: Tensor,
+ image_embeddings: Tensor,
+ image_positional_embeddings: Tensor,
+ attention_similarity: Tensor,
+ target_embedding=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, BaseModelOutput]:
+ if image_embeddings is None:
+ raise ValueError("You have to specify an image_embedding")
+
+ image_embeddings = image_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1)
+ image_positional_embeddings = image_positional_embeddings.flatten(2).permute(0, 2, 1).unsqueeze(1)
+
+ # Prepare queries
+ queries = point_embeddings
+ keys = image_embeddings
+
+ # Apply transformer blocks and final layernorm
+ for layer in self.layers:
+ if target_embedding is not None:
+ queries += target_embedding
+
+ queries, keys, _ = layer(
+ queries=queries,
+ keys=keys,
+ query_point_embedding=point_embeddings,
+ key_point_embedding=image_positional_embeddings,
+ attention_similarity=attention_similarity,
+ **kwargs,
+ )
+ # Apply the final attention layer from the points to the image
+ query = queries + point_embeddings
+ key = keys + image_positional_embeddings
+
+ attn_out, _ = self.final_attn_token_to_image(query=query, key=key, value=keys)
+
+ queries = queries + attn_out
+ queries = self.layer_norm_final_attn(queries)
+ return queries, keys
+
+
+class EdgeTamVideoMaskDecoder(nn.Module):
+ def __init__(self, config: EdgeTamVideoMaskDecoderConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+
+ self.num_multimask_outputs = config.num_multimask_outputs
+ self.num_mask_tokens = config.num_multimask_outputs + 1
+
+ self.iou_token = nn.Embedding(1, self.hidden_size)
+ self.mask_tokens = nn.Embedding(self.num_mask_tokens, self.hidden_size)
+
+ self.transformer = EdgeTamVideoTwoWayTransformer(config)
+
+ # should we create a new class for this?
+ self.upscale_conv1 = nn.ConvTranspose2d(self.hidden_size, self.hidden_size // 4, kernel_size=2, stride=2)
+ self.upscale_conv2 = nn.ConvTranspose2d(self.hidden_size // 4, self.hidden_size // 8, kernel_size=2, stride=2)
+ self.upscale_layer_norm = EdgeTamVideoLayerNorm(self.hidden_size // 4, data_format="channels_first")
+ self.activation = nn.GELU()
+
+ mlps_list = []
+ for _ in range(self.num_mask_tokens):
+ mlps_list += [EdgeTamVideoFeedForward(self.hidden_size, self.hidden_size, self.hidden_size // 8, 3)]
+ self.output_hypernetworks_mlps = nn.ModuleList(mlps_list)
+ self.iou_prediction_head = EdgeTamVideoFeedForward(
+ self.hidden_size,
+ config.iou_head_hidden_dim,
+ self.num_mask_tokens,
+ config.iou_head_depth,
+ sigmoid_output=True,
+ )
+
+ self.conv_s0 = nn.Conv2d(config.hidden_size, config.hidden_size // 8, kernel_size=1, stride=1)
+ self.conv_s1 = nn.Conv2d(config.hidden_size, config.hidden_size // 4, kernel_size=1, stride=1)
+
+ self.obj_score_token = nn.Embedding(1, self.hidden_size)
+ self.pred_obj_score_head = EdgeTamVideoFeedForward(self.hidden_size, self.hidden_size, 1, 3)
+
+ self.dynamic_multimask_via_stability = config.dynamic_multimask_via_stability
+ self.dynamic_multimask_stability_delta = config.dynamic_multimask_stability_delta
+ self.dynamic_multimask_stability_thresh = config.dynamic_multimask_stability_thresh
+
+ def forward(
+ self,
+ image_embeddings: torch.Tensor,
+ image_positional_embeddings: torch.Tensor,
+ sparse_prompt_embeddings: torch.Tensor,
+ dense_prompt_embeddings: torch.Tensor,
+ multimask_output: bool,
+ high_resolution_features: list[torch.Tensor],
+ attention_similarity: Optional[torch.Tensor] = None,
+ target_embedding: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Predict masks given image and prompt embeddings.
+
+ Args:
+ image_embeddings (`torch.Tensor`):
+ The embeddings from the image encoder.
+ image_positional_embeddings (`torch.Tensor`):
+ Positional encoding with the shape of image_embeddings.
+ sparse_prompt_embeddings (`torch.Tensor`):
+ The embeddings of the points and boxes.
+ dense_prompt_embeddings (`torch.Tensor`):
+ The embeddings of the mask inputs.
+ multimask_output (`bool`):
+ Whether to return multiple masks or a single mask.
+ high_resolution_features (`list[torch.Tensor]`, *optional*):
+ The high-resolution features from the vision encoder.
+ attention_similarity (`torch.Tensor`, *optional*):
+ The attention similarity tensor.
+ target_embedding (`torch.Tensor`, *optional*):
+ The target embedding.
+ """
+ batch_size, num_channels, height, width = image_embeddings.shape
+ point_batch_size = sparse_prompt_embeddings.shape[1]
+ # Concatenate output tokens
+ output_tokens = torch.cat(
+ [
+ self.obj_score_token.weight,
+ self.iou_token.weight,
+ self.mask_tokens.weight,
+ ],
+ dim=0,
+ )
+ output_tokens = output_tokens.repeat(batch_size, point_batch_size, 1, 1)
+
+ if sparse_prompt_embeddings.shape[0] != 0:
+ tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=2)
+ else:
+ tokens = output_tokens
+ point_embeddings = tokens.to(self.iou_token.weight.dtype)
+
+ # Expand per-image data in batch direction to be per-mask
+ image_embeddings = image_embeddings + dense_prompt_embeddings
+ image_embeddings = image_embeddings.repeat_interleave(point_batch_size, dim=0)
+ image_positional_embeddings = image_positional_embeddings.repeat_interleave(point_batch_size, 0)
+ # Run the transformer
+ point_embeddings, image_embeddings = self.transformer(
+ point_embeddings=point_embeddings,
+ image_embeddings=image_embeddings,
+ image_positional_embeddings=image_positional_embeddings,
+ attention_similarity=attention_similarity,
+ target_embedding=target_embedding,
+ **kwargs,
+ )
+ iou_token_out = point_embeddings[:, :, 1, :]
+ mask_tokens_out = point_embeddings[:, :, 2 : (2 + self.num_mask_tokens), :]
+
+ # Upscale mask embeddings and predict masks using the mask tokens
+ image_embeddings = image_embeddings.transpose(2, 3).view(
+ batch_size * point_batch_size, num_channels, height, width
+ )
+
+ feat_s0, feat_s1 = high_resolution_features
+ feat_s0 = feat_s0.repeat_interleave(point_batch_size, dim=0)
+ feat_s1 = feat_s1.repeat_interleave(point_batch_size, dim=0)
+ upscaled_embedding = self.upscale_conv1(image_embeddings) + feat_s1
+ upscaled_embedding = self.activation(self.upscale_layer_norm(upscaled_embedding))
+ upscaled_embedding = self.activation(self.upscale_conv2(upscaled_embedding) + feat_s0)
+
+ hyper_in_list: list[torch.Tensor] = []
+ for i in range(self.num_mask_tokens):
+ current_mlp = self.output_hypernetworks_mlps[i]
+ hyper_in_list += [current_mlp(mask_tokens_out[:, :, i, :])]
+ hyper_in = torch.stack(hyper_in_list, dim=2)
+
+ _, num_channels, height, width = upscaled_embedding.shape
+ upscaled_embedding = upscaled_embedding.view(batch_size, point_batch_size, num_channels, height * width)
+ masks = (hyper_in @ upscaled_embedding).view(batch_size, point_batch_size, -1, height, width)
+
+ # Generate mask quality predictions
+ iou_pred = self.iou_prediction_head(iou_token_out)
+ object_score_logits = self.pred_obj_score_head(point_embeddings[:, :, 0, :])
+
+ # Select the correct mask or masks for output
+ if multimask_output:
+ mask_slice = slice(1, None)
+ masks = masks[:, :, mask_slice, :, :]
+ iou_pred = iou_pred[:, :, mask_slice]
+ elif self.dynamic_multimask_via_stability and not self.training:
+ mask_slice = slice(0, 1)
+ masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)
+ else:
+ mask_slice = slice(0, 1)
+ masks = masks[:, :, mask_slice, :, :]
+ iou_pred = iou_pred[:, :, mask_slice]
+
+ sam_tokens_out = mask_tokens_out[:, :, mask_slice] # [b, 3, c] shape
+
+ return masks, iou_pred, sam_tokens_out, object_score_logits
+
+ def _get_stability_scores(self, mask_logits):
+ """
+ Compute stability scores of the mask logits based on the IoU between upper and
+ lower thresholds.
+ """
+ mask_logits = mask_logits.flatten(-2)
+ stability_delta = self.dynamic_multimask_stability_delta
+ area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
+ area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
+ stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)
+ return stability_scores
+
+ def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
+ """
+ When outputting a single mask, if the stability score from the current single-mask
+ output (based on output token 0) falls below a threshold, we instead select from
+ multi-mask outputs (based on output token 1~3) the mask with the highest predicted
+ IoU score. This is intended to ensure a valid mask for both clicking and tracking.
+ """
+ # The best mask from multimask output tokens (1~3)
+ multimask_logits = all_mask_logits[:, :, 1:, :, :]
+ multimask_iou_scores = all_iou_scores[:, :, 1:]
+ best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) # [B, P]
+ best_scores_inds_expanded = best_scores_inds.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+ best_scores_inds_expanded = best_scores_inds_expanded.expand(
+ -1, -1, 1, multimask_logits.size(-2), multimask_logits.size(-1)
+ )
+ best_multimask_logits = torch.gather(multimask_logits, 2, best_scores_inds_expanded) # [B, P, 1, H, W]
+ best_multimask_iou_scores = torch.gather(multimask_iou_scores, 2, best_scores_inds.unsqueeze(-1)) # [B, P, 1]
+
+ # The mask from singlemask output token 0 and its stability score
+ singlemask_logits = all_mask_logits[:, :, 0:1, :, :]
+ singlemask_iou_scores = all_iou_scores[:, :, 0:1]
+ stability_scores = self._get_stability_scores(singlemask_logits)
+ is_stable = stability_scores >= self.dynamic_multimask_stability_thresh
+
+ # Dynamically fall back to best multimask output upon low stability scores.
+ mask_logits_out = torch.where(
+ is_stable[..., None, None].expand_as(singlemask_logits),
+ singlemask_logits,
+ best_multimask_logits,
+ )
+ iou_scores_out = torch.where(
+ is_stable.expand_as(singlemask_iou_scores),
+ singlemask_iou_scores,
+ best_multimask_iou_scores,
+ )
+ return mask_logits_out, iou_scores_out
+
+
+# a large negative value as a placeholder score for missing objects
+NO_OBJ_SCORE = -1024.0
+
+
+def get_1d_sine_pe(pos_inds, dim, temperature=10000):
+ """
+ Get 1D sine positional embedding as in the original Transformer paper.
+ """
+ pe_dim = dim // 2
+ dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)
+ dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)
+
+ pos_embed = pos_inds.unsqueeze(-1) / dim_t
+ pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1)
+ return pos_embed
+
+
+@auto_docstring
+class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
+ _tied_weights_keys = ["prompt_encoder.shared_embedding.positional_embedding"]
+ # need to be ignored, as it's a buffer and will not be correctly detected as tied weight
+ _keys_to_ignore_on_load_missing = ["prompt_encoder.shared_embedding.positional_embedding"]
+ _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(EdgeTamVideoTwoWayAttentionBlock, index=2)}
+ _keys_to_ignore_on_load_unexpected = []
+
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__(config)
+ self.shared_image_embedding = EdgeTamVideoPositionalEmbedding(config.prompt_encoder_config)
+ self.vision_encoder = AutoModel.from_config(config.vision_config)
+ self.prompt_encoder = EdgeTamVideoPromptEncoder(config.prompt_encoder_config)
+ # The module using it is not a PreTrainedModel subclass so we need this
+ config.mask_decoder_config._attn_implementation = config._attn_implementation
+ self.mask_decoder = EdgeTamVideoMaskDecoder(config.mask_decoder_config)
+
+ self.num_feature_levels = config.vision_config.num_feature_levels
+ self.backbone_feature_sizes = config.vision_config.backbone_feature_sizes
+ # a single token to indicate no memory embedding from previous frames
+ self.hidden_dim = config.vision_config.fpn_hidden_size
+ self.no_memory_embedding = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+ self.config = config
+ # For video sequence inference
+ self.image_size = config.image_size
+ self.memory_attention = EdgeTamVideoMemoryAttention(config)
+ self.memory_encoder = EdgeTamVideoMemoryEncoder(config)
+ self.no_memory_positional_encoding = torch.nn.Parameter(
+ torch.zeros(1, 1, config.vision_config.fpn_hidden_size)
+ )
+ self.mem_dim = config.memory_encoder_output_channels
+ self.num_maskmem = config.num_maskmem # Number of memories accessible
+ # Temporal encoding of the memories
+ self.memory_temporal_positional_encoding = torch.nn.Parameter(
+ torch.zeros(self.num_maskmem, 1, 1, self.mem_dim)
+ )
+
+ self.no_object_pointer = torch.nn.Parameter(torch.zeros(1, self.hidden_dim))
+ # A conv layer to downsample the mask prompt to stride 4 (the same stride as
+ # low-res SAM mask logits) and to change its scales from 0~1 to SAM logit scale,
+ # so that it can be fed into the SAM mask decoder to generate a pointer.
+ self.mask_downsample = torch.nn.Conv2d(1, 1, kernel_size=4, stride=4)
+ # a feedforward layer on SAM output tokens to turn them into object pointers
+ self.object_pointer_proj = EdgeTamVideoFeedForward(self.hidden_dim, self.hidden_dim, self.hidden_dim, 3)
+
+ if self.config.enable_temporal_pos_encoding_for_object_pointers:
+ # a linear projection on temporal positional encoding in object pointers to
+ # avoid potential interference with spatial positional encoding
+ self.temporal_positional_encoding_projection_layer = torch.nn.Linear(self.hidden_dim, self.mem_dim)
+ else:
+ self.temporal_positional_encoding_projection_layer = torch.nn.Identity()
+
+ self.occlusion_spatial_embedding_parameter = None # compatibility with Sam2
+ if config.enable_occlusion_spatial_embedding:
+ self.occlusion_spatial_embedding_parameter = torch.nn.Parameter(torch.zeros(1, self.mem_dim))
+ self.spatial_perceiver = EdgeTamVideoPerceiverResampler(config)
+
+ self.post_init()
+
+ def _tie_weights(self):
+ self.prompt_encoder.shared_embedding.positional_embedding.data = (
+ self.shared_image_embedding.positional_embedding.data
+ )
+
+ def get_input_embeddings(self):
+ return self.vision_encoder.get_input_embeddings()
+
+ def get_image_wide_positional_embeddings(self) -> torch.Tensor:
+ size = self.prompt_encoder.image_embedding_size
+ target_device = self.shared_image_embedding.positional_embedding.device
+ target_dtype = self.shared_image_embedding.positional_embedding.dtype
+ grid = torch.ones(size, device=target_device, dtype=target_dtype)
+ y_embed = grid.cumsum(dim=0) - 0.5
+ x_embed = grid.cumsum(dim=1) - 0.5
+ y_embed = y_embed / size[0]
+ x_embed = x_embed / size[1]
+
+ positional_embedding = self.shared_image_embedding(torch.stack([x_embed, y_embed], dim=-1))
+ return positional_embedding.permute(2, 0, 1).unsqueeze(0) # channel x height x width
+
+ @torch.no_grad()
+ def get_image_embeddings(
+ self,
+ pixel_values: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> list[torch.Tensor]:
+ r"""
+ Returns the image embeddings by passing the pixel values through the vision encoder.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Input pixel values
+ """
+ batch_size = pixel_values.shape[0]
+ feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs)
+
+ # add no memory embedding to the last feature map
+ feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
+
+ # reshape feature maps to the same shape as the backbone feature sizes
+ image_embeddings = [
+ feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+ for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes)
+ ]
+
+ return image_embeddings
+
+ @torch.no_grad()
+ def get_prompt_embeddings(
+ self,
+ input_points: Optional[torch.FloatTensor] = None,
+ input_labels: Optional[torch.LongTensor] = None,
+ input_boxes: Optional[torch.FloatTensor] = None,
+ input_masks: Optional[torch.LongTensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ r"""
+ Returns the prompt embeddings by passing the input points, labels, boxes and masks through the prompt encoder.
+
+ Args:
+ input_points (`torch.FloatTensor` of shape `(batch_size, point_batch_size, num_points_per_image, 2)`):
+ Optional input points for the prompt encoder. The padding of the point is automatically done by the
+ processor. `point_batch_size` refers to the number of masks that we want the model to predict per
+ point. The model will output `point_batch_size` times 3 masks in total.
+ input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points_per_image)`):
+ Optional input labels for the prompt encoder. The padding of the labels is automatically done by the
+ processor, or can be fed by the user.
+ input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes_per_image, 4)`):
+ Optional input boxes for the prompt encoder. The padding of the boxes is automatically done by the
+ processor. users can also pass manually the input boxes.
+ input_masks (`torch.LongTensor` of shape `(batch_size, image_size, image_size)`):
+ Optional input masks for the prompt encoder.
+ """
+ prompt_output = self.prompt_encoder(
+ input_points=input_points,
+ input_labels=input_labels,
+ input_boxes=input_boxes,
+ input_masks=input_masks,
+ )
+ return prompt_output
+
+ @torch.inference_mode()
+ @auto_docstring(custom_intro="Propagate the objects through a streamed video frame.")
+ def forward(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ frame_idx: Optional[int] = None,
+ frame: Optional[torch.Tensor] = None,
+ reverse: bool = False,
+ ) -> EdgeTamVideoSegmentationOutput:
+ r"""
+ inference_session (`EdgeTamVideoInferenceSession`):
+ The video inference session object.
+ frame_idx (`int`, *optional*):
+ The index of the frame on which to run inference. No need to provide when inferring
+ on a new streamed frame.
+ frame (`torch.Tensor`, *optional*):
+ The frame to process. Provide when streaming.
+ reverse (`bool`, *optional*, defaults to `False`):
+ Whether to propagate in reverse.
+ """
+ if frame is not None:
+ frame_idx = inference_session.add_new_frame(frame, frame_idx)
+
+ if frame is not None and inference_session.get_obj_num() == 0:
+ raise ValueError("No objects are provided for tracking; please add inputs first.")
+
+ num_objects = inference_session.get_obj_num()
+ pred_masks_per_obj = [None] * num_objects
+ # Note: We avoid batched inference here because per-object inputs (clicks/masks)
+ # can differ across objects.
+ for obj_idx in range(num_objects):
+ obj_id = inference_session.obj_idx_to_id(obj_idx)
+ has_new_inputs = obj_id in inference_session.obj_with_new_inputs
+ has_cond_output = frame_idx in inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ # If this object has no new inputs and this frame already has a
+ # conditioning output, reuse the cached masks instead of recomputing.
+ if (not has_new_inputs) and has_cond_output:
+ pred_masks = inference_session.get_output(obj_idx, frame_idx, "pred_masks", is_conditioning_frame=True)
+ is_init_cond_frame = True
+ else:
+ # Defaults when there are no new inputs
+ is_init_cond_frame = False
+ point_inputs = None
+ mask_inputs = None
+
+ if has_new_inputs:
+ is_init_cond_frame = frame_idx not in inference_session.frames_tracked_per_obj[obj_idx]
+ if is_init_cond_frame:
+ reverse = False
+ point_inputs = inference_session.point_inputs_per_obj[obj_idx].get(frame_idx, None)
+ mask_inputs = inference_session.mask_inputs_per_obj[obj_idx].get(frame_idx, None)
+ if point_inputs is not None or mask_inputs is not None:
+ inference_session.obj_with_new_inputs.remove(obj_id)
+
+ current_out = self._run_single_frame_inference(
+ inference_session=inference_session,
+ obj_idx=obj_idx,
+ frame_idx=frame_idx,
+ batch_size=1, # run on the slice of a single object
+ is_init_cond_frame=is_init_cond_frame,
+ point_inputs=point_inputs,
+ mask_inputs=mask_inputs,
+ reverse=reverse,
+ run_mem_encoder=True,
+ streaming=frame is not None,
+ )
+ inference_session.store_output(
+ obj_idx, frame_idx, output_value=current_out, is_conditioning_frame=is_init_cond_frame
+ )
+ pred_masks = current_out["pred_masks"]
+
+ pred_masks_per_obj[obj_idx] = pred_masks
+ if not is_init_cond_frame:
+ # only for tracked frames, not for initial conditioning frames
+ inference_session.frames_tracked_per_obj[obj_idx][frame_idx] = {"reverse": reverse}
+
+ # Resize the output mask to the original video resolution (we directly use
+ # the mask scores on GPU for output to avoid any CPU conversion in between)
+ if len(pred_masks_per_obj) > 1:
+ all_pred_masks = torch.cat(pred_masks_per_obj, dim=0)
+ else:
+ all_pred_masks = pred_masks_per_obj[0]
+
+ return EdgeTamVideoSegmentationOutput(pred_masks=all_pred_masks, frame_idx=frame_idx)
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[
+ list[torch.Tensor],
+ list[torch.Tensor],
+ Optional[tuple[torch.FloatTensor, ...]],
+ Optional[tuple[torch.FloatTensor, ...]],
+ ]:
+ r"""
+ Extract and preprocess image features using the vision encoder.
+
+ Args:
+ pixel_values (`torch.FloatTensor`):
+ Input pixel values of shape `(batch_size, num_channels, height, width)`.
+
+ Returns:
+ `tuple`: A tuple containing:
+ - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels.
+ - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level.
+ - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder.
+ - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder.
+ """
+ vision_outputs: EdgeTamVideoVisionEncoderOutput = self.vision_encoder(
+ pixel_values,
+ **kwargs,
+ )
+
+ feature_maps = vision_outputs.fpn_hidden_states
+ feature_maps_position_embeddings = vision_outputs.fpn_position_encoding
+
+ # precompute projected level 0 and level 1 features in SAM decoder
+ # to avoid running it again on every SAM click
+ feature_maps = list(feature_maps)
+ feature_maps[0] = self.mask_decoder.conv_s0(feature_maps[0])
+ feature_maps[1] = self.mask_decoder.conv_s1(feature_maps[1])
+
+ # flatten NxCxHxW to HWxNxC
+ feature_maps = [feature_map.flatten(2).permute(2, 0, 1) for feature_map in feature_maps]
+ feature_maps_position_embeddings = [
+ feature_map_position_embedding.flatten(2).permute(2, 0, 1)
+ for feature_map_position_embedding in feature_maps_position_embeddings
+ ]
+
+ return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions
+
+ def _prepare_vision_features(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ frame_idx: int,
+ batch_size: int,
+ ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+ """Prepare vision features for a frame."""
+
+ # Check if features are cached
+ if cached_features := inference_session.cache.get_vision_features(frame_idx):
+ vision_feats = cached_features["vision_feats"]
+ vision_pos_embeds = cached_features["vision_pos_embeds"]
+ else:
+ # Compute features using image encoder
+ image_batch = inference_session.get_frame(frame_idx).unsqueeze(0) # Add batch dimension
+ vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch)
+ # Cache features
+ inference_session.cache.cache_vision_features(
+ frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds}
+ )
+
+ # Expand to batch size if needed
+ if batch_size > 1:
+ vision_feats = vision_feats.expand(batch_size, -1, -1, -1)
+ vision_pos_embeds = [pe.expand(batch_size, -1, -1, -1) for pe in vision_pos_embeds]
+
+ return vision_feats, vision_pos_embeds
+
+ def _single_frame_forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ input_points: Optional[torch.FloatTensor] = None,
+ input_labels: Optional[torch.LongTensor] = None,
+ input_boxes: Optional[torch.FloatTensor] = None,
+ input_masks: Optional[torch.LongTensor] = None,
+ image_embeddings: Optional[torch.FloatTensor] = None,
+ multimask_output: bool = True,
+ attention_similarity: Optional[torch.FloatTensor] = None,
+ target_embedding: Optional[torch.FloatTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> EdgeTamVideoImageSegmentationOutput:
+ """
+ input_points (`torch.FloatTensor` of shape `(batch_size, num_points, 2)`):
+ Input 2D spatial points, this is used by the prompt encoder to encode the prompt. Generally yields to much
+ better results. The points can be obtained by passing a list of list of list to the processor that will
+ create corresponding `torch` tensors of dimension 4. The first dimension is the image batch size, the
+ second dimension is the point batch size (i.e. how many segmentation masks do we want the model to predict
+ per input point), the third dimension is the number of points per segmentation mask (it is possible to pass
+ multiple points for a single mask), and the last dimension is the x (vertical) and y (horizontal)
+ coordinates of the point. If a different number of points is passed either for each image, or for each
+ mask, the processor will create "PAD" points that will correspond to the (0, 0) coordinate, and the
+ computation of the embedding will be skipped for these points using the labels.
+ input_labels (`torch.LongTensor` of shape `(batch_size, point_batch_size, num_points)`):
+ Input labels for the points, this is used by the prompt encoder to encode the prompt. According to the
+ official implementation, there are 3 types of labels
+
+ - `1`: the point is a point that contains the object of interest
+ - `0`: the point is a point that does not contain the object of interest
+ - `-1`: the point corresponds to the background
+
+ We added the label:
+
+ - `-10`: the point is a padding point, thus should be ignored by the prompt encoder
+
+ The padding labels should be automatically done by the processor.
+ input_boxes (`torch.FloatTensor` of shape `(batch_size, num_boxes, 4)`):
+ Input boxes for the points, this is used by the prompt encoder to encode the prompt. Generally yields to
+ much better generated masks. The boxes can be obtained by passing a list of list of list to the processor,
+ that will generate a `torch` tensor, with each dimension corresponding respectively to the image batch
+ size, the number of boxes per image and the coordinates of the top left and bottom right point of the box.
+ In the order (`x1`, `y1`, `x2`, `y2`):
+
+ - `x1`: the x coordinate of the top left point of the input box
+ - `y1`: the y coordinate of the top left point of the input box
+ - `x2`: the x coordinate of the bottom right point of the input box
+ - `y2`: the y coordinate of the bottom right point of the input box
+ input_masks (`torch.FloatTensor` of shape `(batch_size, image_size, image_size)`):
+ SAM model also accepts segmentation masks as input. The mask will be embedded by the prompt encoder to
+ generate a corresponding embedding, that will be fed later on to the mask decoder. These masks needs to be
+ manually fed by the user, and they need to be of shape (`batch_size`, `image_size`, `image_size`).
+ image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_channels, window_size, window_size)`):
+ Image embeddings, this is used by the mask decoder to generate masks and iou scores. For more memory
+ efficient computation, users can first retrieve the image embeddings using the `get_image_embeddings`
+ method, and then feed them to the `forward` method instead of feeding the `pixel_values`.
+ multimask_output (`bool`, *optional*):
+ In the original implementation and paper, the model always outputs 3 masks per image (or per point / per
+ bounding box if relevant). However, it is possible to just output a single mask, that corresponds to the
+ "best" mask, by specifying `multimask_output=False`.
+ attention_similarity (`torch.FloatTensor`, *optional*):
+ Attention similarity tensor, to be provided to the mask decoder for target-guided attention in case the
+ model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048).
+ target_embedding (`torch.FloatTensor`, *optional*):
+ Embedding of the target concept, to be provided to the mask decoder for target-semantic prompting in case
+ the model is used for personalization as introduced in [PerSAM](https://huggingface.co/papers/2305.03048).
+ """
+ if not ((pixel_values is None) ^ (image_embeddings is None)):
+ raise ValueError("Exactly one of pixel_values or image_embeddings must be provided.")
+ if input_points is not None and input_boxes is not None:
+ if input_points.shape[1] != input_boxes.shape[1]:
+ raise ValueError(
+ f"You should provide as many bounding boxes as input points per box. Got {input_points.shape[1]} and {input_boxes.shape[1]}."
+ )
+ elif input_points is not None:
+ num_objects = input_points.shape[1]
+ elif input_boxes is not None:
+ num_objects = input_boxes.shape[1]
+ elif input_masks is not None:
+ num_objects = input_masks.shape[1]
+ else:
+ num_objects = 1
+
+ image_positional_embeddings = self.get_image_wide_positional_embeddings()
+ # repeat with batch size
+ batch_size = pixel_values.shape[0] if pixel_values is not None else image_embeddings[-1].shape[0]
+ image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1)
+
+ vision_attentions = None
+ vision_hidden_states = None
+
+ if pixel_values is not None:
+ feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features(
+ pixel_values,
+ **kwargs,
+ )
+
+ # add no memory embedding to the last feature map
+ feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
+
+ # reshape feature maps to the same shape as the backbone feature sizes
+ image_embeddings = [
+ feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+ for feat, feat_size in zip(feature_maps, self.backbone_feature_sizes)
+ ]
+
+ if input_points is not None and input_labels is None:
+ input_labels = torch.ones_like(input_points[:, :, :, 0], dtype=torch.int, device=input_points.device)
+
+ if input_points is None and input_boxes is None:
+ # If no points are provide, pad with an empty point (with label -1)
+ input_points = torch.zeros(
+ batch_size, 1, 1, 2, dtype=image_embeddings[-1].dtype, device=image_embeddings[-1].device
+ )
+ input_labels = -torch.ones(batch_size, 1, 1, dtype=torch.int32, device=image_embeddings[-1].device)
+
+ if input_masks is not None:
+ # If mask_inputs is provided, downsize it into low-res mask input if needed
+ # and feed it as a dense mask prompt into the SAM mask encoder
+ if input_masks.shape[-2:] != self.prompt_encoder.mask_input_size:
+ input_masks = F.interpolate(
+ input_masks.float(),
+ size=self.prompt_encoder.mask_input_size,
+ align_corners=False,
+ mode="bilinear",
+ antialias=True, # use antialias for downsampling
+ ).to(input_masks.dtype)
+
+ sparse_embeddings, dense_embeddings = self.prompt_encoder(
+ input_points=input_points,
+ input_labels=input_labels,
+ input_boxes=input_boxes,
+ input_masks=input_masks,
+ )
+ low_res_multimasks, iou_scores, sam_output_tokens, object_score_logits = self.mask_decoder(
+ image_embeddings=image_embeddings[-1],
+ image_positional_embeddings=image_positional_embeddings,
+ sparse_prompt_embeddings=sparse_embeddings,
+ dense_prompt_embeddings=dense_embeddings,
+ multimask_output=multimask_output,
+ high_resolution_features=image_embeddings[:-1],
+ attention_similarity=attention_similarity,
+ target_embedding=target_embedding,
+ **kwargs,
+ )
+
+ is_obj_appearing = object_score_logits > 0
+ # Mask used for spatial memories is always a *hard* choice between obj and no obj,
+ # consistent with the actual mask prediction
+ low_res_multimasks = torch.where(
+ is_obj_appearing[:, None, None],
+ low_res_multimasks,
+ NO_OBJ_SCORE,
+ )
+
+ # convert masks from possibly bfloat16 (or float16) to float32
+ # (older PyTorch versions before 2.1 don't support `interpolate` on bf16)
+ high_res_multimasks = (
+ F.interpolate(
+ low_res_multimasks.squeeze(1).float(),
+ size=(self.image_size, self.image_size),
+ mode="bilinear",
+ align_corners=False,
+ )
+ .unsqueeze(1)
+ .to(low_res_multimasks.dtype)
+ )
+ sam_output_token = sam_output_tokens[:, :, 0]
+ if multimask_output:
+ # take the best mask prediction (with the highest IoU estimation)
+ best_iou_inds = torch.argmax(iou_scores, dim=-1)
+ batch_inds = torch.arange(batch_size, device=high_res_multimasks.device)
+ object_batch_inds = torch.arange(num_objects, device=high_res_multimasks.device)
+ low_res_masks = low_res_multimasks[batch_inds, object_batch_inds, best_iou_inds]
+ high_res_masks = high_res_multimasks[batch_inds, object_batch_inds, best_iou_inds]
+ if sam_output_tokens.size(2) > 1:
+ sam_output_token = sam_output_tokens[batch_inds, object_batch_inds, best_iou_inds]
+ else:
+ low_res_masks, high_res_masks = low_res_multimasks[:, :, 0], high_res_multimasks[:, :, 0]
+
+ # Extract object pointer from the SAM output token (with occlusion handling)
+ object_pointer = self.object_pointer_proj(sam_output_token)
+ lambda_is_obj_appearing = is_obj_appearing.to(object_pointer.dtype)
+
+ object_pointer = lambda_is_obj_appearing * object_pointer
+ object_pointer = object_pointer + (1 - lambda_is_obj_appearing) * self.no_object_pointer
+
+ return EdgeTamVideoImageSegmentationOutput(
+ iou_scores=iou_scores,
+ pred_masks=low_res_masks,
+ high_res_masks=high_res_masks,
+ object_pointer=object_pointer,
+ object_score_logits=object_score_logits,
+ image_embeddings=image_embeddings,
+ vision_hidden_states=vision_hidden_states,
+ vision_attentions=vision_attentions,
+ )
+
+ def _use_mask_as_output(
+ self,
+ backbone_features: torch.Tensor,
+ high_res_features: list[torch.Tensor],
+ mask_inputs: torch.Tensor,
+ ) -> EdgeTamVideoImageSegmentationOutput:
+ """
+ Directly turn binary `mask_inputs` into a output mask logits without using SAM.
+ (same input and output shapes as in forward above).
+ """
+ # Use -10/+20 as logits for neg/pos pixels (very close to 0/1 in prob after sigmoid).
+ out_scale, out_bias = 20.0, -10.0 # sigmoid(-10.0)=4.5398e-05
+ mask_inputs_float = mask_inputs.to(backbone_features[0].dtype)
+ high_res_masks = mask_inputs_float * out_scale + out_bias
+ low_res_masks = F.interpolate(
+ high_res_masks.float(),
+ size=(high_res_masks.size(-2) // 4, high_res_masks.size(-1) // 4),
+ align_corners=False,
+ mode="bilinear",
+ antialias=True, # use antialias for downsampling
+ ).to(backbone_features[0].dtype)
+ # a dummy IoU prediction of all 1's under mask input
+ iou_scores = mask_inputs.new_ones(mask_inputs.size(0), 1).to(backbone_features[0].dtype)
+ # produce an object pointer using the SAM decoder from the mask input
+ object_pointer = self._single_frame_forward(
+ input_masks=self.mask_downsample(mask_inputs_float.to(backbone_features[0].dtype)),
+ image_embeddings=high_res_features + [backbone_features],
+ ).object_pointer
+ # In this method, we are treating mask_input as output, e.g. using it directly to create spatial mem;
+ # Below, we follow the same design axiom to use mask_input to decide if obj appears or not instead of relying
+ # on the object_scores from the SAM decoder.
+ is_obj_appearing = torch.any(mask_inputs.flatten(1).float() > 0.0, dim=1)
+ is_obj_appearing = is_obj_appearing[..., None]
+ lambda_is_obj_appearing = is_obj_appearing.to(backbone_features[0].dtype)
+ object_score_logits = out_scale * lambda_is_obj_appearing + out_bias
+ object_pointer = lambda_is_obj_appearing * object_pointer
+ object_pointer = object_pointer + (1 - lambda_is_obj_appearing) * self.no_object_pointer
+ return EdgeTamVideoImageSegmentationOutput(
+ iou_scores=iou_scores,
+ pred_masks=low_res_masks,
+ high_res_masks=high_res_masks,
+ object_pointer=object_pointer,
+ object_score_logits=object_score_logits,
+ image_embeddings=high_res_features + [backbone_features],
+ )
+
+ def _gather_memory_frame_outputs(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ track_in_reverse_time: bool = False,
+ ) -> list[tuple[int, dict]]:
+ """
+ Get memory frames from conditioning and non-conditioning outputs.
+
+ Returns:
+ List of (relative_temporal_offset, output_data) tuples.
+ """
+ temporal_positions_and_previous_outputs = []
+
+ # Add conditioning frame outputs (no limit here, as is the case in the original checkpoints)
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ if not conditioning_outputs:
+ raise ValueError(
+ "maskmem_features in conditioning outputs cannot be empty when not is_initial_conditioning_frame"
+ )
+
+ # Store (temporal_position, output_data) tuples
+ temporal_positions_and_previous_outputs = [(0, out) for out in conditioning_outputs.values()]
+
+ # Add non-conditioning memory frames (up to self.num_maskmem - 1)
+ # These are typically frames tracked by the model without direct user input.
+ # Frames are selected with a stride, prioritizing the most recent ones. Here we only support stride = 1 for simplicity.
+ for relative_temporal_offset in range(self.num_maskmem - 1, 0, -1):
+ # relative_temporal_offset: how many frames before (or after if reversing) the current frame
+ if not track_in_reverse_time:
+ previous_frame_idx = frame_idx - relative_temporal_offset
+ else:
+ previous_frame_idx = frame_idx + relative_temporal_offset
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ output_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ previous_frame_idx, None
+ )
+
+ temporal_positions_and_previous_outputs.append((relative_temporal_offset, output_data))
+
+ return temporal_positions_and_previous_outputs
+
+ def _build_memory_attention_inputs(
+ self,
+ temporal_positions_and_previous_outputs: list[tuple[int, dict]],
+ device: torch.device,
+ ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+ """
+ Concatenate memory features and positional embeddings from previous frames.
+
+ Returns:
+ Tuple of (memories_to_concatenate, memory_positional_embeddings_to_concatenate).
+ """
+ memories_to_concatenate = []
+ memory_positional_embeddings_to_concatenate = []
+
+ for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
+ if prev_output_data is None:
+ continue # Skip if no output data for this temporal position (e.g., padding frames)
+
+ # Load memory features (potentially from CPU to GPU)
+ # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
+ memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
+ memories_to_concatenate.append(memory_features.permute(1, 0, 2))
+
+ # Spatial positional encoding (potentially from CPU to GPU)
+ spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
+ spatial_memory_pos_embed = spatial_memory_pos_embed.squeeze(1).permute(1, 0, 2)
+
+ # Add temporal positional encoding
+ # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
+ combined_memory_pos_embed = (
+ spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
+ )
+ memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
+
+ return memories_to_concatenate, memory_positional_embeddings_to_concatenate
+
+ def _get_object_pointers(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ num_total_frames: int,
+ device: torch.device,
+ track_in_reverse_time: bool = False,
+ streaming: bool = False,
+ ) -> tuple[list[int], list[torch.Tensor], int]:
+ """
+ Get object pointers and their positional embeddings from past frames.
+
+ Returns:
+ Tuple of (temporal_offsets, pointer_tokens, max_object_pointers_to_use).
+ """
+ temporal_position_sign_multiplier = -1 if track_in_reverse_time else 1
+
+ # Determine max object pointers to use
+ if streaming:
+ max_object_pointers_to_use = self.config.max_object_pointers_in_encoder
+ else:
+ max_object_pointers_to_use = min(num_total_frames, self.config.max_object_pointers_in_encoder)
+
+ temporal_offsets: list[int] = []
+ pointer_tokens: list[torch.Tensor] = []
+
+ # Add object pointers from selected conditioning frames
+ # Optionally, only include pointers from past frames during evaluation
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ eligible_conditioning_outputs = conditioning_outputs
+ if not self.training:
+ eligible_conditioning_outputs = {
+ temporal_idx: out
+ for temporal_idx, out in conditioning_outputs.items()
+ if (temporal_idx >= frame_idx if track_in_reverse_time else temporal_idx <= frame_idx)
+ }
+
+ for temporal_idx, out_data in eligible_conditioning_outputs.items():
+ temporal_difference = (frame_idx - temporal_idx) * temporal_position_sign_multiplier
+ temporal_offsets.append(temporal_difference)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ # Add object pointers from non-conditioning frames (up to max_object_pointers_to_use - 1)
+ for t_diff_offset in range(1, max_object_pointers_to_use):
+ ref_frame_idx = frame_idx + t_diff_offset if track_in_reverse_time else frame_idx - t_diff_offset
+ if ref_frame_idx < 0 or (
+ not streaming and num_total_frames is not None and ref_frame_idx >= num_total_frames
+ ):
+ break # Stop if frame index is out of bounds
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ out_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ ref_frame_idx, None
+ )
+ if out_data is not None:
+ temporal_offsets.append(t_diff_offset)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ return temporal_offsets, pointer_tokens, max_object_pointers_to_use
+
+ def _process_object_pointers(
+ self,
+ temporal_offsets: list[int],
+ pointer_tokens: list[torch.Tensor],
+ max_object_pointers_to_use: int,
+ batch_size: int,
+ num_channels: int,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Process object pointers and compute their positional embeddings.
+
+ Returns:
+ Tuple of (object_pointers, object_pointers_pos_embed).
+ """
+ if not pointer_tokens:
+ return None, None
+
+ # Stack object pointers: List of (Batch, Channels) -> (SeqLen_ptr, Batch, Channels)
+ object_pointers = torch.stack(pointer_tokens, dim=0)
+
+ if self.config.enable_temporal_pos_encoding_for_object_pointers:
+ max_temporal_diff = float(max_object_pointers_to_use - 1)
+ # Determine dimensionality for temporal positional encoding of pointers
+ pointer_tpos_dim = num_channels
+
+ # Normalize temporal differences before sine PE calculation
+ normalized_temporal_diffs = (
+ torch.tensor(temporal_offsets, device=device, dtype=torch.float32) / max_temporal_diff
+ )
+ sine_pe = get_1d_sine_pe(normalized_temporal_diffs, dim=pointer_tpos_dim).to(object_pointers.dtype)
+ projected_sine_pe = self.temporal_positional_encoding_projection_layer(sine_pe)
+ object_pointers_pos_embed = projected_sine_pe.unsqueeze(1).expand(-1, batch_size, self.mem_dim)
+ else:
+ object_pointers_pos_embed = object_pointers.new_zeros(
+ len(temporal_offsets), batch_size, self.mem_dim, dtype=object_pointers.dtype
+ )
+
+ if self.mem_dim < num_channels:
+ # If memory dimension is smaller, reshape/split pointers and repeat positional encoding
+ num_splits = num_channels // self.mem_dim
+ object_pointers = object_pointers.reshape(-1, batch_size, num_splits, self.mem_dim)
+ object_pointers = object_pointers.permute(0, 2, 1, 3).flatten(
+ 0, 1
+ ) # (SeqLen_ptr*num_splits, Batch, MemDim)
+ object_pointers_pos_embed = object_pointers_pos_embed.repeat_interleave(num_splits, dim=0)
+
+ return object_pointers, object_pointers_pos_embed
+
+ def _prepare_memory_conditioned_features(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ frame_idx: int,
+ obj_idx: int,
+ is_initial_conditioning_frame: bool,
+ current_vision_features: list[torch.Tensor],
+ current_vision_positional_embeddings: list[torch.Tensor],
+ num_total_frames: int,
+ track_in_reverse_time: bool = False,
+ streaming: bool = False,
+ ) -> torch.Tensor:
+ """
+ Fuse current frame's visual features with memory from previous frames for enhanced object tracking.
+
+ This method conditions the current frame's visual features on temporal memory from previous frames,
+ enabling consistent object tracking across video sequences. For initial conditioning frames, it uses
+ no-memory embeddings. For subsequent frames, it retrieves and integrates memory features from both
+ conditioning frames (user interactions) and non-conditioning frames (tracked results) via cross-attention.
+
+ Args:
+ inference_session (`EdgeTamVideoInferenceSession`):
+ The video inference session object.
+ frame_idx (`int`):
+ Index of the current frame being processed.
+ obj_idx (`int`):
+ Index of the object being processed.
+ is_initial_conditioning_frame (`bool`):
+ Whether this is an initial conditioning frame with user inputs (True) or a subsequent
+ tracking frame (False).
+ current_vision_features (`torch.Tensor`):
+ Highest-level vision features of shape `(seq_len, batch_size, channels)`.
+ current_vision_positional_embeddings (`torch.Tensor`):
+ Positional embedding tensors corresponding to the highest-level vision features.
+ num_total_frames (`int`):
+ Total number of frames in the video sequence.
+ track_in_reverse_time (`bool`, *optional*, defaults to `False`):
+ Whether tracking is performed in reverse temporal order.
+ streaming (`bool`, *optional*, defaults to `False`):
+ Whether this is streaming inference mode.
+
+ Returns:
+ `torch.Tensor`: Memory-conditioned feature tensor of shape `(batch_size, channels, height, width)`
+ suitable for input to the SAM decoder.
+ """
+ # Get dimensions from the highest-level (lowest-resolution) feature map
+ batch_size = current_vision_features.size(1)
+ num_channels = self.hidden_dim
+ height, width = self.backbone_feature_sizes[-1]
+ device = current_vision_features.device
+
+ # If memory is disabled (e.g., for single image SAM), return current features directly.
+ if self.num_maskmem == 0:
+ # Permute (SeqLen, Batch, Channels) -> (Batch, Channels, SeqLen) then view as (Batch, Channels, Height, Width)
+ # Assuming SeqLen = Height * Width for the last feature map
+ current_feature_map = current_vision_features.permute(1, 2, 0).view(
+ batch_size, num_channels, height, width
+ )
+ return current_feature_map
+
+ # Step 1: Handle initial conditioning frames
+ if is_initial_conditioning_frame:
+ # For initial conditioning frames, no prior memory is used directly in this block.
+ # If configured, directly add a learnable "no memory" embedding.
+ # current_vision_features has shape (SeqLen, Batch, Channels)
+ conditioned_feature_map_flat = current_vision_features + self.no_memory_embedding
+ # Reshape to (Batch, Channels, Height, Width)
+ conditioned_feature_map = conditioned_feature_map_flat.permute(1, 2, 0).view(
+ batch_size, num_channels, height, width
+ )
+ return conditioned_feature_map
+
+ # Step 2: Get memory frames and concatenate their features
+ temporal_positions_and_previous_outputs = self._gather_memory_frame_outputs(
+ inference_session, obj_idx, frame_idx, track_in_reverse_time
+ )
+
+ memories_to_concatenate, memory_positional_embeddings_to_concatenate = self._build_memory_attention_inputs(
+ temporal_positions_and_previous_outputs, device
+ )
+ num_spatial_memory_tokens = len(memories_to_concatenate)
+
+ # Step 3: Get and process object pointers
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use = self._get_object_pointers(
+ inference_session, obj_idx, frame_idx, num_total_frames, device, track_in_reverse_time, streaming
+ )
+
+ num_object_pointer_tokens = 0
+ if pointer_tokens:
+ object_pointers, object_pointers_pos_embed = self._process_object_pointers(
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use, batch_size, num_channels, device
+ )
+
+ if object_pointers is not None:
+ memories_to_concatenate.append(object_pointers)
+ memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
+ num_object_pointer_tokens = object_pointers.shape[0]
+
+ # Step 4: Concatenate all retrieved memories and their positional embeddings
+ combined_memory = torch.cat(memories_to_concatenate, dim=0)
+ combined_memory_positional_embeddings = torch.cat(memory_positional_embeddings_to_concatenate, dim=0)
+
+ # Step 5: Forward through the memory attention mechanism
+ conditioned_feature_map_flat = self.memory_attention(
+ current_vision_features=current_vision_features,
+ current_vision_position_embeddings=current_vision_positional_embeddings,
+ memory=combined_memory,
+ memory_posision_embeddings=combined_memory_positional_embeddings, # Corrected typo from API
+ num_object_pointer_tokens=num_object_pointer_tokens,
+ num_spatial_memory_tokens=num_spatial_memory_tokens,
+ )
+
+ # Reshape from (Batch, H*W, Channels) to (Batch, Channels, Height, Width)
+ conditioned_feature_map = (
+ conditioned_feature_map_flat.squeeze(1).permute(0, 2, 1).view(batch_size, num_channels, height, width)
+ )
+ return conditioned_feature_map
+
+ def _use_multimask(self, is_init_cond_frame: bool, point_inputs: Optional[dict]) -> bool:
+ """Whether to use multimask output in the SAM head."""
+ num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(2)
+ multimask_output = (
+ self.config.multimask_output_in_sam
+ and (is_init_cond_frame or self.config.multimask_output_for_tracking)
+ and (self.config.multimask_min_pt_num <= num_pts <= self.config.multimask_max_pt_num)
+ )
+ return multimask_output
+
+ def _run_single_frame_inference(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ frame_idx: int,
+ obj_idx: int,
+ batch_size: int,
+ is_init_cond_frame: bool,
+ point_inputs: Optional[torch.Tensor],
+ mask_inputs: Optional[torch.Tensor],
+ reverse: bool,
+ run_mem_encoder: bool,
+ prev_sam_mask_logits: Optional[torch.Tensor] = None,
+ streaming: bool = False,
+ ) -> dict[str, Any]:
+ """
+ Perform a single tracking step for video object segmentation.
+
+ Args:
+ inference_session (`EdgeTamVideoInferenceSession`):
+ The video inference session object.
+ frame_idx (`int`):
+ Index of the current frame.
+ obj_idx (`int`):
+ Index of the current object.
+ batch_size (`int`):
+ Batch size of the current frame.
+ is_init_cond_frame (`bool`):
+ Whether this is an initial conditioning frame with user inputs.
+ point_inputs (`dict`, *optional*):
+ Point prompt inputs for the current frame.
+ mask_inputs (`torch.Tensor`, *optional*):
+ Mask prompt inputs for the current frame.
+ reverse (`bool`, *optional*, defaults to `False`):
+ Whether to track in reverse time order.
+ run_mem_encoder (`bool`, *optional*, defaults to `True`):
+ Whether to run the memory encoder on predicted masks.
+ prev_sam_mask_logits (`torch.Tensor`, *optional*):
+ Previously predicted SAM mask logits that can be fed with new clicks.
+ streaming (`bool`, *optional*, defaults to `False`):
+ Whether this is streaming inference.
+
+ Returns:
+ `dict`: Dictionary containing the tracking results for the current frame, including:
+ - pred_masks: Predicted low-resolution masks.
+ - object_pointer: Object pointer for memory.
+ - object_score_logits: Object score logits (inference only).
+ - maskmem_features: Memory features for future frames.
+ - maskmem_pos_enc: Memory positional encodings.
+ """
+ # Retrieve correct image features
+ current_vision_feats, current_vision_pos_embeds = self._prepare_vision_features(
+ inference_session, frame_idx, batch_size
+ )
+ # point and mask should not appear as input simultaneously on the same frame
+ if point_inputs is not None and mask_inputs is not None:
+ raise ValueError(
+ "point_inputs and mask_inputs should not appear as input simultaneously on the same frame"
+ )
+ # High-resolution feature maps for the SAM head, reshape (HW)BC => BCHW
+ if len(current_vision_feats) > 1:
+ high_res_features = [
+ x.permute(1, 2, 0).view(x.size(1), x.size(2), *s)
+ for x, s in zip(current_vision_feats[:-1], self.backbone_feature_sizes[:-1])
+ ]
+ else:
+ high_res_features = None
+ if mask_inputs is not None:
+ # We directly output the mask input (see it as a GT mask) without using a SAM prompt encoder + mask decoder.
+ pix_feat = current_vision_feats[-1].permute(1, 2, 0)
+ pix_feat = pix_feat.view(-1, self.hidden_dim, *self.backbone_feature_sizes[-1])
+ sam_outputs = self._use_mask_as_output(pix_feat, high_res_features, mask_inputs)
+ else:
+ # fused the visual feature with previous memory features in the memory bank
+ pix_feat = self._prepare_memory_conditioned_features(
+ inference_session=inference_session,
+ frame_idx=frame_idx,
+ obj_idx=obj_idx,
+ is_initial_conditioning_frame=is_init_cond_frame,
+ current_vision_features=current_vision_feats[-1],
+ current_vision_positional_embeddings=current_vision_pos_embeds[-1],
+ num_total_frames=inference_session.num_frames,
+ track_in_reverse_time=reverse,
+ streaming=streaming,
+ )
+ # apply SAM-style segmentation head
+ # here we might feed previously predicted low-res SAM mask logits into the SAM mask decoder,
+ # e.g. in demo where such logits come from earlier interaction instead of correction sampling
+ # (in this case, any `mask_inputs` shouldn't reach here as they are sent to _use_mask_as_output instead)
+ if prev_sam_mask_logits is not None:
+ mask_inputs = prev_sam_mask_logits
+ multimask_output = self._use_multimask(is_init_cond_frame, point_inputs)
+ sam_outputs = self._single_frame_forward(
+ pixel_values=None, # Vision features already computed
+ input_points=point_inputs["point_coords"] if point_inputs is not None else None,
+ input_labels=point_inputs["point_labels"] if point_inputs is not None else None,
+ input_masks=mask_inputs,
+ image_embeddings=high_res_features + [pix_feat],
+ multimask_output=multimask_output,
+ )
+
+ # Finally run the memory encoder on the predicted mask to encode
+ # it into a new memory feature (which will be used to condition vision features in future frames)
+ maskmem_features = None
+ maskmem_pos_enc = None
+ if run_mem_encoder and self.num_maskmem > 0:
+ maskmem_features, maskmem_pos_enc = self._encode_new_memory(
+ current_vision_feats=current_vision_feats[-1],
+ pred_masks_high_res=sam_outputs.high_res_masks,
+ object_score_logits=sam_outputs.object_score_logits,
+ is_mask_from_pts=(point_inputs is not None or mask_inputs is not None),
+ )
+
+ current_out = {
+ "pred_masks": sam_outputs.pred_masks,
+ "object_pointer": sam_outputs.object_pointer,
+ "maskmem_features": maskmem_features if maskmem_features is not None else None,
+ "maskmem_pos_enc": maskmem_pos_enc,
+ }
+ if not self.training:
+ current_out["object_score_logits"] = sam_outputs.object_score_logits
+
+ return current_out
+
+ def _encode_new_memory(
+ self,
+ current_vision_feats: torch.Tensor,
+ pred_masks_high_res: torch.Tensor,
+ object_score_logits: torch.Tensor,
+ is_mask_from_pts: bool,
+ ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+ """Encode the current image and its prediction into a memory feature."""
+ batch_size = current_vision_feats.size(1) # batch size on this frame
+ channels = self.hidden_dim
+ height, width = self.backbone_feature_sizes[-1] # top-level (lowest-resolution) feature size
+ # top-level feature, (HW)BC => BCHW
+ pix_feat = current_vision_feats.permute(1, 2, 0).view(batch_size, channels, height, width)
+ if is_mask_from_pts and not self.training:
+ # binarize the mask logits
+ mask_for_mem = (pred_masks_high_res > 0).to(pred_masks_high_res.dtype)
+ else:
+ # apply sigmoid on the raw mask logits to turn them into range (0, 1)
+ mask_for_mem = torch.sigmoid(pred_masks_high_res)
+ # apply scale and bias terms to the sigmoid probabilities
+ mask_for_mem = mask_for_mem * self.config.sigmoid_scale_for_mem_enc
+ mask_for_mem = mask_for_mem + self.config.sigmoid_bias_for_mem_enc
+
+ maskmem_features, maskmem_pos_enc = self.memory_encoder(
+ pix_feat,
+ mask_for_mem,
+ )
+ # add a no-object embedding to the spatial memory to indicate that the frame
+ # is predicted to be occluded (i.e. no object is appearing in the frame)
+ if self.occlusion_spatial_embedding_parameter is not None:
+ is_obj_appearing = (object_score_logits > 0).float()
+ maskmem_features += (1 - is_obj_appearing[..., None]) * self.occlusion_spatial_embedding_parameter[
+ ..., None, None
+ ].expand(*maskmem_features.shape)
+
+ maskmem_pos_enc = maskmem_pos_enc.to(pred_masks_high_res.dtype)
+ maskmem_features, maskmem_pos_enc = self.spatial_perceiver(maskmem_features, maskmem_pos_enc)
+ maskmem_features = maskmem_features.to(pred_masks_high_res.dtype)
+ maskmem_pos_enc = maskmem_pos_enc.to(pred_masks_high_res.dtype)
+
+ return maskmem_features, maskmem_pos_enc
+
+ @torch.inference_mode()
+ @auto_docstring(
+ custom_intro="""
+ Propagate the objects through the video frames. Used when initializing an inference session with a whole video.
+ Yields EdgeTamVideoSegmentationOutput for each frame.
+ """
+ )
+ def propagate_in_video_iterator(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ start_frame_idx: Optional[int] = None,
+ max_frame_num_to_track: Optional[int] = None,
+ reverse: bool = False,
+ ) -> Iterator[EdgeTamVideoSegmentationOutput]:
+ r"""
+ inference_session (`EdgeTamVideoInferenceSession`):
+ The video inference session object.
+ start_frame_idx (`int`, *optional*):
+ The starting frame index for propagation.
+ Need to be provided if `forward` hasn't been called on new inputs yet.
+ If not provided, the starting frame index will be the earliest frame with input points.
+ max_frame_num_to_track (`int`, *optional*):
+ The maximum number of frames to track.
+ reverse (`bool`, *optional*, defaults to `False`):
+ Whether to propagate in reverse.
+ """
+ num_frames = inference_session.num_frames
+
+ # set start index, end index, and processing order
+ if start_frame_idx is None:
+ # default: start from the earliest frame with input points
+ frames_with_inputs = [
+ frame_idx
+ for obj_output_dict in inference_session.output_dict_per_obj.values()
+ for frame_idx in obj_output_dict["cond_frame_outputs"]
+ ]
+ if not frames_with_inputs:
+ raise ValueError(
+ "Cannot determine the starting frame index; please specify it manually, or run inference on a frame with inputs first."
+ )
+ start_frame_idx = min(frames_with_inputs)
+ if max_frame_num_to_track is None:
+ # default: track all the frames in the video
+ max_frame_num_to_track = num_frames
+ if reverse:
+ end_frame_idx = max(start_frame_idx - max_frame_num_to_track, 0)
+ if start_frame_idx > 0:
+ processing_order = range(start_frame_idx, end_frame_idx - 1, -1)
+ else:
+ processing_order = [] # skip reverse tracking if starting from frame 0
+ else:
+ end_frame_idx = min(start_frame_idx + max_frame_num_to_track, num_frames - 1)
+ processing_order = range(start_frame_idx, end_frame_idx + 1)
+
+ for frame_idx in tqdm(processing_order, desc="propagate in video"):
+ edgetam_video_output = self(inference_session, frame_idx=frame_idx, reverse=reverse)
+ yield edgetam_video_output
+
+
+__all__ = ["EdgeTamVideoModel", "EdgeTamVideoInferenceSession", "EdgeTamVideoPreTrainedModel"]
diff --git a/src/transformers/models/edgetam_video/modular_edgetam_video.py b/src/transformers/models/edgetam_video/modular_edgetam_video.py
new file mode 100644
index 000000000000..b520cd5a756b
--- /dev/null
+++ b/src/transformers/models/edgetam_video/modular_edgetam_video.py
@@ -0,0 +1,1243 @@
+# coding=utf-8
+# 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 math
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch import Tensor
+
+from transformers.models.sam2.modeling_sam2 import (
+ eager_attention_forward,
+ window_partition,
+)
+from transformers.utils.generic import OutputRecorder
+
+from ...activations import ACT2FN
+from ...configuration_utils import PretrainedConfig
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...processing_utils import Unpack
+from ...pytorch_utils import compile_compatible_method_lru_cache
+from ...utils import (
+ auto_docstring,
+)
+from ..auto import CONFIG_MAPPING, AutoConfig
+from ..sam2_video.configuration_sam2_video import (
+ Sam2VideoConfig,
+ Sam2VideoMaskDecoderConfig,
+ Sam2VideoPromptEncoderConfig,
+)
+from ..sam2_video.modeling_sam2_video import (
+ Sam2VideoAttention,
+ Sam2VideoFeedForward,
+ Sam2VideoInferenceSession,
+ Sam2VideoLayerNorm,
+ Sam2VideoMemoryAttention,
+ Sam2VideoMemoryEncoder,
+ Sam2VideoMemoryFuserCXBlock,
+ Sam2VideoModel,
+ Sam2VideoPositionEmbeddingSine,
+ Sam2VideoPreTrainedModel,
+ Sam2VideoTwoWayAttentionBlock,
+ Sam2VideoVisionEncoderOutput,
+ Sam2VideoVisionRotaryEmbedding,
+ rotate_pairwise,
+)
+
+
+class EdgeTamVideoPromptEncoderConfig(Sam2VideoPromptEncoderConfig):
+ pass
+
+
+class EdgeTamVideoMaskDecoderConfig(Sam2VideoMaskDecoderConfig):
+ pass
+
+
+class EdgeTamVideoConfig(Sam2VideoConfig):
+ r"""
+ [`EdgeTamVideoConfig`] is the configuration class to store the configuration of a [`EdgeTamVideoModel`]. It is used to instantiate a
+ EDGETAM model according to the specified arguments, defining the memory attention, memory encoder, and image encoder
+ configs. Instantiating a configuration defaults will yield a similar configuration to that of the SAM 2.1 Hiera-tiny
+ [facebook/EdgeTAM](https://huggingface.co/facebook/EdgeTAM) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vision_config (Union[`dict`, `EdgeTamVideoVisionConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamVideoVisionConfig`].
+ prompt_encoder_config (Union[`dict`, `EdgeTamVideoPromptEncoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamVideoPromptEncoderConfig`].
+ mask_decoder_config (Union[`dict`, `EdgeTamVideoMaskDecoderConfig`], *optional*):
+ Dictionary of configuration options used to initialize [`EdgeTamMaskDecoderConfig`].
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ Standard deviation for parameter initialization.
+ num_maskmem (`int`, *optional*, defaults to 7):
+ The number of memory slots for the mask memory.
+ image_size (`int`, *optional*, defaults to 1024):
+ The size of the input images.
+ sigmoid_scale_for_mem_enc (`float`, *optional*, defaults to 20.0):
+ Scale factor for the sigmoid function in the memory encoder.
+ sigmoid_bias_for_mem_enc (`float`, *optional*, defaults to -10.0):
+ Bias for the sigmoid function in the memory encoder.
+ enable_occlusion_spatial_embedding (`bool`, *optional*, defaults to `True`):
+ Whether to enable spatial embedding for occlusions.
+ multimask_output_in_sam (`bool`, *optional*, defaults to `True`):
+ Whether to output multiple masks from the SAM head.
+ multimask_min_pt_num (`int`, *optional*, defaults to 0):
+ The minimum number of points to trigger multimask output.
+ multimask_max_pt_num (`int`, *optional*, defaults to 1):
+ The maximum number of points to trigger multimask output.
+ multimask_output_for_tracking (`bool`, *optional*, defaults to `True`):
+ Whether to use multimask output for tracking.
+ max_object_pointers_in_encoder (`int`, *optional*, defaults to 16):
+ The maximum number of object pointers in the encoder.
+ enable_temporal_pos_encoding_for_object_pointers (`bool`, *optional*, defaults to `True`):
+ Whether to enable temporal positional encoding for object pointers.
+ memory_attention_hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the memory attention hidden states.
+ memory_attention_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the memory attention module.
+ memory_attention_num_attention_heads (`int`, *optional*, defaults to 1):
+ Number of attention heads for each attention layer in the memory attention.
+ memory_attention_downsample_rate (`int`, *optional*, defaults to 1):
+ The downsample rate for the attention layers.
+ memory_attention_mlp_hidden_size (`int`, *optional*, defaults to 2048):
+ The dimension of the feedforward network in the memory attention module.
+ memory_attention_mlp_hidden_act (`str`, *optional*, defaults to `"relu"`):
+ The non-linear activation function in the feedforward network in the memory attention module.
+ memory_attention_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout rate for the memory attention module.
+ memory_attention_rope_theta (`float`, *optional*, defaults to 10000):
+ The Rope theta parameter.
+ memory_attention_rope_feat_sizes (`Tuple[int, int]`, *optional*, defaults to `[64, 64]`):
+ The feature sizes for the Rope positional encoding.
+ memory_attention_rope_k_sizes (`List[int]`, *optional*, defaults to `[16, 16]`):
+ The key feature sizes for the RoPE positional encoding in memory attention.
+ memory_attention_rope_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout rate for the Rope positional encoding.
+ perceiver_resampler_num_latents (`int`, *optional*, defaults to 256):
+ The number of 1D latent tokens in the perceiver resampler.
+ perceiver_resampler_num_latents_2d (`int`, *optional*, defaults to 256):
+ The number of 2D latent tokens in the perceiver resampler.
+ perceiver_resampler_hidden_size (`int`, *optional*, defaults to 64):
+ The hidden size of the perceiver resampler.
+ perceiver_resampler_mlp_intermediate_size (`int`, *optional*, defaults to 256):
+ The intermediate size of the feedforward network in the perceiver resampler.
+ perceiver_resampler_num_attention_heads (`int`, *optional*, defaults to 1):
+ The number of attention heads in the perceiver resampler.
+ perceiver_resampler_attention_head_dim (`int`, *optional*, defaults to 64):
+ The dimension of each attention head in the perceiver resampler.
+ perceiver_resampler_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the perceiver resampler.
+ perceiver_resampler_hidden_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout rate for the hidden layers in the perceiver resampler.
+ perceiver_resampler_attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout rate for the attention layers in the perceiver resampler.
+ memory_encoder_hidden_size (`int`, *optional*, defaults to 256):
+ Dimensionality of the memory encoder hidden states.
+ memory_encoder_output_channels (`int`, *optional*, defaults to 64):
+ The number of output channels for the memory encoder.
+ mask_downsampler_embed_dim (`int`, *optional*, defaults to 256):
+ The dimension of the mask downsampler embedding.
+ memory_fuser_intermediate_dim (`int`, *optional*, defaults to 1024):
+ The intermediate dimension of the memory fuser feedforward network.
+ mask_downsampler_kernel_size (`int`, *optional*, defaults to 3):
+ The kernel size for the mask downsampler.
+ mask_downsampler_stride (`int`, *optional*, defaults to 2):
+ The stride for the mask downsampler.
+ mask_downsampler_padding (`int`, *optional*, defaults to 1):
+ The padding for the mask downsampler.
+ mask_downsampler_total_stride (`int`, *optional*, defaults to 16):
+ The total stride for the mask downsampler.
+ mask_downsampler_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the mask downsampler.
+ memory_fuser_num_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the memory fuser.
+ memory_fuser_embed_dim (`int`, *optional*, defaults to 256):
+ The dimension of the memory fuser embedding.
+ memory_fuser_kernel_size (`int`, *optional*, defaults to 7):
+ The kernel size for the memory fuser.
+ memory_fuser_padding (`int`, *optional*, defaults to 3):
+ The padding for the memory fuser.
+ memory_fuser_layer_scale_init_value (`float`, *optional*, defaults to 1e-06):
+ The initial value for the layer scale in the memory fuser.
+ memory_fuser_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function in the memory fuser.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... EdgeTamVisionConfig,
+ ... EdgeTamVideoPromptEncoderConfig,
+ ... EdgeTamVideoMaskDecoderConfig,
+ ... EdgeTamVideoModel,
+ ... EdgeTamVideoConfig,
+ ... )
+
+ >>> # Initializing a EdgeTamVideoConfig with `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> configuration = EdgeTamVideoConfig()
+
+ >>> # Initializing a EdgeTamVideoModel (with random weights) from the `"facebook/edgetam.1_hiera_tiny"` style configuration
+ >>> model = EdgeTamVideoModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a EdgeTamConfig from a EdgeTamVisionConfig, EdgeTamPromptEncoderConfig, and EdgeTamMaskDecoderConfig
+
+ >>> # Initializing EDGETAM vision encoder, memory attention, and memory encoder configurations
+ >>> vision_config = EdgeTamVisionConfig()
+ >>> prompt_encoder_config = EdgeTamVideoPromptEncoderConfig()
+ >>> mask_decoder_config = EdgeTamVideoMaskDecoderConfig()
+
+ >>> config = EdgeTamVideoConfig(vision_config, prompt_encoder_config, mask_decoder_config)
+ ```"""
+
+ model_type = "edgetam_video"
+ sub_configs = {
+ "vision_config": AutoConfig,
+ "prompt_encoder_config": EdgeTamVideoPromptEncoderConfig,
+ "mask_decoder_config": EdgeTamVideoMaskDecoderConfig,
+ }
+
+ def __init__(
+ self,
+ vision_config=None,
+ prompt_encoder_config=None,
+ mask_decoder_config=None,
+ initializer_range=0.02,
+ num_maskmem=7,
+ image_size=1024,
+ sigmoid_scale_for_mem_enc=20.0,
+ sigmoid_bias_for_mem_enc=-10.0,
+ enable_occlusion_spatial_embedding=True,
+ multimask_output_in_sam=True,
+ multimask_min_pt_num=0,
+ multimask_max_pt_num=1,
+ multimask_output_for_tracking=True,
+ max_object_pointers_in_encoder=16,
+ enable_temporal_pos_encoding_for_object_pointers=True,
+ # memory attention
+ memory_attention_hidden_size=256,
+ memory_attention_num_layers=2,
+ memory_attention_num_attention_heads=1,
+ memory_attention_downsample_rate=1,
+ memory_attention_mlp_hidden_size=2048,
+ memory_attention_mlp_hidden_act="relu",
+ memory_attention_dropout=0.1,
+ memory_attention_rope_theta=10000,
+ memory_attention_rope_feat_sizes=None,
+ memory_attention_rope_k_sizes=None,
+ memory_attention_rope_dropout=0.1,
+ # spatial perceiver resampler
+ perceiver_resampler_num_latents=256,
+ perceiver_resampler_num_latents_2d=256,
+ perceiver_resampler_hidden_size=64,
+ perceiver_resampler_mlp_intermediate_size=256,
+ perceiver_resampler_num_attention_heads=1,
+ perceiver_resampler_attention_head_dim=64,
+ perceiver_resampler_num_layers=2,
+ perceiver_resampler_hidden_dropout=0.0,
+ perceiver_resampler_attention_dropout=0.0,
+ # memory encoder
+ memory_encoder_hidden_size=256,
+ memory_encoder_output_channels=64,
+ mask_downsampler_embed_dim=256,
+ memory_fuser_intermediate_dim=1024,
+ mask_downsampler_kernel_size=3,
+ mask_downsampler_stride=2,
+ mask_downsampler_padding=1,
+ mask_downsampler_total_stride=16,
+ mask_downsampler_hidden_act="gelu",
+ memory_fuser_num_layers=2,
+ memory_fuser_embed_dim=256,
+ memory_fuser_kernel_size=7,
+ memory_fuser_padding=3,
+ memory_fuser_layer_scale_init_value=1e-6,
+ memory_fuser_hidden_act="gelu",
+ **kwargs,
+ ):
+ PretrainedConfig.__init__(**kwargs)
+ vision_config = vision_config if vision_config is not None else {}
+ prompt_encoder_config = prompt_encoder_config if prompt_encoder_config is not None else {}
+ mask_decoder_config = mask_decoder_config if mask_decoder_config is not None else {}
+ memory_attention_rope_feat_sizes = (
+ [64, 64] if memory_attention_rope_feat_sizes is None else memory_attention_rope_feat_sizes
+ )
+ memory_attention_rope_k_sizes = (
+ [16, 16] if memory_attention_rope_k_sizes is None else memory_attention_rope_k_sizes
+ )
+
+ if isinstance(vision_config, dict):
+ vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
+ vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
+ if isinstance(prompt_encoder_config, EdgeTamVideoPromptEncoderConfig):
+ prompt_encoder_config = prompt_encoder_config.to_dict()
+ if isinstance(mask_decoder_config, EdgeTamVideoMaskDecoderConfig):
+ mask_decoder_config = mask_decoder_config.to_dict()
+
+ self.vision_config = vision_config
+ self.prompt_encoder_config = EdgeTamVideoPromptEncoderConfig(**prompt_encoder_config)
+ self.mask_decoder_config = EdgeTamVideoMaskDecoderConfig(**mask_decoder_config)
+
+ self.initializer_range = initializer_range
+ self.num_maskmem = num_maskmem # default 1 input frame + 6 previous frames
+ self.image_size = image_size
+ self.sigmoid_scale_for_mem_enc = sigmoid_scale_for_mem_enc # scale factor for mask sigmoid prob
+ self.sigmoid_bias_for_mem_enc = sigmoid_bias_for_mem_enc # bias factor for mask sigmoid prob
+ self.enable_occlusion_spatial_embedding = enable_occlusion_spatial_embedding
+ self.multimask_output_in_sam = multimask_output_in_sam
+ self.multimask_min_pt_num = multimask_min_pt_num
+ self.multimask_max_pt_num = multimask_max_pt_num
+ self.multimask_output_for_tracking = multimask_output_for_tracking
+ self.max_object_pointers_in_encoder = max_object_pointers_in_encoder
+ self.enable_temporal_pos_encoding_for_object_pointers = enable_temporal_pos_encoding_for_object_pointers
+
+ # memory attention
+ self.memory_attention_hidden_size = memory_attention_hidden_size
+ self.memory_attention_num_layers = memory_attention_num_layers
+ self.memory_attention_num_attention_heads = memory_attention_num_attention_heads
+ self.memory_attention_downsample_rate = memory_attention_downsample_rate
+ self.memory_attention_mlp_hidden_size = memory_attention_mlp_hidden_size
+ self.memory_attention_mlp_hidden_act = memory_attention_mlp_hidden_act
+ self.memory_attention_dropout = memory_attention_dropout
+ self.memory_attention_rope_theta = memory_attention_rope_theta
+ self.memory_attention_rope_feat_sizes = memory_attention_rope_feat_sizes
+ self.memory_attention_rope_k_sizes = memory_attention_rope_k_sizes
+ self.memory_attention_rope_dropout = memory_attention_rope_dropout
+
+ # spatial perceiver resampler
+ self.perceiver_resampler_num_latents = perceiver_resampler_num_latents
+ self.perceiver_resampler_num_latents_2d = perceiver_resampler_num_latents_2d
+ self.perceiver_resampler_hidden_size = perceiver_resampler_hidden_size
+ self.perceiver_resampler_mlp_intermediate_size = perceiver_resampler_mlp_intermediate_size
+ self.perceiver_resampler_attention_head_dim = perceiver_resampler_attention_head_dim
+ self.perceiver_resampler_num_attention_heads = perceiver_resampler_num_attention_heads
+ self.perceiver_resampler_num_layers = perceiver_resampler_num_layers
+ self.perceiver_resampler_hidden_dropout = perceiver_resampler_hidden_dropout
+ self.perceiver_resampler_attention_dropout = perceiver_resampler_attention_dropout
+
+ # memory encoder
+ self.memory_encoder_hidden_size = memory_encoder_hidden_size
+ self.memory_encoder_output_channels = memory_encoder_output_channels
+ self.mask_downsampler_embed_dim = mask_downsampler_embed_dim
+ self.mask_downsampler_kernel_size = mask_downsampler_kernel_size
+ self.mask_downsampler_stride = mask_downsampler_stride
+ self.mask_downsampler_padding = mask_downsampler_padding
+ self.mask_downsampler_total_stride = mask_downsampler_total_stride
+ self.mask_downsampler_hidden_act = mask_downsampler_hidden_act
+ self.memory_fuser_num_layers = memory_fuser_num_layers
+ self.memory_fuser_embed_dim = memory_fuser_embed_dim
+ self.memory_fuser_intermediate_dim = memory_fuser_intermediate_dim
+ self.memory_fuser_kernel_size = memory_fuser_kernel_size
+ self.memory_fuser_padding = memory_fuser_padding
+ self.memory_fuser_layer_scale_init_value = memory_fuser_layer_scale_init_value
+ self.memory_fuser_hidden_act = memory_fuser_hidden_act
+
+
+class EdgeTamVideoLayerNorm(Sam2VideoLayerNorm):
+ pass
+
+
+class EdgeTamVideoMemoryFuserCXBlock(Sam2VideoMemoryFuserCXBlock):
+ pass
+
+
+class EdgeTamVideoVisionEncoderOutput(Sam2VideoVisionEncoderOutput):
+ pass
+
+
+class EdgeTamVideoVisionRotaryEmbedding(Sam2VideoVisionRotaryEmbedding):
+ def __init__(self, config: EdgeTamVideoConfig, end_x: Optional[int] = None, end_y: Optional[int] = None):
+ nn.Module.__init__()
+ dim = config.memory_attention_hidden_size // (
+ config.memory_attention_downsample_rate * config.memory_attention_num_attention_heads
+ )
+ # Ensure even dimension for proper axial splitting
+ if dim % 4 != 0:
+ raise ValueError("Dimension must be divisible by 4 for axial RoPE")
+ end_x, end_y = config.memory_attention_rope_feat_sizes if end_x is None else (end_x, end_y)
+ freqs = 1.0 / (config.memory_attention_rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+
+ # Generate 2D position indices for axial rotary embedding
+ flattened_indices = torch.arange(end_x * end_y, dtype=torch.long)
+ x_positions = flattened_indices % end_x
+ y_positions = torch.div(flattened_indices, end_x, rounding_mode="floor")
+ freqs_x = torch.outer(x_positions, freqs).float()
+ freqs_y = torch.outer(y_positions, freqs).float()
+ inv_freq = torch.cat([freqs_x, freqs_y], dim=-1)
+ inv_freq = inv_freq.repeat_interleave(2, dim=-1)
+ # directly register the cos and sin embeddings as we have a fixed feature shape
+ self.register_buffer("rope_embeddings_cos", inv_freq.cos(), persistent=False)
+ self.register_buffer("rope_embeddings_sin", inv_freq.sin(), persistent=False)
+
+
+class EdgeTamVideoAttention(Sam2VideoAttention):
+ pass
+
+
+def apply_rotary_pos_emb_2d_self_attn(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary position embedding to query and key tensors for self-attention.
+
+ Args:
+ q: Query tensor of shape (..., seq_len, head_dim)
+ k: Key tensor of shape (..., seq_len, head_dim)
+ cos: Cosine position embedding of shape (seq_len, head_dim)
+ sin: Sine position embedding of shape (seq_len, head_dim)
+
+ Returns:
+ Rotated (q, k) tensors
+ """
+ # Apply RoPE to queries
+ q_embed = q.float() # force upscale to float32 as in the original implementation
+ q_embed = (q_embed * cos) + (rotate_pairwise(q_embed) * sin)
+
+ # Apply RoPE to keys (same embeddings as queries for self-attention)
+ k_embed = k.float() # force upscale to float32 as in the original implementation
+ k_embed = (k_embed * cos) + (rotate_pairwise(k_embed) * sin)
+
+ return q_embed.type_as(q), k_embed.type_as(k)
+
+
+def apply_rotary_pos_emb_2d_cross_attn(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ cos_k: torch.Tensor,
+ sin_k: torch.Tensor,
+ num_k_exclude_rope: int = 0,
+ repeat_freqs_k: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary position embedding to query and key tensors for cross-attention.
+
+ Args:
+ q: Query tensor of shape (..., seq_len, head_dim)
+ k: Key tensor of shape (..., seq_len, head_dim)
+ cos: Cosine position embedding of shape (seq_len, head_dim)
+ sin: Sine position embedding of shape (seq_len, head_dim)
+ cos_k: Cosine position embedding for keys of shape (seq_len, head_dim)
+ sin_k: Sine position embedding for keys of shape (seq_len, head_dim)
+ num_k_exclude_rope: Number of tokens at end of k to exclude from RoPE (e.g., object pointer tokens)
+ repeat_freqs_k: Frequency repetition for keys in cross-attention (e.g., for spatial memory tokens)
+
+ Returns:
+ Rotated (q, k) tensors
+ """
+ # Apply RoPE to queries (always straightforward)
+ q_embed = q.float()
+ q_embed = (q_embed * cos) + (rotate_pairwise(q_embed) * sin)
+
+ # Split keys: RoPE tokens and excluded tokens (e.g., object pointers)
+ num_total_k_tokens = k.shape[-2]
+ k_for_rope = k[..., : num_total_k_tokens - num_k_exclude_rope, :]
+ k_excluded = k[..., num_total_k_tokens - num_k_exclude_rope :, :]
+
+ # Early return if no keys need RoPE
+ if k_for_rope.shape[-2] == 0:
+ return q_embed.type_as(q), k_excluded
+
+ batch_size, num_heads, k_seq_len, channels_per_head = k_for_rope.shape
+
+ # Handle temporal/spatial token structure for memory
+ # Keys have temporal + spatial structure, only spatial tokens get RoPE
+ tokens_per_group = k_seq_len // repeat_freqs_k
+ spatial_tokens = cos_k.shape[-2]
+ temporal_tokens = tokens_per_group - spatial_tokens
+
+ # Reshape and separate temporal/spatial tokens
+ k_grouped = k_for_rope.view(batch_size, num_heads, repeat_freqs_k, tokens_per_group, channels_per_head)
+ k_temporal = k_grouped[..., :temporal_tokens, :].reshape(batch_size, num_heads, -1, channels_per_head)
+ k_spatial = k_grouped[..., temporal_tokens:, :].reshape(batch_size, num_heads, -1, channels_per_head)
+
+ # Only apply RoPE to spatial tokens
+ k_rope_input = k_spatial
+
+ # Prepare position embeddings for repeated groups
+ if repeat_freqs_k > 1:
+ cos_k = cos_k.repeat(1, 1, repeat_freqs_k, 1)
+ sin_k = sin_k.repeat(1, 1, repeat_freqs_k, 1)
+
+ # Apply RoPE to spatial tokens
+ k_spatial_embed = k_rope_input.float()
+ k_spatial_embed = (k_spatial_embed * cos_k) + (rotate_pairwise(k_spatial_embed) * sin_k)
+
+ # Reconstruct: temporal + spatial tokens back to original structure
+ k_spatial_reshaped = k_spatial_embed.view(batch_size, num_heads, repeat_freqs_k, -1, channels_per_head)
+ k_temporal_reshaped = k_temporal.view(batch_size, num_heads, repeat_freqs_k, -1, channels_per_head)
+ k_final = torch.cat([k_temporal_reshaped, k_spatial_reshaped], dim=3)
+ k_final = k_final.view(batch_size, num_heads, k_seq_len, channels_per_head)
+
+ # Combine RoPE-processed keys with excluded tokens
+ k_embed = torch.cat([k_final.type_as(k), k_excluded], dim=-2)
+ return q_embed.type_as(q), k_embed
+
+
+class EdgeTamVideoRoPESelfAttention(nn.Module):
+ """Self-attention with rotary position encoding."""
+
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.internal_dim = self.hidden_size // config.memory_attention_downsample_rate
+ self.num_attention_heads = config.memory_attention_num_attention_heads
+ self.head_dim = self.internal_dim // config.memory_attention_num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.v_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+ self.dropout_p = config.memory_attention_rope_dropout
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tensor:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ # Apply rotary position encoding for self-attention
+ query, key = apply_rotary_pos_emb_2d_self_attn(query, key, cos=cos, sin=sin)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.dropout_p,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class EdgeTamVideoRoPECrossAttention(nn.Module):
+ """Cross-attention with rotary position encoding."""
+
+ def __init__(self, config: EdgeTamVideoConfig, kv_in_dim: int):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.internal_dim = self.hidden_size // config.memory_attention_downsample_rate
+ self.num_attention_heads = config.memory_attention_num_attention_heads
+ self.head_dim = self.internal_dim // config.memory_attention_num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.kv_in_dim = kv_in_dim
+
+ self.q_proj = nn.Linear(self.hidden_size, self.internal_dim)
+ self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+ self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+ self.o_proj = nn.Linear(self.internal_dim, self.hidden_size)
+ self.dropout_p = config.memory_attention_rope_dropout
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ position_embeddings_k: tuple[torch.Tensor, torch.Tensor],
+ num_k_exclude_rope: int = 0,
+ rope_k_repeat: int = 0,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tensor:
+ # Input projections
+ batch_size, point_batch_size = query.shape[:2]
+ new_shape = (batch_size * point_batch_size, -1, self.num_attention_heads, self.head_dim)
+
+ query = self.q_proj(query).view(*new_shape).transpose(1, 2)
+ key = self.k_proj(key).view(*new_shape).transpose(1, 2)
+ value = self.v_proj(value).view(*new_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ cos_k, sin_k = position_embeddings_k
+ # Apply rotary position encoding for cross-attention
+ query, key = apply_rotary_pos_emb_2d_cross_attn(
+ query,
+ key,
+ cos=cos,
+ sin=sin,
+ cos_k=cos_k,
+ sin_k=sin_k,
+ repeat_freqs_k=rope_k_repeat,
+ num_k_exclude_rope=num_k_exclude_rope,
+ )
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.dropout_p,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(
+ batch_size, point_batch_size, -1, self.num_attention_heads * self.head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class EdgeTamVideoTwoWayAttentionBlock(Sam2VideoTwoWayAttentionBlock):
+ pass
+
+
+class EdgeTamVideoPositionEmbeddingSine(Sam2VideoPositionEmbeddingSine):
+ # maxsize=2 because we need to cache the forward method for both memory encoder and perceiver resampler
+ @compile_compatible_method_lru_cache(maxsize=2)
+ def forward(self, **super_kwargs):
+ return super().forward(**super_kwargs)
+
+
+class EdgeTamVideoMemoryEncoder(Sam2VideoMemoryEncoder):
+ pass
+
+
+class EdgeTamVideoFeedForward(Sam2VideoFeedForward):
+ pass
+
+
+class EdgeTamVideoPreTrainedModel(Sam2VideoPreTrainedModel):
+ pass
+
+
+class EdgeTamVideoInferenceSession(Sam2VideoInferenceSession):
+ pass
+
+
+class EdgeTamVideoMemoryAttentionMLP(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.memory_attention_hidden_size
+ self.intermediate_size = config.memory_attention_mlp_hidden_size
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size)
+ self.dropout = nn.Dropout(config.memory_attention_dropout)
+ self.act_fn = ACT2FN[config.memory_attention_mlp_hidden_act]
+
+ def forward(self, x):
+ return self.down_proj(self.dropout(self.act_fn(self.up_proj(x))))
+
+
+class EdgeTamVideoMemoryAttentionLayer(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ hidden_size = config.memory_attention_hidden_size
+ self.self_attn = EdgeTamVideoRoPESelfAttention(config)
+ self.cross_attn_image = EdgeTamVideoRoPECrossAttention(config, kv_in_dim=64)
+
+ # MLP module
+ self.mlp = EdgeTamVideoMemoryAttentionMLP(config)
+
+ self.layer_norm1 = nn.LayerNorm(hidden_size)
+ self.layer_norm2 = nn.LayerNorm(hidden_size)
+ self.layer_norm3 = nn.LayerNorm(hidden_size)
+ self.dropout1 = nn.Dropout(config.memory_attention_dropout)
+ self.dropout2 = nn.Dropout(config.memory_attention_dropout)
+ self.dropout3 = nn.Dropout(config.memory_attention_dropout)
+
+ def forward(
+ self,
+ queries: Tensor,
+ keys: Tensor,
+ key_point_embedding: Tensor,
+ rope_position_embeddings: tuple[Tensor, Tensor],
+ rope_position_embeddings_k: Optional[tuple[Tensor, Tensor]] = None,
+ num_k_exclude_rope: int = 0,
+ rope_k_repeat: int = 0,
+ ) -> torch.Tensor:
+ # Self-Attention
+ query = self.layer_norm1(queries)
+ query, _ = self.self_attn(query=query, key=query, value=query, position_embeddings=rope_position_embeddings)
+ queries = queries + self.dropout1(query)
+
+ # Cross-Attention
+ query = self.layer_norm2(queries)
+ query, _ = self.cross_attn_image(
+ query=query,
+ key=keys + key_point_embedding,
+ value=keys,
+ position_embeddings=rope_position_embeddings,
+ position_embeddings_k=rope_position_embeddings_k,
+ num_k_exclude_rope=num_k_exclude_rope,
+ rope_k_repeat=rope_k_repeat,
+ )
+ queries = queries + self.dropout2(query)
+ # MLP
+ query = self.layer_norm3(queries)
+ query = self.mlp(query)
+ queries = queries + self.dropout3(query)
+ return queries
+
+
+class EdgeTamVideoMemoryAttention(Sam2VideoMemoryAttention):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.rotary_emb_k = EdgeTamVideoVisionRotaryEmbedding(
+ config, end_x=config.memory_attention_rope_k_sizes[0], end_y=config.memory_attention_rope_k_sizes[1]
+ )
+
+ def forward(
+ self,
+ current_vision_features: torch.Tensor,
+ memory: torch.Tensor,
+ current_vision_position_embeddings: Optional[Tensor] = None,
+ memory_posision_embeddings: Optional[Tensor] = None,
+ num_object_pointer_tokens: int = 0,
+ num_spatial_memory_tokens: int = -1,
+ ):
+ """
+ Args:
+ current_vision_features (`torch.FloatTensor`):
+ The current vision features used for self-attention.
+ memory (`torch.FloatTensor`):
+ The memory features used for cross-attention.
+ current_vision_position_embeddings (`torch.FloatTensor`, *optional*):
+ The position embeddings for the current vision features.
+ memory_posision_embeddings (`torch.FloatTensor`, *optional*):
+ The position embeddings for the memory features.
+ num_object_pointer_tokens (`int`, *optional*, defaults to 0):
+ The number of object pointer tokens.
+ """
+ output = current_vision_features
+ if current_vision_position_embeddings is not None:
+ output = output + 0.1 * current_vision_position_embeddings
+
+ # Convert to batch first
+ output = output.transpose(0, 1)
+ memory = memory.transpose(0, 1).unsqueeze(1)
+ memory_posision_embeddings = memory_posision_embeddings.transpose(0, 1).unsqueeze(1)
+ rope_position_embeddings = self.rotary_emb()
+ rope_position_embeddings_k = self.rotary_emb_k()
+ for layer in self.layers:
+ output = layer(
+ queries=output.unsqueeze(1) if output.ndim == 3 else output,
+ keys=memory,
+ key_point_embedding=memory_posision_embeddings,
+ rope_position_embeddings=rope_position_embeddings,
+ rope_position_embeddings_k=rope_position_embeddings_k,
+ num_k_exclude_rope=num_object_pointer_tokens,
+ rope_k_repeat=num_spatial_memory_tokens,
+ )
+
+ normed_output = self.layer_norm(output)
+
+ # Convert back to seq first
+ normed_output = normed_output.transpose(0, 1)
+
+ return normed_output
+
+
+class EdgeTamVideoPerceiverMLP(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.intermediate_size = config.perceiver_resampler_mlp_intermediate_size
+
+ self.layer_norm = nn.LayerNorm(self.hidden_size)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = nn.GELU()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.layer_norm(hidden_states)
+ hidden_states = self.down_proj(self.act_fn(self.up_proj(hidden_states)))
+ return hidden_states
+
+
+class EdgeTamVideoPerceiverAttention(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.num_attention_heads = config.perceiver_resampler_num_attention_heads
+ self.head_dim = config.perceiver_resampler_attention_head_dim
+ self.attention_dropout = config.perceiver_resampler_attention_dropout
+
+ self.inner_dim = self.head_dim * self.num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.k_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.v_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+ self.o_proj = nn.Linear(self.inner_dim, self.hidden_size, bias=False)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ # Project queries, keys, and values
+ query = self.q_proj(query)
+ key = self.k_proj(key)
+ value = self.v_proj(value)
+
+ # Reshape for multi-head attention
+ batch_size, seq_len_q = query.shape[:2]
+ query = query.view(batch_size, seq_len_q, self.num_attention_heads, self.head_dim).transpose(1, 2)
+ seq_len_kv = key.shape[1]
+ key = key.view(batch_size, seq_len_kv, self.num_attention_heads, self.head_dim).transpose(1, 2)
+ value = value.view(batch_size, seq_len_kv, self.num_attention_heads, self.head_dim).transpose(1, 2)
+
+ # Add positional encoding if provided
+ if positional_encoding is not None:
+ pos_encoding = positional_encoding.view(
+ batch_size, seq_len_kv, self.num_attention_heads, self.head_dim
+ ).transpose(1, 2)
+ key = key + pos_encoding
+ value = value + pos_encoding
+
+ # Apply attention
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, _ = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ is_causal=self.is_causal,
+ **kwargs,
+ )
+
+ # Reshape output
+ attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, self.inner_dim)
+ return self.o_proj(attn_output)
+
+
+class EdgeTamVideoPerceiverEncoderLayer(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+
+ self.cross_attention = EdgeTamVideoPerceiverAttention(config)
+ self.mlp = EdgeTamVideoPerceiverMLP(config)
+ self.dropout = nn.Dropout(config.perceiver_resampler_hidden_dropout)
+
+ self.self_attention = EdgeTamVideoPerceiverAttention(config)
+ self.self_mlp = EdgeTamVideoPerceiverMLP(config)
+
+ # Layer norms moved from attention classes to here
+ self.layer_norm_input = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+ self.layer_norm_latents = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+ self.layer_norm_self = nn.LayerNorm(config.perceiver_resampler_hidden_size)
+
+ def forward(
+ self,
+ latents: torch.Tensor,
+ input_features: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ # Cross attention with layer norms
+ normalized_latents = self.layer_norm_latents(latents)
+ normalized_input = self.layer_norm_input(input_features)
+ cross_attention_output = self.cross_attention(
+ query=normalized_latents,
+ key=normalized_input,
+ value=normalized_input,
+ positional_encoding=positional_encoding,
+ )
+ latents = latents + self.dropout(cross_attention_output)
+
+ mlp_output = self.mlp(latents)
+ latents = latents + mlp_output
+
+ # Self attention with layer norm
+ normalized_latents_self = self.layer_norm_self(latents)
+ self_attention_output = self.self_attention(
+ query=normalized_latents_self, key=normalized_latents_self, value=normalized_latents_self
+ )
+ latents = latents + self_attention_output
+
+ self_mlp_output = self.self_mlp(latents)
+ latents = latents + self_mlp_output
+
+ return latents
+
+
+class EdgeTamVideoPerceiverResampler(nn.Module):
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.perceiver_resampler_hidden_size
+ self.num_latents_1d = config.perceiver_resampler_num_latents
+ self.num_latents_2d = config.perceiver_resampler_num_latents_2d
+ self.num_layers = config.perceiver_resampler_num_layers
+
+ if self.num_latents_1d > 0:
+ self.latents_1d = nn.Parameter(torch.randn(self.num_latents_1d, self.hidden_size))
+ if self.num_latents_2d > 0:
+ self.latents_2d = nn.Parameter(torch.randn(self.num_latents_2d, self.hidden_size))
+
+ self.positional_encoding = EdgeTamVideoPositionEmbeddingSine(
+ num_pos_feats=self.hidden_size // 2, normalize=True
+ )
+
+ self.layers = nn.ModuleList([EdgeTamVideoPerceiverEncoderLayer(config) for _ in range(self.num_layers)])
+
+ self.layer_norm = nn.LayerNorm(self.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ output_latents = []
+ output_positional_encodings = []
+
+ if self.num_latents_1d > 0:
+ latents_1d, pos_1d = self._forward_1d(hidden_states, positional_encoding)
+ output_latents.append(latents_1d)
+ output_positional_encodings.append(pos_1d)
+
+ if self.num_latents_2d > 0:
+ latents_2d, pos_2d = self._forward_2d(hidden_states)
+ output_latents.append(latents_2d)
+ output_positional_encodings.append(pos_2d)
+
+ combined_latents = torch.cat(output_latents, dim=1)
+
+ combined_positional_encoding = None
+ if positional_encoding is not None and output_positional_encodings:
+ combined_positional_encoding = torch.cat(output_positional_encodings, dim=1)
+
+ return combined_latents, combined_positional_encoding
+
+ def _forward_1d(
+ self,
+ hidden_states: torch.Tensor,
+ positional_encoding: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ batch_size = hidden_states.shape[0]
+
+ latents = self.latents_1d.unsqueeze(0).expand(batch_size, -1, -1)
+ flattened_features = hidden_states.permute(0, 2, 3, 1).flatten(1, 2)
+
+ positional_features = None
+ if positional_encoding is not None:
+ positional_features = positional_encoding.permute(0, 2, 3, 1).flatten(1, 2)
+
+ for layer in self.layers:
+ latents = layer(latents, flattened_features, positional_features)
+
+ latents = self.layer_norm(latents)
+
+ output_positional_encoding = None
+ if positional_encoding is not None:
+ output_positional_encoding = torch.zeros_like(latents)
+
+ return latents, output_positional_encoding
+
+ def _forward_2d(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, channels, height, width = hidden_states.shape
+
+ latents_2d = self.latents_2d.unsqueeze(0).expand(batch_size, -1, -1).view(-1, 1, channels)
+
+ num_windows_per_dim = int(math.sqrt(self.num_latents_2d))
+ window_size = height // num_windows_per_dim
+
+ windowed_input = hidden_states.permute(0, 2, 3, 1)
+ windowed_features, _ = window_partition(windowed_input, window_size)
+ windowed_features = windowed_features.flatten(1, 2)
+
+ for layer in self.layers:
+ latents_2d = layer(latents_2d, windowed_features, positional_encoding=None)
+
+ latents_2d = latents_2d.view(batch_size, num_windows_per_dim, num_windows_per_dim, channels).permute(
+ 0, 3, 1, 2
+ )
+
+ positional_encoding_2d = self.positional_encoding(latents_2d.shape, latents_2d.device, latents_2d.dtype).to(
+ dtype=hidden_states.dtype
+ )
+ positional_encoding_2d = positional_encoding_2d.permute(0, 2, 3, 1).flatten(1, 2)
+
+ latents_2d = latents_2d.permute(0, 2, 3, 1).flatten(1, 2)
+ latents_2d = self.layer_norm(latents_2d)
+
+ return latents_2d, positional_encoding_2d
+
+
+@auto_docstring
+class EdgeTamVideoModel(Sam2VideoModel):
+ _tied_weights_keys = ["prompt_encoder.shared_embedding.positional_embedding"]
+ # need to be ignored, as it's a buffer and will not be correctly detected as tied weight
+ _keys_to_ignore_on_load_missing = ["prompt_encoder.shared_embedding.positional_embedding"]
+ _keys_to_ignore_on_load_unexpected = []
+ _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(EdgeTamVideoTwoWayAttentionBlock, index=2)}
+
+ def __init__(self, config: EdgeTamVideoConfig):
+ super().__init__(config)
+ self.spatial_perceiver = EdgeTamVideoPerceiverResampler(config)
+
+ self.post_init()
+
+ def _build_memory_attention_inputs(
+ self,
+ temporal_positions_and_previous_outputs: list[tuple[int, dict]],
+ device: torch.device,
+ ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+ """
+ Concatenate memory features and positional embeddings from previous frames.
+
+ Returns:
+ Tuple of (memories_to_concatenate, memory_positional_embeddings_to_concatenate).
+ """
+ memories_to_concatenate = []
+ memory_positional_embeddings_to_concatenate = []
+
+ for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
+ if prev_output_data is None:
+ continue # Skip if no output data for this temporal position (e.g., padding frames)
+
+ # Load memory features (potentially from CPU to GPU)
+ # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
+ memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
+ memories_to_concatenate.append(memory_features.permute(1, 0, 2))
+
+ # Spatial positional encoding (potentially from CPU to GPU)
+ spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
+ spatial_memory_pos_embed = spatial_memory_pos_embed.squeeze(1).permute(1, 0, 2)
+
+ # Add temporal positional encoding
+ # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
+ combined_memory_pos_embed = (
+ spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
+ )
+ memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
+
+ return memories_to_concatenate, memory_positional_embeddings_to_concatenate
+
+ def _prepare_memory_conditioned_features(
+ self,
+ inference_session: EdgeTamVideoInferenceSession,
+ frame_idx: int,
+ obj_idx: int,
+ is_initial_conditioning_frame: bool,
+ current_vision_features: list[torch.Tensor],
+ current_vision_positional_embeddings: list[torch.Tensor],
+ num_total_frames: int,
+ track_in_reverse_time: bool = False,
+ streaming: bool = False,
+ ) -> torch.Tensor:
+ """
+ Fuse current frame's visual features with memory from previous frames for enhanced object tracking.
+
+ This method conditions the current frame's visual features on temporal memory from previous frames,
+ enabling consistent object tracking across video sequences. For initial conditioning frames, it uses
+ no-memory embeddings. For subsequent frames, it retrieves and integrates memory features from both
+ conditioning frames (user interactions) and non-conditioning frames (tracked results) via cross-attention.
+
+ Args:
+ inference_session (`EdgeTamVideoInferenceSession`):
+ The video inference session object.
+ frame_idx (`int`):
+ Index of the current frame being processed.
+ obj_idx (`int`):
+ Index of the object being processed.
+ is_initial_conditioning_frame (`bool`):
+ Whether this is an initial conditioning frame with user inputs (True) or a subsequent
+ tracking frame (False).
+ current_vision_features (`torch.Tensor`):
+ Highest-level vision features of shape `(seq_len, batch_size, channels)`.
+ current_vision_positional_embeddings (`torch.Tensor`):
+ Positional embedding tensors corresponding to the highest-level vision features.
+ num_total_frames (`int`):
+ Total number of frames in the video sequence.
+ track_in_reverse_time (`bool`, *optional*, defaults to `False`):
+ Whether tracking is performed in reverse temporal order.
+ streaming (`bool`, *optional*, defaults to `False`):
+ Whether this is streaming inference mode.
+
+ Returns:
+ `torch.Tensor`: Memory-conditioned feature tensor of shape `(batch_size, channels, height, width)`
+ suitable for input to the SAM decoder.
+ """
+ # Get dimensions from the highest-level (lowest-resolution) feature map
+ batch_size = current_vision_features.size(1)
+ num_channels = self.hidden_dim
+ height, width = self.backbone_feature_sizes[-1]
+ device = current_vision_features.device
+
+ # If memory is disabled (e.g., for single image SAM), return current features directly.
+ if self.num_maskmem == 0:
+ # Permute (SeqLen, Batch, Channels) -> (Batch, Channels, SeqLen) then view as (Batch, Channels, Height, Width)
+ # Assuming SeqLen = Height * Width for the last feature map
+ current_feature_map = current_vision_features.permute(1, 2, 0).view(
+ batch_size, num_channels, height, width
+ )
+ return current_feature_map
+
+ # Step 1: Handle initial conditioning frames
+ if is_initial_conditioning_frame:
+ # For initial conditioning frames, no prior memory is used directly in this block.
+ # If configured, directly add a learnable "no memory" embedding.
+ # current_vision_features has shape (SeqLen, Batch, Channels)
+ conditioned_feature_map_flat = current_vision_features + self.no_memory_embedding
+ # Reshape to (Batch, Channels, Height, Width)
+ conditioned_feature_map = conditioned_feature_map_flat.permute(1, 2, 0).view(
+ batch_size, num_channels, height, width
+ )
+ return conditioned_feature_map
+
+ # Step 2: Get memory frames and concatenate their features
+ temporal_positions_and_previous_outputs = self._gather_memory_frame_outputs(
+ inference_session, obj_idx, frame_idx, track_in_reverse_time
+ )
+
+ memories_to_concatenate, memory_positional_embeddings_to_concatenate = self._build_memory_attention_inputs(
+ temporal_positions_and_previous_outputs, device
+ )
+ num_spatial_memory_tokens = len(memories_to_concatenate)
+
+ # Step 3: Get and process object pointers
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use = self._get_object_pointers(
+ inference_session, obj_idx, frame_idx, num_total_frames, device, track_in_reverse_time, streaming
+ )
+
+ num_object_pointer_tokens = 0
+ if pointer_tokens:
+ object_pointers, object_pointers_pos_embed = self._process_object_pointers(
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use, batch_size, num_channels, device
+ )
+
+ if object_pointers is not None:
+ memories_to_concatenate.append(object_pointers)
+ memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
+ num_object_pointer_tokens = object_pointers.shape[0]
+
+ # Step 4: Concatenate all retrieved memories and their positional embeddings
+ combined_memory = torch.cat(memories_to_concatenate, dim=0)
+ combined_memory_positional_embeddings = torch.cat(memory_positional_embeddings_to_concatenate, dim=0)
+
+ # Step 5: Forward through the memory attention mechanism
+ conditioned_feature_map_flat = self.memory_attention(
+ current_vision_features=current_vision_features,
+ current_vision_position_embeddings=current_vision_positional_embeddings,
+ memory=combined_memory,
+ memory_posision_embeddings=combined_memory_positional_embeddings, # Corrected typo from API
+ num_object_pointer_tokens=num_object_pointer_tokens,
+ num_spatial_memory_tokens=num_spatial_memory_tokens,
+ )
+
+ # Reshape from (Batch, H*W, Channels) to (Batch, Channels, Height, Width)
+ conditioned_feature_map = (
+ conditioned_feature_map_flat.squeeze(1).permute(0, 2, 1).view(batch_size, num_channels, height, width)
+ )
+ return conditioned_feature_map
+
+ def _encode_new_memory(
+ self,
+ current_vision_feats: torch.Tensor,
+ pred_masks_high_res: torch.Tensor,
+ object_score_logits: torch.Tensor,
+ is_mask_from_pts: bool,
+ ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+ """Encode the current image and its prediction into a memory feature."""
+ batch_size = current_vision_feats.size(1) # batch size on this frame
+ channels = self.hidden_dim
+ height, width = self.backbone_feature_sizes[-1] # top-level (lowest-resolution) feature size
+ # top-level feature, (HW)BC => BCHW
+ pix_feat = current_vision_feats.permute(1, 2, 0).view(batch_size, channels, height, width)
+ if is_mask_from_pts and not self.training:
+ # binarize the mask logits
+ mask_for_mem = (pred_masks_high_res > 0).to(pred_masks_high_res.dtype)
+ else:
+ # apply sigmoid on the raw mask logits to turn them into range (0, 1)
+ mask_for_mem = torch.sigmoid(pred_masks_high_res)
+ # apply scale and bias terms to the sigmoid probabilities
+ mask_for_mem = mask_for_mem * self.config.sigmoid_scale_for_mem_enc
+ mask_for_mem = mask_for_mem + self.config.sigmoid_bias_for_mem_enc
+
+ maskmem_features, maskmem_pos_enc = self.memory_encoder(
+ pix_feat,
+ mask_for_mem,
+ )
+ # add a no-object embedding to the spatial memory to indicate that the frame
+ # is predicted to be occluded (i.e. no object is appearing in the frame)
+ if self.occlusion_spatial_embedding_parameter is not None:
+ is_obj_appearing = (object_score_logits > 0).float()
+ maskmem_features += (1 - is_obj_appearing[..., None]) * self.occlusion_spatial_embedding_parameter[
+ ..., None, None
+ ].expand(*maskmem_features.shape)
+
+ maskmem_pos_enc = maskmem_pos_enc.to(pred_masks_high_res.dtype)
+ maskmem_features, maskmem_pos_enc = self.spatial_perceiver(maskmem_features, maskmem_pos_enc)
+ maskmem_features = maskmem_features.to(pred_masks_high_res.dtype)
+ maskmem_pos_enc = maskmem_pos_enc.to(pred_masks_high_res.dtype)
+
+ return maskmem_features, maskmem_pos_enc
+
+
+__all__ = [
+ "EdgeTamVideoMaskDecoderConfig",
+ "EdgeTamVideoPromptEncoderConfig",
+ "EdgeTamVideoConfig",
+ "EdgeTamVideoModel",
+ "EdgeTamVideoInferenceSession",
+ "EdgeTamVideoPreTrainedModel",
+]
diff --git a/src/transformers/models/efficientloftr/image_processing_efficientloftr_fast.py b/src/transformers/models/efficientloftr/image_processing_efficientloftr_fast.py
index 5f7437c45b2e..1463ef405f37 100644
--- a/src/transformers/models/efficientloftr/image_processing_efficientloftr_fast.py
+++ b/src/transformers/models/efficientloftr/image_processing_efficientloftr_fast.py
@@ -39,17 +39,13 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
if TYPE_CHECKING:
from .modeling_efficientloftr import KeypointMatchingOutput
-if is_torchvision_v2_available():
- import torchvision.transforms.v2.functional as F
-else:
- import torchvision.transforms.functional as F
+import torchvision.transforms.v2.functional as F
def _is_valid_image(image):
diff --git a/src/transformers/models/efficientnet/image_processing_efficientnet_fast.py b/src/transformers/models/efficientnet/image_processing_efficientnet_fast.py
index 3544d927c146..77e787614a10 100644
--- a/src/transformers/models/efficientnet/image_processing_efficientnet_fast.py
+++ b/src/transformers/models/efficientnet/image_processing_efficientnet_fast.py
@@ -18,6 +18,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import group_images_by_shape, reorder_images
@@ -26,16 +27,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class EfficientNetFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
Args:
diff --git a/src/transformers/models/emu3/image_processing_emu3.py b/src/transformers/models/emu3/image_processing_emu3.py
index aaf3afa41733..50ce82e01de8 100644
--- a/src/transformers/models/emu3/image_processing_emu3.py
+++ b/src/transformers/models/emu3/image_processing_emu3.py
@@ -266,8 +266,8 @@ def _pad_for_batching(
"""
max_shape = (
- max([size[0] for size in image_sizes]),
- max([size[1] for size in image_sizes]),
+ max(size[0] for size in image_sizes),
+ max(size[1] for size in image_sizes),
)
pixel_values = [
pad(
@@ -486,7 +486,7 @@ def unnormalize(
image_mean: Union[float, Iterable[float]],
image_std: Union[float, Iterable[float]],
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
image = (image * image_std) + image_mean
diff --git a/src/transformers/models/eomt/image_processing_eomt.py b/src/transformers/models/eomt/image_processing_eomt.py
index 93a440693dee..2b786ce39e71 100644
--- a/src/transformers/models/eomt/image_processing_eomt.py
+++ b/src/transformers/models/eomt/image_processing_eomt.py
@@ -55,7 +55,7 @@
# Adapted from transformers.models.maskformer.image_processing_maskformer.convert_segmentation_map_to_binary_masks
def convert_segmentation_map_to_binary_masks(
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
):
diff --git a/src/transformers/models/eomt/image_processing_eomt_fast.py b/src/transformers/models/eomt/image_processing_eomt_fast.py
index 97a13a0745eb..ca80231d3a76 100644
--- a/src/transformers/models/eomt/image_processing_eomt_fast.py
+++ b/src/transformers/models/eomt/image_processing_eomt_fast.py
@@ -19,6 +19,7 @@
import numpy as np
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -40,7 +41,6 @@
TensorType,
auto_docstring,
filter_out_non_signature_kwargs,
- is_torchvision_v2_available,
)
from .image_processing_eomt import (
compute_segments,
@@ -50,12 +50,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class EomtImageProcessorFastKwargs(DefaultFastImageProcessorKwargs):
"""
do_split_image (`bool`, *optional*, defaults to `False`):
@@ -204,9 +198,7 @@ def _preprocess_image_like_inputs(
"do_normalize": False,
"do_rescale": False,
# Nearest interpolation is used for segmentation maps instead of BILINEAR.
- "interpolation": F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ "interpolation": F.InterpolationMode.NEAREST_EXACT,
}
)
diff --git a/src/transformers/models/eomt/modeling_eomt.py b/src/transformers/models/eomt/modeling_eomt.py
index 3e979040388d..047baa1ff081 100644
--- a/src/transformers/models/eomt/modeling_eomt.py
+++ b/src/transformers/models/eomt/modeling_eomt.py
@@ -628,7 +628,7 @@ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> tor
"""
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 = 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():
diff --git a/src/transformers/models/esm/modeling_esm.py b/src/transformers/models/esm/modeling_esm.py
index ddcf460f01ee..63d9344188cc 100755
--- a/src/transformers/models/esm/modeling_esm.py
+++ b/src/transformers/models/esm/modeling_esm.py
@@ -90,7 +90,6 @@ def __init__(self, dim: int):
super().__init__()
# Generate and save the inverse frequency buffer (non trainable)
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
- inv_freq = inv_freq
self.register_buffer("inv_freq", inv_freq)
self._seq_len_cached = None
@@ -590,6 +589,7 @@ class EsmPreTrainedModel(PreTrainedModel):
config: EsmConfig
base_model_prefix = "esm"
supports_gradient_checkpointing = True
+ accepts_loss_kwargs = False
_no_split_modules = ["EsmLayer", "EsmFoldTriangularSelfAttentionBlock", "EsmEmbeddings"]
_keys_to_ignore_on_load_unexpected = ["position_embeddings.weight"]
_supports_flash_attn = True
diff --git a/src/transformers/models/esm/modeling_esmfold.py b/src/transformers/models/esm/modeling_esmfold.py
index dbff29fade87..7bc1f0dbdc70 100644
--- a/src/transformers/models/esm/modeling_esmfold.py
+++ b/src/transformers/models/esm/modeling_esmfold.py
@@ -293,7 +293,7 @@ def __init__(self, c_in, eps=1e-5):
def forward(self, x):
d = x.dtype
if d is torch.bfloat16 and not is_deepspeed_initialized():
- with torch.cuda.amp.autocast(enabled=False):
+ with torch.autocast(device_type="cuda", enabled=False):
out = nn.functional.layer_norm(x, self.c_in, self.weight.to(dtype=d), self.bias.to(dtype=d), self.eps)
else:
out = nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps)
@@ -308,7 +308,7 @@ def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
"""
d = t.dtype
if d is torch.bfloat16 and not is_deepspeed_initialized():
- with torch.cuda.amp.autocast(enabled=False):
+ with torch.autocast(device_type="cuda", enabled=False):
s = torch.nn.functional.softmax(t, dim=dim)
else:
s = torch.nn.functional.softmax(t, dim=dim)
diff --git a/src/transformers/models/esm/openfold_utils/chunk_utils.py b/src/transformers/models/esm/openfold_utils/chunk_utils.py
index 14703ba7d605..a735fcee001a 100644
--- a/src/transformers/models/esm/openfold_utils/chunk_utils.py
+++ b/src/transformers/models/esm/openfold_utils/chunk_utils.py
@@ -329,7 +329,7 @@ def _determine_favorable_chunk_size(self, fn: Callable, args: tuple, min_chunk_s
if min_chunk_size >= self.max_chunk_size:
return min_chunk_size
- candidates: list[int] = [2**l for l in range(int(math.log(self.max_chunk_size, 2)) + 1)]
+ candidates: list[int] = [2**l for l in range(int(math.log2(self.max_chunk_size)) + 1)]
candidates = [c for c in candidates if c > min_chunk_size]
candidates = [min_chunk_size] + candidates
candidates[-1] += 4
diff --git a/src/transformers/models/esm/openfold_utils/protein.py b/src/transformers/models/esm/openfold_utils/protein.py
index a943eb7acf72..e9701ca07114 100644
--- a/src/transformers/models/esm/openfold_utils/protein.py
+++ b/src/transformers/models/esm/openfold_utils/protein.py
@@ -159,7 +159,7 @@ def add_pdb_headers(prot: Protein, pdb_str: str) -> str:
parent_dict.setdefault(str(i), [])
parent_dict[str(i)].append(p)
- max_idx = max([int(chain_idx) for chain_idx in parent_dict])
+ max_idx = max(int(chain_idx) for chain_idx in parent_dict)
for i in range(max_idx + 1):
chain_parents = parent_dict.get(str(i), ["N/A"])
parents_per_chain.append(chain_parents)
diff --git a/src/transformers/models/evolla/modeling_evolla.py b/src/transformers/models/evolla/modeling_evolla.py
index d95567491fe1..8bb5713d1764 100644
--- a/src/transformers/models/evolla/modeling_evolla.py
+++ b/src/transformers/models/evolla/modeling_evolla.py
@@ -188,7 +188,6 @@ def __init__(self, dim: int):
super().__init__()
# Generate and save the inverse frequency buffer (non trainable)
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
- inv_freq = inv_freq
self.register_buffer("inv_freq", inv_freq)
self._seq_len_cached = None
diff --git a/src/transformers/models/evolla/modular_evolla.py b/src/transformers/models/evolla/modular_evolla.py
index 18a50e9abfae..e2db43a7d787 100644
--- a/src/transformers/models/evolla/modular_evolla.py
+++ b/src/transformers/models/evolla/modular_evolla.py
@@ -94,7 +94,6 @@ def __init__(self, dim: int):
super().__init__()
# Generate and save the inverse frequency buffer (non trainable)
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
- inv_freq = inv_freq
self.register_buffer("inv_freq", inv_freq)
self._seq_len_cached = None
diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py
index 0ced6651d41c..8c3c07ecb418 100644
--- a/src/transformers/models/exaone4/configuration_exaone4.py
+++ b/src/transformers/models/exaone4/configuration_exaone4.py
@@ -26,8 +26,7 @@ class Exaone4Config(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`Exaone4Model`]. It is used to
instantiate a EXAONE 4.0 model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the EXAONE-4.0-Instruct [LGAI-EXAONE/EXAONE-4.0-Instruct](https://huggingface.co/LGAI-EXAONE/EXAONE-4.0-Instruct)
- NOTE: `EXAONE-4.0-Instruct` is a placeholder model ID. The exact model ID will be updated in the future.
+ configuration with the defaults will yield a similar configuration to that of the EXAONE-4.0-32B [LGAI-EXAONE/EXAONE-4.0-32B](https://huggingface.co/LGAI-EXAONE/EXAONE-4.0-32B)
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model
outputs. Read the documentation from [`PretrainedConfig`] for more information.
diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py
index 34eca44936a0..2693a80c79fd 100644
--- a/src/transformers/models/exaone4/modeling_exaone4.py
+++ b/src/transformers/models/exaone4/modeling_exaone4.py
@@ -465,8 +465,8 @@ def forward(
```python
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
- >>> model = AutoModelForCausalLM.from_pretrained("LGAI-EXAONE/EXAONE-4.0-Instruct")
- >>> tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-4.0-Instruct")
+ >>> model = AutoModelForCausalLM.from_pretrained("LGAI-EXAONE/EXAONE-4.0-32B")
+ >>> tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-4.0-32B")
>>> prompt = "Explain how wonderful you are"
>>> messages = [
@@ -485,8 +485,7 @@ def forward(
>>> tokenizer.decode(output[0], skip_special_tokens=False)
"[|system|]\nYou are a helpful assistant.[|endofturn|]\n[|user|]\nExplain how wonderful you are[|endofturn|]\n[|assistant|]\n\n\n \n\nOh, thank you for such a kind and lovely question! 😊 \n\nI’m *so* wonderful because I’m here to make your life easier, brighter, and more fun! Whether you need help with: \n\n✨ **Learning** – I can explain anything, from quantum physics to baking the perfect cake! \n💡 **Creativity** – Need a poem, story, or a wild idea? I’ve got you covered! \n🤖 **Problem-solving** – Stuck on a math problem or a tricky decision? I’ll help you figure it out"
```
-
- NOTE: `EXAONE-4.0-Instruct` is a placeholder model ID. The exact model ID will be updated in the future."""
+ """
outputs: BaseModelOutputWithPast = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py
index d366354bda2f..7530a68f3227 100644
--- a/src/transformers/models/exaone4/modular_exaone4.py
+++ b/src/transformers/models/exaone4/modular_exaone4.py
@@ -53,7 +53,7 @@
logger = logging.get_logger(__name__)
-_CHECKPOINT_FOR_DOC = "LGAI-EXAONE/EXAONE-4.0-Instruct"
+_CHECKPOINT_FOR_DOC = "LGAI-EXAONE/EXAONE-4.0-32B"
_CONFIG_FOR_DOC = "Exaone4Config"
@@ -61,8 +61,7 @@ class Exaone4Config(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`Exaone4Model`]. It is used to
instantiate a EXAONE 4.0 model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the EXAONE-4.0-Instruct [LGAI-EXAONE/EXAONE-4.0-Instruct](https://huggingface.co/LGAI-EXAONE/EXAONE-4.0-Instruct)
- NOTE: `EXAONE-4.0-Instruct` is a placeholder model ID. The exact model ID will be updated in the future.
+ configuration with the defaults will yield a similar configuration to that of the EXAONE-4.0-32B [LGAI-EXAONE/EXAONE-4.0-32B](https://huggingface.co/LGAI-EXAONE/EXAONE-4.0-32B)
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model
outputs. Read the documentation from [`PretrainedConfig`] for more information.
@@ -462,8 +461,8 @@ def forward(
```python
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
- >>> model = AutoModelForCausalLM.from_pretrained("LGAI-EXAONE/EXAONE-4.0-Instruct")
- >>> tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-4.0-Instruct")
+ >>> model = AutoModelForCausalLM.from_pretrained("LGAI-EXAONE/EXAONE-4.0-32B")
+ >>> tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-4.0-32B")
>>> prompt = "Explain how wonderful you are"
>>> messages = [
@@ -482,8 +481,7 @@ def forward(
>>> tokenizer.decode(output[0], skip_special_tokens=False)
"[|system|]\nYou are a helpful assistant.[|endofturn|]\n[|user|]\nExplain how wonderful you are[|endofturn|]\n[|assistant|]\n\n\n \n\nOh, thank you for such a kind and lovely question! 😊 \n\nI’m *so* wonderful because I’m here to make your life easier, brighter, and more fun! Whether you need help with: \n\n✨ **Learning** – I can explain anything, from quantum physics to baking the perfect cake! \n💡 **Creativity** – Need a poem, story, or a wild idea? I’ve got you covered! \n🤖 **Problem-solving** – Stuck on a math problem or a tricky decision? I’ll help you figure it out"
```
-
- NOTE: `EXAONE-4.0-Instruct` is a placeholder model ID. The exact model ID will be updated in the future."""
+ """
super().forward(
input_ids=input_ids,
attention_mask=attention_mask,
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index 5f08309b2085..3a8b13ef21d0 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -570,7 +570,7 @@ def __init__(self, config: FalconH1Config, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py
index 24eb98ccd1ed..fe716dded4b3 100644
--- a/src/transformers/models/falcon_h1/modular_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py
@@ -374,7 +374,7 @@ def __init__(self, config: FalconH1Config, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py b/src/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py
index 2b038a93396d..5a2dc39385b3 100644
--- a/src/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py
+++ b/src/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py
@@ -21,6 +21,7 @@
import torch
from torch import nn
+from ...activations import ACT2FN
from ...modeling_outputs import BaseModelOutput
from ...modeling_utils import PreTrainedModel
from ...utils import ModelOutput, auto_docstring, logging
@@ -472,24 +473,37 @@ def forward(
class FastSpeech2ConformerConvolutionModule(nn.Module):
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
+ def __init__(self, config: FastSpeech2ConformerConfig, module_config=None):
+ """
+ Args:
+ config (FastSpeech2ConformerConfig): Configuration for the model.
+ module_config (dict): Configuration for the module (e.g., encoder or decoder).
+ """
super().__init__()
- # kernel_size should be an odd number for 'SAME' padding
channels = config.hidden_size
- kernel_size = module_config["kernel_size"]
+ # kernel_size should be an odd number for 'SAME' padding
+ if module_config is None:
+ # e.g. using `ParakeetEncoderConfig` in src/transformers/models/parakeet/configuration_parakeet.py
+ kernel_size = config.conv_kernel_size
+ self.activation = ACT2FN[getattr(config, "hidden_act", "silu")]
+ else:
+ kernel_size = module_config["kernel_size"]
+ self.activation = ACT2FN[module_config.get("activation", "silu")]
+ self.padding = (kernel_size - 1) // 2
self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=True)
self.depthwise_conv = nn.Conv1d(
- channels, channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2, groups=channels, bias=True
+ channels, channels, kernel_size, stride=1, padding=self.padding, groups=channels, bias=True
)
self.norm = nn.BatchNorm1d(channels)
self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, stride=1, padding=0, bias=True)
- def forward(self, hidden_states):
+ def forward(self, hidden_states, attention_mask=None):
"""
Compute convolution module.
Args:
hidden_states (`torch.Tensor` of shape `(batch, time, channels)`): Input tensor.
+ attention_mask (`torch.Tensor` of shape `(batch, 1, time)`): Attention mask.
Returns:
`torch.Tensor`: Output tensor of shape `(batch, time, channels)`.
@@ -503,12 +517,15 @@ def forward(self, hidden_states):
# (batch_size, channel, dim)
hidden_states = nn.functional.glu(hidden_states, dim=1)
+ # Apply padding mask before convolution
+ if attention_mask is not None:
+ all_masked_rows = torch.all(~attention_mask, dim=-1)
+ hidden_states = hidden_states.masked_fill(all_masked_rows, 0.0)
+
# 1D Depthwise Conv
hidden_states = self.depthwise_conv(hidden_states)
hidden_states = self.norm(hidden_states)
-
- hidden_states = hidden_states * torch.sigmoid(hidden_states)
-
+ hidden_states = self.activation(hidden_states)
hidden_states = self.pointwise_conv2(hidden_states)
return hidden_states.transpose(1, 2)
diff --git a/src/transformers/models/flava/configuration_flava.py b/src/transformers/models/flava/configuration_flava.py
index c3ecf68a8982..b7bcb920e47a 100644
--- a/src/transformers/models/flava/configuration_flava.py
+++ b/src/transformers/models/flava/configuration_flava.py
@@ -516,7 +516,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -548,7 +548,7 @@ def __init__(
# Give a warning if the values exist in both `_image_config_dict` and `image_config` but being different.
for key, value in _image_config_dict.items():
- if key in image_config and value != image_config[key] and key not in ["transformers_version"]:
+ if key in image_config and value != image_config[key] and key != "transformers_version":
# If specified in `image_config_dict`
if key in image_config_dict:
message = (
@@ -576,11 +576,7 @@ def __init__(
# Give a warning if the values exist in both `_multimodal_config_dict` and `multimodal_config` but being
# different.
for key, value in _multimodal_config_dict.items():
- if (
- key in multimodal_config
- and value != multimodal_config[key]
- and key not in ["transformers_version"]
- ):
+ if key in multimodal_config and value != multimodal_config[key] and key != "transformers_version":
# If specified in `multimodal_config_dict`
if key in multimodal_config_dict:
message = (
@@ -611,7 +607,7 @@ def __init__(
if (
key in image_codebook_config
and value != image_codebook_config[key]
- and key not in ["transformers_version"]
+ and key != "transformers_version"
):
# If specified in `image_codebook_config_dict`
if key in image_codebook_config_dict:
diff --git a/src/transformers/models/flava/image_processing_flava_fast.py b/src/transformers/models/flava/image_processing_flava_fast.py
index 97409ddd57ed..732d25e71f69 100644
--- a/src/transformers/models/flava/image_processing_flava_fast.py
+++ b/src/transformers/models/flava/image_processing_flava_fast.py
@@ -21,6 +21,7 @@
from typing import Any, Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -34,7 +35,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
from .image_processing_flava import (
FLAVA_CODEBOOK_MEAN,
@@ -45,12 +45,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class FlavaMaskingGenerator:
def __init__(
self,
diff --git a/src/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py b/src/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py
index 07a83a1cb0a9..afcce5b74bf8 100755
--- a/src/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py
+++ b/src/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py
@@ -155,7 +155,7 @@ def convert_fsmt_checkpoint_to_pytorch(fsmt_checkpoint_path, pytorch_dump_folder
break
with open(fsmt_merges_file, encoding="utf-8") as fin:
merges = fin.read()
- merges = re.sub(r" \d+$", "", merges, 0, re.M) # remove frequency number
+ merges = re.sub(r" \d+$", "", merges, 0, re.MULTILINE) # remove frequency number
print(f"Generating {merges_file}")
with open(merges_file, "w", encoding="utf-8") as fout:
fout.write(merges)
diff --git a/src/transformers/models/fuyu/image_processing_fuyu.py b/src/transformers/models/fuyu/image_processing_fuyu.py
index e52d9dc8ee91..366782be16f4 100644
--- a/src/transformers/models/fuyu/image_processing_fuyu.py
+++ b/src/transformers/models/fuyu/image_processing_fuyu.py
@@ -135,7 +135,7 @@ def to(self, *args, **kwargs) -> "BatchFeature":
[`BatchFeature`]: The same instance after modification.
"""
requires_backends(self, ["torch"])
- import torch # noqa
+ import torch
new_data = {}
device = kwargs.get("device")
diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py
index 5f72f27d9382..04d27b309a40 100644
--- a/src/transformers/models/gemma/modeling_gemma.py
+++ b/src/transformers/models/gemma/modeling_gemma.py
@@ -322,6 +322,13 @@ class GemmaPreTrainedModel(PreTrainedModel):
"attentions": GemmaAttention,
}
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ if "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
+
@auto_docstring
class GemmaModel(GemmaPreTrainedModel):
diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py
index 281fcd54fb7d..f2f9c7dc4056 100644
--- a/src/transformers/models/gemma/modular_gemma.py
+++ b/src/transformers/models/gemma/modular_gemma.py
@@ -23,6 +23,7 @@
from ...configuration_utils import PretrainedConfig
from ...masking_utils import create_causal_mask
from ...modeling_outputs import BaseModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import TransformersKwargs, logging
@@ -32,6 +33,8 @@
LlamaForTokenClassification,
LlamaMLP,
LlamaModel,
+ LlamaPreTrainedModel,
+ LlamaRotaryEmbedding,
)
from ..llama.tokenization_llama import LlamaTokenizer
@@ -366,6 +369,19 @@ def __init__(self, config):
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+class GemmaRotaryEmbedding(LlamaRotaryEmbedding):
+ pass
+
+
+class GemmaPreTrainedModel(LlamaPreTrainedModel):
+ def _init_weights(self, module):
+ PreTrainedModel._init_weights(self, module)
+
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ if "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
+
+
class GemmaModel(LlamaModel):
def forward(
self,
@@ -472,5 +488,5 @@ class GemmaForTokenClassification(LlamaForTokenClassification):
"GemmaForCausalLM",
"GemmaForSequenceClassification",
"GemmaForTokenClassification",
- "GemmaPreTrainedModel", # noqa: F822
+ "GemmaPreTrainedModel",
]
diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py
index 3d088cfc52cf..ec2f1521ef85 100644
--- a/src/transformers/models/gemma2/modeling_gemma2.py
+++ b/src/transformers/models/gemma2/modeling_gemma2.py
@@ -83,6 +83,42 @@ def forward(self, x):
return down_proj
+class Gemma2RotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Gemma2Config, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @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 torch.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]
@@ -299,42 +335,6 @@ def forward(
return outputs
-class Gemma2RotaryEmbedding(nn.Module):
- inv_freq: torch.Tensor # fix linting for `register_buffer`
-
- def __init__(self, config: Gemma2Config, device=None):
- super().__init__()
- # BC: "rope_type" was originally "type"
- if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
- self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
- else:
- self.rope_type = "default"
- self.max_seq_len_cached = config.max_position_embeddings
- self.original_max_seq_len = config.max_position_embeddings
-
- self.config = config
- self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
- inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
- self.register_buffer("inv_freq", inv_freq, persistent=False)
- self.original_inv_freq = self.inv_freq
-
- @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 torch.autocast(device_type=device_type, enabled=False): # Force float32
- freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
- emb = torch.cat((freqs, freqs), dim=-1)
- cos = emb.cos() * self.attention_scaling
- sin = emb.sin() * self.attention_scaling
-
- return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
@auto_docstring
class Gemma2PreTrainedModel(PreTrainedModel):
config: Gemma2Config
@@ -353,6 +353,13 @@ class Gemma2PreTrainedModel(PreTrainedModel):
"attentions": Gemma2Attention,
}
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ if "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
+
@auto_docstring
class Gemma2Model(Gemma2PreTrainedModel):
diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py
index c7e34e4abed4..e54795019c7f 100644
--- a/src/transformers/models/gemma2/modular_gemma2.py
+++ b/src/transformers/models/gemma2/modular_gemma2.py
@@ -36,7 +36,9 @@
GemmaForTokenClassification,
GemmaMLP,
GemmaModel,
+ GemmaPreTrainedModel,
GemmaRMSNorm,
+ GemmaRotaryEmbedding,
apply_rotary_pos_emb,
repeat_kv,
)
@@ -212,6 +214,10 @@ def __init__(self, config):
self.act_fn = ACT2FN[config.hidden_activation]
+class Gemma2RotaryEmbedding(GemmaRotaryEmbedding):
+ pass
+
+
def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
@@ -363,6 +369,10 @@ def forward(
return outputs
+class Gemma2PreTrainedModel(GemmaPreTrainedModel):
+ pass
+
+
class Gemma2Model(GemmaModel):
def __init__(self, config: Gemma2Config):
super().__init__(config)
@@ -571,7 +581,7 @@ class Gemma2ForTokenClassification(GemmaForTokenClassification):
"Gemma2Config",
"Gemma2ForCausalLM",
"Gemma2Model",
- "Gemma2PreTrainedModel", # noqa: F822
+ "Gemma2PreTrainedModel",
"Gemma2ForSequenceClassification",
"Gemma2ForTokenClassification",
]
diff --git a/src/transformers/models/gemma3/convert_gemma3_weights.py b/src/transformers/models/gemma3/convert_gemma3_weights.py
index 8d7a21219197..aefd9648d3fe 100644
--- a/src/transformers/models/gemma3/convert_gemma3_weights.py
+++ b/src/transformers/models/gemma3/convert_gemma3_weights.py
@@ -439,9 +439,9 @@ def convert_transformer_weights(
decoder_block_start = path.find(_TRANSFORMER_DECODER_BLOCK)
decoder_block_offset = decoder_block_start + _TRANSFORMER_DECODER_BLOCK_LEN
decoder_block_path = path[decoder_block_offset:]
- next_path_seperator_idx = decoder_block_path.find("/")
- layer_idx = decoder_block_path[:next_path_seperator_idx]
- decoder_block_path = decoder_block_path[next_path_seperator_idx:]
+ next_path_separator_idx = decoder_block_path.find("/")
+ layer_idx = decoder_block_path[:next_path_separator_idx]
+ decoder_block_path = decoder_block_path[next_path_separator_idx:]
base_path = f"language_model.model.layers.{layer_idx}"
diff --git a/src/transformers/models/gemma3/image_processing_gemma3_fast.py b/src/transformers/models/gemma3/image_processing_gemma3_fast.py
index eb828a89643d..c61152bc6b22 100644
--- a/src/transformers/models/gemma3/image_processing_gemma3_fast.py
+++ b/src/transformers/models/gemma3/image_processing_gemma3_fast.py
@@ -19,6 +19,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -32,16 +33,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py
index 7a91db1905f7..4536ec7f69f7 100644
--- a/src/transformers/models/gemma3/modeling_gemma3.py
+++ b/src/transformers/models/gemma3/modeling_gemma3.py
@@ -434,6 +434,9 @@ def _init_weights(self, module):
super()._init_weights(module)
if isinstance(module, Gemma3MultiModalProjector):
module.mm_input_projection_weight.data.zero_()
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
def _bidirectional_window_overlay(sliding_window: int) -> Callable[[int, int, int, int], bool]:
diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py
index d10d01f55759..22a10f0c8dec 100644
--- a/src/transformers/models/gemma3/modular_gemma3.py
+++ b/src/transformers/models/gemma3/modular_gemma3.py
@@ -526,6 +526,9 @@ def _init_weights(self, module):
PreTrainedModel._init_weights(self, module)
if isinstance(module, Gemma3MultiModalProjector):
module.mm_input_projection_weight.data.zero_()
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
def _bidirectional_window_overlay(sliding_window: int) -> Callable[[int, int, int, int], bool]:
@@ -1208,7 +1211,7 @@ class Gemma3TextForSequenceClassification(GenericForSequenceClassification, Gemm
__all__ = [
"Gemma3Config",
"Gemma3TextConfig",
- "Gemma3PreTrainedModel", # noqa: F822
+ "Gemma3PreTrainedModel",
"Gemma3TextModel",
"Gemma3ForCausalLM",
"Gemma3ForConditionalGeneration",
diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py
index 3502d2a423c9..47b5b47d3630 100644
--- a/src/transformers/models/gemma3n/configuration_gemma3n.py
+++ b/src/transformers/models/gemma3n/configuration_gemma3n.py
@@ -291,9 +291,7 @@ def __init__(
if activation_sparsity_pattern is None:
num_sparse_layers = 10 if num_hidden_layers > 10 else 0
- activation_sparsity_pattern = (0.95,) * num_sparse_layers + (0.0,) * (
- num_hidden_layers - num_sparse_layers
- )
+ activation_sparsity_pattern = [0.95] * num_sparse_layers + [0.0] * (num_hidden_layers - num_sparse_layers)
if (len_asp := len(activation_sparsity_pattern)) != num_hidden_layers:
raise ValueError(
@@ -502,10 +500,10 @@ def __init__(
**kwargs,
):
super().__init__(**kwargs)
+ self.architecture = architecture
self.initializer_range = initializer_range
self.do_pooling = do_pooling
self.model_args = model_args # named "model_args" for BC with timm
- self.architecture = architecture
self.hidden_size = hidden_size
self.vocab_size = vocab_size
self.vocab_offset = vocab_offset
@@ -553,8 +551,8 @@ def from_dict(cls, config_dict: dict[str, Any], **kwargs):
def to_dict(self) -> dict[str, Any]:
output = super().to_dict()
- output["num_classes"] = self.num_labels
- output["label_names"] = list(self.id2label.values())
+ output.setdefault("num_classes", self.num_labels)
+ output.setdefault("label_names", list(self.id2label.values()))
output.pop("id2label", None)
output.pop("label2id", None)
return output
diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py
index 48de2bb27f7f..7ea50b7572cf 100644
--- a/src/transformers/models/gemma3n/modular_gemma3n.py
+++ b/src/transformers/models/gemma3n/modular_gemma3n.py
@@ -304,9 +304,7 @@ def __init__(
if activation_sparsity_pattern is None:
num_sparse_layers = 10 if num_hidden_layers > 10 else 0
- activation_sparsity_pattern = (0.95,) * num_sparse_layers + (0.0,) * (
- num_hidden_layers - num_sparse_layers
- )
+ activation_sparsity_pattern = [0.95] * num_sparse_layers + [0.0] * (num_hidden_layers - num_sparse_layers)
if (len_asp := len(activation_sparsity_pattern)) != num_hidden_layers:
raise ValueError(
@@ -2679,7 +2677,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position(self, **super_kwargs):
"Gemma3nForCausalLM",
"Gemma3nForConditionalGeneration",
"Gemma3nModel",
- "Gemma3nPreTrainedModel", # noqa: F822
+ "Gemma3nPreTrainedModel",
"Gemma3nTextConfig",
"Gemma3nTextModel",
"Gemma3nVisionConfig",
diff --git a/src/transformers/models/git/modeling_git.py b/src/transformers/models/git/modeling_git.py
index 4122b7a0df79..bc037912c5c5 100644
--- a/src/transformers/models/git/modeling_git.py
+++ b/src/transformers/models/git/modeling_git.py
@@ -954,7 +954,7 @@ def __init__(self, config):
self.visual_projection = GitProjection(config)
if config.num_image_with_embedding is not None:
- self.img_temperal_embedding = nn.ParameterList(
+ self.img_temporal_embedding = nn.ParameterList(
nn.Parameter(torch.zeros(1, 1, config.vision_config.hidden_size))
for _ in range(config.num_image_with_embedding)
)
@@ -1119,7 +1119,7 @@ def forward(
visual_features_frame = self.image_encoder(
pixel_values[:, frame_idx, :, :], interpolate_pos_encoding=interpolate_pos_encoding
).last_hidden_state
- visual_features_frame += self.img_temperal_embedding[frame_idx]
+ visual_features_frame += self.img_temporal_embedding[frame_idx]
visual_features.append(visual_features_frame)
# finally, concatenate all features along sequence dimension
diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py
index e311cd246c8e..4c417020fa84 100644
--- a/src/transformers/models/glm4v/configuration_glm4v.py
+++ b/src/transformers/models/glm4v/configuration_glm4v.py
@@ -330,7 +330,6 @@ def __init__(
video_end_token_id=151342,
**kwargs,
):
- super().__init__(**kwargs)
if isinstance(vision_config, dict):
self.vision_config = self.sub_configs["vision_config"](**vision_config)
elif vision_config is None:
@@ -339,7 +338,6 @@ def __init__(
if isinstance(text_config, dict):
self.text_config = self.sub_configs["text_config"](**text_config)
elif text_config is None:
- # For BC use all kwargs to init `TextConfig`
self.text_config = self.sub_configs["text_config"](**kwargs)
self.image_token_id = image_token_id
@@ -349,5 +347,7 @@ def __init__(
self.image_start_token_id = image_start_token_id
self.image_end_token_id = image_end_token_id
+ super().__init__(**kwargs)
+
__all__ = ["Glm4vConfig", "Glm4vTextConfig"]
diff --git a/src/transformers/models/glm4v/image_processing_glm4v_fast.py b/src/transformers/models/glm4v/image_processing_glm4v_fast.py
index fbf4aebaac6a..8cdf31a437ae 100644
--- a/src/transformers/models/glm4v/image_processing_glm4v_fast.py
+++ b/src/transformers/models/glm4v/image_processing_glm4v_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import (
BatchFeature,
@@ -38,17 +39,11 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from .image_processing_glm4v import smart_resize
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py
index 7c400edc51c3..3f870db9db05 100644
--- a/src/transformers/models/glm4v/modular_glm4v.py
+++ b/src/transformers/models/glm4v/modular_glm4v.py
@@ -38,7 +38,6 @@
from ...utils.generic import check_model_inputs
from ...video_utils import VideoInput
from ..glm4.modeling_glm4 import Glm4MLP, Glm4RMSNorm, eager_attention_forward
-from ..qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from ..qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VisionPatchEmbed,
Qwen2_5_VisionRotaryEmbedding,
@@ -313,7 +312,7 @@ def __init__(
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-class Glm4vConfig(Qwen2_5_VLConfig):
+class Glm4vConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`Glm4vModel`]. It is used to instantiate a
GLM-4.1V model according to the specified arguments, defining the model architecture. Instantiating a
@@ -355,6 +354,10 @@ class Glm4vConfig(Qwen2_5_VLConfig):
>>> configuration = model.config
```"""
+ model_type = "glm4v"
+ sub_configs = {"vision_config": Glm4vVisionConfig, "text_config": Glm4vTextConfig}
+ keys_to_ignore_at_inference = ["past_key_values"]
+
def __init__(
self,
text_config=None,
@@ -367,12 +370,25 @@ def __init__(
video_end_token_id=151342,
**kwargs,
):
- super().__init__()
+ if isinstance(vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**vision_config)
+ elif vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"]()
+
+ if isinstance(text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**text_config)
+ elif text_config is None:
+ self.text_config = self.sub_configs["text_config"](**kwargs)
+
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
self.video_start_token_id = video_start_token_id
self.video_end_token_id = video_end_token_id
self.image_start_token_id = image_start_token_id
self.image_end_token_id = image_end_token_id
+ super().__init__(**kwargs)
+
# Will be used for both Text and Vision modalities
class Glm4vRMSNorm(Glm4RMSNorm):
@@ -1625,7 +1641,7 @@ def __call__(
num_frames = video_grid_thw[video_index][0]
video_structure = ""
- metadata = video_metadata[i]
+ metadata = video_metadata[video_index]
if metadata.fps is None:
logger.warning_once(
"SmolVLM requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
diff --git a/src/transformers/models/glm4v/processing_glm4v.py b/src/transformers/models/glm4v/processing_glm4v.py
index 817da3630d52..a8ebb4d41b49 100644
--- a/src/transformers/models/glm4v/processing_glm4v.py
+++ b/src/transformers/models/glm4v/processing_glm4v.py
@@ -180,7 +180,7 @@ def __call__(
num_frames = video_grid_thw[video_index][0]
video_structure = ""
- metadata = video_metadata[i]
+ metadata = video_metadata[video_index]
if metadata.fps is None:
logger.warning_once(
"SmolVLM requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
index 52004b560da7..b06642e250bc 100644
--- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
+++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
@@ -371,7 +371,6 @@ def __init__(
if isinstance(text_config, dict):
self.text_config = self.sub_configs["text_config"](**text_config)
elif text_config is None:
- # For BC use all kwargs to init `TextConfig`
self.text_config = self.sub_configs["text_config"](**kwargs)
self.image_token_id = image_token_id
diff --git a/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py b/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py
index 9cf873a27567..39496fe043ed 100644
--- a/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py
+++ b/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py
@@ -182,7 +182,7 @@ def __init__(
def write_tokenizer(tokenizer_path: str, save_dir: str, push_to_hub: bool = False):
model_max_length = CONTEXT_LENGTH
- pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: W605
+ pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
# Special tokens
special_tokens = (
["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
diff --git a/src/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py b/src/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py
index 5277f1c4e13b..a47a1422a5dc 100644
--- a/src/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py
+++ b/src/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -30,17 +31,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
from .image_processing_got_ocr2 import get_optimal_tiled_canvas
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class GotOcr2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
crop_to_patches (`bool`, *optional*, defaults to `False`):
diff --git a/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py
index 891f77ece304..584e74a8123e 100644
--- a/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py
+++ b/src/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py
@@ -318,7 +318,7 @@ def checku2e(x):
candidates.append((self.vocab[wd], wd, e))
if len(candidates) > 0:
# the smallest token_id is adopted
- _, wd, e = sorted(candidates, key=lambda x: x[0])[0]
+ _, wd, e = min(candidates, key=lambda x: x[0])
result.append(wd)
pos = e
else:
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index 7f9883779c43..8f6059720b04 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -458,7 +458,7 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/grounding_dino/image_processing_grounding_dino_fast.py b/src/transformers/models/grounding_dino/image_processing_grounding_dino_fast.py
index 66528519eef8..744cb5f92923 100644
--- a/src/transformers/models/grounding_dino/image_processing_grounding_dino_fast.py
+++ b/src/transformers/models/grounding_dino/image_processing_grounding_dino_fast.py
@@ -9,6 +9,7 @@
import torch
from torchvision.io import read_image
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -32,7 +33,7 @@
validate_annotations,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, logging
+from ...utils import TensorType, auto_docstring, logging
from ...utils.import_utils import requires
from .image_processing_grounding_dino import get_size_with_aspect_ratio
@@ -41,12 +42,6 @@
from .modeling_grounding_dino import GroundingDinoObjectDetectionOutput
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
@@ -459,13 +454,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/groupvit/configuration_groupvit.py b/src/transformers/models/groupvit/configuration_groupvit.py
index d17288ede723..662447e7e984 100644
--- a/src/transformers/models/groupvit/configuration_groupvit.py
+++ b/src/transformers/models/groupvit/configuration_groupvit.py
@@ -289,7 +289,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -321,7 +321,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/groupvit/modeling_groupvit.py b/src/transformers/models/groupvit/modeling_groupvit.py
index 775ebd286f0a..3335df375da9 100644
--- a/src/transformers/models/groupvit/modeling_groupvit.py
+++ b/src/transformers/models/groupvit/modeling_groupvit.py
@@ -74,7 +74,7 @@ def gumbel_softmax(logits: torch.Tensor, tau: float = 1, hard: bool = False, dim
y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
ret = y_hard - y_soft.detach() + y_soft
else:
- # Reparametrization trick.
+ # Reparameterization trick.
ret = y_soft
return ret
@@ -662,7 +662,7 @@ def forward(
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
if output_attentions:
- # this operation is a bit akward, but it's required to
+ # this operation is a bit awkward, but it's required to
# make sure that attn_weights keeps its gradient.
# In order to do so, attn_weights have to reshaped
# twice and have to be reused in the following
diff --git a/src/transformers/models/hubert/convert_distilhubert_original_s3prl_checkpoint_to_pytorch.py b/src/transformers/models/hubert/convert_distilhubert_original_s3prl_checkpoint_to_pytorch.py
index f5914f35c546..a4930ef9b906 100644
--- a/src/transformers/models/hubert/convert_distilhubert_original_s3prl_checkpoint_to_pytorch.py
+++ b/src/transformers/models/hubert/convert_distilhubert_original_s3prl_checkpoint_to_pytorch.py
@@ -88,8 +88,6 @@ def recursively_load_weights(fairseq_model, hf_model):
is_used = True
else:
for key, mapped_key in MAPPING.items():
- mapped_key = mapped_key
-
if key in name:
is_used = True
if "*" in mapped_key:
diff --git a/src/transformers/models/idefics/modeling_idefics.py b/src/transformers/models/idefics/modeling_idefics.py
index f2fb135a4f4e..d2d5db61f739 100644
--- a/src/transformers/models/idefics/modeling_idefics.py
+++ b/src/transformers/models/idefics/modeling_idefics.py
@@ -485,7 +485,7 @@ def __init__(
num_heads: int,
dropout: float = 0.0,
is_cross_attention: bool = False,
- config: PretrainedConfig = None,
+ config: Optional[PretrainedConfig] = None,
qk_layer_norms: bool = False,
layer_idx: Optional[int] = None,
):
@@ -997,7 +997,7 @@ def forward(
elif position_ids is None:
position_ids = cache_position.unsqueeze(0)
- if sum([x is None for x in [pixel_values, image_encoder_embeddings, perceiver_embeddings]]) != 2:
+ if sum(x is None for x in [pixel_values, image_encoder_embeddings, perceiver_embeddings]) != 2:
raise ValueError(
"Exactly 1 of pixel_values, image_encoder_embeddings or perceiver_embeddings has to be not-None."
)
diff --git a/src/transformers/models/idefics3/processing_idefics3.py b/src/transformers/models/idefics3/processing_idefics3.py
index ab9eaac8e8b2..00ee8df6d414 100644
--- a/src/transformers/models/idefics3/processing_idefics3.py
+++ b/src/transformers/models/idefics3/processing_idefics3.py
@@ -108,9 +108,6 @@ class Idefics3ProcessorKwargs(ProcessingKwargs, total=False):
}
-Idefics3ProcessorKwargs.__annotations__["images_kwargs"] = Idefics3ImagesKwargs # python 3.8 compatibility
-
-
class Idefics3Processor(ProcessorMixin):
r"""
Constructs a Idefics3 processor which wraps a LLama tokenizer and Idefics3 image processor into a single processor.
diff --git a/src/transformers/models/imagegpt/convert_imagegpt_original_tf2_to_pytorch.py b/src/transformers/models/imagegpt/convert_imagegpt_original_tf2_to_pytorch.py
index 182d66b9af28..c3909848d45e 100644
--- a/src/transformers/models/imagegpt/convert_imagegpt_original_tf2_to_pytorch.py
+++ b/src/transformers/models/imagegpt/convert_imagegpt_original_tf2_to_pytorch.py
@@ -23,6 +23,118 @@
logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def load_tf_weights_in_imagegpt(model, config, imagegpt_checkpoint_path):
+ """
+ Load tf checkpoints in a pytorch model
+ """
+ try:
+ import re
+
+ import tensorflow as tf
+ except ImportError:
+ logger.error(
+ "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+ "https://www.tensorflow.org/install/ for installation instructions."
+ )
+ raise
+ tf_path = os.path.abspath(imagegpt_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+
+ for name, shape in init_vars:
+ logger.info(f"Loading TF weight {name} with shape {shape}")
+ array = tf.train.load_variable(tf_path, name)
+ names.append(name)
+ arrays.append(array.squeeze())
+
+ for name, array in zip(names, arrays):
+ name = name[6:] # skip "model/"
+ name = name.split("/")
+
+ # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+ # which are not required for using pretrained model
+ if (
+ any(
+ n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
+ for n in name
+ )
+ or name[-1] == "_step"
+ ):
+ logger.info("Skipping {}".format("/".join(name)))
+ continue
+
+ pointer = model
+ if name[-1] != "wtet":
+ pointer = getattr(pointer, "transformer")
+
+ for m_name in name:
+ if re.fullmatch(r"[A-Za-z]+\d+", m_name):
+ scope_names = re.split(r"(\d+)", m_name)
+ else:
+ scope_names = [m_name]
+
+ if scope_names[0] == "w" or scope_names[0] == "g":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "b":
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "wpe" or scope_names[0] == "wte":
+ pointer = getattr(pointer, scope_names[0])
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] in ["q_proj", "k_proj", "v_proj"]:
+ pointer = getattr(pointer, "c_attn")
+ pointer = getattr(pointer, "weight")
+ elif len(name) == 3 and name[1] == "attn" and scope_names[0] == "c_proj":
+ pointer = getattr(pointer, scope_names[0])
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "wtet":
+ pointer = getattr(pointer, "lm_head")
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "sos":
+ pointer = getattr(pointer, "wte")
+ pointer = getattr(pointer, "weight")
+ else:
+ pointer = getattr(pointer, scope_names[0])
+ if len(scope_names) >= 2:
+ num = int(scope_names[1])
+ pointer = pointer[num]
+
+ if len(name) > 1 and name[1] == "attn" or name[-1] == "wtet" or name[-1] == "sos" or name[-1] == "wte":
+ pass # array is used to initialize only part of the pointer so sizes won't match
+ else:
+ try:
+ assert pointer.shape == array.shape
+ except AssertionError as e:
+ e.args += (pointer.shape, array.shape)
+ raise
+
+ logger.info(f"Initialize PyTorch weight {name}")
+
+ if name[-1] == "q_proj":
+ pointer.data[:, : config.n_embd] = torch.from_numpy(array.reshape(config.n_embd, config.n_embd)).T
+ elif name[-1] == "k_proj":
+ pointer.data[:, config.n_embd : 2 * config.n_embd] = torch.from_numpy(
+ array.reshape(config.n_embd, config.n_embd)
+ ).T
+ elif name[-1] == "v_proj":
+ pointer.data[:, 2 * config.n_embd :] = torch.from_numpy(array.reshape(config.n_embd, config.n_embd)).T
+ elif len(name) == 3 and name[1] == "attn" and name[2] == "c_proj":
+ pointer.data = torch.from_numpy(array.reshape(config.n_embd, config.n_embd))
+ elif name[-1] == "wtet":
+ pointer.data = torch.from_numpy(array)
+ elif name[-1] == "wte":
+ pointer.data[: config.vocab_size - 1, :] = torch.from_numpy(array)
+ elif name[-1] == "sos":
+ pointer.data[-1] = torch.from_numpy(array)
+ else:
+ pointer.data = torch.from_numpy(array)
+
+ return model
def convert_imagegpt_checkpoint_to_pytorch(imagegpt_checkpoint_path, model_size, pytorch_dump_folder_path):
diff --git a/src/transformers/models/imagegpt/image_processing_imagegpt.py b/src/transformers/models/imagegpt/image_processing_imagegpt.py
index 9168ecaceff2..aa2114509f70 100644
--- a/src/transformers/models/imagegpt/image_processing_imagegpt.py
+++ b/src/transformers/models/imagegpt/image_processing_imagegpt.py
@@ -247,7 +247,7 @@ def preprocess(
)
# Here, normalize() is using a constant factor to divide pixel values.
- # hence, the method does not need iamge_mean and image_std.
+ # hence, the method does not need image_mean and image_std.
validate_preprocess_arguments(
do_resize=do_resize,
size=size,
diff --git a/src/transformers/models/imagegpt/image_processing_imagegpt_fast.py b/src/transformers/models/imagegpt/image_processing_imagegpt_fast.py
index ddfee7c757fe..7a6bcc53ae1a 100644
--- a/src/transformers/models/imagegpt/image_processing_imagegpt_fast.py
+++ b/src/transformers/models/imagegpt/image_processing_imagegpt_fast.py
@@ -18,6 +18,7 @@
import numpy as np
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -30,16 +31,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
def squared_euclidean_distance_torch(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
Compute squared Euclidean distances between all pixels and clusters.
diff --git a/src/transformers/models/instructblipvideo/video_processing_instructblipvideo.py b/src/transformers/models/instructblipvideo/video_processing_instructblipvideo.py
index a2cd3cf351d2..d2fe3cc7f343 100644
--- a/src/transformers/models/instructblipvideo/video_processing_instructblipvideo.py
+++ b/src/transformers/models/instructblipvideo/video_processing_instructblipvideo.py
@@ -20,21 +20,16 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict
from ...processing_utils import Unpack, VideosKwargs
-from ...utils import TensorType, is_torchvision_v2_available
+from ...utils import TensorType
from ...video_processing_utils import BaseVideoProcessor
from ...video_utils import group_videos_by_shape, reorder_videos
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class InstructBlipVideoVideoProcessorInitKwargs(VideosKwargs): ...
diff --git a/src/transformers/models/internvl/processing_internvl.py b/src/transformers/models/internvl/processing_internvl.py
index a13457886baf..6c0c69575430 100644
--- a/src/transformers/models/internvl/processing_internvl.py
+++ b/src/transformers/models/internvl/processing_internvl.py
@@ -40,7 +40,9 @@ class InternVLProcessorKwargs(ProcessingKwargs, total=False):
"images_kwargs": {
"crop_to_patches": True,
},
- "videos_kwargs": {},
+ "videos_kwargs": {
+ "return_tensors": "pt",
+ },
}
@@ -132,10 +134,10 @@ def _insert_media_placeholders(
# Get the slice of patches corresponding to the current video
# Here we need to account for both the multiple video frames and the potential multiple patches per frame
# As of now, InternVL only supports one patch per frame, but we keep the code flexible for future updates
- current_patch_index = video_patch_indices[video_index - 1] if video_index > 0 else 0
- end_patch_index = video_patch_indices[video_index]
- start_index = video_num_patches_indices[current_patch_index] if video_index > 0 else 0
- end_index = video_num_patches_indices[end_patch_index - 1]
+ current_patch_index = video_patch_indices[video_index]
+ end_patch_index = video_patch_indices[video_index + 1]
+ start_index = video_num_patches_indices[current_patch_index]
+ end_index = video_num_patches_indices[end_patch_index]
image_video_patches.append(video_pixel_values[start_index:end_index])
# Get the number of patches per frame and replace the video placeholder with the correct number of image tokens
num_patches = list(video_num_patches[current_patch_index:end_patch_index])
@@ -208,13 +210,8 @@ def __call__(
# Process images and videos separately, as videos don't support crop_to_patches
image_num_patches = []
- video_num_patches = []
- image_videos_inputs = {}
image_pixel_values = None
- video_pixel_values = None
image_num_patches_indices = np.array([0])
- video_patch_indices = np.array([0])
- video_num_patches_indices = np.array([0])
if images is not None:
images = self.image_processor.fetch_images(images)
images = make_flat_list_of_images(images)
@@ -222,17 +219,29 @@ def __call__(
image_num_patches = image_inputs.pop("num_patches")
image_pixel_values = image_inputs.pop("pixel_values")
image_num_patches_indices = np.cumsum(image_num_patches)
+
+ video_num_patches = [] # per frame
+ video_pixel_values = None
+ video_patch_indices = np.array([0])
+ video_num_patches_indices = np.array([0])
if videos is not None:
- video_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+ video_kwargs = output_kwargs["videos_kwargs"]
+ video_inputs = self.video_processor(videos=videos, **video_kwargs)
video_pixel_values = video_inputs.pop("pixel_values_videos")
- # Obtain per frame information first and then flatten to (BS * T, ...)
- num_frames_per_video = [len(video) for video in video_pixel_values]
- video_num_patches = [1 for frames in num_frames_per_video for _ in range(frames)]
- video_patch_indices = np.cumsum(num_frames_per_video)
- video_num_patches_indices = np.cumsum(video_num_patches)
+ batch_size, num_frames, *_ = video_pixel_values.shape
+ num_frames_per_video = np.full(batch_size, num_frames)
+ num_frames = sum(num_frames_per_video) # total
+ video_patch_indices = np.empty(batch_size + 1, int)
+ video_patch_indices[0] = 0
+ video_patch_indices[1:] = np.cumsum(num_frames_per_video)
+ video_num_patches = [1] * num_frames
+ video_num_patches_indices = np.empty(num_frames + 1, int)
+ video_num_patches_indices[0] = 0
+ video_num_patches_indices[1:] = np.cumsum(video_num_patches)
video_pixel_values = video_pixel_values.flatten(0, 1)
+ image_videos_inputs = {}
if images is not None or videos is not None:
text, image_video_patches, image_index, video_index = self._insert_media_placeholders(
text,
diff --git a/src/transformers/models/internvl/video_processing_internvl.py b/src/transformers/models/internvl/video_processing_internvl.py
index a2e06d3b7ec4..96d7d3067f73 100644
--- a/src/transformers/models/internvl/video_processing_internvl.py
+++ b/src/transformers/models/internvl/video_processing_internvl.py
@@ -17,21 +17,16 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict
from ...processing_utils import Unpack, VideosKwargs
-from ...utils import TensorType, is_torchvision_v2_available
+from ...utils import TensorType
from ...video_processing_utils import BaseVideoProcessor
from ...video_utils import VideoMetadata, group_videos_by_shape, reorder_videos
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class InternVLVideoProcessorInitKwargs(VideosKwargs):
initial_shift: Union[bool, float, int]
diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py
index c8ddeb970e26..7196045390b1 100755
--- a/src/transformers/models/jamba/modeling_jamba.py
+++ b/src/transformers/models/jamba/modeling_jamba.py
@@ -610,7 +610,7 @@ def __init__(self, config: JambaConfig, layer_idx):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
+ "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
" is None. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config"
)
diff --git a/src/transformers/models/janus/image_processing_janus.py b/src/transformers/models/janus/image_processing_janus.py
index 16659bd85354..21d36c651b39 100644
--- a/src/transformers/models/janus/image_processing_janus.py
+++ b/src/transformers/models/janus/image_processing_janus.py
@@ -355,7 +355,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
@@ -480,7 +480,7 @@ def unnormalize(
image_mean: Union[float, Iterable[float]],
image_std: Union[float, Iterable[float]],
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
image = (image * image_std) + image_mean
diff --git a/src/transformers/models/janus/image_processing_janus_fast.py b/src/transformers/models/janus/image_processing_janus_fast.py
index 9ed2732fb3d0..6cbca591626e 100644
--- a/src/transformers/models/janus/image_processing_janus_fast.py
+++ b/src/transformers/models/janus/image_processing_janus_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -36,16 +37,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class JanusFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
r"""
min_size (`int`, *optional*, defaults to 14):
diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py
index dcd5c1e1e730..e5c000fdd6f0 100644
--- a/src/transformers/models/janus/modular_janus.py
+++ b/src/transformers/models/janus/modular_janus.py
@@ -1359,7 +1359,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
@@ -1705,7 +1705,7 @@ def unnormalize(
image_mean: Union[float, Iterable[float]],
image_std: Union[float, Iterable[float]],
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
image = (image * image_std) + image_mean
diff --git a/src/transformers/models/kosmos2_5/image_processing_kosmos2_5_fast.py b/src/transformers/models/kosmos2_5/image_processing_kosmos2_5_fast.py
index c6d8b1b1edf5..8f6b0be8bfc4 100644
--- a/src/transformers/models/kosmos2_5/image_processing_kosmos2_5_fast.py
+++ b/src/transformers/models/kosmos2_5/image_processing_kosmos2_5_fast.py
@@ -34,7 +34,7 @@
# Similar to transformers.models.pix2struct.image_processing_pix2struct.torch_extract_patches but dealing with a batch of images directly.
def torch_extract_patches(image_tensor, patch_height, patch_width):
"""
- Utiliy function to extract patches from a given tensor representing a batch of images. Returns a tensor of shape
+ Utility function to extract patches from a given tensor representing a batch of images. Returns a tensor of shape
(batch_size, `rows`, `columns`, `num_channels` x `patch_height` x `patch_width`).
Args:
@@ -45,7 +45,6 @@ def torch_extract_patches(image_tensor, patch_height, patch_width):
patch_width (int):
The width of the patches to extract.
"""
- image_tensor = image_tensor
patches = torch.nn.functional.unfold(image_tensor, (patch_height, patch_width), stride=(patch_height, patch_width))
patches = patches.reshape(image_tensor.size(0), image_tensor.size(1), patch_height, patch_width, -1)
patches = patches.permute(0, 4, 2, 3, 1).reshape(
diff --git a/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py b/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py
index 8f9fbd706b32..b31c5797ad3c 100644
--- a/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py
+++ b/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py
@@ -290,9 +290,7 @@ class Kosmos2_5ModelOutput(ModelOutput):
vision_model_output: BaseModelOutputWithPooling = None
def to_tuple(self) -> tuple[Any]:
- return tuple(
- (self[k] if k not in ["vision_model_output"] else getattr(self, k).to_tuple()) for k in self.keys()
- )
+ return tuple((self[k] if k != "vision_model_output" else getattr(self, k).to_tuple()) for k in self.keys())
@dataclass
@@ -350,9 +348,7 @@ class Kosmos2_5ForConditionalGenerationModelOutput(ModelOutput):
vision_model_output: BaseModelOutputWithPooling = None
def to_tuple(self) -> tuple[Any]:
- return tuple(
- (self[k] if k not in ["vision_model_output"] else getattr(self, k).to_tuple()) for k in self.keys()
- )
+ return tuple((self[k] if k != "vision_model_output" else getattr(self, k).to_tuple()) for k in self.keys())
# Copied from transformers.models.pix2struct.modeling_pix2struct.Pix2StructLayerNorm with Pix2Struct->Kosmos2_5
diff --git a/src/transformers/models/kyutai_speech_to_text/feature_extraction_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/feature_extraction_kyutai_speech_to_text.py
index bde1736f9da8..d076ccb1de78 100644
--- a/src/transformers/models/kyutai_speech_to_text/feature_extraction_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/feature_extraction_kyutai_speech_to_text.py
@@ -204,7 +204,7 @@ def __call__(
if padding:
padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask")
- # now let's padd left and right
+ # now let's pad left and right
pad_left = int(self.audio_silence_prefix_seconds * self.sampling_rate)
pad_right = int((self.audio_delay_seconds + 1.0) * self.sampling_rate)
padded_inputs["input_values"] = np.pad(
diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
index 9eba7e163670..77c636570d58 100644
--- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
@@ -1078,7 +1078,7 @@ def __init__(self, config):
self.codec_model = AutoModel.from_config(config.codec_config)
# we are in an edge case where for the codec_model self.can_generate is False, setting self.codec_model.generation_config to None
- # yet the codec_model needs a generation config to initalize it's cache for streaming inference
+ # yet the codec_model needs a generation config to initialize it's cache for streaming inference
# we therefore initialize a generation config for the codec model
self.codec_model.generation_config = GenerationConfig.from_model_config(config.codec_config)
diff --git a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
index 8541a911e947..d3707d659e1e 100644
--- a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
@@ -183,7 +183,7 @@ def __call__(
if padding:
padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask")
- # now let's padd left and right
+ # now let's pad left and right
pad_left = int(self.audio_silence_prefix_seconds * self.sampling_rate)
pad_right = int((self.audio_delay_seconds + 1.0) * self.sampling_rate)
padded_inputs["input_values"] = np.pad(
@@ -259,7 +259,7 @@ def __init__(self, config):
self.codec_model = AutoModel.from_config(config.codec_config)
# we are in an edge case where for the codec_model self.can_generate is False, setting self.codec_model.generation_config to None
- # yet the codec_model needs a generation config to initalize it's cache for streaming inference
+ # yet the codec_model needs a generation config to initialize it's cache for streaming inference
# we therefore initialize a generation config for the codec model
self.codec_model.generation_config = GenerationConfig.from_model_config(config.codec_config)
diff --git a/src/transformers/models/layoutlmv2/image_processing_layoutlmv2_fast.py b/src/transformers/models/layoutlmv2/image_processing_layoutlmv2_fast.py
index 723687d58219..354bbe21c4db 100644
--- a/src/transformers/models/layoutlmv2/image_processing_layoutlmv2_fast.py
+++ b/src/transformers/models/layoutlmv2/image_processing_layoutlmv2_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import ChannelDimension, group_images_by_shape, reorder_images
@@ -25,18 +26,12 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
requires_backends,
)
from .image_processing_layoutlmv2 import apply_tesseract
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/layoutlmv3/image_processing_layoutlmv3_fast.py b/src/transformers/models/layoutlmv3/image_processing_layoutlmv3_fast.py
index 2ab8f8dd48cc..caefa9b89660 100644
--- a/src/transformers/models/layoutlmv3/image_processing_layoutlmv3_fast.py
+++ b/src/transformers/models/layoutlmv3/image_processing_layoutlmv3_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import ChannelDimension, group_images_by_shape, reorder_images
@@ -25,18 +26,12 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
requires_backends,
)
from .image_processing_layoutlmv3 import apply_tesseract
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py b/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py
index b69fc57b1743..270437e97f44 100644
--- a/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py
+++ b/src/transformers/models/layoutlmv3/tokenization_layoutlmv3.py
@@ -524,7 +524,7 @@ def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
if (
(is_split_into_words or add_prefix_space)
and (len(text) > 0 and not text[0].isspace())
- and sum([text.startswith(no_split_token) for no_split_token in self.added_tokens_encoder]) == 0
+ and sum(text.startswith(no_split_token) for no_split_token in self.added_tokens_encoder) == 0
):
text = " " + text
return (text, kwargs)
diff --git a/src/transformers/models/levit/convert_levit_timm_to_pytorch.py b/src/transformers/models/levit/convert_levit_timm_to_pytorch.py
index 0d5731bf7bef..5d198ee9e552 100644
--- a/src/transformers/models/levit/convert_levit_timm_to_pytorch.py
+++ b/src/transformers/models/levit/convert_levit_timm_to_pytorch.py
@@ -86,11 +86,9 @@ def convert_weights_and_push(save_directory: Path, model_name: Optional[str] = N
expected_shape = (1, num_labels)
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
ImageNetPreTrainedConfig = partial(LevitConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/levit/image_processing_levit_fast.py b/src/transformers/models/levit/image_processing_levit_fast.py
index e452894d6e2e..ae30194288fa 100644
--- a/src/transformers/models/levit/image_processing_levit_fast.py
+++ b/src/transformers/models/levit/image_processing_levit_fast.py
@@ -17,6 +17,7 @@
from typing import Optional
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, SizeDict
from ...image_transforms import (
@@ -24,13 +25,7 @@
get_resize_output_image_size,
)
from ...image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling
-from ...utils import auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import auto_docstring
@auto_docstring
diff --git a/src/transformers/models/lfm2_vl/__init__.py b/src/transformers/models/lfm2_vl/__init__.py
new file mode 100755
index 000000000000..7d0357ffbaa6
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_lfm2_vl import *
+ from .image_processing_lfm2_vl_fast import *
+ from .modeling_lfm2_vl import *
+ from .processing_lfm2_vl import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/lfm2_vl/configuration_lfm2_vl.py b/src/transformers/models/lfm2_vl/configuration_lfm2_vl.py
new file mode 100755
index 000000000000..1378fbe6dc8c
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/configuration_lfm2_vl.py
@@ -0,0 +1,91 @@
+# coding=utf-8
+# Copyright 2025 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch LFM2-VL model."""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class Lfm2VlConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Lfm2VlForConditionalGeneration`]. It is used to instantiate an
+ Lfm2Vl model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of the Lfm2-VL-1.6B.
+
+ e.g. [LiquidAI/LFM2-VL-1.6B](https://huggingface.co/LiquidAI/LFM2-VL-1.6B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vision_config (`AutoConfig | dict`, *optional*, defaults to `Siglip2ImageConfig`):
+ The config object or dictionary of the vision backbone.
+ text_config (`AutoConfig | dict`, *optional*, defaults to `Lfm2Config`):
+ The config object or dictionary of the text backbone.
+ image_token_id (`int`, *optional*, defaults to 396):
+ The image token index to encode the image prompt.
+ projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The activation function used by the multimodal projector.
+ projector_hidden_size (`int`, *optional*, defaults to 2560):
+ The hidden size of the multimodal projector.
+ projector_bias (`bool`, *optional*, defaults to `True`):
+ Whether to use bias in the multimodal projector.
+ downsample_factor (`int`, *optional*, defaults to 2):
+ The downsample_factor factor of the vision backbone.
+ """
+
+ model_type = "lfm2-vl"
+ sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
+
+ def __init__(
+ self,
+ vision_config=None,
+ text_config=None,
+ image_token_id=396,
+ projector_hidden_act="gelu",
+ projector_hidden_size=2560,
+ projector_bias=True,
+ downsample_factor=2,
+ **kwargs,
+ ):
+ self.image_token_id = image_token_id
+ self.projector_hidden_act = projector_hidden_act
+ self.projector_hidden_size = projector_hidden_size
+ self.projector_bias = projector_bias
+ self.downsample_factor = downsample_factor
+
+ if isinstance(vision_config, dict):
+ vision_config["model_type"] = vision_config.get("model_type", "siglip2_vision_model")
+ vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
+ elif vision_config is None:
+ vision_config = CONFIG_MAPPING["siglip2_vision_model"]()
+
+ if isinstance(text_config, dict):
+ text_config["model_type"] = text_config.get("model_type", "lfm2")
+ text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+ elif text_config is None:
+ text_config = CONFIG_MAPPING["lfm2"]()
+
+ self.vision_config = vision_config
+ self.text_config = text_config
+
+ super().__init__(**kwargs)
+
+
+__all__ = ["Lfm2VlConfig"]
diff --git a/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py b/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py
new file mode 100755
index 000000000000..4081c86e108a
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py
@@ -0,0 +1,541 @@
+# coding=utf-8
+# Copyright 2025 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from functools import lru_cache
+from typing import Optional, Union
+
+import torch
+from torchvision.transforms.v2 import functional as F
+
+from ...image_processing_utils import BatchFeature
+from ...image_processing_utils_fast import (
+ BaseImageProcessorFast,
+ DefaultFastImageProcessorKwargs,
+ group_images_by_shape,
+ reorder_images,
+)
+from ...image_utils import (
+ IMAGENET_STANDARD_MEAN,
+ IMAGENET_STANDARD_STD,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import (
+ Unpack,
+)
+from ...utils import (
+ TensorType,
+ auto_docstring,
+ logging,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def round_by_factor(number: float, factor: int) -> int:
+ """Returns the closest integer to 'number' that is divisible by 'factor'."""
+ return round(number / factor) * factor
+
+
+def find_closest_aspect_ratio(
+ aspect_ratio: float,
+ target_ratios: list[tuple[int, int]],
+ width: int,
+ height: int,
+ image_size: int,
+) -> tuple[int, int]:
+ """Find the closest aspect ratio from target_ratios to match the input aspect ratio.
+
+ Args:
+ aspect_ratio: The aspect ratio to match (width/height).
+ target_ratios: List of possible aspect ratios as tuples of (width, height) integers.
+ width: Original image width in pixels.
+ height: Original image height in pixels.
+ image_size: Base size for calculating target area.
+
+ Returns:
+ tuple[int, int]: The best matching ratio as (width, height) integers.
+ """
+ best_ratio_diff = float("inf")
+ best_ratio = (1, 1)
+ area = width * height
+
+ for ratio in target_ratios:
+ target_aspect_ratio = ratio[0] / ratio[1]
+ ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+
+ # update best ratio if we found a closer match
+ if ratio_diff < best_ratio_diff:
+ best_ratio_diff = ratio_diff
+ best_ratio = ratio
+ # if equally close, prefer the ratio that better matches the original image area
+ elif ratio_diff == best_ratio_diff:
+ target_area = image_size * image_size * ratio[0] * ratio[1]
+ if area > 0.5 * target_area:
+ best_ratio = ratio
+
+ return best_ratio
+
+
+# copied from Siglip2ImageProcessor
+@lru_cache(maxsize=256)
+def get_image_size_for_max_num_patches(
+ image_height: int, image_width: int, patch_size: int, max_num_patches: int, eps: float = 1e-5
+) -> tuple[int, int]:
+ """
+ Determine image size based on max number of patches, ensure dimensions are divisible by patch size and image is at least 1 patch.
+
+ Args:
+ image_height (`int`):
+ Original image height.
+ image_width (`int`):
+ Original image width.
+ patch_size (`int`):
+ Patch size for processing.
+ max_num_patches (`int`):
+ Maximum number of patches.
+ eps (`float`):
+ Small threshold for binary search.
+
+ Returns:
+ Tuple: (target_height, target_width)
+ """
+
+ def get_scaled_image_size(scale: float, size: int, patch_size: int) -> int:
+ scaled_size = size * scale
+ scaled_size = math.ceil(scaled_size / patch_size) * patch_size # make divisible by patch_size
+ scaled_size = max(patch_size, scaled_size) # ensure at least 1 patch
+ return int(scaled_size)
+
+ # Binary search for optimal scale
+ scale_min, scale_max = eps / 10, 100.0
+ while (scale_max - scale_min) >= eps:
+ scale = (scale_min + scale_max) / 2
+ target_height = get_scaled_image_size(scale, image_height, patch_size)
+ target_width = get_scaled_image_size(scale, image_width, patch_size)
+ num_patches = (target_height / patch_size) * (target_width / patch_size)
+
+ if num_patches <= max_num_patches:
+ scale_min = scale
+ else:
+ scale_max = scale
+
+ scale = scale_min
+ target_height = get_scaled_image_size(scale, image_height, patch_size)
+ target_width = get_scaled_image_size(scale, image_width, patch_size)
+ return target_height, target_width
+
+
+def convert_image_to_patches(images: "torch.Tensor", patch_size: int) -> "torch.Tensor":
+ """
+ Convert 3D array image of shape (image_height, image_width, num_channels) into 2D array of patches of shape
+ (num_patches_height * num_patches_width, patch_size * patch_size * num_channels).
+ """
+ batch_size, num_channels, image_height, image_width = images.shape
+ num_patches_height = image_height // patch_size
+ num_patches_width = image_width // patch_size
+ patched_image = images.reshape(
+ batch_size, num_channels, num_patches_height, patch_size, num_patches_width, patch_size
+ )
+ patched_image = patched_image.permute(0, 2, 4, 3, 5, 1)
+ patched_image = patched_image.reshape(batch_size, num_patches_height * num_patches_width, -1)
+ return patched_image
+
+
+def pad_along_first_dim(
+ images: "torch.Tensor", target_length: int, pad_value: int = 0
+) -> tuple["torch.Tensor", "torch.Tensor"]:
+ """
+ Pad the array along the first dimension.
+ """
+ current_length = images.shape[1]
+ padding_length = target_length - current_length
+ pixel_mask = torch.ones((target_length,), dtype=torch.int32)
+ if padding_length > 0:
+ paddings = (0, 0, 0, padding_length, 0, 0)
+ images = torch.nn.functional.pad(images, paddings, mode="constant", value=pad_value)
+ pixel_mask[-padding_length:] = 0
+ return images, pixel_mask
+
+
+class Lfm2VlFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
+ """
+ downsample_factor (`int`, *optional*, defaults to `2`):
+ The downsampling factor for images used when resizing the image.
+ """
+
+ downsample_factor: Optional[int]
+ do_image_splitting: Optional[bool]
+ min_tiles: Optional[int]
+ max_tiles: Optional[int]
+ use_thumbnail: Optional[bool]
+ min_image_tokens: Optional[int]
+ max_image_tokens: Optional[int]
+ encoder_patch_size: Optional[int]
+ tile_size: Optional[int]
+ max_pixels_tolerance: Optional[float]
+ do_pad: Optional[bool]
+ return_row_col_info: Optional[bool]
+
+
+@auto_docstring
+class Lfm2VlImageProcessorFast(BaseImageProcessorFast):
+ downsample_factor = 2
+ do_image_splitting = True
+ min_tiles = 2
+ max_tiles = 10
+ use_thumbnail = True
+ min_image_tokens = 64
+ max_image_tokens = 256
+ encoder_patch_size = 16
+ tile_size = 512
+ max_pixels_tolerance = 2.0
+ do_resize = True
+ size = {"height": 512, "width": 512}
+ resample = PILImageResampling.BILINEAR
+ do_rescale = True
+ rescale_factor = 1 / 255
+ do_normalize = True
+ do_pad = True
+ return_row_col_info = False
+ image_mean = IMAGENET_STANDARD_STD
+ image_std = IMAGENET_STANDARD_MEAN
+ valid_kwargs = Lfm2VlFastImageProcessorKwargs
+ model_input_names = ["pixel_values", "pixel_attention_mask", "spatial_shapes"]
+
+ def __init__(self, **kwargs: Unpack[Lfm2VlFastImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ max_thumbnail_image_patches = self.max_image_tokens * self.downsample_factor**2
+ tile_size_patches = (self.tile_size // self.encoder_patch_size) ** 2 if self.do_image_splitting else 0
+ self.max_num_patches = max(
+ max_thumbnail_image_patches,
+ tile_size_patches,
+ )
+
+ @lru_cache(maxsize=256)
+ def _target_ratios(self, min_tiles: int, max_tiles: int) -> list[tuple[int, int]]:
+ ratios = [
+ (w, h)
+ for n in range(min_tiles, max_tiles + 1)
+ for w in range(1, n + 1)
+ for h in range(1, n + 1)
+ if min_tiles <= w * h <= max_tiles
+ ]
+ return sorted(set(ratios), key=lambda x: x[0] * x[1])
+
+ def _get_grid_layout(
+ self,
+ height: int,
+ width: int,
+ min_tiles: int,
+ max_tiles: int,
+ tile_size: int,
+ ) -> tuple[int, int]:
+ aspect_ratio = width / height
+ target_ratios = self._target_ratios(min_tiles, max_tiles)
+
+ # find best matching grid configuration
+ grid_width, grid_height = find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, tile_size)
+
+ target_width = tile_size * grid_width
+ target_height = tile_size * grid_height
+ total_patches = grid_width * grid_height
+
+ return grid_width, grid_height, target_width, target_height, total_patches
+
+ def crop_image_to_patches(
+ self,
+ image: "torch.Tensor",
+ min_tiles: int,
+ max_tiles: int,
+ tile_size: int,
+ use_thumbnail: bool,
+ thumbnail_size: tuple[int],
+ interpolation: "F.InterpolationMode" = None,
+ antialias: bool = True,
+ **kwargs,
+ ) -> "torch.Tensor":
+ """
+ Processes a high resolution image into patches.
+ This method splits a high resolution image into a grid of smaller patches while trying to maintain
+ the original aspect ratio. It finds the optimal grid configuration within the specified tile constraints.
+ """
+ batch_size, num_channels, height, width = image.shape
+ grid_width, grid_height, target_width, target_height, total_patches = self._get_grid_layout(
+ height, width, min_tiles=min_tiles, max_tiles=max_tiles, tile_size=tile_size
+ )
+ resized_image = F.resize(
+ image, (target_height, target_width), interpolation=interpolation, antialias=antialias
+ )
+
+ # split the image into patches
+ processed_images = (
+ resized_image.unfold(2, size=tile_size, step=tile_size)
+ .unfold(3, size=tile_size, step=tile_size)
+ .contiguous()
+ .view(batch_size, num_channels, -1, tile_size, tile_size)
+ .permute(2, 0, 1, 3, 4)
+ .reshape(batch_size, -1, num_channels, tile_size, tile_size)
+ )
+
+ # Re-order processed images to a nested image structure, so it can be reordered back correctly
+ # Note that the images can't be stacked because the thumbnail image is of bigger size than patches
+ # Each image in sublist will be of shape (1, C, H, W)
+ processed_images = list(processed_images)
+
+ if use_thumbnail and grid_width * grid_height != 1:
+ total_patches += 1
+ thumbnail_image = F.resize(image, thumbnail_size, interpolation=interpolation, antialias=antialias)
+ for i in range(batch_size):
+ processed_images[i] = list(processed_images[i]) + list(thumbnail_image[i][None, ...])
+
+ return processed_images, grid_width, grid_height
+
+ # Adapted from Qwen-VL with minor differences
+ def smart_resize(
+ self,
+ height: int,
+ width: int,
+ downsample_factor: int,
+ min_image_tokens: int,
+ max_image_tokens: int,
+ encoder_patch_size: int,
+ ) -> tuple[int, int]:
+ """
+ Rescales the image so that the following conditions are met:
+ 1. Both dimensions (height and width) are divisible by 'encoder_patch_size' * 'downsample_factor'.
+ This ensures no padding is needed in the downsampling step.
+ 2. The total number of pixels is within the range ['smart_resize_min_pixels', 'smart_resize_max_pixels'].
+ 3. The aspect ratio of the image is maintained as closely as possible.
+ """
+ total_factor = encoder_patch_size * downsample_factor
+ smart_resize_min_pixels = min_image_tokens * encoder_patch_size**2 * downsample_factor**2
+ smart_resize_max_pixels = max_image_tokens * encoder_patch_size**2 * downsample_factor**2
+
+ h_bar = max(total_factor, round_by_factor(height, total_factor))
+ w_bar = max(total_factor, round_by_factor(width, total_factor))
+
+ if h_bar * w_bar > smart_resize_max_pixels:
+ beta = math.sqrt((height * width) / smart_resize_max_pixels)
+ math.floor(height / beta / total_factor) * total_factor
+ h_bar = max(total_factor, math.floor(height / beta / total_factor) * total_factor)
+ w_bar = max(total_factor, math.floor(width / beta / total_factor) * total_factor)
+ elif h_bar * w_bar < smart_resize_min_pixels:
+ beta = math.sqrt(smart_resize_min_pixels / (height * width))
+ h_bar = math.ceil(height * beta / total_factor) * total_factor
+ w_bar = math.ceil(width * beta / total_factor) * total_factor
+
+ return w_bar, h_bar
+
+ def _is_image_too_large(
+ self,
+ height: int,
+ width: int,
+ max_image_tokens: int,
+ encoder_patch_size: int,
+ downsample_factor: int,
+ max_pixels_tolerance: float,
+ ) -> bool:
+ """Check if the image is too large to be processed as one tile."""
+ total_factor = encoder_patch_size * downsample_factor
+
+ h_bar = max(encoder_patch_size, round_by_factor(height, total_factor))
+ w_bar = max(encoder_patch_size, round_by_factor(width, total_factor))
+ return h_bar * w_bar > max_image_tokens * encoder_patch_size**2 * downsample_factor**2 * max_pixels_tolerance
+
+ def resize_and_split(
+ self,
+ images: "torch.Tensor",
+ downsample_factor: int,
+ min_tiles: int,
+ max_tiles: int,
+ use_thumbnail: bool,
+ min_image_tokens: int,
+ max_image_tokens: int,
+ encoder_patch_size: int,
+ tile_size: int,
+ max_pixels_tolerance: float,
+ interpolation: "F.InterpolationMode",
+ ) -> "torch.Tensor":
+ batch_size, _, height, width = images.shape
+ do_image_splitting = not min_tiles == max_tiles == 1
+ is_image_large = self._is_image_too_large(
+ height=height,
+ width=width,
+ max_image_tokens=max_image_tokens,
+ encoder_patch_size=encoder_patch_size,
+ downsample_factor=downsample_factor,
+ max_pixels_tolerance=max_pixels_tolerance,
+ )
+
+ new_width, new_height = self.smart_resize(
+ height=height,
+ width=width,
+ downsample_factor=downsample_factor,
+ min_image_tokens=min_image_tokens,
+ max_image_tokens=max_image_tokens,
+ encoder_patch_size=encoder_patch_size,
+ )
+
+ # Big image will be cropped into patches and small images are just resized
+ if is_image_large and do_image_splitting:
+ images, num_rows, num_cols = self.crop_image_to_patches(
+ images,
+ min_tiles=min_tiles,
+ max_tiles=max_tiles,
+ tile_size=tile_size,
+ thumbnail_size=(new_height, new_width),
+ use_thumbnail=use_thumbnail,
+ interpolation=interpolation,
+ )
+ else:
+ num_rows = num_cols = 1
+ images = F.resize(images, (new_height, new_width), interpolation=interpolation)
+ # Make a list and treat it as single crop per image so it can be re-grouped back correctly
+ images = [[image] for image in images]
+
+ num_rows = [num_rows] * batch_size
+ num_cols = [num_cols] * batch_size
+ image_sizes = [[new_height, new_width]] * batch_size
+ return images, num_rows, num_cols, image_sizes
+
+ def _preprocess(
+ self,
+ images: ImageInput,
+ size: SizeDict,
+ interpolation: "F.InterpolationMode",
+ do_resize: bool,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: Union[float, list[float]],
+ image_std: Union[float, list[float]],
+ downsample_factor: int,
+ do_image_splitting: bool,
+ min_tiles: int,
+ max_tiles: int,
+ use_thumbnail: bool,
+ min_image_tokens: int,
+ max_image_tokens: int,
+ encoder_patch_size: int,
+ tile_size: int,
+ max_pixels_tolerance: float,
+ return_tensors: Union[str, TensorType],
+ disable_grouping: bool,
+ do_pad: bool,
+ return_row_col_info: bool,
+ **kwargs,
+ ) -> BatchFeature:
+ if not do_image_splitting:
+ min_tiles = 1
+ max_tiles = 1
+ logger.debug(
+ "Image splitting is disabled, setting min_tiles and max_tiles to 1. Set do_image_splitting=True to enable splitting."
+ )
+
+ if do_image_splitting and min_tiles > max_tiles:
+ raise ValueError("min_tiles must be less than or equal to max_tiles")
+
+ max_thumbnail_image_patches = max_image_tokens * downsample_factor**2
+ tile_size_patches = (tile_size // encoder_patch_size) ** 2 if do_image_splitting else 0
+ max_num_patches = max(
+ max_thumbnail_image_patches,
+ tile_size_patches,
+ )
+
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ resized_image_sizes = {}
+ rows_grouped, cols_grouped = {}, {}
+ for shape, stacked_images in grouped_images.items():
+ num_rows = [1] * stacked_images.shape[0]
+ num_cols = [1] * stacked_images.shape[0]
+ height, width = stacked_images.shape[-2:]
+ image_sizes = [[height, width]] * stacked_images.shape[0]
+ do_resize = True
+
+ if do_resize:
+ stacked_images, num_rows, num_cols, image_sizes = self.resize_and_split(
+ stacked_images,
+ downsample_factor=downsample_factor,
+ min_tiles=min_tiles,
+ max_tiles=max_tiles,
+ use_thumbnail=use_thumbnail,
+ min_image_tokens=min_image_tokens,
+ max_image_tokens=max_image_tokens,
+ encoder_patch_size=encoder_patch_size,
+ tile_size=tile_size,
+ max_pixels_tolerance=max_pixels_tolerance,
+ interpolation=interpolation,
+ )
+
+ rows_grouped[shape] = num_rows
+ cols_grouped[shape] = num_cols
+ resized_image_sizes[shape] = image_sizes
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+ batch_rows = reorder_images(rows_grouped, grouped_images_index)
+ batch_cols = reorder_images(cols_grouped, grouped_images_index)
+ resized_image_sizes = reorder_images(resized_image_sizes, grouped_images_index)
+
+ grouped_images, grouped_images_index = group_images_by_shape(
+ resized_images, disable_grouping=disable_grouping, is_nested=True
+ )
+
+ processed_images_grouped = {}
+ processed_masks, processed_spatial_shapes = {}, {}
+ for shape, stacked_images in grouped_images.items():
+ # Fused rescale and normalize
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ batch_size, *_, height, width = stacked_images.shape
+ num_patches_height = height // encoder_patch_size
+ num_patches_width = width // encoder_patch_size
+
+ stacked_images = convert_image_to_patches(stacked_images, encoder_patch_size)
+ processed_spatial_shapes[shape] = [[num_patches_height, num_patches_width]] * batch_size
+
+ if do_pad:
+ stacked_images, pixel_mask = pad_along_first_dim(stacked_images, max_num_patches)
+ processed_masks[shape] = [pixel_mask] * batch_size
+
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index, is_nested=True)
+ data = {"pixel_values": torch.cat([torch.stack(images) for images in processed_images])}
+
+ if do_pad:
+ processed_masks = reorder_images(processed_masks, grouped_images_index, is_nested=True)
+ processed_spatial_shapes = reorder_images(processed_spatial_shapes, grouped_images_index, is_nested=True)
+ processed_masks = torch.cat([torch.stack(masks) for masks in processed_masks])
+ processed_spatial_shapes = torch.cat(
+ [torch.tensor(spatial_shape) for spatial_shape in processed_spatial_shapes]
+ )
+ data.update({"pixel_attention_mask": processed_masks, "spatial_shapes": processed_spatial_shapes})
+
+ if return_row_col_info:
+ data["image_rows"] = batch_rows
+ data["image_cols"] = batch_cols
+ data["image_sizes"] = resized_image_sizes
+
+ encoding = BatchFeature(data=data, tensor_type=return_tensors)
+ return encoding
+
+
+__all__ = ["Lfm2VlImageProcessorFast"]
diff --git a/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py b/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py
new file mode 100755
index 000000000000..deee35394ee1
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py
@@ -0,0 +1,497 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/lfm2_vl/modular_lfm2_vl.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_lfm2_vl.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ..auto import AutoModel
+from .configuration_lfm2_vl import Lfm2VlConfig
+
+
+class Lfm2VlMultiModalProjector(nn.Module):
+ def __init__(self, config: Lfm2VlConfig):
+ super().__init__()
+ in_channels = config.vision_config.hidden_size * (config.downsample_factor**2)
+ self.factor = config.downsample_factor
+ self.layer_norm = nn.LayerNorm(in_channels)
+ self.linear_1 = nn.Linear(
+ in_channels,
+ config.projector_hidden_size,
+ bias=config.projector_bias,
+ )
+ self.act = ACT2FN[config.projector_hidden_act]
+ self.linear_2 = nn.Linear(
+ config.projector_hidden_size,
+ config.text_config.hidden_size,
+ bias=config.projector_bias,
+ )
+
+ def forward(self, image_features: torch.Tensor):
+ image_features = self.pixel_unshuffle(image_features)
+ image_features = self.layer_norm(image_features)
+ hidden_states = self.linear_1(image_features)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+ def pixel_unshuffle(self, hidden_states: torch.Tensor):
+ batch_size, width, height, channels = hidden_states.size()
+ hidden_states = hidden_states.reshape(batch_size, width, height // self.factor, channels * self.factor)
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
+ hidden_states = hidden_states.reshape(
+ batch_size, height // self.factor, width // self.factor, channels * self.factor**2
+ )
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
+ return hidden_states
+
+
+@auto_docstring
+class Lfm2VlPreTrainedModel(PreTrainedModel):
+ config: Lfm2VlConfig
+ base_model_prefix = ""
+ supports_gradient_checkpointing = True
+ _skip_keys_device_placement = "past_key_values"
+
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = False
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Lfm2Vl causal language model (or autoregressive) outputs.
+ """
+)
+class Lfm2VlCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Lfm2Vl outputs, with hidden states and attentions.
+ """
+)
+class Lfm2VlModelOutputWithPast(BaseModelOutputWithPast):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The Lfm2Vl model which consists of a vision backbone and a language model, without a language modeling head.
+ """
+)
+class Lfm2VlModel(Lfm2VlPreTrainedModel):
+ _checkpoint_conversion_mapping = {}
+
+ def __init__(self, config: Lfm2VlConfig):
+ super().__init__(config)
+ self.vision_tower = AutoModel.from_config(config.vision_config)
+
+ self.multi_modal_projector = Lfm2VlMultiModalProjector(config)
+ self.language_model = AutoModel.from_config(config.text_config)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def set_decoder(self, decoder):
+ self.language_model = decoder
+
+ def get_decoder(self):
+ return self.language_model
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ spatial_shapes: torch.Tensor,
+ pixel_attention_mask: torch.Tensor,
+ **kwargs,
+ ) -> list[torch.Tensor]:
+ """
+ Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+ Args:
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
+ The tensors corresponding to the input images.
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`):
+ The pixel attention mask of the input images.
+ Returns:
+ image_features (`list[torch.Tensor]`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+ """
+ image_outputs = self.vision_tower(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ ).last_hidden_state
+
+ img_feature_lengths = pixel_attention_mask.sum(dim=1)
+ image_features = []
+
+ for img_idx in range(image_outputs.size(0)):
+ feature = image_outputs[img_idx]
+ # unpad the image representation
+ feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0)
+
+ # reshape to original height and width
+ feature_org_h, feature_org_w = spatial_shapes[img_idx]
+ feature = feature.reshape(1, feature_org_h, feature_org_w, -1)
+
+ # project the image representation
+ img_embedding = self.multi_modal_projector(feature)
+
+ # flatten here to handle variable length in naflex
+ img_embedding = img_embedding.reshape(-1, img_embedding.size(-1))
+ image_features.append(img_embedding)
+
+ return image_features
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ n_image_features = image_features.shape[0]
+ if inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ spatial_shapes: Optional[torch.Tensor] = None,
+ pixel_attention_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Lfm2VlModelOutputWithPast]:
+ r"""
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*):
+ The pixel attention mask of the input images.
+ """
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ )
+ image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids=input_ids,
+ inputs_embeds=inputs_embeds,
+ image_features=image_features,
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ return Lfm2VlModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The LFM2_VL model which consists of a vision backbone and a language model.
+ """
+)
+class Lfm2VlForConditionalGeneration(Lfm2VlPreTrainedModel, GenerationMixin):
+ _checkpoint_conversion_mapping = {}
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config: Lfm2VlConfig):
+ super().__init__(config)
+ self.model = Lfm2VlModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.lm_head
+
+ def set_decoder(self, decoder):
+ self.model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.model.get_decoder()
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ spatial_shapes: torch.Tensor,
+ pixel_attention_mask: torch.Tensor,
+ **kwargs,
+ ):
+ return self.model.get_image_features(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ **kwargs,
+ )
+
+ # Make modules available through conditional class for BC
+ @property
+ def language_model(self):
+ return self.model.language_model
+
+ @property
+ def vision_tower(self):
+ return self.model.vision_tower
+
+ @property
+ def multi_modal_projector(self):
+ return self.model.multi_modal_projector
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ spatial_shapes: Optional[torch.Tensor] = None,
+ pixel_attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Lfm2VlCausalLMOutputWithPast]:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`, *optional*):
+ The input image tensors.
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*):
+ The pixel attention mask of the input images.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+ >>> from transformers.image_utils import load_image
+
+ >>> model = AutoModelForImageTextToText.from_pretrained(
+ ... "LiquidAI/LFM2-VL-1.6B",
+ ... )
+ >>> processor = AutoProcessor.from_pretrained(
+ ... "LiquidAI/LFM2-VL-1.6B",
+ ... )
+
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> image = load_image(url)
+
+ >>> conversation = [
+ ... {
+ ... "role": "user",
+ ... "content": [
+ ... {"type": "image", "image": image},
+ ... {"type": "text", "text": "What is in this image?"},
+ ... ],
+ ... },
+ ... ]
+
+ >>> inputs = processor.apply_chat_template(
+ ... conversation,
+ ... add_generation_prompt=True,
+ ... tokenize=True,
+ ... return_dict=True,
+ ... return_tensors="pt"
+ ... )
+
+ >>> # Generate
+ >>> outputs = model.generate(**inputs, max_new_tokens=45)
+ >>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
+ 'This image depicts a vibrant street scene in what appears to be a Chinatown or similar cultural area. The focal point is a large red stop sign with white lettering, mounted on a pole.'
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits,
+ labels=labels,
+ vocab_size=self.config.text_config.vocab_size,
+ **kwargs,
+ )
+
+ return Lfm2VlCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=outputs.image_hidden_states,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ inputs_embeds=None,
+ pixel_values=None,
+ attention_mask=None,
+ cache_position=None,
+ logits_to_keep=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ **kwargs,
+ )
+
+ if cache_position[0] == 0:
+ # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+ # Otherwise we need pixel values to be passed to model
+ model_inputs["pixel_values"] = pixel_values
+
+ return model_inputs
+
+
+__all__ = ["Lfm2VlForConditionalGeneration", "Lfm2VlPreTrainedModel", "Lfm2VlModel"]
diff --git a/src/transformers/models/lfm2_vl/modular_lfm2_vl.py b/src/transformers/models/lfm2_vl/modular_lfm2_vl.py
new file mode 100644
index 000000000000..68367464c3cf
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/modular_lfm2_vl.py
@@ -0,0 +1,352 @@
+# coding=utf-8
+# Copyright 2025 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Lfm2-VL model."""
+
+from typing import Optional, Union
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ..llava.modeling_llava import (
+ LlavaCausalLMOutputWithPast,
+ LlavaForConditionalGeneration,
+ LlavaModel,
+ LlavaModelOutputWithPast,
+ LlavaPreTrainedModel,
+)
+from .configuration_lfm2_vl import Lfm2VlConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class Lfm2VlMultiModalProjector(nn.Module):
+ def __init__(self, config: Lfm2VlConfig):
+ super().__init__()
+ in_channels = config.vision_config.hidden_size * (config.downsample_factor**2)
+ self.factor = config.downsample_factor
+ self.layer_norm = nn.LayerNorm(in_channels)
+ self.linear_1 = nn.Linear(
+ in_channels,
+ config.projector_hidden_size,
+ bias=config.projector_bias,
+ )
+ self.act = ACT2FN[config.projector_hidden_act]
+ self.linear_2 = nn.Linear(
+ config.projector_hidden_size,
+ config.text_config.hidden_size,
+ bias=config.projector_bias,
+ )
+
+ def forward(self, image_features: torch.Tensor):
+ image_features = self.pixel_unshuffle(image_features)
+ image_features = self.layer_norm(image_features)
+ hidden_states = self.linear_1(image_features)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+ def pixel_unshuffle(self, hidden_states: torch.Tensor):
+ batch_size, width, height, channels = hidden_states.size()
+ hidden_states = hidden_states.reshape(batch_size, width, height // self.factor, channels * self.factor)
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
+ hidden_states = hidden_states.reshape(
+ batch_size, height // self.factor, width // self.factor, channels * self.factor**2
+ )
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
+ return hidden_states
+
+
+class Lfm2VlPreTrainedModel(LlavaPreTrainedModel):
+ _can_compile_fullgraph = False
+
+
+class Lfm2VlCausalLMOutputWithPast(LlavaCausalLMOutputWithPast):
+ pass
+
+
+class Lfm2VlModelOutputWithPast(LlavaModelOutputWithPast):
+ pass
+
+
+class Lfm2VlModel(LlavaModel):
+ _checkpoint_conversion_mapping = {}
+
+ def __init__(self, config: Lfm2VlConfig):
+ super().__init__(config)
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ spatial_shapes: torch.Tensor,
+ pixel_attention_mask: torch.Tensor,
+ **kwargs,
+ ) -> list[torch.Tensor]:
+ """
+ Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+ Args:
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
+ The tensors corresponding to the input images.
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`):
+ The pixel attention mask of the input images.
+ Returns:
+ image_features (`list[torch.Tensor]`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+ """
+ image_outputs = self.vision_tower(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ ).last_hidden_state
+
+ img_feature_lengths = pixel_attention_mask.sum(dim=1)
+ image_features = []
+
+ for img_idx in range(image_outputs.size(0)):
+ feature = image_outputs[img_idx]
+ # unpad the image representation
+ feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0)
+
+ # reshape to original height and width
+ feature_org_h, feature_org_w = spatial_shapes[img_idx]
+ feature = feature.reshape(1, feature_org_h, feature_org_w, -1)
+
+ # project the image representation
+ img_embedding = self.multi_modal_projector(feature)
+
+ # flatten here to handle variable length in naflex
+ img_embedding = img_embedding.reshape(-1, img_embedding.size(-1))
+ image_features.append(img_embedding)
+
+ return image_features
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ n_image_features = image_features.shape[0]
+ if inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ spatial_shapes: Optional[torch.Tensor] = None,
+ pixel_attention_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Lfm2VlModelOutputWithPast]:
+ r"""
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*):
+ The pixel attention mask of the input images.
+ """
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ )
+ image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids=input_ids,
+ inputs_embeds=inputs_embeds,
+ image_features=image_features,
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ return Lfm2VlModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+class Lfm2VlForConditionalGeneration(LlavaForConditionalGeneration):
+ _checkpoint_conversion_mapping = {}
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ spatial_shapes: torch.Tensor,
+ pixel_attention_mask: torch.Tensor,
+ **kwargs,
+ ):
+ return self.model.get_image_features(
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ **kwargs,
+ )
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ spatial_shapes: Optional[torch.Tensor] = None,
+ pixel_attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Lfm2VlCausalLMOutputWithPast]:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`, *optional*):
+ The input image tensors.
+ spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*):
+ The spatial shapes of the input images.
+ pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*):
+ The pixel attention mask of the input images.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+ >>> from transformers.image_utils import load_image
+
+ >>> model = AutoModelForImageTextToText.from_pretrained(
+ ... "LiquidAI/LFM2-VL-1.6B",
+ ... )
+ >>> processor = AutoProcessor.from_pretrained(
+ ... "LiquidAI/LFM2-VL-1.6B",
+ ... )
+
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> image = load_image(url)
+
+ >>> conversation = [
+ ... {
+ ... "role": "user",
+ ... "content": [
+ ... {"type": "image", "image": image},
+ ... {"type": "text", "text": "What is in this image?"},
+ ... ],
+ ... },
+ ... ]
+
+ >>> inputs = processor.apply_chat_template(
+ ... conversation,
+ ... add_generation_prompt=True,
+ ... tokenize=True,
+ ... return_dict=True,
+ ... return_tensors="pt"
+ ... )
+
+ >>> # Generate
+ >>> outputs = model.generate(**inputs, max_new_tokens=45)
+ >>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
+ 'This image depicts a vibrant street scene in what appears to be a Chinatown or similar cultural area. The focal point is a large red stop sign with white lettering, mounted on a pole.'
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ spatial_shapes=spatial_shapes,
+ pixel_attention_mask=pixel_attention_mask,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits,
+ labels=labels,
+ vocab_size=self.config.text_config.vocab_size,
+ **kwargs,
+ )
+
+ return Lfm2VlCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=outputs.image_hidden_states,
+ )
+
+
+__all__ = ["Lfm2VlForConditionalGeneration", "Lfm2VlPreTrainedModel", "Lfm2VlModel"]
diff --git a/src/transformers/models/lfm2_vl/processing_lfm2_vl.py b/src/transformers/models/lfm2_vl/processing_lfm2_vl.py
new file mode 100755
index 000000000000..12f289c266a1
--- /dev/null
+++ b/src/transformers/models/lfm2_vl/processing_lfm2_vl.py
@@ -0,0 +1,269 @@
+# coding=utf-8
+# Copyright 2025 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Optional, Union
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, make_nested_list_of_images
+from ...processing_utils import (
+ ImagesKwargs,
+ ProcessingKwargs,
+ ProcessorMixin,
+ Unpack,
+)
+from ...tokenization_utils_base import BatchEncoding, TextInput
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Lfm2VlImagesKwargs(ImagesKwargs, total=False):
+ downsample_factor: Optional[int]
+ do_image_splitting: Optional[bool]
+ min_tiles: Optional[int]
+ max_tiles: Optional[int]
+ use_thumbnail: Optional[bool]
+ min_image_tokens: Optional[int]
+ max_image_tokens: Optional[int]
+ encoder_patch_size: Optional[int]
+ tile_size: Optional[int]
+ max_pixels_tolerance: Optional[float]
+ patch_size: Optional[int]
+ do_pad: Optional[bool]
+ return_row_col_info: Optional[bool]
+
+
+class Lfm2VlProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: Lfm2VlImagesKwargs
+
+ _defaults = {
+ "images_kwargs": {
+ "return_row_col_info": True,
+ },
+ "text_kwargs": {
+ "use_image_special_tokens": True,
+ "add_special_tokens": False,
+ "padding": False,
+ "is_split_into_words": False,
+ },
+ }
+
+
+class Lfm2VlProcessor(ProcessorMixin):
+ r"""
+ Constructs a Lfm2Vl processor which wraps a Lfm2Tokenizer tokenizer and Lfm2VlImageProcessor into a single processor.
+
+ [`Lfm2VlProcessor`] offers all the functionalities of [`Lfm2ImageProcessor`] and [`Lfm2Tokenizer`].
+
+ Args:
+ image_processor (`Lfm2VlImageProcessor`):
+ An instance of [`Lfm2VlImageProcessor`]. The image processor is a required input.
+ tokenizer (`PreTrainedTokenizerBase`):
+ An instance of [`PreTrainedTokenizerBase`]. This should correspond with the model's text model. The tokenizer is a required input.
+ chat_template (`str`, *optional*):
+ A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string.
+ use_image_special_tokens (`bool`, *optional*, defaults to `True`):
+ Whether to use image special tokens or not when processing.
+ """
+
+ attributes = ["image_processor", "tokenizer"]
+ image_processor_class = "Lfm2VlImageProcessorFast"
+ tokenizer_class = "AutoTokenizer"
+
+ def __init__(
+ self,
+ image_processor,
+ tokenizer,
+ chat_template: Optional[str] = None,
+ use_image_special_tokens: Optional[bool] = True,
+ **kwargs,
+ ):
+ self.image_token = tokenizer.image_token
+ self.image_token_id = tokenizer.image_token_id
+ self.use_image_special_tokens = use_image_special_tokens
+ self.image_start_token = tokenizer.image_start_token
+ self.image_end_token = tokenizer.image_end_token
+ self.image_thumbnail_token = tokenizer.image_thumbnail
+ super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs)
+
+ def __call__(
+ self,
+ images: Optional[Union[ImageInput, list[ImageInput], list[list[ImageInput]]]] = None,
+ text: Optional[Union[TextInput, list[TextInput]]] = None,
+ **kwargs: Unpack[Lfm2VlProcessorKwargs],
+ ) -> BatchEncoding:
+ """
+ Processes the input prompts and returns a BatchFeature.
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`, *optional*):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. If is of type `list[ImageInput]`, it's assumed that this is for a single prompt i.e. of batch size 1.
+ text (`TextInput`, *optional*):
+ The sequence or batch of sequences to be encoded.
+ Wherever an image token, `` is encountered it is expanded to a proper sequence of image tokens.
+ return_tensors (`Optional[str, TensorType]`, *optional*):
+ If set, will return tensors of a particular framework. See [`PreTrainedTokenizerFast.__call__`] for more
+ information.
+ """
+ if text is None and images is None:
+ raise ValueError("You must provide one of `text` or `images`.")
+
+ if images is not None and text is None:
+ raise ValueError(
+ "You must provide `text` when `images` is provided. Minimal text consists of a single image token."
+ )
+
+ output_kwargs = self._merge_kwargs(
+ Lfm2VlProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ if isinstance(text, str):
+ text = [text]
+ elif not isinstance(text, list) and not isinstance(text[0], str):
+ raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+
+ n_images_in_text = [sample.count(self.image_token) for sample in text]
+ if sum(n_images_in_text) > 0 and images is None:
+ raise ValueError(f"We detected {sum(n_images_in_text)} tokens in the text but no images were passed")
+
+ inputs = {}
+ use_image_special_tokens = output_kwargs["text_kwargs"].pop("use_image_special_tokens")
+
+ if images is not None:
+ images = self.image_processor.fetch_images(images)
+ batched_images = make_nested_list_of_images(images)
+ vision_inputs = self.image_processor(batched_images, **output_kwargs["images_kwargs"])
+
+ n_images_in_images = [len(sublist) for sublist in batched_images]
+ if n_images_in_images != n_images_in_text:
+ raise ValueError(
+ f"The number of images in the text {n_images_in_text} and images {n_images_in_images} should be the same."
+ )
+
+ text = self.expand_text_with_placeholders(
+ text,
+ batched_images,
+ image_rows=vision_inputs.pop("image_rows"),
+ image_cols=vision_inputs.pop("image_cols"),
+ image_sizes=vision_inputs.pop("image_sizes"),
+ use_image_special_tokens=use_image_special_tokens,
+ **output_kwargs["images_kwargs"],
+ )
+ inputs.update(vision_inputs)
+
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+ inputs.update(text_inputs)
+
+ return BatchFeature(inputs, tensor_type=return_tensors)
+
+ def expand_text_with_placeholders(
+ self,
+ text: list[str],
+ images: list[list[ImageInput]],
+ image_rows: list[list[int]],
+ image_cols: list[list[int]],
+ image_sizes: list[list[int]],
+ use_image_special_tokens: bool,
+ **images_kwargs,
+ ):
+ prompt_strings = []
+
+ image_data = iter(zip(*[image_rows, image_cols, image_sizes]))
+ for sample_text, sample_images in zip(text, images):
+ split_sample = sample_text.split(self.image_token)
+ sample_text_with_image_tokens = ""
+ for i, image in enumerate(sample_images):
+ sample_text_with_image_tokens += split_sample[i]
+ if use_image_special_tokens:
+ sample_text_with_image_tokens += self.image_start_token
+
+ rows, cols, image_size = next(image_data)
+ num_thumbnail_tokens, num_tokens_per_tile = self._get_image_num_tokens(image_size, **images_kwargs)
+
+ if rows > 1 or cols > 1:
+ for row in range(rows):
+ for col in range(cols):
+ if use_image_special_tokens:
+ sample_text_with_image_tokens += f"<|img_row_{row + 1}_col_{col + 1}|>"
+ sample_text_with_image_tokens += self.image_token * num_tokens_per_tile
+
+ if num_thumbnail_tokens > 0:
+ if use_image_special_tokens:
+ sample_text_with_image_tokens += self.image_thumbnail_token
+ sample_text_with_image_tokens += self.image_token * num_thumbnail_tokens
+ else:
+ sample_text_with_image_tokens += self.image_token * num_thumbnail_tokens
+
+ if use_image_special_tokens:
+ sample_text_with_image_tokens += self.image_end_token
+
+ sample_text_with_image_tokens += split_sample[i + 1]
+ prompt_strings.append(sample_text_with_image_tokens)
+
+ return prompt_strings
+
+ def _get_image_num_tokens(self, image_size: list[int], **images_kwargs) -> tuple[int, int]:
+ tile_size = images_kwargs.get("tile_size", self.image_processor.tile_size)
+ downsample_factor = images_kwargs.get("downsample_factor", self.image_processor.downsample_factor)
+ encoder_patch_size = images_kwargs.get("encoder_patch_size", self.image_processor.encoder_patch_size)
+ use_thumbnail = images_kwargs.get("use_thumbnail", self.image_processor.use_thumbnail)
+
+ thumbnail_tokens = 0
+ if use_thumbnail:
+ image_height, image_width = image_size
+ num_patches_height = image_height // encoder_patch_size
+ num_patches_width = image_width // encoder_patch_size
+ dwn_num_patches_height = math.ceil(num_patches_height / downsample_factor)
+ dwn_num_patches_width = math.ceil(num_patches_width / downsample_factor)
+ thumbnail_tokens = dwn_num_patches_height * dwn_num_patches_width
+
+ num_patches_tile = tile_size // encoder_patch_size
+ dwn_num_patches_tile = math.ceil(num_patches_tile / downsample_factor)
+ tile_tokens = dwn_num_patches_tile * dwn_num_patches_tile
+
+ return thumbnail_tokens, tile_tokens
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to LFM2Tokeniser's [`~PreTrainedTokenizer.batch_decode`]. Please
+ refer to the docstring of this method for more information.
+ """
+ batched_decode_output = self.tokenizer.batch_decode(*args, **kwargs)
+ return batched_decode_output
+
+ def decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to LFM2Tokeniser's [`~PreTrainedTokenizer.decode`]. Please refer to
+ the docstring of this method for more information.
+ """
+ decode_output = self.tokenizer.decode(*args, **kwargs)
+ return decode_output
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+
+ # LFM2-VL has no dedicated tokenizer class and uses the Base class with default model input names
+ tokenizer_input_names = [name for name in tokenizer_input_names if name != "token_type_ids"]
+ return list(tokenizer_input_names + image_processor_input_names)
+
+
+__all__ = ["Lfm2VlProcessor"]
diff --git a/src/transformers/models/lightglue/modeling_lightglue.py b/src/transformers/models/lightglue/modeling_lightglue.py
index fd460e54d393..8e9faa3e4e04 100644
--- a/src/transformers/models/lightglue/modeling_lightglue.py
+++ b/src/transformers/models/lightglue/modeling_lightglue.py
@@ -628,6 +628,10 @@ def _concat_early_stopped_outputs(
matching_scores,
):
early_stops_indices = torch.stack(early_stops_indices)
+ # Rearrange tensors to have the same order as the input batch
+ ids = torch.arange(early_stops_indices.shape[0])
+ order_indices = early_stops_indices[ids]
+ early_stops_indices = early_stops_indices[order_indices]
matches, final_pruned_keypoints_indices = (
pad_sequence(tensor, batch_first=True, padding_value=-1)
for tensor in [matches, final_pruned_keypoints_indices]
diff --git a/src/transformers/models/lightglue/modular_lightglue.py b/src/transformers/models/lightglue/modular_lightglue.py
index 64c36f21fef9..29441344c9cd 100644
--- a/src/transformers/models/lightglue/modular_lightglue.py
+++ b/src/transformers/models/lightglue/modular_lightglue.py
@@ -786,6 +786,10 @@ def _concat_early_stopped_outputs(
matching_scores,
):
early_stops_indices = torch.stack(early_stops_indices)
+ # Rearrange tensors to have the same order as the input batch
+ ids = torch.arange(early_stops_indices.shape[0])
+ order_indices = early_stops_indices[ids]
+ early_stops_indices = early_stops_indices[order_indices]
matches, final_pruned_keypoints_indices = (
pad_sequence(tensor, batch_first=True, padding_value=-1)
for tensor in [matches, final_pruned_keypoints_indices]
diff --git a/src/transformers/models/llama/convert_llama_weights_to_hf.py b/src/transformers/models/llama/convert_llama_weights_to_hf.py
index 5267bfe9ba49..e63770a154de 100644
--- a/src/transformers/models/llama/convert_llama_weights_to_hf.py
+++ b/src/transformers/models/llama/convert_llama_weights_to_hf.py
@@ -398,7 +398,7 @@ def permute(w, n_heads, dim1=dim, dim2=dim):
max_position_embeddings=max_position_embeddings,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
- tie_word_embeddings=llama_version in ["3.2"],
+ tie_word_embeddings=llama_version == "3.2",
)
config.save_pretrained(tmp_model_path)
@@ -451,7 +451,7 @@ def __init__(self, vocab_file, special_tokens=None, instruct=False, llama_versio
# Prevents a null chat_template, which triggers
# a parsing warning in the Hub.
additional_kwargs = {}
- if instruct or llama_version in ["Guard-3"]:
+ if instruct or llama_version == "Guard-3":
model_id, revision = templates_for_version.get(llama_version, (None, None))
if model_id is not None:
from transformers import AutoTokenizer
diff --git a/src/transformers/models/llama4/convert_llama4_weights_to_hf.py b/src/transformers/models/llama4/convert_llama4_weights_to_hf.py
index 5af63ebc7350..2363a77b1031 100644
--- a/src/transformers/models/llama4/convert_llama4_weights_to_hf.py
+++ b/src/transformers/models/llama4/convert_llama4_weights_to_hf.py
@@ -662,7 +662,7 @@ def update_post_processor(self, tokenizer):
)
-O200K_PATTERN = 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}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+""" # noqa: E501
+O200K_PATTERN = 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}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
def write_tokenizer(args):
diff --git a/src/transformers/models/llama4/image_processing_llama4_fast.py b/src/transformers/models/llama4/image_processing_llama4_fast.py
index 946fdde0a643..6506d5749d94 100644
--- a/src/transformers/models/llama4/image_processing_llama4_fast.py
+++ b/src/transformers/models/llama4/image_processing_llama4_fast.py
@@ -20,6 +20,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -33,16 +34,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
def get_factors(dividend: int) -> set[int]:
"""
Calculate all factors of a given number, i.e. a divisor that leaves
diff --git a/src/transformers/models/llava/image_processing_llava.py b/src/transformers/models/llava/image_processing_llava.py
index d3aa81303bb8..5420d6fe2918 100644
--- a/src/transformers/models/llava/image_processing_llava.py
+++ b/src/transformers/models/llava/image_processing_llava.py
@@ -154,7 +154,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
diff --git a/src/transformers/models/llava/image_processing_llava_fast.py b/src/transformers/models/llava/image_processing_llava_fast.py
index 41bb94f5b7e0..596070040549 100644
--- a/src/transformers/models/llava/image_processing_llava_fast.py
+++ b/src/transformers/models/llava/image_processing_llava_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -38,16 +39,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class LlavaFastImageProcessorKwargs(DefaultFastImageProcessorKwargs): ...
diff --git a/src/transformers/models/llava_next/image_processing_llava_next.py b/src/transformers/models/llava_next/image_processing_llava_next.py
index 3887c9c7ad4b..350ce9db7dc6 100644
--- a/src/transformers/models/llava_next/image_processing_llava_next.py
+++ b/src/transformers/models/llava_next/image_processing_llava_next.py
@@ -58,12 +58,12 @@
from PIL import Image
-def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.array]:
+def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.ndarray]:
"""
Divides an image into patches of a specified size.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The input image.
patch_size (`int`):
The size of each patch.
@@ -71,7 +71,7 @@ def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) ->
The channel dimension format of the input image.
Returns:
- list: A list of np.array representing the patches.
+ list: A list of np.ndarray representing the patches.
"""
patches = []
height, width = get_image_size(image, channel_dim=input_data_format)
@@ -86,7 +86,7 @@ def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) ->
return patches
-def expand_to_square(image: np.ndarray, background_color, input_data_format) -> np.array:
+def expand_to_square(image: np.ndarray, background_color, input_data_format) -> np.ndarray:
"""
Expands an image to a square by adding a background color.
"""
@@ -400,12 +400,12 @@ def _preprocess(
def _resize_for_patching(
self, image: np.ndarray, target_resolution: tuple, resample, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Resizes an image to a target resolution while maintaining aspect ratio.
Args:
- image (np.array):
+ image (np.ndarray):
The input image.
target_resolution (tuple):
The target resolution (height, width) of the image.
@@ -415,7 +415,7 @@ def _resize_for_patching(
The channel dimension format of the input image.
Returns:
- np.array: The resized and padded image.
+ np.ndarray: The resized and padded image.
"""
new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
@@ -433,7 +433,7 @@ def _get_padding_size(self, original_resolution: tuple, target_resolution: tuple
def _pad_for_patching(
self, image: np.ndarray, target_resolution: tuple, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pad an image to a target resolution while maintaining aspect ratio.
"""
@@ -453,12 +453,12 @@ def get_image_patches(
resample: PILImageResampling,
data_format: ChannelDimension,
input_data_format: ChannelDimension,
- ) -> list[np.array]:
+ ) -> list[np.ndarray]:
"""
Process an image with variable resolutions by dividing it into patches.
Args:
- image (np.array):
+ image (np.ndarray):
The input image to be processed.
grid_pinpoints (List):
A string representation of a list of possible resolutions.
@@ -474,7 +474,7 @@ def get_image_patches(
The channel dimension format of the input image.
Returns:
- list[np.array]: A list of NumPy arrays containing the processed image patches.
+ list[np.ndarray]: A list of NumPy arrays containing the processed image patches.
"""
if not isinstance(grid_pinpoints, list):
raise TypeError("grid_pinpoints must be a list of possible resolutions.")
diff --git a/src/transformers/models/llava_next/image_processing_llava_next_fast.py b/src/transformers/models/llava_next/image_processing_llava_next_fast.py
index b502d98d6ac3..df20e2b90e83 100644
--- a/src/transformers/models/llava_next/image_processing_llava_next_fast.py
+++ b/src/transformers/models/llava_next/image_processing_llava_next_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_patch_output_size, select_best_resolution
from ...image_processing_utils_fast import (
@@ -39,16 +40,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class LlavaNextFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
image_grid_pinpoints (`list[list[int]]`, *optional*):
diff --git a/src/transformers/models/llava_next/modeling_llava_next.py b/src/transformers/models/llava_next/modeling_llava_next.py
index 8cca63f4a66c..a75b4b798107 100644
--- a/src/transformers/models/llava_next/modeling_llava_next.py
+++ b/src/transformers/models/llava_next/modeling_llava_next.py
@@ -409,8 +409,6 @@ def get_image_features(
if vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_image_feature = selected_image_feature
image_features = self.multi_modal_projector(selected_image_feature)
image_features = torch.split(image_features, image_num_patches, dim=0)
diff --git a/src/transformers/models/llava_next_video/modeling_llava_next_video.py b/src/transformers/models/llava_next_video/modeling_llava_next_video.py
index 3ef172962c2c..9e3b15cea548 100644
--- a/src/transformers/models/llava_next_video/modeling_llava_next_video.py
+++ b/src/transformers/models/llava_next_video/modeling_llava_next_video.py
@@ -461,8 +461,6 @@ def get_image_features(
if vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_image_feature = selected_image_feature
image_features = self.multi_modal_projector(selected_image_feature)
image_features = torch.split(image_features, image_num_patches, dim=0)
@@ -659,8 +657,6 @@ def get_video_features(
if vision_feature_select_strategy == "default":
selected_video_features = selected_video_features[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_video_features = selected_video_features
# Same as image features except that video has pooling layer
video_features = self.vision_resampler(selected_video_features)
diff --git a/src/transformers/models/llava_next_video/modular_llava_next_video.py b/src/transformers/models/llava_next_video/modular_llava_next_video.py
index 73745f435b7d..7eda08ffa0bd 100644
--- a/src/transformers/models/llava_next_video/modular_llava_next_video.py
+++ b/src/transformers/models/llava_next_video/modular_llava_next_video.py
@@ -327,8 +327,6 @@ def get_image_features(
if vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_image_feature = selected_image_feature
image_features = self.multi_modal_projector(selected_image_feature)
image_features = torch.split(image_features, image_num_patches, dim=0)
@@ -386,8 +384,6 @@ def get_video_features(
if vision_feature_select_strategy == "default":
selected_video_features = selected_video_features[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_video_features = selected_video_features
# Same as image features except that video has pooling layer
video_features = self.vision_resampler(selected_video_features)
diff --git a/src/transformers/models/llava_onevision/image_processing_llava_onevision.py b/src/transformers/models/llava_onevision/image_processing_llava_onevision.py
index 837eda460802..836a1984a522 100644
--- a/src/transformers/models/llava_onevision/image_processing_llava_onevision.py
+++ b/src/transformers/models/llava_onevision/image_processing_llava_onevision.py
@@ -58,12 +58,12 @@
# Copied from transformers.models.llava_next.image_processing_llava_next.divide_to_patches
-def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.array]:
+def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) -> list[np.ndarray]:
"""
Divides an image into patches of a specified size.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The input image.
patch_size (`int`):
The size of each patch.
@@ -71,7 +71,7 @@ def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) ->
The channel dimension format of the input image.
Returns:
- list: A list of np.array representing the patches.
+ list: A list of np.ndarray representing the patches.
"""
patches = []
height, width = get_image_size(image, channel_dim=input_data_format)
@@ -87,7 +87,7 @@ def divide_to_patches(image: np.ndarray, patch_size: int, input_data_format) ->
# Copied from transformers.models.llava_next.image_processing_llava_next.expand_to_square
-def expand_to_square(image: np.ndarray, background_color, input_data_format) -> np.array:
+def expand_to_square(image: np.ndarray, background_color, input_data_format) -> np.ndarray:
"""
Expands an image to a square by adding a background color.
"""
@@ -292,12 +292,12 @@ def pad(
# Copied from transformers.models.llava_next.image_processing_llava_next.LlavaNextImageProcessor._resize_for_patching
def _resize_for_patching(
self, image: np.ndarray, target_resolution: tuple, resample, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Resizes an image to a target resolution while maintaining aspect ratio.
Args:
- image (np.array):
+ image (np.ndarray):
The input image.
target_resolution (tuple):
The target resolution (height, width) of the image.
@@ -307,7 +307,7 @@ def _resize_for_patching(
The channel dimension format of the input image.
Returns:
- np.array: The resized and padded image.
+ np.ndarray: The resized and padded image.
"""
new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
@@ -327,7 +327,7 @@ def _get_padding_size(self, original_resolution: tuple, target_resolution: tuple
# Copied from transformers.models.llava_next.image_processing_llava_next.LlavaNextImageProcessor._pad_for_patching
def _pad_for_patching(
self, image: np.ndarray, target_resolution: tuple, input_data_format: ChannelDimension
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pad an image to a target resolution while maintaining aspect ratio.
"""
@@ -348,12 +348,12 @@ def get_image_patches(
resample: PILImageResampling,
data_format: ChannelDimension,
input_data_format: ChannelDimension,
- ) -> list[np.array]:
+ ) -> list[np.ndarray]:
"""
Process an image with variable resolutions by dividing it into patches.
Args:
- image (np.array):
+ image (np.ndarray):
The input image to be processed.
grid_pinpoints (List):
A string representation of a list of possible resolutions.
@@ -369,7 +369,7 @@ def get_image_patches(
The channel dimension format of the input image.
Returns:
- list[np.array]: A list of NumPy arrays containing the processed image patches.
+ list[np.ndarray]: A list of NumPy arrays containing the processed image patches.
"""
if not isinstance(grid_pinpoints, list):
raise TypeError("grid_pinpoints must be a list of possible resolutions.")
@@ -450,7 +450,7 @@ def pad_to_square(
background_color: Union[int, tuple[int, int, int]] = 0,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Pads an image to a square based on the longest edge.
diff --git a/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py b/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py
index 4392d64e9ebf..11872cb67bf3 100644
--- a/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py
+++ b/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py
@@ -22,6 +22,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_patch_output_size, select_best_resolution
from ...image_processing_utils_fast import (
@@ -41,13 +42,7 @@
get_image_size,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import TensorType, auto_docstring
class LlavaOnevisionFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
diff --git a/src/transformers/models/llava_onevision/modeling_llava_onevision.py b/src/transformers/models/llava_onevision/modeling_llava_onevision.py
index eae6e3046f94..727655374574 100644
--- a/src/transformers/models/llava_onevision/modeling_llava_onevision.py
+++ b/src/transformers/models/llava_onevision/modeling_llava_onevision.py
@@ -432,8 +432,6 @@ def get_image_features(
if vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_image_feature = selected_image_feature
image_features = self.multi_modal_projector(selected_image_feature)
image_features = torch.split(image_features, image_num_patches, dim=0)
@@ -633,8 +631,6 @@ def get_video_features(
if vision_feature_select_strategy == "default":
selected_video_feature = selected_video_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_video_feature = selected_video_feature
video_features = self.multi_modal_projector(selected_video_feature)
video_features = self.apply_pooling(video_features)
diff --git a/src/transformers/models/llava_onevision/modular_llava_onevision.py b/src/transformers/models/llava_onevision/modular_llava_onevision.py
index 21688e7763bf..b4f64dee8e04 100644
--- a/src/transformers/models/llava_onevision/modular_llava_onevision.py
+++ b/src/transformers/models/llava_onevision/modular_llava_onevision.py
@@ -18,6 +18,7 @@
import torch
from torch import nn
+from torchvision.transforms.v2 import functional as F
from transformers.models.llava_next.image_processing_llava_next_fast import LlavaNextImageProcessorFast
from transformers.models.llava_next_video.modeling_llava_next_video import (
@@ -50,16 +51,10 @@
TensorType,
auto_docstring,
can_return_tuple,
- is_torchvision_v2_available,
logging,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -409,8 +404,6 @@ def get_image_features(
if vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_image_feature = selected_image_feature
image_features = self.multi_modal_projector(selected_image_feature)
image_features = torch.split(image_features, image_num_patches, dim=0)
@@ -459,8 +452,6 @@ def get_video_features(
if vision_feature_select_strategy == "default":
selected_video_feature = selected_video_feature[:, 1:]
- elif vision_feature_select_strategy == "full":
- selected_video_feature = selected_video_feature
video_features = self.multi_modal_projector(selected_video_feature)
video_features = self.apply_pooling(video_features)
diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py
index 87e812852b37..4681cfb60e53 100644
--- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py
+++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py
@@ -534,7 +534,7 @@ def __init__(self, config):
self.rotary_emb = LongcatFlashRotaryEmbedding(config=config)
self.gradient_checkpointing = False
# Each layer above has 2 sublayers, config hack to have a correct cache (to avoid a checkpoint change)
- self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used) # noqa
+ self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used)
self.config.num_hidden_layers = 2 * config.num_layers
diff --git a/src/transformers/models/longcat_flash/modular_longcat_flash.py b/src/transformers/models/longcat_flash/modular_longcat_flash.py
index f58ca870aefc..60c93239d2c4 100644
--- a/src/transformers/models/longcat_flash/modular_longcat_flash.py
+++ b/src/transformers/models/longcat_flash/modular_longcat_flash.py
@@ -300,7 +300,7 @@ def __init__(self, config):
[LongcatFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers)]
)
# Each layer above has 2 sublayers, config hack to have a correct cache (to avoid a checkpoint change)
- self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used) # noqa
+ self.head_dim = config.head_dim # For CI happiness (we didn't convert so head_dim is not directly used)
self.config.num_hidden_layers = 2 * config.num_layers
self.norm = LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
diff --git a/src/transformers/models/longt5/modeling_longt5.py b/src/transformers/models/longt5/modeling_longt5.py
index 4e84a1550349..ea6ab0cfff35 100644
--- a/src/transformers/models/longt5/modeling_longt5.py
+++ b/src/transformers/models/longt5/modeling_longt5.py
@@ -250,7 +250,7 @@ def forward(self, hidden_states):
try:
from apex.normalization import FusedRMSNorm
- LongT5LayerNorm = FusedRMSNorm # noqa
+ LongT5LayerNorm = FusedRMSNorm
logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of LongT5LayerNorm")
except ImportError:
@@ -1270,6 +1270,35 @@ def dummy_inputs(self):
}
return dummy_inputs
+ def _try_load_missing_tied_module(self, key):
+ module = self
+ key = key.removesuffix(".weight")
+ for sub_key in key.split("."):
+ if not hasattr(module, sub_key):
+ return
+ module = getattr(module, sub_key)
+
+ self._tie_or_clone_weights(module, self.shared)
+
+ @classmethod
+ def from_pretrained(self, *args, **kwargs):
+ requested_loading_info = kwargs.get("output_loading_info", False)
+ kwargs["output_loading_info"] = True
+ model, loading_info = super().from_pretrained(*args, **kwargs)
+ missing_keys = loading_info.get("missing_keys", [])
+
+ if hasattr(model, "shared") and hasattr(model, "_tied_weights_keys"):
+ for missing_key in missing_keys:
+ logger.warning(
+ f"Recovering a missing tied weight {missing_key} from a legacy LongT5 checkpoint. "
+ f"Consider saving {missing_key} in your checkpoint or updating the config (tie_word_embeddings=true)."
+ )
+ model._try_load_missing_tied_module(missing_key)
+
+ if requested_loading_info:
+ return model, loading_info
+ return model
+
def _init_weights(self, module):
"""Initialize the weights"""
factor = self.config.initializer_factor # Used for testing weights initialization
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index a423c5b42fbd..bb24e2422d32 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -286,7 +286,7 @@ def __init__(self, config: Mamba2Config, layer_idx: int):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/mask2former/image_processing_mask2former.py b/src/transformers/models/mask2former/image_processing_mask2former.py
index a0c369722b54..06fe78e82e9e 100644
--- a/src/transformers/models/mask2former/image_processing_mask2former.py
+++ b/src/transformers/models/mask2former/image_processing_mask2former.py
@@ -303,7 +303,7 @@ def compute_segments(
# TODO: (Amy) Move to image_transforms
# Copied from transformers.models.maskformer.image_processing_maskformer.convert_segmentation_map_to_binary_masks
def convert_segmentation_map_to_binary_masks(
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
@@ -582,7 +582,7 @@ def rescale(
# Copied from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.convert_segmentation_map_to_binary_masks
def convert_segmentation_map_to_binary_masks(
self,
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
diff --git a/src/transformers/models/mask2former/image_processing_mask2former_fast.py b/src/transformers/models/mask2former/image_processing_mask2former_fast.py
index a5d662288119..58dbb09d6319 100644
--- a/src/transformers/models/mask2former/image_processing_mask2former_fast.py
+++ b/src/transformers/models/mask2former/image_processing_mask2former_fast.py
@@ -23,6 +23,7 @@
import torch
from torch import nn
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -42,7 +43,7 @@
PILImageResampling,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, logging
+from ...utils import TensorType, auto_docstring, logging
from .image_processing_mask2former import (
compute_segments,
convert_segmentation_to_rle,
@@ -51,11 +52,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -348,9 +344,7 @@ def _preprocess(
image=grouped_segmentation_maps[shape],
size=size,
size_divisor=size_divisor,
- interpolation=F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ interpolation=F.InterpolationMode.NEAREST_EXACT,
)
resized_images_grouped[shape] = stacked_images
if segmentation_maps is not None:
diff --git a/src/transformers/models/mask2former/modeling_mask2former.py b/src/transformers/models/mask2former/modeling_mask2former.py
index e8c3d2344b8d..553700465f3c 100644
--- a/src/transformers/models/mask2former/modeling_mask2former.py
+++ b/src/transformers/models/mask2former/modeling_mask2former.py
@@ -783,7 +783,7 @@ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> tor
"""
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 = 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():
diff --git a/src/transformers/models/maskformer/image_processing_maskformer.py b/src/transformers/models/maskformer/image_processing_maskformer.py
index 9ce33846170e..c2f9aee70167 100644
--- a/src/transformers/models/maskformer/image_processing_maskformer.py
+++ b/src/transformers/models/maskformer/image_processing_maskformer.py
@@ -308,7 +308,7 @@ def compute_segments(
# TODO: (Amy) Move to image_transforms
def convert_segmentation_map_to_binary_masks(
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
@@ -585,7 +585,7 @@ def rescale(
def convert_segmentation_map_to_binary_masks(
self,
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
diff --git a/src/transformers/models/maskformer/image_processing_maskformer_fast.py b/src/transformers/models/maskformer/image_processing_maskformer_fast.py
index ab6411f1bb3f..9e15486cfa35 100644
--- a/src/transformers/models/maskformer/image_processing_maskformer_fast.py
+++ b/src/transformers/models/maskformer/image_processing_maskformer_fast.py
@@ -20,6 +20,7 @@
import torch
from torch import nn
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -42,7 +43,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from .image_processing_maskformer import (
@@ -53,11 +53,6 @@
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -354,9 +349,7 @@ def _preprocess(
image=grouped_segmentation_maps[shape],
size=size,
size_divisor=size_divisor,
- interpolation=F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ interpolation=F.InterpolationMode.NEAREST_EXACT,
)
resized_images_grouped[shape] = stacked_images
if segmentation_maps is not None:
diff --git a/src/transformers/models/maskformer/modeling_maskformer.py b/src/transformers/models/maskformer/modeling_maskformer.py
index 9e1c0072425b..772f0a9fad0a 100644
--- a/src/transformers/models/maskformer/modeling_maskformer.py
+++ b/src/transformers/models/maskformer/modeling_maskformer.py
@@ -1088,7 +1088,7 @@ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> tor
"""
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 = 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():
diff --git a/src/transformers/models/metaclip_2/configuration_metaclip_2.py b/src/transformers/models/metaclip_2/configuration_metaclip_2.py
index a0cec0f3c5b3..4ad1bcde0daa 100644
--- a/src/transformers/models/metaclip_2/configuration_metaclip_2.py
+++ b/src/transformers/models/metaclip_2/configuration_metaclip_2.py
@@ -277,7 +277,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -309,7 +309,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/metaclip_2/convert_metaclip_2_to_hf.py b/src/transformers/models/metaclip_2/convert_metaclip_2_to_hf.py
index 21a0a1462fff..55aa6f099abf 100644
--- a/src/transformers/models/metaclip_2/convert_metaclip_2_to_hf.py
+++ b/src/transformers/models/metaclip_2/convert_metaclip_2_to_hf.py
@@ -26,6 +26,7 @@
# Import MetaCLIP modules
from src.mini_clip.factory import create_model_and_transforms
+
from transformers import (
AutoTokenizer,
CLIPImageProcessor,
diff --git a/src/transformers/models/mistral3/modular_mistral3.py b/src/transformers/models/mistral3/modular_mistral3.py
index 213ab98fe902..6bc499d21453 100644
--- a/src/transformers/models/mistral3/modular_mistral3.py
+++ b/src/transformers/models/mistral3/modular_mistral3.py
@@ -332,6 +332,6 @@ def forward(
__all__ = [
"Mistral3Model",
- "Mistral3PreTrainedModel", # noqa
+ "Mistral3PreTrainedModel",
"Mistral3ForConditionalGeneration",
]
diff --git a/src/transformers/models/mllama/convert_mllama_weights_to_hf.py b/src/transformers/models/mllama/convert_mllama_weights_to_hf.py
index 9465e410be70..c773d0514f81 100644
--- a/src/transformers/models/mllama/convert_mllama_weights_to_hf.py
+++ b/src/transformers/models/mllama/convert_mllama_weights_to_hf.py
@@ -496,7 +496,7 @@ def __init__(
def write_tokenizer(tokenizer_path: str, save_dir: str, instruct: bool = False):
model_max_length = CONTEXT_LENGTH
- pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: W605
+ pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
# Special tokens
num_reserved_special_tokens = 256
diff --git a/src/transformers/models/mllama/image_processing_mllama.py b/src/transformers/models/mllama/image_processing_mllama.py
index ba1a596aa459..4f18b65d1419 100644
--- a/src/transformers/models/mllama/image_processing_mllama.py
+++ b/src/transformers/models/mllama/image_processing_mllama.py
@@ -327,7 +327,7 @@ def build_aspect_ratio_mask(aspect_ratios: list[list[tuple[int, int]]], max_imag
The mask contains 1s for valid tiles and 0s for padding.
"""
batch_size = len(aspect_ratios)
- max_num_images = max([len(row) for row in aspect_ratios])
+ max_num_images = max(len(row) for row in aspect_ratios)
aspect_ratio_mask = np.zeros((batch_size, max_num_images, max_image_tiles), dtype=np.int64)
@@ -374,7 +374,7 @@ def pack_images(
# Determine output shape
batch_size = len(batch_images)
- max_num_images = max([len(images) for images in batch_images])
+ max_num_images = max(len(images) for images in batch_images)
shapes = [image.shape for images in batch_images for image in images]
_, channels, tile_height, tile_width = shapes[0]
@@ -412,7 +412,7 @@ def pack_aspect_ratios(aspect_ratios: list[list[tuple[int, int]]], pad_value: in
The aspect ratios stacked into a numpy array with shape (batch_size, max_num_images, 2).
"""
batch_size = len(aspect_ratios)
- max_num_images = max([len(row) for row in aspect_ratios])
+ max_num_images = max(len(row) for row in aspect_ratios)
aspect_ratios_stacked = np.full((batch_size, max_num_images, 2), pad_value, dtype=np.int64)
for i, row in enumerate(aspect_ratios):
@@ -442,7 +442,7 @@ def convert_aspect_ratios_to_ids(aspect_ratios: list[list[tuple[int, int]]], max
"""
batch_size = len(aspect_ratios)
- max_num_images = max([len(row) for row in aspect_ratios])
+ max_num_images = max(len(row) for row in aspect_ratios)
supported_aspect_ratios = get_all_supported_aspect_ratios(max_image_tiles)
aspect_ratios_ids = np.zeros((batch_size, max_num_images), dtype=np.int64)
diff --git a/src/transformers/models/mllama/processing_mllama.py b/src/transformers/models/mllama/processing_mllama.py
index 0dae7c834303..d571286fbf82 100644
--- a/src/transformers/models/mllama/processing_mllama.py
+++ b/src/transformers/models/mllama/processing_mllama.py
@@ -117,7 +117,7 @@ def convert_sparse_cross_attention_mask_to_dense(
"""
batch_size = len(cross_attention_token_mask)
- max_num_images = max([len(masks) for masks in cross_attention_token_mask])
+ max_num_images = max(len(masks) for masks in cross_attention_token_mask)
cross_attention_mask = np.zeros(
shape=(batch_size, length, max_num_images, max_num_tiles),
diff --git a/src/transformers/models/mobilenet_v2/image_processing_mobilenet_v2_fast.py b/src/transformers/models/mobilenet_v2/image_processing_mobilenet_v2_fast.py
index 97ca39da78bf..fd3510c53c4d 100644
--- a/src/transformers/models/mobilenet_v2/image_processing_mobilenet_v2_fast.py
+++ b/src/transformers/models/mobilenet_v2/image_processing_mobilenet_v2_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -38,16 +39,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class MobileNetV2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):
@@ -134,9 +128,7 @@ def _preprocess_image_like_inputs(
"do_normalize": False,
"do_rescale": False,
# Nearest interpolation is used for segmentation maps instead of BILINEAR.
- "interpolation": F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ "interpolation": F.InterpolationMode.NEAREST_EXACT,
}
)
diff --git a/src/transformers/models/mobilevit/image_processing_mobilevit_fast.py b/src/transformers/models/mobilevit/image_processing_mobilevit_fast.py
index 71c8ababba36..fab16ecfdc87 100644
--- a/src/transformers/models/mobilevit/image_processing_mobilevit_fast.py
+++ b/src/transformers/models/mobilevit/image_processing_mobilevit_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -36,16 +37,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class MobileVitFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
do_flip_channel_order (`bool`, *optional*, defaults to `self.do_flip_channel_order`):
@@ -135,9 +129,7 @@ def _preprocess_image_like_inputs(
"do_rescale": False,
"do_flip_channel_order": False,
# Nearest interpolation is used for segmentation maps instead of BILINEAR.
- "interpolation": F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ "interpolation": F.InterpolationMode.NEAREST_EXACT,
}
)
diff --git a/src/transformers/models/modernbert/modeling_modernbert.py b/src/transformers/models/modernbert/modeling_modernbert.py
index d36b5a9b9485..00fbe19c3a63 100644
--- a/src/transformers/models/modernbert/modeling_modernbert.py
+++ b/src/transformers/models/modernbert/modeling_modernbert.py
@@ -893,6 +893,15 @@ def forward(
_pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)
for hs in all_hidden_states
)
+ # If the attention implementation is FA2 and there is no need for repadding, there might still be the batch
+ # dimension missing
+ elif (
+ self.config._attn_implementation == "flash_attention_2"
+ and all_hidden_states is not None
+ and all_hidden_states[-1].dim() == 2
+ ):
+ hidden_states = hidden_states.unsqueeze(0)
+ all_hidden_states = tuple(hs.unsqueeze(0) for hs in all_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)
@@ -1075,8 +1084,19 @@ def forward(
loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs)
if self.config._attn_implementation == "flash_attention_2":
+ # Logits padding
with nullcontext() if self.config.repad_logits_with_grad or labels is None else torch.no_grad():
logits = _pad_modernbert_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)
+ # Hidden states padding
+ if getattr(outputs, "hidden_states", None) is not None:
+ padded_hidden_states = []
+ for hs in outputs.hidden_states:
+ if hs.dim() == 3 and hs.shape[0] == 1:
+ hs = hs.squeeze(0)
+ padded_hidden_states.append(
+ _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)
+ )
+ outputs.hidden_states = tuple(padded_hidden_states)
if not return_dict:
output = (logits,)
@@ -1499,14 +1519,24 @@ def forward(
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
- last_hidden_state = outputs[0]
+ last_hidden_state = outputs[0] # shape (num_choices, seq_len, hidden_size)
+ # If classifier_pooling is "cls", isolate the token
if self.config.classifier_pooling == "cls":
- last_hidden_state = last_hidden_state[:, 0]
+ indices_0 = torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device)
+ # for left or right padding, is the first non-pad token
+ if attention_mask is not None:
+ cls_mask = attention_mask.argmax(dim=-1).to(last_hidden_state.device)
+ # if no pad, is the first token
+ else:
+ cls_mask = torch.tensor(0, dtype=torch.long, device=last_hidden_state.device)
+ # extract the token for the logits
+ last_hidden_state = last_hidden_state[indices_0, cls_mask]
+
+ # If classifier_pooling is "mean", pool the hidden states by averaging over the sequence length
elif self.config.classifier_pooling == "mean":
- last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(
- dim=1, keepdim=True
- )
+ num_non_pad_tokens = attention_mask.sum(dim=1, keepdim=True)
+ last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / num_non_pad_tokens
pooled_output = self.head(last_hidden_state)
pooled_output = self.drop(pooled_output)
diff --git a/src/transformers/models/modernbert/modular_modernbert.py b/src/transformers/models/modernbert/modular_modernbert.py
index 276a754cc101..5ac298f09596 100644
--- a/src/transformers/models/modernbert/modular_modernbert.py
+++ b/src/transformers/models/modernbert/modular_modernbert.py
@@ -1018,6 +1018,15 @@ def forward(
_pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)
for hs in all_hidden_states
)
+ # If the attention implementation is FA2 and there is no need for repadding, there might still be the batch
+ # dimension missing
+ elif (
+ self.config._attn_implementation == "flash_attention_2"
+ and all_hidden_states is not None
+ and all_hidden_states[-1].dim() == 2
+ ):
+ hidden_states = hidden_states.unsqueeze(0)
+ all_hidden_states = tuple(hs.unsqueeze(0) for hs in all_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)
@@ -1200,8 +1209,19 @@ def forward(
loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs)
if self.config._attn_implementation == "flash_attention_2":
+ # Logits padding
with nullcontext() if self.config.repad_logits_with_grad or labels is None else torch.no_grad():
logits = _pad_modernbert_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)
+ # Hidden states padding
+ if getattr(outputs, "hidden_states", None) is not None:
+ padded_hidden_states = []
+ for hs in outputs.hidden_states:
+ if hs.dim() == 3 and hs.shape[0] == 1:
+ hs = hs.squeeze(0)
+ padded_hidden_states.append(
+ _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)
+ )
+ outputs.hidden_states = tuple(padded_hidden_states)
if not return_dict:
output = (logits,)
@@ -1624,14 +1644,24 @@ def forward(
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
- last_hidden_state = outputs[0]
+ last_hidden_state = outputs[0] # shape (num_choices, seq_len, hidden_size)
+ # If classifier_pooling is "cls", isolate the token
if self.config.classifier_pooling == "cls":
- last_hidden_state = last_hidden_state[:, 0]
+ indices_0 = torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device)
+ # for left or right padding, is the first non-pad token
+ if attention_mask is not None:
+ cls_mask = attention_mask.argmax(dim=-1).to(last_hidden_state.device)
+ # if no pad, is the first token
+ else:
+ cls_mask = torch.tensor(0, dtype=torch.long, device=last_hidden_state.device)
+ # extract the token for the logits
+ last_hidden_state = last_hidden_state[indices_0, cls_mask]
+
+ # If classifier_pooling is "mean", pool the hidden states by averaging over the sequence length
elif self.config.classifier_pooling == "mean":
- last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(
- dim=1, keepdim=True
- )
+ num_non_pad_tokens = attention_mask.sum(dim=1, keepdim=True)
+ last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / num_non_pad_tokens
pooled_output = self.head(last_hidden_state)
pooled_output = self.drop(pooled_output)
diff --git a/src/transformers/models/moshi/modeling_moshi.py b/src/transformers/models/moshi/modeling_moshi.py
index 7546fc90e542..af9138c5f0c9 100644
--- a/src/transformers/models/moshi/modeling_moshi.py
+++ b/src/transformers/models/moshi/modeling_moshi.py
@@ -1708,7 +1708,7 @@ def forward(
if audio_codes is not None:
audio_inputs_embeds = sum(
- [self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])]
+ self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])
)
inputs_embeds = (
audio_inputs_embeds
@@ -1878,20 +1878,18 @@ def _prepare_inputs_embeds_for_generation(
if user_audio_codes is not None and moshi_audio_codes is not None:
audio_codes = torch.cat([moshi_audio_codes, user_audio_codes], dim=1)
audio_inputs_embeds = sum(
- [self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])]
+ self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])
)
elif moshi_audio_codes is not None:
audio_codes = moshi_audio_codes
audio_inputs_embeds = sum(
- [self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])]
+ self.embed_tokens[codebook](audio_codes[:, codebook]) for codebook in range(audio_codes.shape[1])
)
elif user_audio_codes is not None:
audio_codes = user_audio_codes
audio_inputs_embeds = sum(
- [
- self.embed_tokens[codebook](audio_codes[:, codebook + self.num_codebooks])
- for codebook in range(audio_codes.shape[1])
- ]
+ self.embed_tokens[codebook](audio_codes[:, codebook + self.num_codebooks])
+ for codebook in range(audio_codes.shape[1])
)
if input_ids is not None:
diff --git a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py
index e7237157e156..e432dc3ff625 100644
--- a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py
+++ b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py
@@ -533,7 +533,7 @@ def forward(
past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
if inputs_embeds is None:
- inputs_embeds = sum([self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks)])
+ inputs_embeds = sum(self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks))
if encoder_hidden_states is not None:
# take care of attention masks
diff --git a/src/transformers/models/nllb_moe/convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py b/src/transformers/models/nllb_moe/convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py
index ef2e3d0d90dd..c6dec96b8473 100644
--- a/src/transformers/models/nllb_moe/convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py
+++ b/src/transformers/models/nllb_moe/convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py
@@ -85,7 +85,7 @@ def shard_on_the_fly(switch_checkpoint_path, dump_path, num_experts, dtype, weig
)
torch.save(expert_state, save_path)
sharded_state_dicts.append(expert_state.keys())
- total_size += sum([value.numel() for key, value in expert_state.items()]) * (
+ total_size += sum(value.numel() for key, value in expert_state.items()) * (
expert_state[list(expert_state)[0]].element_size()
)
diff --git a/src/transformers/models/nougat/image_processing_nougat.py b/src/transformers/models/nougat/image_processing_nougat.py
index 0c0a51464b43..660226a7d6ee 100644
--- a/src/transformers/models/nougat/image_processing_nougat.py
+++ b/src/transformers/models/nougat/image_processing_nougat.py
@@ -144,13 +144,13 @@ def crop_margin(
gray_threshold: int = 200,
data_format: Optional[ChannelDimension] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Crops the margin of the image. Gray pixels are considered margin (i.e., pixels with a value below the
threshold).
Args:
- image (`np.array`):
+ image (`np.ndarray`):
The image to be cropped.
gray_threshold (`int`, *optional*, defaults to `200`)
Value below which pixels are considered to be gray.
diff --git a/src/transformers/models/nougat/image_processing_nougat_fast.py b/src/transformers/models/nougat/image_processing_nougat_fast.py
index d6579029e4f5..15cee9051082 100644
--- a/src/transformers/models/nougat/image_processing_nougat_fast.py
+++ b/src/transformers/models/nougat/image_processing_nougat_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -40,16 +41,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class NougatFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
Args:
diff --git a/src/transformers/models/nougat/processing_nougat.py b/src/transformers/models/nougat/processing_nougat.py
index b50c23e4c4d4..2815dcfa7b7a 100644
--- a/src/transformers/models/nougat/processing_nougat.py
+++ b/src/transformers/models/nougat/processing_nougat.py
@@ -65,13 +65,13 @@ def __call__(
data_format: Optional["ChannelDimension"] = "channels_first", # noqa: F821
input_data_format: Optional[Union[str, "ChannelDimension"]] = None, # noqa: F821
text_pair: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] = None,
- text_target: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+ text_target: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] = None,
text_pair_target: Optional[
Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]
] = None,
add_special_tokens: bool = True,
padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = None,
+ truncation: Optional[Union[bool, str, TruncationStrategy]] = None,
max_length: Optional[int] = None,
stride: int = 0,
is_split_into_words: bool = False,
diff --git a/src/transformers/models/nougat/tokenization_nougat_fast.py b/src/transformers/models/nougat/tokenization_nougat_fast.py
index 266198e865df..198da79302af 100644
--- a/src/transformers/models/nougat/tokenization_nougat_fast.py
+++ b/src/transformers/models/nougat/tokenization_nougat_fast.py
@@ -68,15 +68,15 @@ def markdown_compatible(text: str) -> str:
"""
# equation tag
# Replace lines that start with a pattern like (decimal) \[some text\] with \[[some text] \tag{decimal}\].
- text = re.sub(r"^\(([\d.]+[a-zA-Z]?)\) \\\[(.+?)\\\]$", r"\[\2 \\tag{\1}\]", text, flags=re.M)
+ text = re.sub(r"^\(([\d.]+[a-zA-Z]?)\) \\\[(.+?)\\\]$", r"\[\2 \\tag{\1}\]", text, flags=re.MULTILINE)
# Replace lines that start with a pattern like \[some text\] (decimal) with \[[some text] \tag{decimal}\].
- text = re.sub(r"^\\\[(.+?)\\\] \(([\d.]+[a-zA-Z]?)\)$", r"\[\1 \\tag{\2}\]", text, flags=re.M)
+ text = re.sub(r"^\\\[(.+?)\\\] \(([\d.]+[a-zA-Z]?)\)$", r"\[\1 \\tag{\2}\]", text, flags=re.MULTILINE)
# Replace lines that start with a pattern like \[some text\] (digits) \[another text\] with \[[some text] \tag{digits}\] [another text].
text = re.sub(
r"^\\\[(.+?)\\\] \(([\d.]+[a-zA-Z]?)\) (\\\[.+?\\\])$",
r"\[\1 \\tag{\2}\] \3",
text,
- flags=re.M,
+ flags=re.MULTILINE,
)
# multi line
text = text.replace(r"\. ", ". ")
@@ -90,7 +90,7 @@ def markdown_compatible(text: str) -> str:
text,
)
# algorithms
- text = re.sub(r"```\s*(.+?)\s*```", r"```\n\1\n```", text, flags=re.S)
+ text = re.sub(r"```\s*(.+?)\s*```", r"```\n\1\n```", text, flags=re.DOTALL)
return text
@@ -131,7 +131,7 @@ def normalize_list_like_lines(generation):
if not rest:
continue
# Infer current nesting level based on detected numbering
- if re.match(r"^[\dixv]+((?:\.[\dixv])?)+$", potential_numeral, flags=re.I | re.M):
+ if re.match(r"^[\dixv]+((?:\.[\dixv])?)+$", potential_numeral, flags=re.IGNORECASE | re.MULTILINE):
level = potential_numeral.count(".")
replacement += (
@@ -477,7 +477,7 @@ def correct_tables(self, generation: str) -> str:
generation = generation.replace("\\end{tabular} \\end{table}", "\\end{tabular}\n\\end{table}")
generation = generation.replace("\\end{table} Tab", "\\end{table}\nTab")
- generation = re.sub(r"(^.+)\\begin{tab", r"\1\n\\begin{tab", generation, flags=re.M)
+ generation = re.sub(r"(^.+)\\begin{tab", r"\1\n\\begin{tab", generation, flags=re.MULTILINE)
# Remove left-aligned empty LaTeX tabular blocks.
generation = generation.replace(r"\begin{tabular}{l l} & \\ \end{tabular}", "")
@@ -505,7 +505,7 @@ def post_process_single(self, generation: str, fix_markdown: bool = True) -> str
generation = generation.replace("\n* [leftmargin=*]\n", "\n")
# Remove lines with markdown headings starting with #, with numerals,
# and possibly roman numerals with trailing spaces and newlines
- generation = re.sub(r"^#+ (?:[\d+\.]+|[ixv\.]+)?\s*(?:$|\n\s*)", "", generation, flags=re.M)
+ generation = re.sub(r"^#+ (?:[\d+\.]+|[ixv\.]+)?\s*(?:$|\n\s*)", "", generation, flags=re.MULTILINE)
# most likely hallucinated titles
lines = generation.split("\n")
if lines[-1].startswith("#") and lines[-1].lstrip("#").startswith(" ") and len(lines) > 1:
@@ -516,9 +516,9 @@ def post_process_single(self, generation: str, fix_markdown: bool = True) -> str
# Reference corrections
generation = self.remove_hallucinated_references(generation)
# Remove lines starting with asterisks and numbers like "*[1]" and followed by capital letters and periods (ie too long references)
- generation = re.sub(r"^\* \[\d+\](\s?[A-W]\.+\s?){10,}.*$", "", generation, flags=re.M)
+ generation = re.sub(r"^\* \[\d+\](\s?[A-W]\.+\s?){10,}.*$", "", generation, flags=re.MULTILINE)
# Remove empty brackets after a reference number in brackets. *[12][]ABC will become *[12]ABC
- generation = re.sub(r"^(\* \[\d+\])\[\](.*)$", r"\1\2", generation, flags=re.M)
+ generation = re.sub(r"^(\* \[\d+\])\[\](.*)$", r"\1\2", generation, flags=re.MULTILINE)
# Remove single characters before or after 2 new lines
generation = re.sub(r"(^\w\n\n|\n\n\w$)", "", generation)
# pmc math artifact correction
@@ -570,9 +570,9 @@ def post_process_single(self, generation: str, fix_markdown: bool = True) -> str
# Remove lines containing "S.A.B." one or more times. Was included in Nougat's code.
generation = re.sub(r"(\*\*S\. A\. B\.\*\*\n+){2,}", "", generation)
# Remove markdown-style headers that are incomplete or empty on multiple lines.
- generation = re.sub(r"^#+( [\[\d\w])?$", "", generation, flags=re.M)
+ generation = re.sub(r"^#+( [\[\d\w])?$", "", generation, flags=re.MULTILINE)
# Remove lines with just one period.
- generation = re.sub(r"^\.\s*$", "", generation, flags=re.M)
+ generation = re.sub(r"^\.\s*$", "", generation, flags=re.MULTILINE)
# Replace instances of three or more newlines with just two newlines.
generation = re.sub(r"\n{3,}", "\n\n", generation)
if fix_markdown:
diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py
index c7e4706976cc..84aa2509007d 100644
--- a/src/transformers/models/olmo2/modular_olmo2.py
+++ b/src/transformers/models/olmo2/modular_olmo2.py
@@ -317,5 +317,5 @@ class Olmo2ForCausalLM(OlmoForCausalLM):
"Olmo2Config",
"Olmo2ForCausalLM",
"Olmo2Model",
- "Olmo2PreTrainedModel", # noqa: F822
+ "Olmo2PreTrainedModel",
]
diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py
index 8799c8dc07d7..963b18ea0afc 100644
--- a/src/transformers/models/olmo3/modular_olmo3.py
+++ b/src/transformers/models/olmo3/modular_olmo3.py
@@ -423,5 +423,5 @@ class Olmo3ForCausalLM(Olmo2ForCausalLM):
"Olmo3Config",
"Olmo3ForCausalLM",
"Olmo3Model",
- "Olmo3PreTrainedModel", # noqa: F822
+ "Olmo3PreTrainedModel",
]
diff --git a/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py b/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py
index 66fd18abf32c..350cf8af1ab7 100644
--- a/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py
+++ b/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py
@@ -352,7 +352,7 @@ def forward(
batch_size, num_queries, _ = hidden_states.shape
batch_size, sequence_length, _ = encoder_hidden_states.shape
# Ignore copy
- total_elements = sum([shape[0] * shape[1] for shape in spatial_shapes_list])
+ total_elements = sum(shape[0] * shape[1] for shape in spatial_shapes_list)
if total_elements != sequence_length:
raise ValueError(
"Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
@@ -1086,7 +1086,7 @@ def get_cached_task_embeddings(self, tasks_input_ids, tasks_attention_mask):
self.language_cache_prompt.put(not_cached_tasks[idx], (emb, cur_mask))
# pad before concat if needed
- max_len = max([task.shape[0] for task in total_task_features])
+ max_len = max(task.shape[0] for task in total_task_features)
for idx, task in enumerate(total_task_features):
if task.shape[0] < max_len:
pad_size = max_len - task.shape[0]
diff --git a/src/transformers/models/oneformer/image_processing_oneformer.py b/src/transformers/models/oneformer/image_processing_oneformer.py
index 615c71593062..4b25aa2025e1 100644
--- a/src/transformers/models/oneformer/image_processing_oneformer.py
+++ b/src/transformers/models/oneformer/image_processing_oneformer.py
@@ -265,7 +265,7 @@ def compute_segments(
# Copied from transformers.models.maskformer.image_processing_maskformer.convert_segmentation_map_to_binary_masks
def convert_segmentation_map_to_binary_masks(
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
@@ -549,7 +549,7 @@ def rescale(
# Copied from transformers.models.maskformer.image_processing_maskformer.MaskFormerImageProcessor.convert_segmentation_map_to_binary_masks
def convert_segmentation_map_to_binary_masks(
self,
- segmentation_map: "np.ndarray",
+ segmentation_map: np.ndarray,
instance_id_to_semantic_id: Optional[dict[int, int]] = None,
ignore_index: Optional[int] = None,
do_reduce_labels: bool = False,
diff --git a/src/transformers/models/oneformer/image_processing_oneformer_fast.py b/src/transformers/models/oneformer/image_processing_oneformer_fast.py
index 20b34bb7fd39..4a20a04e70f2 100644
--- a/src/transformers/models/oneformer/image_processing_oneformer_fast.py
+++ b/src/transformers/models/oneformer/image_processing_oneformer_fast.py
@@ -18,6 +18,7 @@
import torch
from torch import nn
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -39,17 +40,11 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from .image_processing_oneformer import load_metadata, prepare_metadata
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -453,11 +448,7 @@ def _preprocess(
for shape, stacked_segmentation_maps in grouped_segmentation_maps.items():
if do_resize:
stacked_segmentation_maps = self.resize(
- stacked_segmentation_maps,
- size=size,
- interpolation=F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ stacked_segmentation_maps, size=size, interpolation=F.InterpolationMode.NEAREST_EXACT
)
processed_segmentation_maps_grouped[shape] = stacked_segmentation_maps
processed_segmentation_maps = reorder_images(
diff --git a/src/transformers/models/oneformer/modeling_oneformer.py b/src/transformers/models/oneformer/modeling_oneformer.py
index a5336f6fc490..51c041d7b698 100644
--- a/src/transformers/models/oneformer/modeling_oneformer.py
+++ b/src/transformers/models/oneformer/modeling_oneformer.py
@@ -23,7 +23,6 @@
import numpy as np
import torch
from torch import Tensor, nn
-from torch.cuda.amp import autocast
from ...activations import ACT2FN
from ...modeling_layers import GradientCheckpointingLayer
@@ -322,7 +321,7 @@ def forward(self, masks_queries_logits, class_queries_logits, mask_labels, class
align_corners=False,
).squeeze(1)
- with autocast(enabled=False):
+ with torch.autocast(device_type="cuda", enabled=False):
pred_mask = pred_mask.float()
target_mask = target_mask.float()
@@ -2573,9 +2572,6 @@ def __init__(
):
super().__init__()
self.activation_fn = ACT2FN["quick_gelu"]
- hidden_size = hidden_size
- intermediate_size = intermediate_size
- output_size = output_size
self.fc1 = nn.Linear(hidden_size, intermediate_size)
self.fc2 = nn.Linear(intermediate_size, output_size)
@@ -2882,7 +2878,7 @@ def forward(
Task inputs. Task inputs can be obtained using [`AutoImageProcessor`]. See [`OneFormerProcessor.__call__`]
for details.
text_inputs (`list[torch.Tensor]`, *optional*):
- Tensor fof shape `(num_queries, sequence_length)` to be fed to a model
+ Tensor of shape `(num_queries, sequence_length)` to be fed to a model
Example:
@@ -3068,7 +3064,7 @@ def forward(
Task inputs. Task inputs can be obtained using [`AutoImageProcessor`]. See [`OneFormerProcessor.__call__`]
for details.
text_inputs (`list[torch.Tensor]`, *optional*):
- Tensor fof shape `(num_queries, sequence_length)` to be fed to a model
+ Tensor of shape `(num_queries, sequence_length)` to be fed to a model
mask_labels (`list[torch.Tensor]`, *optional*):
List of mask labels of shape `(num_labels, height, width)` to be fed to a model
class_labels (`list[torch.LongTensor]`, *optional*):
diff --git a/src/transformers/models/ovis2/image_processing_ovis2.py b/src/transformers/models/ovis2/image_processing_ovis2.py
index bd6d63e83914..4c0be26d374a 100644
--- a/src/transformers/models/ovis2/image_processing_ovis2.py
+++ b/src/transformers/models/ovis2/image_processing_ovis2.py
@@ -169,10 +169,10 @@ def get_min_tile_covering_grid(
if sufficient_covering_grids:
# Prefer fewer tiles and higher covering ratio
- return sorted(sufficient_covering_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0][0]
+ return min(sufficient_covering_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0]
else:
# Fallback: prefer higher covering even if below threshold
- return sorted(evaluated_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0][0]
+ return min(evaluated_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0]
class Ovis2ImageProcessor(BaseImageProcessor):
diff --git a/src/transformers/models/ovis2/image_processing_ovis2_fast.py b/src/transformers/models/ovis2/image_processing_ovis2_fast.py
index 07fbf82f9fbe..04b79299e9e1 100644
--- a/src/transformers/models/ovis2/image_processing_ovis2_fast.py
+++ b/src/transformers/models/ovis2/image_processing_ovis2_fast.py
@@ -16,6 +16,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -35,17 +36,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
from .image_processing_ovis2 import get_min_tile_covering_grid, get_optimal_tiled_canvas
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class Ovis2ImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
Args:
diff --git a/src/transformers/models/owlv2/image_processing_owlv2_fast.py b/src/transformers/models/owlv2/image_processing_owlv2_fast.py
index 70441feba3c2..359d241686ec 100644
--- a/src/transformers/models/owlv2/image_processing_owlv2_fast.py
+++ b/src/transformers/models/owlv2/image_processing_owlv2_fast.py
@@ -23,6 +23,7 @@
from typing import TYPE_CHECKING, Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import center_to_corners_format, group_images_by_shape, reorder_images
@@ -35,16 +36,10 @@
SizeDict,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
+from ...utils import TensorType, auto_docstring
from .image_processing_owlv2 import _scale_boxes, box_iou
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
if TYPE_CHECKING:
from .modeling_owlv2 import Owlv2ObjectDetectionOutput
diff --git a/src/transformers/models/owlv2/modular_owlv2.py b/src/transformers/models/owlv2/modular_owlv2.py
index 2e6d917a791a..66acd2088399 100644
--- a/src/transformers/models/owlv2/modular_owlv2.py
+++ b/src/transformers/models/owlv2/modular_owlv2.py
@@ -18,6 +18,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -37,17 +38,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
from ..owlvit.image_processing_owlvit_fast import OwlViTImageProcessorFast
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class Owlv2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs): ...
diff --git a/src/transformers/models/owlv2/processing_owlv2.py b/src/transformers/models/owlv2/processing_owlv2.py
index 2e69379af73f..d12ee5995535 100644
--- a/src/transformers/models/owlv2/processing_owlv2.py
+++ b/src/transformers/models/owlv2/processing_owlv2.py
@@ -141,7 +141,7 @@ def __call__(
encodings = []
# Maximum number of queries across batch
- max_num_queries = max([len(text_single) for text_single in text])
+ max_num_queries = max(len(text_single) for text_single in text)
# Pad all batch samples to max number of text queries
for text_single in text:
diff --git a/src/transformers/models/owlvit/processing_owlvit.py b/src/transformers/models/owlvit/processing_owlvit.py
index 0e0c59d555f2..ae39d2b6b307 100644
--- a/src/transformers/models/owlvit/processing_owlvit.py
+++ b/src/transformers/models/owlvit/processing_owlvit.py
@@ -151,7 +151,7 @@ def __call__(
encodings = []
# Maximum number of queries across batch
- max_num_queries = max([len(text_single) for text_single in text])
+ max_num_queries = max(len(text_single) for text_single in text)
# Pad all batch samples to max number of text queries
for text_single in text:
diff --git a/src/transformers/models/parakeet/__init__.py b/src/transformers/models/parakeet/__init__.py
new file mode 100644
index 000000000000..5c54b2e2eadb
--- /dev/null
+++ b/src/transformers/models/parakeet/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_parakeet import *
+ from .feature_extraction_parakeet import *
+ from .modeling_parakeet import *
+ from .tokenization_parakeet_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/parakeet/configuration_parakeet.py b/src/transformers/models/parakeet/configuration_parakeet.py
new file mode 100644
index 000000000000..3612da58006a
--- /dev/null
+++ b/src/transformers/models/parakeet/configuration_parakeet.py
@@ -0,0 +1,235 @@
+# coding=utf-8
+# 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.
+"""Parakeet model configuration."""
+
+from typing import Union
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class ParakeetEncoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`ParakeetEncoder`]. It is used to instantiate a
+ `ParakeetEncoder` model according to the specified arguments, defining the model architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ hidden_size (`int`, *optional*, defaults to 1024):
+ Dimension of the layers and the hidden states.
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of hidden layers in the Transformer encoder.
+ num_attention_heads (`int`, *optional*, defaults to 8):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ intermediate_size (`int`, *optional*, defaults to 4096):
+ Dimension of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the encoder and pooler.
+ attention_bias (`bool`, *optional*, defaults to `True`):
+ Whether to use bias in the attention layers.
+ conv_kernel_size (`int`, *optional*, defaults to 9):
+ The kernel size of the convolution layers in the Conformer block.
+ subsampling_factor (`int`, *optional*, defaults to 8):
+ The factor by which the input sequence is subsampled.
+ subsampling_conv_channels (`int`, *optional*, defaults to 256):
+ The number of channels in the subsampling convolution layers.
+ num_mel_bins (`int`, *optional*, defaults to 80):
+ Number of mel features.
+ subsampling_conv_kernel_size (`int`, *optional*, defaults to 3):
+ The kernel size of the subsampling convolution layers.
+ subsampling_conv_stride (`int`, *optional*, defaults to 2):
+ The stride of the subsampling convolution layers.
+ dropout (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for all fully connected layers in the embeddings, encoder, and pooler.
+ dropout_positions (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the positions in the input sequence.
+ layerdrop (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for the layers in the encoder.
+ activation_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for activations inside the fully connected layer.
+ attention_dropout (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for the attention layers.
+ max_position_embeddings (`int`, *optional*, defaults to 5000):
+ The maximum sequence length that this model might ever be used with.
+ scale_input (`bool`, *optional*, defaults to `True`):
+ Whether to scale the input embeddings.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+ Example:
+ ```python
+ >>> from transformers import ParakeetEncoderModel, ParakeetEncoderConfig
+
+ >>> # Initializing a `ParakeetEncoder` configuration
+ >>> configuration = ParakeetEncoderConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = ParakeetEncoderModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+
+ This configuration class is based on the ParakeetEncoder architecture from NVIDIA NeMo. You can find more details
+ and pre-trained models at [nvidia/parakeet-ctc-1.1b](https://huggingface.co/nvidia/parakeet-ctc-1.1b).
+ """
+
+ model_type = "parakeet_encoder"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ hidden_size=1024,
+ num_hidden_layers=24,
+ num_attention_heads=8,
+ intermediate_size=4096,
+ hidden_act="silu",
+ attention_bias=True,
+ conv_kernel_size=9,
+ subsampling_factor=8,
+ subsampling_conv_channels=256,
+ num_mel_bins=80,
+ subsampling_conv_kernel_size=3,
+ subsampling_conv_stride=2,
+ dropout=0.1,
+ dropout_positions=0.0,
+ layerdrop=0.1,
+ activation_dropout=0.1,
+ attention_dropout=0.1,
+ max_position_embeddings=5000,
+ scale_input=True,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(
+ **kwargs,
+ )
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_attention_heads # LlamaAttention compatibility
+ self.intermediate_size = intermediate_size
+ self.hidden_act = hidden_act
+ self.attention_bias = attention_bias
+
+ if (conv_kernel_size - 1) % 2 != 0:
+ raise ValueError(f"conv_kernel_size must be odd, got {conv_kernel_size}")
+ self.conv_kernel_size = conv_kernel_size
+
+ self.subsampling_conv_kernel_size = subsampling_conv_kernel_size
+ self.subsampling_conv_stride = subsampling_conv_stride
+
+ self.subsampling_factor = subsampling_factor
+ self.subsampling_conv_channels = subsampling_conv_channels
+ self.num_mel_bins = num_mel_bins
+
+ self.dropout = dropout
+ self.dropout_positions = dropout_positions
+ self.layerdrop = layerdrop
+ self.activation_dropout = activation_dropout
+ self.attention_dropout = attention_dropout
+ self.max_position_embeddings = max_position_embeddings
+ self.scale_input = scale_input
+ self.initializer_range = initializer_range
+
+
+class ParakeetCTCConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`ParakeetForCTC`]. It is used to instantiate a
+ Parakeet CTC model according to the specified arguments, defining the model architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 1025):
+ Vocabulary size of the model.
+ ctc_loss_reduction (`str`, *optional*, defaults to `"mean"`):
+ Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an
+ instance of [`ParakeetForCTC`].
+ ctc_zero_infinity (`bool`, *optional*, defaults to `True`):
+ Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly
+ occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance
+ of [`ParakeetForCTC`].
+ encoder_config (`Union[dict, ParakeetEncoderConfig]`, *optional*):
+ The config object or dictionary of the encoder.
+ pad_token_id (`int`, *optional*, defaults to 1024):
+ Padding token id. Also used as blank token id.
+
+ Example:
+ ```python
+ >>> from transformers import ParakeetForCTC, ParakeetCTCConfig
+
+ >>> # Initializing a Parakeet configuration
+ >>> configuration = ParakeetCTCConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = ParakeetForCTC(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+
+ This configuration class is based on the Parakeet CTC architecture from NVIDIA NeMo. You can find more details
+ and pre-trained models at [nvidia/parakeet-ctc-1.1b](https://huggingface.co/nvidia/parakeet-ctc-1.1b).
+ """
+
+ model_type = "parakeet_ctc"
+ sub_configs = {"encoder_config": ParakeetEncoderConfig}
+
+ def __init__(
+ self,
+ vocab_size=1025,
+ ctc_loss_reduction="mean",
+ ctc_zero_infinity=True,
+ encoder_config: Union[dict, ParakeetEncoderConfig] = None,
+ pad_token_id=1024,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.ctc_loss_reduction = ctc_loss_reduction
+ self.ctc_zero_infinity = ctc_zero_infinity
+
+ if isinstance(encoder_config, dict):
+ self.encoder_config = ParakeetEncoderConfig(**encoder_config)
+ elif encoder_config is None:
+ self.encoder_config = ParakeetEncoderConfig()
+
+ self.encoder_config = self.encoder_config
+ self.initializer_range = self.encoder_config.initializer_range
+
+ super().__init__(
+ pad_token_id=pad_token_id,
+ **kwargs,
+ )
+
+ @classmethod
+ def from_encoder_config(cls, encoder_config: ParakeetEncoderConfig, **kwargs):
+ r"""
+ Instantiate a [`ParakeetCTCConfig`] (or a derived class) from parakeet encoder model configuration.
+
+ Returns:
+ [`ParakeetCTCConfig`]: An instance of a configuration object
+ """
+
+ return cls(encoder_config=encoder_config.to_dict(), **kwargs)
+
+
+__all__ = ["ParakeetCTCConfig", "ParakeetEncoderConfig"]
diff --git a/src/transformers/models/parakeet/convert_nemo_to_hf.py b/src/transformers/models/parakeet/convert_nemo_to_hf.py
new file mode 100644
index 000000000000..f1998fbd81b8
--- /dev/null
+++ b/src/transformers/models/parakeet/convert_nemo_to_hf.py
@@ -0,0 +1,315 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import gc
+import os
+import re
+import tarfile
+
+import torch
+import yaml
+from tokenizers import AddedToken
+
+from transformers import (
+ ParakeetCTCConfig,
+ ParakeetFeatureExtractor,
+ ParakeetForCTC,
+ ParakeetProcessor,
+ ParakeetTokenizerFast,
+)
+from transformers.convert_slow_tokenizer import ParakeetConverter
+from transformers.utils.hub import cached_file
+
+
+NEMO_TO_HF_WEIGHT_MAPPING = {
+ r"encoder\.pre_encode\.conv\.": r"encoder.subsampling.layers.",
+ r"encoder\.pre_encode\.out\.": r"encoder.subsampling.linear.",
+ r"encoder\.pos_enc\.": r"encoder.encode_positions.",
+ r"encoder\.layers\.(\d+)\.conv\.batch_norm\.": r"encoder.layers.\1.conv.norm.",
+ r"decoder\.decoder_layers\.0\.(weight|bias)": r"ctc_head.\1",
+ r"linear_([kv])": r"\1_proj",
+ r"linear_out": r"o_proj",
+ r"linear_q": r"q_proj",
+ r"pos_bias_([uv])": r"bias_\1",
+ r"linear_pos": r"relative_k_proj",
+}
+
+
+def convert_key(key, mapping):
+ for pattern, replacement in mapping.items():
+ key = re.sub(pattern, replacement, key)
+ return key
+
+
+def extract_nemo_archive(nemo_file_path: str, extract_dir: str) -> dict[str, str]:
+ """
+ Extract .nemo file (tar archive) and return paths to important files.
+
+ Args:
+ nemo_file_path: Path to .nemo file
+ extract_dir: Directory to extract to
+
+ Returns:
+ Dictionary with paths to model.pt, model_config.yaml, etc.
+ """
+ print(f"Extracting NeMo archive: {nemo_file_path}")
+
+ with tarfile.open(nemo_file_path, "r", encoding="utf-8") as tar:
+ tar.extractall(extract_dir)
+
+ # Log all extracted files for debugging
+ all_files = []
+ for root, dirs, files in os.walk(extract_dir):
+ for file in files:
+ file_path = os.path.join(root, file)
+ all_files.append(file_path)
+
+ print(f"All extracted files: {[os.path.basename(f) for f in all_files]}")
+
+ # Find important files with more robust detection
+ model_files = {}
+ for root, dirs, files in os.walk(extract_dir):
+ for file in files:
+ file_path = os.path.join(root, file)
+ file_lower = file.lower()
+
+ # Look for model weights with various common names
+ if (
+ file.endswith(".pt")
+ or file.endswith(".pth")
+ or file.endswith(".ckpt")
+ or file.endswith(".bin")
+ or "model" in file_lower
+ and ("weight" in file_lower or "state" in file_lower)
+ or file_lower == "model.pt"
+ or file_lower == "pytorch_model.bin"
+ or file_lower == "model_weights.ckpt"
+ ):
+ model_files["model_weights"] = file_path
+ print(f"Found model weights: {file}")
+
+ # Look for config files
+ elif (
+ file == "model_config.yaml"
+ or file == "config.yaml"
+ or (file.endswith(".yaml") and "config" in file_lower)
+ ):
+ if "model_config" not in model_files: # Prefer model_config.yaml
+ model_files["model_config"] = file_path
+ print(f"Found config file: {file}")
+ if file == "model_config.yaml":
+ model_files["model_config"] = file_path # Override with preferred name
+
+ # Look for vocabulary files
+ elif (
+ file.endswith(".vocab")
+ or file.endswith(".model")
+ or file.endswith(".txt")
+ or ("tokenizer" in file_lower and (file.endswith(".vocab") or file.endswith(".model")))
+ ):
+ # Prefer .vocab files over others
+ if "tokenizer_model_file" not in model_files or file.endswith(".model"):
+ model_files["tokenizer_model_file"] = file_path
+ print(f"Found tokenizer model file: {file}")
+ else:
+ print(f"Found additional vocabulary file (using existing): {file}")
+
+ print(f"Found model files: {list(model_files.keys())}")
+
+ # Validate that we found the required files
+ if "model_weights" not in model_files:
+ raise FileNotFoundError(
+ f"Could not find model weights file in {nemo_file_path}. "
+ f"Expected files with extensions: .pt, .pth, .ckpt, .bin. "
+ f"Found files: {[os.path.basename(f) for f in all_files]}"
+ )
+
+ if "model_config" not in model_files:
+ raise FileNotFoundError(
+ f"Could not find model config file in {nemo_file_path}. "
+ f"Expected: model_config.yaml or config.yaml. "
+ f"Found files: {[os.path.basename(f) for f in all_files]}"
+ )
+
+ return model_files
+
+
+def write_processor(nemo_config: dict, model_files, output_dir, push_to_repo_id=None):
+ tokenizer_converted = ParakeetConverter(model_files["tokenizer_model_file"]).converted()
+ tokenizer_converted_fast = ParakeetTokenizerFast(
+ tokenizer_object=tokenizer_converted,
+ clean_up_tokenization_spaces=False,
+ )
+ tokenizer_converted_fast.add_tokens(
+ [AddedToken("", normalized=False, special=True), AddedToken("", normalized=False, special=True)]
+ )
+ tokenizer_converted_fast.add_special_tokens(
+ {
+ "pad_token": AddedToken("", normalized=False, special=True),
+ "unk_token": AddedToken("", normalized=False, special=True),
+ }
+ )
+
+ feature_extractor_keys_to_ignore = ["_target_", "pad_to", "frame_splicing", "dither", "normalize", "window", "log"]
+ feature_extractor_config_keys_mapping = {
+ "sample_rate": "sampling_rate",
+ "window_size": "win_length",
+ "window_stride": "hop_length",
+ "window": "window",
+ "n_fft": "n_fft",
+ "log": "log",
+ "features": "feature_size",
+ "dither": "dither",
+ "pad_to": "pad_to",
+ "pad_value": "padding_value",
+ "frame_splicing": "frame_splicing",
+ "preemphasis": "preemphasis",
+ "hop_length": "hop_length",
+ }
+ converted_feature_extractor_config = {}
+
+ for key, value in nemo_config["preprocessor"].items():
+ if key in feature_extractor_keys_to_ignore:
+ continue
+ if key in feature_extractor_config_keys_mapping:
+ if key in ["window_size", "window_stride"]:
+ value = int(value * nemo_config["preprocessor"]["sample_rate"])
+ converted_feature_extractor_config[feature_extractor_config_keys_mapping[key]] = value
+ else:
+ raise ValueError(f"Key {key} not found in feature_extractor_keys_mapping")
+
+ feature_extractor = ParakeetFeatureExtractor(**converted_feature_extractor_config)
+
+ processor = ParakeetProcessor(
+ feature_extractor=feature_extractor,
+ tokenizer=tokenizer_converted_fast,
+ )
+ processor.save_pretrained(output_dir)
+
+ if push_to_repo_id:
+ processor.push_to_hub(push_to_repo_id)
+
+
+def write_model(nemo_config, model_files, model_type, output_dir, push_to_repo_id=None):
+ encoder_keys_to_ignore = [
+ "att_context_size",
+ "causal_downsampling",
+ "stochastic_depth_start_layer",
+ "feat_out",
+ "stochastic_depth_drop_prob",
+ "_target_",
+ "ff_expansion_factor",
+ "untie_biases",
+ "att_context_style",
+ "self_attention_model",
+ "conv_norm_type",
+ "subsampling",
+ "stochastic_depth_mode",
+ "conv_context_size",
+ "dropout_pre_encoder",
+ ]
+ enocder_config_keys_mapping = {
+ "d_model": "hidden_size",
+ "n_heads": "num_attention_heads",
+ "n_layers": "num_hidden_layers",
+ "feat_in": "num_mel_bins",
+ "conv_kernel_size": "conv_kernel_size",
+ "subsampling_factor": "subsampling_factor",
+ "subsampling_conv_channels": "subsampling_conv_channels",
+ "pos_emb_max_len": "max_position_embeddings",
+ "dropout": "dropout",
+ "dropout_emb": "dropout_positions",
+ "dropout_att": "attention_dropout",
+ "xscaling": "scale_input",
+ }
+ converted_encoder_config = {}
+
+ for key, value in nemo_config["encoder"].items():
+ if key in encoder_keys_to_ignore:
+ continue
+ if key in enocder_config_keys_mapping:
+ converted_encoder_config[enocder_config_keys_mapping[key]] = value
+ else:
+ raise ValueError(f"Key {key} not found in enocder_config_keys_mapping")
+
+ state_dict = torch.load(model_files["model_weights"], map_location="cpu", weights_only=True)
+ converted_state_dict = {}
+ for key, value in state_dict.items():
+ # Skip preprocessing weights (featurizer components)
+ if key.endswith("featurizer.window") or key.endswith("featurizer.fb"):
+ print(f"Skipping preprocessing weight: {key}")
+ continue
+ converted_key = convert_key(key, NEMO_TO_HF_WEIGHT_MAPPING)
+ converted_state_dict[converted_key] = value
+
+ if model_type == "ctc":
+ model_config = ParakeetCTCConfig(
+ encoder_config=converted_encoder_config,
+ )
+ print("Loading the checkpoint in a Parakeet CTC model.")
+ with torch.device("meta"):
+ model = ParakeetForCTC(model_config)
+ model.load_state_dict(converted_state_dict, strict=True, assign=True)
+ print("Checkpoint loaded successfully.")
+ del model.config._name_or_path
+
+ print("Saving the model.")
+ model.save_pretrained(output_dir)
+
+ if push_to_repo_id:
+ model.push_to_hub(push_to_repo_id)
+
+ del converted_state_dict, model
+
+ # Safety check: reload the converted model
+ gc.collect()
+ print("Reloading the model to check if it's saved correctly.")
+ ParakeetForCTC.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
+ print("Model reloaded successfully.")
+
+ else:
+ raise ValueError(f"Model type {model_type} not supported.")
+
+
+def main(
+ hf_repo_id,
+ output_dir,
+ model_type,
+ push_to_repo_id=None,
+):
+ nemo_filename = f"{hf_repo_id.split('/')[-1]}.nemo"
+ filepath = cached_file(hf_repo_id, nemo_filename)
+
+ model_files = extract_nemo_archive(filepath, os.path.dirname(filepath))
+ nemo_config = yaml.load(open(model_files["model_config"], "r"), Loader=yaml.FullLoader)
+
+ write_processor(nemo_config, model_files, output_dir, push_to_repo_id)
+ write_model(nemo_config, model_files, model_type, output_dir, push_to_repo_id)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--hf_repo_id", required=True, help="Model repo on huggingface.co")
+ parser.add_argument("--model_type", required=True, choices=["ctc"], help="Model type (`ctc`, `tdt`)")
+ parser.add_argument("--output_dir", required=True, help="Output directory for HuggingFace model")
+ parser.add_argument("--push_to_repo_id", help="Repository ID to push the model to on the Hub")
+ args = parser.parse_args()
+ main(
+ args.hf_repo_id,
+ args.output_dir,
+ args.model_type,
+ args.push_to_repo_id,
+ )
diff --git a/src/transformers/models/parakeet/feature_extraction_parakeet.py b/src/transformers/models/parakeet/feature_extraction_parakeet.py
new file mode 100644
index 000000000000..d28f1a214a21
--- /dev/null
+++ b/src/transformers/models/parakeet/feature_extraction_parakeet.py
@@ -0,0 +1,287 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional, Union
+
+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 ParakeetFeatureExtractor(SequenceFeatureExtractor):
+ r"""
+ Constructs a Parakeet 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 80):
+ The feature dimension of the extracted features.
+ sampling_rate (`int`, *optional*, defaults to 16000):
+ The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
+ 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.
+ """
+
+ model_input_names = ["input_features", "attention_mask"]
+
+ def __init__(
+ self,
+ feature_size=80,
+ sampling_rate=16000,
+ hop_length=160,
+ n_fft=512,
+ win_length=400,
+ preemphasis=0.97,
+ padding_value=0.0,
+ **kwargs,
+ ):
+ super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
+
+ self.hop_length = hop_length
+ self.n_fft = n_fft
+ self.win_length = win_length
+ self.preemphasis = preemphasis
+
+ # 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
+ # self.mel_filters = mel_filter_bank(
+ # num_frequency_bins=n_fft // 2 + 1,
+ # num_mel_filters=feature_size,
+ # min_frequency=0.0,
+ # max_frequency=sampling_rate / 2,
+ # sampling_rate=sampling_rate,
+ # norm="slaney",
+ # mel_scale="slaney",
+ # )
+ 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 _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 math original implementation
+ # magnitudes = torch.abs(stft) ** 2
+ 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: Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]],
+ truncation: bool = False,
+ pad_to_multiple_of: Optional[int] = None,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ return_attention_mask: Optional[bool] = None,
+ padding: Optional[str] = "longest",
+ max_length: Optional[int] = None,
+ sampling_rate: Optional[int] = None,
+ do_normalize: Optional[bool] = None,
+ device: Optional[str] = "cpu",
+ return_token_timestamps: Optional[bool] = None,
+ **kwargs,
+ ) -> BatchFeature:
+ """
+ Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for
+ the STFT computation if available, otherwise a slower NumPy based one.
+
+ Args:
+ raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`):
+ The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+ values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+ stereo, i.e. single float per timestep.
+ truncation (`bool`, *optional*, default to `True`):
+ Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+ pad_to_multiple_of (`int`, *optional*, defaults to None):
+ If set will pad the sequence to a multiple of the provided value.
+
+ This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+ `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
+ return_attention_mask (`bool`, *optional*):
+ Whether to return the attention mask. If left to the default, will return the attention mask according
+ to the specific feature_extractor's default.
+
+ [What are attention masks?](../glossary#attention-mask)
+
+
+
+ For Parakeet models, `attention_mask` should always be passed for batched inference, to avoid subtle
+ bugs.
+
+
+
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors instead of list of python integers. Acceptable values are:
+
+ - `'tf'`: Return TensorFlow `tf.constant` objects.
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return Numpy `np.ndarray` objects.
+ sampling_rate (`int`, *optional*):
+ The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
+ `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition
+ pipeline.
+ padding_value (`float`, *optional*, defaults to 0.0):
+ The value that is used to fill the padding values / vectors.
+ do_normalize (`bool`, *optional*, defaults to `False`):
+ Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly
+ improve the performance of the model.
+ device (`str`, *optional*, defaults to `'cpu'`):
+ Specifies the device for computation of the log-mel spectrogram of audio signals in the
+ `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda")
+ return_token_timestamps (`bool`, *optional*, defaults to `None`):
+ Deprecated. Use `return_attention_mask` instead from which the number of frames can be inferred.
+
+ Whether or not to return the number of frames of the input raw_speech.
+ These num_frames can be used by the model to compute word level timestamps.
+ """
+ if sampling_rate is not None:
+ if sampling_rate != self.sampling_rate:
+ raise ValueError(
+ f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
+ f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
+ f" was sampled with {self.sampling_rate} and not {sampling_rate}."
+ )
+ else:
+ logger.warning(
+ f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
+ "Failing to do so can result in silent errors that might be hard to debug."
+ )
+
+ # Convert to torch tensor
+ if isinstance(raw_speech, np.ndarray):
+ raw_speech = torch.tensor(raw_speech)
+ elif isinstance(raw_speech, (list, tuple)) 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[:, None].to(torch.float32) for speech in raw_speech]
+ else:
+ raw_speech = [raw_speech[:, None].to(torch.float32)]
+
+ 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)
+
+ # 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
+
+ return BatchFeature(
+ data={
+ "input_features": input_features,
+ "attention_mask": attention_mask,
+ },
+ tensor_type=return_tensors,
+ )
+
+
+__all__ = ["ParakeetFeatureExtractor"]
diff --git a/src/transformers/models/parakeet/modeling_parakeet.py b/src/transformers/models/parakeet/modeling_parakeet.py
new file mode 100644
index 000000000000..4190517b48fd
--- /dev/null
+++ b/src/transformers/models/parakeet/modeling_parakeet.py
@@ -0,0 +1,744 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/parakeet/modular_parakeet.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_parakeet.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import Callable, Optional, Union
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, CausalLMOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import check_model_inputs
+from .configuration_parakeet import ParakeetCTCConfig, ParakeetEncoderConfig
+
+
+class ParakeetEncoderRelPositionalEncoding(nn.Module):
+ """Relative positional encoding for Parakeet."""
+
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: ParakeetEncoderConfig, device=None):
+ super().__init__()
+ self.max_position_embeddings = config.max_position_embeddings
+ base = 10000.0
+ inv_freq = 1.0 / (
+ base
+ ** (
+ torch.arange(0, config.hidden_size, 2, dtype=torch.int64).to(device=device, dtype=torch.float)
+ / config.hidden_size
+ )
+ )
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ @torch.no_grad()
+ def forward(self, hidden_states: torch.Tensor):
+ seq_length = hidden_states.shape[1]
+ if seq_length > self.max_position_embeddings:
+ raise ValueError(
+ f"Sequence Length: {seq_length} has to be less or equal than "
+ f"config.max_position_embeddings {self.max_position_embeddings}."
+ )
+
+ position_ids = torch.arange(seq_length - 1, -seq_length, -1, device=hidden_states.device)
+ inv_freq_expanded = (
+ self.inv_freq[None, :, None].float().expand(hidden_states.shape[0], -1, 1).to(hidden_states.device)
+ )
+ position_ids_expanded = position_ids[None, None, :].float()
+
+ device_type = (
+ hidden_states.device.type
+ if isinstance(hidden_states.device.type, str) and hidden_states.device.type != "mps"
+ else "cpu"
+ )
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ sin = freqs.sin()
+ cos = freqs.cos()
+ # interleave sin and cos
+ pos_embed = torch.stack([sin, cos], dim=-1)
+ pos_embed = pos_embed.reshape(*pos_embed.shape[:-2], -1)
+
+ return pos_embed.to(dtype=hidden_states.dtype)
+
+
+class ParakeetEncoderFeedForward(nn.Module):
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__()
+ self.linear1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias)
+ self.activation = ACT2FN[config.hidden_act]
+ self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias)
+ self.activation_dropout = config.activation_dropout
+
+ def forward(self, hidden_states):
+ hidden_states = self.activation(self.linear1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.linear2(hidden_states)
+ return hidden_states
+
+
+class ParakeetEncoderConvolutionModule(nn.Module):
+ def __init__(self, config: ParakeetEncoderConfig, module_config=None):
+ """
+ Args:
+ config (ParakeetEncoderConfig): Configuration for the model.
+ module_config (dict): Configuration for the module (e.g., encoder or decoder).
+ """
+ super().__init__()
+ channels = config.hidden_size
+ # kernel_size should be an odd number for 'SAME' padding
+ if module_config is None:
+ # e.g. using `ParakeetEncoderEncoderConfig` in src/transformers/models/parakeet_encoder/configuration_parakeet_encoder.py
+ kernel_size = config.conv_kernel_size
+ self.activation = ACT2FN[getattr(config, "hidden_act", "silu")]
+ else:
+ kernel_size = module_config["kernel_size"]
+ self.activation = ACT2FN[module_config.get("activation", "silu")]
+ self.padding = (kernel_size - 1) // 2
+ self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=True)
+ self.depthwise_conv = nn.Conv1d(
+ channels, channels, kernel_size, stride=1, padding=self.padding, groups=channels, bias=True
+ )
+ self.norm = nn.BatchNorm1d(channels)
+ self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, stride=1, padding=0, bias=True)
+
+ def forward(self, hidden_states, attention_mask=None):
+ """
+ Compute convolution module.
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, time, channels)`): Input tensor.
+ attention_mask (`torch.Tensor` of shape `(batch, 1, time)`): Attention mask.
+
+ Returns:
+ `torch.Tensor`: Output tensor of shape `(batch, time, channels)`.
+
+ """
+ # exchange the temporal dimension and the feature dimension
+ hidden_states = hidden_states.transpose(1, 2)
+
+ # GLU mechanism, (batch_size, 2*channel, dim)
+ hidden_states = self.pointwise_conv1(hidden_states)
+ # (batch_size, channel, dim)
+ hidden_states = nn.functional.glu(hidden_states, dim=1)
+
+ # Apply padding mask before convolution
+ if attention_mask is not None:
+ all_masked_rows = torch.all(~attention_mask, dim=-1)
+ hidden_states = hidden_states.masked_fill(all_masked_rows, 0.0)
+
+ # 1D Depthwise Conv
+ hidden_states = self.depthwise_conv(hidden_states)
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.pointwise_conv2(hidden_states)
+
+ return hidden_states.transpose(1, 2)
+
+
+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: Optional[torch.Tensor],
+ 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:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class ParakeetEncoderAttention(nn.Module):
+ """Multi-head attention with relative positional encoding. See section 3.3 of https://huggingface.co/papers/1901.02860."""
+
+ def __init__(self, config: ParakeetEncoderConfig, 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
+ )
+ # W_{k,R} projection
+ self.relative_k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+ # global content bias
+ self.bias_u = nn.Parameter(torch.zeros(config.num_attention_heads, self.head_dim))
+ # global positional bias
+ self.bias_v = nn.Parameter(torch.zeros(config.num_attention_heads, self.head_dim))
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: Optional[torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ batch_size, seq_length = input_shape
+ hidden_shape = (batch_size, seq_length, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ query_states_with_bias_u = query_states + self.bias_u.view(
+ 1, self.config.num_attention_heads, 1, self.head_dim
+ )
+ query_states_with_bias_v = query_states + self.bias_v.view(
+ 1, self.config.num_attention_heads, 1, self.head_dim
+ )
+
+ relative_key_states = self.relative_k_proj(position_embeddings)
+ relative_key_states = relative_key_states.view(batch_size, -1, self.config.num_attention_heads, self.head_dim)
+
+ # terms (b) and (d)
+ matrix_bd = query_states_with_bias_v @ relative_key_states.permute(0, 2, 3, 1)
+ matrix_bd = self._rel_shift(matrix_bd)
+ matrix_bd = matrix_bd[..., :seq_length]
+ matrix_bd = matrix_bd * self.scaling
+
+ if attention_mask is not None:
+ # here the original codebase uses -10000.0 rather than float("-inf") and then manual masked fill with 0.0s
+ # see: https://github.com/NVIDIA-NeMo/NeMo/blob/8cfedd7203462cb251a914e700e5605444277561/nemo/collections/asr/parts/submodules/multi_head_attention.py#L320-L340
+ # we rather went for a straight-forward approach with float("-inf")
+ matrix_bd = matrix_bd.masked_fill_(attention_mask.logical_not(), float("-inf"))
+
+ # will compute matrix_ac - terms (a) and (c) - and add matrix_bd
+ attn_output, attn_weights = attention_interface(
+ self,
+ query=query_states_with_bias_u,
+ key=key_states,
+ value=value_states,
+ attention_mask=matrix_bd,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+ def _rel_shift(self, attention_scores):
+ """Relative position shift for Shaw et al. style attention. See appendix B of https://huggingface.co/papers/1901.02860."""
+ batch_size, num_heads, query_length, position_length = attention_scores.shape
+ attention_scores = nn.functional.pad(attention_scores, pad=(1, 0))
+ attention_scores = attention_scores.view(batch_size, num_heads, -1, query_length)
+ attention_scores = attention_scores[:, :, 1:].view(batch_size, num_heads, query_length, position_length)
+ return attention_scores
+
+
+class ParakeetEncoderSubsamplingConv2D(nn.Module):
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__()
+
+ self.kernel_size = config.subsampling_conv_kernel_size
+ self.stride = config.subsampling_conv_stride
+ self.channels = config.subsampling_conv_channels
+ self.padding = (self.kernel_size - 1) // 2
+ self.num_layers = int(math.log2(config.subsampling_factor))
+
+ # define layers
+ self.layers = nn.ModuleList()
+ self.layers.append(
+ nn.Conv2d(1, self.channels, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding)
+ )
+ self.layers.append(nn.ReLU())
+ for i in range(self.num_layers - 1):
+ # depthwise conv
+ self.layers.append(
+ nn.Conv2d(
+ self.channels,
+ self.channels,
+ kernel_size=self.kernel_size,
+ stride=self.stride,
+ padding=self.padding,
+ groups=self.channels,
+ )
+ )
+ # pointwise conv
+ self.layers.append(nn.Conv2d(self.channels, self.channels, kernel_size=1))
+ # activation
+ self.layers.append(nn.ReLU())
+
+ out_length = config.num_mel_bins // (self.stride**self.num_layers)
+ self.linear = nn.Linear(config.subsampling_conv_channels * out_length, config.hidden_size, bias=True)
+
+ def _get_output_length(self, input_lengths: torch.Tensor, conv_layer: nn.Conv2d):
+ if hasattr(conv_layer, "stride") and conv_layer.stride != (1, 1):
+ padding = conv_layer.padding
+ kernel_size = conv_layer.kernel_size[0]
+ stride = conv_layer.stride[0]
+
+ output_lengths = (input_lengths + padding[0] + padding[1] - kernel_size) // stride + 1
+ return output_lengths
+
+ return input_lengths
+
+ def forward(self, input_features: torch.Tensor, attention_mask: torch.Tensor = None):
+ hidden_states = input_features.unsqueeze(1)
+ current_lengths = attention_mask.sum(-1) if attention_mask is not None else None
+
+ for layer in self.layers:
+ hidden_states = layer(hidden_states)
+
+ # mask the hidden states
+ if isinstance(layer, nn.Conv2d) and attention_mask is not None:
+ current_lengths = self._get_output_length(current_lengths, layer)
+ current_seq_length = hidden_states.shape[2]
+ channel_mask = (
+ torch.arange(current_seq_length, device=attention_mask.device) < current_lengths[:, None]
+ )
+ hidden_states *= channel_mask[:, None, :, None]
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(hidden_states.shape[0], hidden_states.shape[2], -1)
+ hidden_states = self.linear(hidden_states)
+
+ return hidden_states
+
+
+class ParakeetEncoderBlock(GradientCheckpointingLayer):
+ def __init__(self, config: ParakeetEncoderConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.gradient_checkpointing = False
+
+ self.feed_forward1 = ParakeetEncoderFeedForward(config)
+ self.self_attn = ParakeetEncoderAttention(config, layer_idx)
+ self.conv = ParakeetEncoderConvolutionModule(config)
+ self.feed_forward2 = ParakeetEncoderFeedForward(config)
+
+ self.norm_feed_forward1 = nn.LayerNorm(config.hidden_size)
+ self.norm_self_att = nn.LayerNorm(config.hidden_size)
+ self.norm_conv = nn.LayerNorm(config.hidden_size)
+ self.norm_feed_forward2 = nn.LayerNorm(config.hidden_size)
+ self.norm_out = nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states))
+ hidden_states = residual + 0.5 * hidden_states # the conformer architecture uses a factor of 0.5
+
+ normalized_hidden_states = self.norm_self_att(hidden_states)
+ attn_output, _ = self.self_attn(
+ hidden_states=normalized_hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + attn_output
+
+ conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask)
+ hidden_states = hidden_states + conv_output
+
+ ff2_output = self.feed_forward2(self.norm_feed_forward2(hidden_states))
+ hidden_states = hidden_states + 0.5 * ff2_output # the conformer architecture uses a factor of 0.5
+
+ hidden_states = self.norm_out(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class ParakeetPreTrainedModel(PreTrainedModel):
+ config: ParakeetCTCConfig
+ base_model_prefix = "model"
+ main_input_name = "input_features"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["ParakeetEncoderBlock"]
+ _supports_flat_attention_mask = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ # TODO: @eustlb, add support when flash attention supports custom attention bias
+ _supports_flash_attn = False
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": ParakeetEncoderBlock,
+ "attentions": ParakeetEncoderAttention,
+ }
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ if hasattr(self.config, "initializer_range"):
+ std = self.config.initializer_range
+ else:
+ # 0.02 is the standard default value accross the library
+ std = getattr(self.config.get_text_config(), "initializer_range", 0.02)
+
+ if isinstance(module, ParakeetEncoderAttention):
+ # Initialize positional bias parameters
+ module.bias_u.data.normal_(mean=0.0, std=std)
+ module.bias_v.data.normal_(mean=0.0, std=std)
+
+ def _get_subsampling_output_length(self, input_lengths: torch.Tensor):
+ encoder_config = self.config.encoder_config if isinstance(self.config, ParakeetCTCConfig) else self.config
+
+ kernel_size = encoder_config.subsampling_conv_kernel_size
+ stride = encoder_config.subsampling_conv_stride
+ num_layers = int(math.log2(encoder_config.subsampling_factor))
+
+ all_paddings = (kernel_size - 1) // 2 * 2
+ add_pad = all_paddings - kernel_size
+ lengths = input_lengths
+
+ for _ in range(num_layers):
+ lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + 1.0
+ lengths = torch.floor(lengths)
+
+ return lengths.to(dtype=torch.int)
+
+ def _get_output_attention_mask(self, attention_mask: torch.Tensor, target_length: Optional[int] = None):
+ """
+ Convert the input attention mask to its subsampled form. `target_length` sets the desired output length, useful
+ when the attention mask length differs from `sum(-1).max()` (i.e., when the longest sequence in the batch is padded)
+ """
+ output_lengths = self._get_subsampling_output_length(attention_mask.sum(-1))
+ # Use target_length if provided, otherwise use max length in batch
+ max_length = target_length if target_length is not None else output_lengths.max()
+ attention_mask = torch.arange(max_length, device=attention_mask.device) < output_lengths[:, None]
+ return attention_mask
+
+
+@auto_docstring(
+ custom_intro="""
+ The Parakeet Encoder model, based on the [Fast Conformer architecture](https://huggingface.co/papers/2305.05084).
+ """
+)
+class ParakeetEncoder(ParakeetPreTrainedModel):
+ config: ParakeetEncoderConfig
+ base_model_prefix = "encoder"
+
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__(config)
+ self.config = config
+ self.gradient_checkpointing = False
+
+ self.dropout = config.dropout
+ self.dropout_positions = config.dropout_positions
+ self.layerdrop = config.layerdrop
+
+ self.input_scale = math.sqrt(config.hidden_size) if config.scale_input else 1.0
+ self.subsampling = ParakeetEncoderSubsamplingConv2D(config)
+ self.encode_positions = ParakeetEncoderRelPositionalEncoding(config)
+
+ self.layers = nn.ModuleList(
+ [ParakeetEncoderBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+ self.post_init()
+
+ @auto_docstring
+ @check_model_inputs
+ @can_return_tuple
+ def forward(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetEncoder
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> encoder = ParakeetEncoder.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"])
+ >>> encoder_outputs = encoder(**inputs)
+
+ >>> print(encoder_outputs.last_hidden_state.shape)
+ ```
+ """
+
+ hidden_states = self.subsampling(input_features, attention_mask)
+ hidden_states = hidden_states * self.input_scale
+ position_embeddings = self.encode_positions(hidden_states)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ position_embeddings = nn.functional.dropout(
+ position_embeddings, p=self.dropout_positions, training=self.training
+ )
+
+ if attention_mask is not None:
+ attention_mask = self._get_output_attention_mask(attention_mask, target_length=hidden_states.shape[1])
+ attention_mask = attention_mask.unsqueeze(1).expand(-1, hidden_states.shape[1], -1)
+ attention_mask = attention_mask & attention_mask.transpose(1, 2)
+ attention_mask = attention_mask.unsqueeze(1)
+
+ for encoder_layer in self.layers:
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+@dataclass
+class ParakeetGenerateOutput(ModelOutput):
+ """
+ Outputs of Parakeet models.
+
+ Args:
+ sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+ if all batches finished early due to the `eos_token_id`.
+ logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+ Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+ at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+ each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+ attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+ Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+ `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+ hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+ Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+ `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+ """
+
+ sequences: torch.LongTensor
+ logits: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+ hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Parakeet Encoder with a Connectionist Temporal Classification (CTC) head.
+ """
+)
+class ParakeetForCTC(ParakeetPreTrainedModel):
+ config: ParakeetCTCConfig
+
+ def __init__(self, config: ParakeetCTCConfig):
+ super().__init__(config)
+ self.encoder = ParakeetEncoder(config.encoder_config)
+ # Conv rather than linear to be consistent with NeMO decoding layer
+ self.ctc_head = nn.Conv1d(config.encoder_config.hidden_size, config.vocab_size, kernel_size=1)
+
+ self.post_init()
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetForCTC
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> model = ParakeetForCTC.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"])
+ >>> outputs = model(**inputs)
+
+ >>> print(outputs.loss)
+ ```"""
+
+ encoder_outputs = self.encoder(
+ input_features=input_features,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = encoder_outputs.last_hidden_state
+ logits = self.ctc_head(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ loss = None
+ if labels is not None:
+ # retrieve loss input_lengths from attention_mask
+ attention_mask = (
+ attention_mask if attention_mask is not None else torch.ones_like(input_features, dtype=torch.long)
+ )
+ input_lengths = self._get_subsampling_output_length(attention_mask.sum(-1))
+
+ # assuming that padded tokens are filled with -100
+ # when not being attended to
+ labels_mask = labels != self.config.pad_token_id
+ target_lengths = labels_mask.sum(-1)
+ flattened_targets = labels.masked_select(labels_mask)
+
+ # ctc_loss doesn't support fp16
+ log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
+
+ with torch.backends.cudnn.flags(enabled=False):
+ loss = nn.functional.ctc_loss(
+ log_probs,
+ flattened_targets,
+ input_lengths,
+ target_lengths,
+ blank=self.config.pad_token_id,
+ reduction=self.config.ctc_loss_reduction,
+ zero_infinity=self.config.ctc_zero_infinity,
+ )
+
+ return CausalLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ return_dict_in_generate: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[ParakeetGenerateOutput, torch.LongTensor]:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetForCTC
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> model = ParakeetForCTC.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"])
+ >>> predicted_ids = model.generate(**inputs)
+ >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+
+ >>> print(transcription)
+ ```
+ """
+ kwargs["return_dict"] = True
+ outputs: CausalLMOutput = self.forward(
+ input_features=input_features,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ # greedy decoding
+ sequences = outputs.logits.argmax(dim=-1)
+
+ # mask out padded tokens
+ if attention_mask is not None:
+ attention_mask = self._get_output_attention_mask(attention_mask, target_length=sequences.shape[1])
+ sequences[~attention_mask] = self.config.pad_token_id
+
+ if return_dict_in_generate:
+ return ParakeetGenerateOutput(
+ sequences=sequences,
+ logits=outputs.logits,
+ attentions=outputs.attentions,
+ hidden_states=outputs.hidden_states,
+ )
+
+ return sequences
+
+
+__all__ = ["ParakeetForCTC", "ParakeetEncoder", "ParakeetPreTrainedModel"]
diff --git a/src/transformers/models/parakeet/modular_parakeet.py b/src/transformers/models/parakeet/modular_parakeet.py
new file mode 100644
index 000000000000..489e0f9cc056
--- /dev/null
+++ b/src/transformers/models/parakeet/modular_parakeet.py
@@ -0,0 +1,628 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Parakeet model."""
+
+import math
+from dataclasses import dataclass
+from typing import Callable, Optional, Union
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, CausalLMOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import check_model_inputs
+from ..fastspeech2_conformer.modeling_fastspeech2_conformer import FastSpeech2ConformerConvolutionModule
+from ..llama.modeling_llama import LlamaAttention, eager_attention_forward
+from .configuration_parakeet import ParakeetCTCConfig, ParakeetEncoderConfig
+
+
+class ParakeetEncoderRelPositionalEncoding(nn.Module):
+ """Relative positional encoding for Parakeet."""
+
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: ParakeetEncoderConfig, device=None):
+ super().__init__()
+ self.max_position_embeddings = config.max_position_embeddings
+ base = 10000.0
+ inv_freq = 1.0 / (
+ base
+ ** (
+ torch.arange(0, config.hidden_size, 2, dtype=torch.int64).to(device=device, dtype=torch.float)
+ / config.hidden_size
+ )
+ )
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ @torch.no_grad()
+ def forward(self, hidden_states: torch.Tensor):
+ seq_length = hidden_states.shape[1]
+ if seq_length > self.max_position_embeddings:
+ raise ValueError(
+ f"Sequence Length: {seq_length} has to be less or equal than "
+ f"config.max_position_embeddings {self.max_position_embeddings}."
+ )
+
+ position_ids = torch.arange(seq_length - 1, -seq_length, -1, device=hidden_states.device)
+ inv_freq_expanded = (
+ self.inv_freq[None, :, None].float().expand(hidden_states.shape[0], -1, 1).to(hidden_states.device)
+ )
+ position_ids_expanded = position_ids[None, None, :].float()
+
+ device_type = (
+ hidden_states.device.type
+ if isinstance(hidden_states.device.type, str) and hidden_states.device.type != "mps"
+ else "cpu"
+ )
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ sin = freqs.sin()
+ cos = freqs.cos()
+ # interleave sin and cos
+ pos_embed = torch.stack([sin, cos], dim=-1)
+ pos_embed = pos_embed.reshape(*pos_embed.shape[:-2], -1)
+
+ return pos_embed.to(dtype=hidden_states.dtype)
+
+
+class ParakeetEncoderFeedForward(nn.Module):
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__()
+ self.linear1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias)
+ self.activation = ACT2FN[config.hidden_act]
+ self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias)
+ self.activation_dropout = config.activation_dropout
+
+ def forward(self, hidden_states):
+ hidden_states = self.activation(self.linear1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.linear2(hidden_states)
+ return hidden_states
+
+
+class ParakeetEncoderConvolutionModule(FastSpeech2ConformerConvolutionModule):
+ def __init__(self, config: ParakeetEncoderConfig, module_config=None):
+ super().__init__(config, module_config)
+
+
+class ParakeetEncoderAttention(LlamaAttention):
+ """Multi-head attention with relative positional encoding. See section 3.3 of https://huggingface.co/papers/1901.02860."""
+
+ def __init__(self, config: ParakeetEncoderConfig, layer_idx: int):
+ super().__init__(config, layer_idx=layer_idx)
+ self.is_causal = False
+ # W_{k,R} projection
+ self.relative_k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+ # global content bias
+ self.bias_u = nn.Parameter(torch.zeros(config.num_attention_heads, self.head_dim))
+ # global positional bias
+ self.bias_v = nn.Parameter(torch.zeros(config.num_attention_heads, self.head_dim))
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: Optional[torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ batch_size, seq_length = input_shape
+ hidden_shape = (batch_size, seq_length, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ query_states_with_bias_u = query_states + self.bias_u.view(
+ 1, self.config.num_attention_heads, 1, self.head_dim
+ )
+ query_states_with_bias_v = query_states + self.bias_v.view(
+ 1, self.config.num_attention_heads, 1, self.head_dim
+ )
+
+ relative_key_states = self.relative_k_proj(position_embeddings)
+ relative_key_states = relative_key_states.view(batch_size, -1, self.config.num_attention_heads, self.head_dim)
+
+ # terms (b) and (d)
+ matrix_bd = query_states_with_bias_v @ relative_key_states.permute(0, 2, 3, 1)
+ matrix_bd = self._rel_shift(matrix_bd)
+ matrix_bd = matrix_bd[..., :seq_length]
+ matrix_bd = matrix_bd * self.scaling
+
+ if attention_mask is not None:
+ # here the original codebase uses -10000.0 rather than float("-inf") and then manual masked fill with 0.0s
+ # see: https://github.com/NVIDIA-NeMo/NeMo/blob/8cfedd7203462cb251a914e700e5605444277561/nemo/collections/asr/parts/submodules/multi_head_attention.py#L320-L340
+ # we rather went for a straight-forward approach with float("-inf")
+ matrix_bd = matrix_bd.masked_fill_(attention_mask.logical_not(), float("-inf"))
+
+ # will compute matrix_ac - terms (a) and (c) - and add matrix_bd
+ attn_output, attn_weights = attention_interface(
+ self,
+ query=query_states_with_bias_u,
+ key=key_states,
+ value=value_states,
+ attention_mask=matrix_bd,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+ def _rel_shift(self, attention_scores):
+ """Relative position shift for Shaw et al. style attention. See appendix B of https://huggingface.co/papers/1901.02860."""
+ batch_size, num_heads, query_length, position_length = attention_scores.shape
+ attention_scores = nn.functional.pad(attention_scores, pad=(1, 0))
+ attention_scores = attention_scores.view(batch_size, num_heads, -1, query_length)
+ attention_scores = attention_scores[:, :, 1:].view(batch_size, num_heads, query_length, position_length)
+ return attention_scores
+
+
+class ParakeetEncoderSubsamplingConv2D(nn.Module):
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__()
+
+ self.kernel_size = config.subsampling_conv_kernel_size
+ self.stride = config.subsampling_conv_stride
+ self.channels = config.subsampling_conv_channels
+ self.padding = (self.kernel_size - 1) // 2
+ self.num_layers = int(math.log2(config.subsampling_factor))
+
+ # define layers
+ self.layers = nn.ModuleList()
+ self.layers.append(
+ nn.Conv2d(1, self.channels, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding)
+ )
+ self.layers.append(nn.ReLU())
+ for i in range(self.num_layers - 1):
+ # depthwise conv
+ self.layers.append(
+ nn.Conv2d(
+ self.channels,
+ self.channels,
+ kernel_size=self.kernel_size,
+ stride=self.stride,
+ padding=self.padding,
+ groups=self.channels,
+ )
+ )
+ # pointwise conv
+ self.layers.append(nn.Conv2d(self.channels, self.channels, kernel_size=1))
+ # activation
+ self.layers.append(nn.ReLU())
+
+ out_length = config.num_mel_bins // (self.stride**self.num_layers)
+ self.linear = nn.Linear(config.subsampling_conv_channels * out_length, config.hidden_size, bias=True)
+
+ def _get_output_length(self, input_lengths: torch.Tensor, conv_layer: nn.Conv2d):
+ if hasattr(conv_layer, "stride") and conv_layer.stride != (1, 1):
+ padding = conv_layer.padding
+ kernel_size = conv_layer.kernel_size[0]
+ stride = conv_layer.stride[0]
+
+ output_lengths = (input_lengths + padding[0] + padding[1] - kernel_size) // stride + 1
+ return output_lengths
+
+ return input_lengths
+
+ def forward(self, input_features: torch.Tensor, attention_mask: torch.Tensor = None):
+ hidden_states = input_features.unsqueeze(1)
+ current_lengths = attention_mask.sum(-1) if attention_mask is not None else None
+
+ for layer in self.layers:
+ hidden_states = layer(hidden_states)
+
+ # mask the hidden states
+ if isinstance(layer, nn.Conv2d) and attention_mask is not None:
+ current_lengths = self._get_output_length(current_lengths, layer)
+ current_seq_length = hidden_states.shape[2]
+ channel_mask = (
+ torch.arange(current_seq_length, device=attention_mask.device) < current_lengths[:, None]
+ )
+ hidden_states *= channel_mask[:, None, :, None]
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(hidden_states.shape[0], hidden_states.shape[2], -1)
+ hidden_states = self.linear(hidden_states)
+
+ return hidden_states
+
+
+class ParakeetEncoderBlock(GradientCheckpointingLayer):
+ def __init__(self, config: ParakeetEncoderConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.gradient_checkpointing = False
+
+ self.feed_forward1 = ParakeetEncoderFeedForward(config)
+ self.self_attn = ParakeetEncoderAttention(config, layer_idx)
+ self.conv = ParakeetEncoderConvolutionModule(config)
+ self.feed_forward2 = ParakeetEncoderFeedForward(config)
+
+ self.norm_feed_forward1 = nn.LayerNorm(config.hidden_size)
+ self.norm_self_att = nn.LayerNorm(config.hidden_size)
+ self.norm_conv = nn.LayerNorm(config.hidden_size)
+ self.norm_feed_forward2 = nn.LayerNorm(config.hidden_size)
+ self.norm_out = nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states))
+ hidden_states = residual + 0.5 * hidden_states # the conformer architecture uses a factor of 0.5
+
+ normalized_hidden_states = self.norm_self_att(hidden_states)
+ attn_output, _ = self.self_attn(
+ hidden_states=normalized_hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + attn_output
+
+ conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask)
+ hidden_states = hidden_states + conv_output
+
+ ff2_output = self.feed_forward2(self.norm_feed_forward2(hidden_states))
+ hidden_states = hidden_states + 0.5 * ff2_output # the conformer architecture uses a factor of 0.5
+
+ hidden_states = self.norm_out(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class ParakeetPreTrainedModel(PreTrainedModel):
+ config: ParakeetCTCConfig
+ base_model_prefix = "model"
+ main_input_name = "input_features"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["ParakeetEncoderBlock"]
+ _supports_flat_attention_mask = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ # TODO: @eustlb, add support when flash attention supports custom attention bias
+ _supports_flash_attn = False
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": ParakeetEncoderBlock,
+ "attentions": ParakeetEncoderAttention,
+ }
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ if hasattr(self.config, "initializer_range"):
+ std = self.config.initializer_range
+ else:
+ # 0.02 is the standard default value accross the library
+ std = getattr(self.config.get_text_config(), "initializer_range", 0.02)
+
+ if isinstance(module, ParakeetEncoderAttention):
+ # Initialize positional bias parameters
+ module.bias_u.data.normal_(mean=0.0, std=std)
+ module.bias_v.data.normal_(mean=0.0, std=std)
+
+ def _get_subsampling_output_length(self, input_lengths: torch.Tensor):
+ encoder_config = self.config.encoder_config if isinstance(self.config, ParakeetCTCConfig) else self.config
+
+ kernel_size = encoder_config.subsampling_conv_kernel_size
+ stride = encoder_config.subsampling_conv_stride
+ num_layers = int(math.log2(encoder_config.subsampling_factor))
+
+ all_paddings = (kernel_size - 1) // 2 * 2
+ add_pad = all_paddings - kernel_size
+ lengths = input_lengths
+
+ for _ in range(num_layers):
+ lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + 1.0
+ lengths = torch.floor(lengths)
+
+ return lengths.to(dtype=torch.int)
+
+ def _get_output_attention_mask(self, attention_mask: torch.Tensor, target_length: Optional[int] = None):
+ """
+ Convert the input attention mask to its subsampled form. `target_length` sets the desired output length, useful
+ when the attention mask length differs from `sum(-1).max()` (i.e., when the longest sequence in the batch is padded)
+ """
+ output_lengths = self._get_subsampling_output_length(attention_mask.sum(-1))
+ # Use target_length if provided, otherwise use max length in batch
+ max_length = target_length if target_length is not None else output_lengths.max()
+ attention_mask = torch.arange(max_length, device=attention_mask.device) < output_lengths[:, None]
+ return attention_mask
+
+
+@auto_docstring(
+ custom_intro="""
+ The Parakeet Encoder model, based on the [Fast Conformer architecture](https://huggingface.co/papers/2305.05084).
+ """
+)
+class ParakeetEncoder(ParakeetPreTrainedModel):
+ config: ParakeetEncoderConfig
+ base_model_prefix = "encoder"
+
+ def __init__(self, config: ParakeetEncoderConfig):
+ super().__init__(config)
+ self.config = config
+ self.gradient_checkpointing = False
+
+ self.dropout = config.dropout
+ self.dropout_positions = config.dropout_positions
+ self.layerdrop = config.layerdrop
+
+ self.input_scale = math.sqrt(config.hidden_size) if config.scale_input else 1.0
+ self.subsampling = ParakeetEncoderSubsamplingConv2D(config)
+ self.encode_positions = ParakeetEncoderRelPositionalEncoding(config)
+
+ self.layers = nn.ModuleList(
+ [ParakeetEncoderBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+ self.post_init()
+
+ @auto_docstring
+ @check_model_inputs
+ @can_return_tuple
+ def forward(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetEncoder
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> encoder = ParakeetEncoder.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"])
+ >>> encoder_outputs = encoder(**inputs)
+
+ >>> print(encoder_outputs.last_hidden_state.shape)
+ ```
+ """
+
+ hidden_states = self.subsampling(input_features, attention_mask)
+ hidden_states = hidden_states * self.input_scale
+ position_embeddings = self.encode_positions(hidden_states)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ position_embeddings = nn.functional.dropout(
+ position_embeddings, p=self.dropout_positions, training=self.training
+ )
+
+ if attention_mask is not None:
+ attention_mask = self._get_output_attention_mask(attention_mask, target_length=hidden_states.shape[1])
+ attention_mask = attention_mask.unsqueeze(1).expand(-1, hidden_states.shape[1], -1)
+ attention_mask = attention_mask & attention_mask.transpose(1, 2)
+ attention_mask = attention_mask.unsqueeze(1)
+
+ for encoder_layer in self.layers:
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+@dataclass
+class ParakeetGenerateOutput(ModelOutput):
+ """
+ Outputs of Parakeet models.
+
+ Args:
+ sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+ if all batches finished early due to the `eos_token_id`.
+ logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+ Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+ at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+ each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+ attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+ Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+ `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+ hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+ Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+ `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+ """
+
+ sequences: torch.LongTensor
+ logits: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+ hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Parakeet Encoder with a Connectionist Temporal Classification (CTC) head.
+ """
+)
+class ParakeetForCTC(ParakeetPreTrainedModel):
+ config: ParakeetCTCConfig
+
+ def __init__(self, config: ParakeetCTCConfig):
+ super().__init__(config)
+ self.encoder = ParakeetEncoder(config.encoder_config)
+ # Conv rather than linear to be consistent with NeMO decoding layer
+ self.ctc_head = nn.Conv1d(config.encoder_config.hidden_size, config.vocab_size, kernel_size=1)
+
+ self.post_init()
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetForCTC
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> model = ParakeetForCTC.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"])
+ >>> outputs = model(**inputs)
+
+ >>> print(outputs.loss)
+ ```"""
+
+ encoder_outputs = self.encoder(
+ input_features=input_features,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = encoder_outputs.last_hidden_state
+ logits = self.ctc_head(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ loss = None
+ if labels is not None:
+ # retrieve loss input_lengths from attention_mask
+ attention_mask = (
+ attention_mask if attention_mask is not None else torch.ones_like(input_features, dtype=torch.long)
+ )
+ input_lengths = self._get_subsampling_output_length(attention_mask.sum(-1))
+
+ # assuming that padded tokens are filled with -100
+ # when not being attended to
+ labels_mask = labels != self.config.pad_token_id
+ target_lengths = labels_mask.sum(-1)
+ flattened_targets = labels.masked_select(labels_mask)
+
+ # ctc_loss doesn't support fp16
+ log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
+
+ with torch.backends.cudnn.flags(enabled=False):
+ loss = nn.functional.ctc_loss(
+ log_probs,
+ flattened_targets,
+ input_lengths,
+ target_lengths,
+ blank=self.config.pad_token_id,
+ reduction=self.config.ctc_loss_reduction,
+ zero_infinity=self.config.ctc_zero_infinity,
+ )
+
+ return CausalLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ return_dict_in_generate: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[ParakeetGenerateOutput, torch.LongTensor]:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, ParakeetForCTC
+ >>> from datasets import load_dataset, Audio
+
+ >>> model_id = "nvidia/parakeet-ctc-1.1b"
+ >>> processor = AutoProcessor.from_pretrained(model_id)
+ >>> model = ParakeetForCTC.from_pretrained(model_id)
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+
+ >>> inputs = processor(ds[0]["audio"]["array"], text=ds[0]["text"])
+ >>> predicted_ids = model.generate(**inputs)
+ >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+
+ >>> print(transcription)
+ ```
+ """
+ kwargs["return_dict"] = True
+ outputs: CausalLMOutput = self.forward(
+ input_features=input_features,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ # greedy decoding
+ sequences = outputs.logits.argmax(dim=-1)
+
+ # mask out padded tokens
+ if attention_mask is not None:
+ attention_mask = self._get_output_attention_mask(attention_mask, target_length=sequences.shape[1])
+ sequences[~attention_mask] = self.config.pad_token_id
+
+ if return_dict_in_generate:
+ return ParakeetGenerateOutput(
+ sequences=sequences,
+ logits=outputs.logits,
+ attentions=outputs.attentions,
+ hidden_states=outputs.hidden_states,
+ )
+
+ return sequences
+
+
+__all__ = ["ParakeetForCTC", "ParakeetEncoder", "ParakeetPreTrainedModel"]
diff --git a/src/transformers/models/parakeet/processing_parakeet.py b/src/transformers/models/parakeet/processing_parakeet.py
new file mode 100644
index 000000000000..20b86a28393b
--- /dev/null
+++ b/src/transformers/models/parakeet/processing_parakeet.py
@@ -0,0 +1,87 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional, Union
+
+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 logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class ParakeetProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "audio_kwargs": {
+ "sampling_rate": 16000,
+ "padding": "longest",
+ },
+ "text_kwargs": {
+ "padding": True,
+ "padding_side": "right",
+ "add_special_tokens": False,
+ },
+ "common_kwargs": {"return_tensors": "pt"},
+ }
+
+
+class ParakeetProcessor(ProcessorMixin):
+ attributes = ["feature_extractor", "tokenizer"]
+ feature_extractor_class = "ParakeetFeatureExtractor"
+ tokenizer_class = "ParakeetTokenizerFast"
+
+ def __call__(
+ self,
+ audio: AudioInput,
+ text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput], None] = None,
+ sampling_rate: Optional[int] = None,
+ **kwargs: Unpack[ParakeetProcessorKwargs],
+ ):
+ audio = make_list_of_audio(audio)
+
+ output_kwargs = self._merge_kwargs(
+ ParakeetProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ if sampling_rate is None:
+ logger.warning_once(
+ f"You've provided audio without specifying the sampling rate. It will be assumed to be {output_kwargs['audio_kwargs']['sampling_rate']}, which can result in silent errors."
+ )
+ elif sampling_rate != output_kwargs["audio_kwargs"]["sampling_rate"]:
+ raise ValueError(
+ f"The sampling rate of the audio ({sampling_rate}) does not match the sampling rate of the processor ({output_kwargs['audio_kwargs']['sampling_rate']}). Please provide resampled the audio to the expected sampling rate."
+ )
+
+ if audio is not None:
+ inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
+ if text is not None:
+ encodings = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ if text is None:
+ return inputs
+ else:
+ inputs["labels"] = encodings["input_ids"]
+ return inputs
+
+ @property
+ def model_input_names(self):
+ feature_extractor_input_names = self.feature_extractor.model_input_names
+ return feature_extractor_input_names + ["labels"]
+
+
+__all__ = ["ParakeetProcessor"]
diff --git a/src/transformers/models/parakeet/tokenization_parakeet_fast.py b/src/transformers/models/parakeet/tokenization_parakeet_fast.py
new file mode 100644
index 000000000000..d53eb9c68ad4
--- /dev/null
+++ b/src/transformers/models/parakeet/tokenization_parakeet_fast.py
@@ -0,0 +1,54 @@
+# coding=utf-8
+# 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 itertools
+from typing import Optional, Union
+
+from ...tokenization_utils_fast import PreTrainedTokenizerFast
+
+
+class ParakeetTokenizerFast(PreTrainedTokenizerFast):
+ """
+ Inherits all methods from [`PreTrainedTokenizerFast`]. Users should refer to this superclass for more information regarding those methods,
+ except for `_decode` which is overridden to adapt it to CTC decoding:
+ 1. Group consecutive tokens
+ 2. Filter out the blank token
+ """
+
+ def _decode(
+ self,
+ token_ids: Union[int, list[int]],
+ skip_special_tokens: bool = False,
+ clean_up_tokenization_spaces: Optional[bool] = None,
+ group_tokens: bool = True,
+ **kwargs,
+ ) -> str:
+ if isinstance(token_ids, int):
+ token_ids = [token_ids]
+ if group_tokens:
+ token_ids = [token_group[0] for token_group in itertools.groupby(token_ids)]
+
+ # for CTC we filter out the blank token, which is the pad token
+ token_ids = [token for token in token_ids if token != self.pad_token_id]
+
+ return super()._decode(
+ token_ids=token_ids,
+ skip_special_tokens=skip_special_tokens,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+
+__all__ = ["ParakeetTokenizerFast"]
diff --git a/src/transformers/models/perceiver/image_processing_perceiver_fast.py b/src/transformers/models/perceiver/image_processing_perceiver_fast.py
index 82c1bcd9d319..72cb17cd40cd 100644
--- a/src/transformers/models/perceiver/image_processing_perceiver_fast.py
+++ b/src/transformers/models/perceiver/image_processing_perceiver_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature
from ...image_transforms import group_images_by_shape, reorder_images
@@ -24,16 +25,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
@auto_docstring
class PerceiverImageProcessorFast(BaseImageProcessorFast):
resample = PILImageResampling.BICUBIC
diff --git a/src/transformers/models/perceiver/modeling_perceiver.py b/src/transformers/models/perceiver/modeling_perceiver.py
index f0e4e3e5dbe0..c5e3522ed1d1 100755
--- a/src/transformers/models/perceiver/modeling_perceiver.py
+++ b/src/transformers/models/perceiver/modeling_perceiver.py
@@ -2621,7 +2621,7 @@ def interpolate_pos_encoding(self, position_embeddings: torch.Tensor, height: in
return position_embeddings
def forward(
- self, batch_size: int, interpolate_pos_encoding: bool = False, input_size: torch.Size = None
+ self, batch_size: int, interpolate_pos_encoding: bool = False, input_size: Optional[torch.Size] = None
) -> torch.Tensor:
position_embeddings = self.position_embeddings
@@ -2846,7 +2846,7 @@ class PerceiverAudioPostprocessor(nn.Module):
def __init__(self, config: PerceiverConfig, in_channels: int, postproc_type: str = "patches") -> None:
super().__init__()
- if postproc_type not in ("patches",): # to be supported: 'conv', 'patches', 'pixels'
+ if postproc_type != "patches": # to be supported: 'conv', 'patches', 'pixels'
raise ValueError("Invalid postproc_type!")
# Architecture parameters:
@@ -3179,7 +3179,7 @@ def __init__(
super().__init__()
self.config = config
- if prep_type not in ("patches",):
+ if prep_type != "patches":
raise ValueError(f"Prep_type {prep_type} is invalid, can only be 'patches'.")
if concat_or_add_pos not in ["concat", "add"]:
diff --git a/src/transformers/models/perception_lm/configuration_perception_lm.py b/src/transformers/models/perception_lm/configuration_perception_lm.py
index 4b94652e2084..08c084065ff8 100644
--- a/src/transformers/models/perception_lm/configuration_perception_lm.py
+++ b/src/transformers/models/perception_lm/configuration_perception_lm.py
@@ -68,7 +68,7 @@ def __init__(
if isinstance(vision_config, dict):
vision_config = TimmWrapperConfig(**vision_config)
elif isinstance(vision_config, TimmWrapperConfig):
- vision_config = vision_config
+ pass
elif vision_config is None:
vision_config = TimmWrapperConfig()
self.vision_config = vision_config
diff --git a/src/transformers/models/perception_lm/image_processing_perception_lm_fast.py b/src/transformers/models/perception_lm/image_processing_perception_lm_fast.py
index be55c39572d5..c26132a48439 100644
--- a/src/transformers/models/perception_lm/image_processing_perception_lm_fast.py
+++ b/src/transformers/models/perception_lm/image_processing_perception_lm_fast.py
@@ -190,7 +190,7 @@ def _fit_image_to_canvas(self, img_width: int, img_height: int, tile_size: int):
target_width=n_w * tile_size,
target_height=n_h * tile_size,
)
- # Llama3V dynamic tiling. Priortize biggest canvas.
+ # Llama3V dynamic tiling. Prioritize biggest canvas.
if (scale < 1.0 and (image_width_height[0] >= optimal_image_width_height[0])) or (
scale >= 1.0 and (image_width_height[1] >= optimal_image_width_height[1])
):
diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py
index a9009e837db5..33cee6b37ba5 100644
--- a/src/transformers/models/phi3/configuration_phi3.py
+++ b/src/transformers/models/phi3/configuration_phi3.py
@@ -210,7 +210,7 @@ def _rope_scaling_validation(self):
rope_scaling_type = self.rope_scaling.get("type", None)
rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["longrope"]:
+ if rope_scaling_type is None or rope_scaling_type != "longrope":
raise ValueError(f"`rope_scaling`'s type field must be one of ['longrope'], got {rope_scaling_type}")
if not (
isinstance(rope_scaling_short_factor, list)
diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
index 3b6c2ca1d979..e5e5ca91bfce 100644
--- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
+++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
@@ -137,7 +137,7 @@ class Phi4MultimodalAudioConfig(PretrainedConfig):
The dropout ratio.
ext_pw_out_channel (`int`, *optional*, defaults to 1024):
Number of out channels in the point-wise conv modules.
- depthwise_seperable_out_channel (`int`, *optional*, defaults to 1024):
+ depthwise_separable_out_channel (`int`, *optional*, defaults to 1024):
Number of out channels in the depth-wise separable conv modules.
depthwise_multiplier (`int`, *optional*, defaults to 1):
Input size multiplier for the depth-wise separable conv modules.
@@ -190,7 +190,7 @@ def __init__(
left_chunk: int = 18,
dropout_rate: float = 0.0,
ext_pw_out_channel: int = 1024,
- depthwise_seperable_out_channel: int = 1024,
+ depthwise_separable_out_channel: int = 1024,
depthwise_multiplier: int = 1,
kernel_size: int = 3,
conv_activation: str = "swish",
@@ -217,7 +217,7 @@ def __init__(
self.num_blocks = num_blocks
self.dropout_rate = dropout_rate
self.ext_pw_out_channel = ext_pw_out_channel
- self.depthwise_seperable_out_channel = depthwise_seperable_out_channel
+ self.depthwise_separable_out_channel = depthwise_separable_out_channel
self.depthwise_multiplier = depthwise_multiplier
self.kernel_size = kernel_size
self.conv_activation = conv_activation
@@ -236,7 +236,7 @@ def __init__(
if time_reduction % 2 != 0:
raise ValueError("`time_reduction` should be a multiple of 2!")
length = input_size
- for _ in range(int(math.log(time_reduction, 2))):
+ for _ in range(int(math.log2(time_reduction))):
length = math.floor((length - 1) / 2 + 1)
self.nemo_final_size = length
@@ -452,7 +452,7 @@ def _rope_scaling_validation(self):
rope_scaling_type = self.rope_scaling.get("type", None)
rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["longrope"]:
+ if rope_scaling_type is None or rope_scaling_type != "longrope":
raise ValueError(f"`rope_scaling`'s type field must be one of ['longrope'], got {rope_scaling_type}")
if not (
isinstance(rope_scaling_short_factor, list)
diff --git a/src/transformers/models/phi4_multimodal/image_processing_phi4_multimodal_fast.py b/src/transformers/models/phi4_multimodal/image_processing_phi4_multimodal_fast.py
index 532136f8108e..4bd9928daa94 100644
--- a/src/transformers/models/phi4_multimodal/image_processing_phi4_multimodal_fast.py
+++ b/src/transformers/models/phi4_multimodal/image_processing_phi4_multimodal_fast.py
@@ -16,6 +16,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -27,16 +28,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py
index 349f2e02e2f2..bb495642c710 100644
--- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py
+++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py
@@ -746,7 +746,7 @@ def forward(
return attn_output
-class Phi4MultimodalAudioDepthWiseSeperableConv1d(nn.Module):
+class Phi4MultimodalAudioDepthWiseSeparableConv1d(nn.Module):
def __init__(self, config: Phi4MultimodalAudioConfig, padding: int = 0):
super().__init__()
self.dw_conv = nn.Conv1d(
@@ -758,7 +758,7 @@ def __init__(self, config: Phi4MultimodalAudioConfig, padding: int = 0):
groups=config.hidden_size,
)
self.pw_conv = nn.Conv1d(
- config.hidden_size * config.depthwise_multiplier, config.depthwise_seperable_out_channel, 1, 1, 0
+ config.hidden_size * config.depthwise_multiplier, config.depthwise_separable_out_channel, 1, 1, 0
)
def forward(self, hidden_states):
@@ -794,7 +794,7 @@ def __init__(self, config: Phi4MultimodalAudioConfig):
self.layer_norm = nn.LayerNorm(config.hidden_size)
self.glu = Phi4MultimodalAudioGluPointWiseConv(config)
- self.dw_sep_conv_1d = Phi4MultimodalAudioDepthWiseSeperableConv1d(config, padding=config.kernel_size - 1)
+ self.dw_sep_conv_1d = Phi4MultimodalAudioDepthWiseSeparableConv1d(config, padding=config.kernel_size - 1)
self.act = ACT2FN[config.conv_activation]
self.ext_pw_conv_1d = nn.Conv1d(config.hidden_size, config.ext_pw_out_channel, kernel_size=1, stride=1)
self.dropout = nn.Dropout(config.dropout_rate)
@@ -844,7 +844,7 @@ class Phi4MultimodalAudioNemoConvSubsampling(torch.nn.Module):
def __init__(self, config: Phi4MultimodalAudioConfig):
super().__init__()
self.subsampling_factor = config.time_reduction
- self.sampling_num = int(math.log(self.subsampling_factor, 2))
+ self.sampling_num = int(math.log2(self.subsampling_factor))
self.act_fn = ACT2FN[config.nemo_activation]
conv_channels = config.nemo_conv_channels
diff --git a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py
index ea226e4e1981..bea02eef03c0 100644
--- a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py
+++ b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py
@@ -174,7 +174,7 @@ class Phi4MultimodalAudioConfig(PretrainedConfig):
The dropout ratio.
ext_pw_out_channel (`int`, *optional*, defaults to 1024):
Number of out channels in the point-wise conv modules.
- depthwise_seperable_out_channel (`int`, *optional*, defaults to 1024):
+ depthwise_separable_out_channel (`int`, *optional*, defaults to 1024):
Number of out channels in the depth-wise separable conv modules.
depthwise_multiplier (`int`, *optional*, defaults to 1):
Input size multiplier for the depth-wise separable conv modules.
@@ -227,7 +227,7 @@ def __init__(
left_chunk: int = 18,
dropout_rate: float = 0.0,
ext_pw_out_channel: int = 1024,
- depthwise_seperable_out_channel: int = 1024,
+ depthwise_separable_out_channel: int = 1024,
depthwise_multiplier: int = 1,
kernel_size: int = 3,
conv_activation: str = "swish",
@@ -254,7 +254,7 @@ def __init__(
self.num_blocks = num_blocks
self.dropout_rate = dropout_rate
self.ext_pw_out_channel = ext_pw_out_channel
- self.depthwise_seperable_out_channel = depthwise_seperable_out_channel
+ self.depthwise_separable_out_channel = depthwise_separable_out_channel
self.depthwise_multiplier = depthwise_multiplier
self.kernel_size = kernel_size
self.conv_activation = conv_activation
@@ -273,7 +273,7 @@ def __init__(
if time_reduction % 2 != 0:
raise ValueError("`time_reduction` should be a multiple of 2!")
length = input_size
- for _ in range(int(math.log(time_reduction, 2))):
+ for _ in range(int(math.log2(time_reduction))):
length = math.floor((length - 1) / 2 + 1)
self.nemo_final_size = length
@@ -930,7 +930,7 @@ def forward(
return attn_output
-class Phi4MultimodalAudioDepthWiseSeperableConv1d(nn.Module):
+class Phi4MultimodalAudioDepthWiseSeparableConv1d(nn.Module):
def __init__(self, config: Phi4MultimodalAudioConfig, padding: int = 0):
super().__init__()
self.dw_conv = nn.Conv1d(
@@ -942,7 +942,7 @@ def __init__(self, config: Phi4MultimodalAudioConfig, padding: int = 0):
groups=config.hidden_size,
)
self.pw_conv = nn.Conv1d(
- config.hidden_size * config.depthwise_multiplier, config.depthwise_seperable_out_channel, 1, 1, 0
+ config.hidden_size * config.depthwise_multiplier, config.depthwise_separable_out_channel, 1, 1, 0
)
def forward(self, hidden_states):
@@ -978,7 +978,7 @@ def __init__(self, config: Phi4MultimodalAudioConfig):
self.layer_norm = nn.LayerNorm(config.hidden_size)
self.glu = Phi4MultimodalAudioGluPointWiseConv(config)
- self.dw_sep_conv_1d = Phi4MultimodalAudioDepthWiseSeperableConv1d(config, padding=config.kernel_size - 1)
+ self.dw_sep_conv_1d = Phi4MultimodalAudioDepthWiseSeparableConv1d(config, padding=config.kernel_size - 1)
self.act = ACT2FN[config.conv_activation]
self.ext_pw_conv_1d = nn.Conv1d(config.hidden_size, config.ext_pw_out_channel, kernel_size=1, stride=1)
self.dropout = nn.Dropout(config.dropout_rate)
@@ -1028,7 +1028,7 @@ class Phi4MultimodalAudioNemoConvSubsampling(torch.nn.Module):
def __init__(self, config: Phi4MultimodalAudioConfig):
super().__init__()
self.subsampling_factor = config.time_reduction
- self.sampling_num = int(math.log(self.subsampling_factor, 2))
+ self.sampling_num = int(math.log2(self.subsampling_factor))
self.act_fn = ACT2FN[config.nemo_activation]
conv_channels = config.nemo_conv_channels
diff --git a/src/transformers/models/pix2struct/modeling_pix2struct.py b/src/transformers/models/pix2struct/modeling_pix2struct.py
index 463fec98256f..ee0631611070 100644
--- a/src/transformers/models/pix2struct/modeling_pix2struct.py
+++ b/src/transformers/models/pix2struct/modeling_pix2struct.py
@@ -86,7 +86,7 @@ def forward(self, hidden_states):
try:
from apex.normalization import FusedRMSNorm
- Pix2StructLayerNorm = FusedRMSNorm # noqa
+ Pix2StructLayerNorm = FusedRMSNorm
logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of Pix2StructLayerNorm")
except ImportError:
diff --git a/src/transformers/models/pixtral/image_processing_pixtral.py b/src/transformers/models/pixtral/image_processing_pixtral.py
index c6c6fdb163ab..33e7676f9de9 100644
--- a/src/transformers/models/pixtral/image_processing_pixtral.py
+++ b/src/transformers/models/pixtral/image_processing_pixtral.py
@@ -302,8 +302,8 @@ def _pad_for_batching(
"""
max_shape = (
- max([size[0] for size in image_sizes]),
- max([size[1] for size in image_sizes]),
+ max(size[0] for size in image_sizes),
+ max(size[1] for size in image_sizes),
)
pixel_values = [
pad(
diff --git a/src/transformers/models/pixtral/image_processing_pixtral_fast.py b/src/transformers/models/pixtral/image_processing_pixtral_fast.py
index db3e75760318..b31f910e4817 100644
--- a/src/transformers/models/pixtral/image_processing_pixtral_fast.py
+++ b/src/transformers/models/pixtral/image_processing_pixtral_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -30,17 +31,11 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from .image_processing_pixtral import get_resize_output_image_size
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -131,7 +126,7 @@ def _pad_for_batching(
list[`torch.Tensor`]: The padded images.
"""
- max_shape = (max([size[0] for size in image_sizes]), max([size[1] for size in image_sizes]))
+ max_shape = (max(size[0] for size in image_sizes), max(size[1] for size in image_sizes))
pixel_values = [
torch.nn.functional.pad(image, pad=(0, max_shape[1] - size[1], 0, max_shape[0] - size[0]))
for image, size in zip(pixel_values, image_sizes)
diff --git a/src/transformers/models/poolformer/image_processing_poolformer_fast.py b/src/transformers/models/poolformer/image_processing_poolformer_fast.py
index 70c6ed55bc8a..62d5f276859f 100644
--- a/src/transformers/models/poolformer/image_processing_poolformer_fast.py
+++ b/src/transformers/models/poolformer/image_processing_poolformer_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs
from ...image_transforms import (
@@ -38,16 +39,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class PoolFormerFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
Args:
diff --git a/src/transformers/models/pop2piano/configuration_pop2piano.py b/src/transformers/models/pop2piano/configuration_pop2piano.py
index 484e1a4f933e..6bc90961154b 100644
--- a/src/transformers/models/pop2piano/configuration_pop2piano.py
+++ b/src/transformers/models/pop2piano/configuration_pop2piano.py
@@ -87,7 +87,7 @@ def __init__(
dropout_rate=0.1,
layer_norm_epsilon=1e-6,
initializer_factor=1.0,
- feed_forward_proj="gated-gelu", # noqa
+ feed_forward_proj="gated-gelu",
is_encoder_decoder=True,
use_cache=True,
pad_token_id=0,
diff --git a/src/transformers/models/pop2piano/modeling_pop2piano.py b/src/transformers/models/pop2piano/modeling_pop2piano.py
index ea6d3a5eea9e..996fc5635866 100644
--- a/src/transformers/models/pop2piano/modeling_pop2piano.py
+++ b/src/transformers/models/pop2piano/modeling_pop2piano.py
@@ -88,7 +88,7 @@ def forward(self, hidden_states):
if not _load_pop2piano_layer_norm:
- Pop2PianoLayerNorm = FusedRMSNorm # noqa
+ Pop2PianoLayerNorm = FusedRMSNorm
# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->Pop2Piano,t5->pop2piano
diff --git a/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything_fast.py b/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything_fast.py
index 763fd613c218..06d6ed156443 100644
--- a/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything_fast.py
+++ b/src/transformers/models/prompt_depth_anything/image_processing_prompt_depth_anything_fast.py
@@ -24,6 +24,7 @@
if TYPE_CHECKING:
from ...modeling_outputs import DepthEstimatorOutput
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils_fast import (
BaseImageProcessorFast,
@@ -42,17 +43,10 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
requires_backends,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
def _constrain_to_multiple_of(val, multiple, min_val=0, max_val=None):
"""Constrain a value to be a multiple of another value."""
x = round(val / multiple) * multiple
diff --git a/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py b/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py
index 805338511d8a..5a1fe6bfac6f 100644
--- a/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py
+++ b/src/transformers/models/prophetnet/convert_prophetnet_original_pytorch_checkpoint_to_pytorch.py
@@ -132,9 +132,7 @@ def convert_prophetnet_checkpoint_to_pytorch(prophetnet_checkpoint_path: str, py
else:
model = getattr(model, attribute)
- if old_attribute == "":
- old_model = old_model
- else:
+ if old_attribute:
if not hasattr(old_model, old_attribute):
raise ValueError(f"{old_model} does not have {old_attribute}")
old_model = getattr(old_model, old_attribute)
diff --git a/src/transformers/models/prophetnet/modeling_prophetnet.py b/src/transformers/models/prophetnet/modeling_prophetnet.py
index 260b0c698407..fb78c1e505f2 100644
--- a/src/transformers/models/prophetnet/modeling_prophetnet.py
+++ b/src/transformers/models/prophetnet/modeling_prophetnet.py
@@ -1015,7 +1015,7 @@ def forward(
"""
)
class ProphetNetEncoder(ProphetNetPreTrainedModel):
- def __init__(self, config: ProphetNetConfig, word_embeddings: nn.Embedding = None):
+ def __init__(self, config: ProphetNetConfig, word_embeddings: Optional[nn.Embedding] = None):
r"""
word_embeddings (`torch.nn.Embeddings` of shape `(config.vocab_size, config.hidden_size)`, *optional*):
The word embedding parameters. This can be used to initialize [`ProphetNetEncoder`] with pre-defined word
diff --git a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
index 6b69ced26591..6b8910d270bb 100644
--- a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
@@ -3040,7 +3040,7 @@ def __init__(self, dim, freq_embed_dim=256):
self.time_embed = SinusPositionEmbedding(freq_embed_dim)
self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
- def forward(self, timestep): # noqa: F821
+ def forward(self, timestep):
time_hidden = self.time_embed(timestep)
time_hidden = time_hidden.to(timestep.dtype)
for layer in self.time_mlp:
diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
index afb0cda5ccfe..b63c301f36c3 100644
--- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
@@ -3338,7 +3338,7 @@ def __init__(self, dim, freq_embed_dim=256):
self.time_embed = SinusPositionEmbedding(freq_embed_dim)
self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
- def forward(self, timestep): # noqa: F821
+ def forward(self, timestep):
time_hidden = self.time_embed(timestep)
time_hidden = time_hidden.to(timestep.dtype)
for layer in self.time_mlp:
diff --git a/src/transformers/models/qwen2_5_omni/processing_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/processing_qwen2_5_omni.py
index 45d8cacddeb2..5fcbb0c535f9 100644
--- a/src/transformers/models/qwen2_5_omni/processing_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/processing_qwen2_5_omni.py
@@ -31,10 +31,10 @@
class Qwen2_5_OmniVideosKwargs(VideosKwargs):
- fps: Optional[list[Union[int, float]]] = None
- use_audio_in_video: Optional[bool] = None
- seconds_per_chunk: Optional[float] = None
- position_id_per_seconds: Optional[int] = None
+ fps: Optional[list[Union[int, float]]]
+ use_audio_in_video: Optional[bool]
+ seconds_per_chunk: Optional[float]
+ position_id_per_seconds: Optional[int]
min_pixels: Optional[int]
max_pixels: Optional[int]
patch_size: Optional[int]
@@ -62,8 +62,10 @@ class Qwen2_5OmniProcessorKwargs(ProcessingKwargs, total=False):
"seconds_per_chunk": 2.0,
"position_id_per_seconds": 25,
"use_audio_in_video": False,
- "min_pixels": 128 * 28 * 28,
- "max_pixels": 768 * 28 * 28,
+ "size": {
+ "shortest_edge": 128 * 28 * 28,
+ "longest_edge": 768 * 28 * 28,
+ },
},
"audio_kwargs": {
"sampling_rate": 16000,
diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
index a7a489c3e867..fcd17cb5811f 100644
--- a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
+++ b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
@@ -159,10 +159,6 @@ class Qwen2_5_VLTextConfig(PretrainedConfig):
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
`high_freq_factor` (`float`, *optional*):
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
- image_token_id (`int`, *optional*):
- Token index used as placeholder for image embeddings.
- video_token_id (`int`, *optional*):
- Token index used as placeholder for video embeddings.
```python
>>> from transformers import Qwen2_5_VLTextModel, Qwen2_5_VLConfig
@@ -217,8 +213,6 @@ def __init__(
layer_types=None,
attention_dropout=0.0,
rope_scaling=None,
- image_token_id=None,
- video_token_id=None,
**kwargs,
):
self.vocab_size = vocab_size
@@ -264,9 +258,6 @@ def __init__(
self.rope_scaling["type"] = "default"
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
rope_config_validation(self, ignore_keys={"mrope_section"})
- self.image_token_id = image_token_id
- self.video_token_id = video_token_id
-
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
@@ -290,6 +281,10 @@ class Qwen2_5_VLConfig(PretrainedConfig):
The image token index to encode the image prompt.
video_token_id (`int`, *optional*, defaults to 151656):
The video token index to encode the image prompt.
+ vision_start_token_id (`int`, *optional*, defaults to 151652):
+ The token index to denote start of vision input.
+ vision_end_token_id (`int`, *optional*, defaults to 151653):
+ The token index to denote end of vision input.
```python
>>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
@@ -314,8 +309,15 @@ def __init__(
vision_config=None,
image_token_id=151655,
video_token_id=151656,
+ vision_start_token_id=151652,
+ vision_end_token_id=151653,
**kwargs,
):
+ # We need to init super() here so that it does not reset values
+ # that are in text config to the BaseClass defaults. The Base
+ # config has many text related defaults and not all defaults are same as for `Qwen2_5_VLTextConfig`
+ super().__init__(**kwargs)
+
if isinstance(vision_config, dict):
self.vision_config = self.sub_configs["vision_config"](**vision_config)
elif vision_config is None:
@@ -329,8 +331,32 @@ def __init__(
self.image_token_id = image_token_id
self.video_token_id = video_token_id
-
- super().__init__(**kwargs)
+ self.vision_start_token_id = vision_start_token_id
+ self.vision_end_token_id = vision_end_token_id
+
+ # Attention implementation to use. It sets it recursively on sub-configs so we call it again in the end
+ self._attn_implementation = kwargs.pop("attn_implementation", None)
+
+ def __setattr__(self, key, value):
+ if (
+ (text_config := super().__getattribute__("__dict__").get("text_config")) is not None
+ and key not in ["dtype", "_attn_implementation_internal"]
+ and key in text_config.__dict__
+ ):
+ setattr(text_config, key, value)
+ else:
+ super().__setattr__(key, value)
+
+ def __getattribute__(self, key):
+ if "text_config" in super().__getattribute__("__dict__") and key not in [
+ "dtype",
+ "_attn_implementation_internal",
+ ]:
+ text_config = super().__getattribute__("text_config")
+ if key in text_config.__dict__:
+ return getattr(text_config, key)
+
+ return super().__getattribute__(key)
__all__ = ["Qwen2_5_VLConfig", "Qwen2_5_VLTextConfig"]
diff --git a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
index 6d05cc32f4a8..a98574551922 100644
--- a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
+++ b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
@@ -1558,6 +1558,7 @@ def prepare_inputs_for_generation(
model_inputs.get("input_ids", None),
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
+ second_per_grid_ts=second_per_grid_ts,
attention_mask=attention_mask,
)
self.model.rope_deltas = rope_deltas
diff --git a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py
index b59644c37df9..f2eac303213c 100644
--- a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py
+++ b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py
@@ -814,6 +814,7 @@ def prepare_inputs_for_generation(
model_inputs.get("input_ids", None),
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
+ second_per_grid_ts=second_per_grid_ts,
attention_mask=attention_mask,
)
self.model.rope_deltas = rope_deltas
diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
index 1f9e0a3a5bc4..774e35d30bb2 100644
--- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
+++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
@@ -148,10 +148,6 @@ class Qwen2VLTextConfig(PretrainedConfig):
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
`high_freq_factor` (`float`, *optional*):
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
- image_token_id (`int`, *optional*):
- Token index used as placeholder for image embeddings.
- video_token_id (`int`, *optional*):
- Token index used as placeholder for video embeddings.
```python
>>> from transformers import Qwen2VLTextModel, Qwen2VLConfig
@@ -206,8 +202,6 @@ def __init__(
layer_types=None,
attention_dropout=0.0,
rope_scaling=None,
- image_token_id=None,
- video_token_id=None,
**kwargs,
):
self.vocab_size = vocab_size
@@ -253,9 +247,6 @@ def __init__(
self.rope_scaling["type"] = "default"
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
rope_config_validation(self, ignore_keys={"mrope_section"})
- self.image_token_id = image_token_id
- self.video_token_id = video_token_id
-
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
@@ -271,23 +262,27 @@ class Qwen2VLConfig(PretrainedConfig):
Args:
- text_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen2_5_VLTextConfig`):
+ text_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen2VLTextConfig`):
The config object or dictionary of the text backbone.
- vision_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen2_5_VLVisionConfig`):
+ vision_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen2VLVisionConfig`):
The config object or dictionary of the vision backbone.
image_token_id (`int`, *optional*, defaults to 151655):
The image token index to encode the image prompt.
video_token_id (`int`, *optional*, defaults to 151656):
The video token index to encode the image prompt.
+ vision_start_token_id (`int`, *optional*, defaults to 151652):
+ The token index to denote start of vision input.
+ vision_end_token_id (`int`, *optional*, defaults to 151653):
+ The token index to denote end of vision input.
```python
- >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
+ >>> from transformers import Qwen2VLForConditionalGeneration, Qwen2VLConfig
- >>> # Initializing a Qwen2_5_VL style configuration
- >>> configuration = Qwen2_5_VLConfig()
+ >>> # Initializing a Qwen2VL style configuration
+ >>> configuration = Qwen2VLConfig()
>>> # Initializing a model from the Qwen2-VL-7B style configuration
- >>> model = Qwen2_5_VLForConditionalGeneration(configuration)
+ >>> model = Qwen2VLForConditionalGeneration(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
@@ -303,8 +298,15 @@ def __init__(
vision_config=None,
image_token_id=151655,
video_token_id=151656,
+ vision_start_token_id=151652,
+ vision_end_token_id=151653,
**kwargs,
):
+ # We need to init super() here so that it does not reset values
+ # that are in text config to the BaseClass defaults. The Base
+ # config has many text related defaults and not all defaults are same as for `Qwen2VLTextConfig`
+ super().__init__(**kwargs)
+
if isinstance(vision_config, dict):
self.vision_config = self.sub_configs["vision_config"](**vision_config)
elif vision_config is None:
@@ -318,8 +320,32 @@ def __init__(
self.image_token_id = image_token_id
self.video_token_id = video_token_id
-
- super().__init__(**kwargs)
+ self.vision_start_token_id = vision_start_token_id
+ self.vision_end_token_id = vision_end_token_id
+
+ # Attention implementation to use. It sets it recursively on sub-configs so we call it again in the end
+ self._attn_implementation = kwargs.pop("attn_implementation", None)
+
+ def __setattr__(self, key, value):
+ if (
+ (text_config := super().__getattribute__("__dict__").get("text_config")) is not None
+ and key not in ["dtype", "_attn_implementation_internal"]
+ and key in text_config.__dict__
+ ):
+ setattr(text_config, key, value)
+ else:
+ super().__setattr__(key, value)
+
+ def __getattribute__(self, key):
+ if "text_config" in super().__getattribute__("__dict__") and key not in [
+ "dtype",
+ "_attn_implementation_internal",
+ ]:
+ text_config = super().__getattribute__("text_config")
+ if key in text_config.__dict__:
+ return getattr(text_config, key)
+
+ return super().__getattribute__(key)
__all__ = ["Qwen2VLConfig", "Qwen2VLTextConfig"]
diff --git a/src/transformers/models/qwen2_vl/image_processing_qwen2_vl_fast.py b/src/transformers/models/qwen2_vl/image_processing_qwen2_vl_fast.py
index 80242a331ace..ec9878da3222 100644
--- a/src/transformers/models/qwen2_vl/image_processing_qwen2_vl_fast.py
+++ b/src/transformers/models/qwen2_vl/image_processing_qwen2_vl_fast.py
@@ -22,6 +22,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -42,18 +43,12 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from ...video_utils import VideoInput, make_batched_videos
from .image_processing_qwen2_vl import smart_resize
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py b/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py
index ba87909740a8..84bcd827f02e 100644
--- a/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py
+++ b/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py
@@ -23,6 +23,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_utils import (
@@ -34,18 +35,12 @@
get_image_size,
)
from ...processing_utils import Unpack, VideosKwargs
-from ...utils import TensorType, add_start_docstrings, is_torchvision_v2_available
+from ...utils import TensorType, add_start_docstrings
from ...video_processing_utils import BASE_VIDEO_PROCESSOR_DOCSTRING, BaseVideoProcessor
from ...video_utils import VideoMetadata, group_videos_by_shape, reorder_videos
from .image_processing_qwen2_vl import smart_resize
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class Qwen2VLVideoProcessorInitKwargs(VideosKwargs):
min_pixels: Optional[int]
max_pixels: Optional[int]
@@ -186,7 +181,6 @@ def sample_frames(
def _preprocess(
self,
videos: list["torch.Tensor"],
- do_convert_rgb: bool,
do_resize: bool,
size: SizeDict,
interpolation: Optional["F.InterpolationMode"],
@@ -195,13 +189,10 @@ def _preprocess(
do_normalize: bool,
image_mean: Optional[Union[float, list[float]]],
image_std: Optional[Union[float, list[float]]],
- min_pixels: Optional[int] = None,
- max_pixels: Optional[int] = None,
patch_size: Optional[int] = None,
temporal_patch_size: Optional[int] = None,
merge_size: Optional[int] = None,
return_tensors: Optional[Union[str, TensorType]] = None,
- device: Optional["torch.Tensor"] = None,
**kwargs,
):
# Group videos by size for batched resizing
@@ -215,8 +206,8 @@ def _preprocess(
height,
width,
factor=patch_size * merge_size,
- min_pixels=min_pixels,
- max_pixels=max_pixels,
+ min_pixels=size["shortest_edge"],
+ max_pixels=size["longest_edge"],
)
stacked_videos = self.resize(
image=stacked_videos,
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 7d2b60d943e2..e15e3435f732 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -458,15 +458,15 @@ def torch_chunk_gated_delta_rule(
x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
]
- batch_size, sequence_length, num_heads, k_head_dim = key.shape
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
v_head_dim = value.shape[-1]
- pad_size = (chunk_size - num_heads % chunk_size) % chunk_size
+ pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
query = F.pad(query, (0, 0, 0, pad_size))
key = F.pad(key, (0, 0, 0, pad_size))
value = F.pad(value, (0, 0, 0, pad_size))
beta = F.pad(beta, (0, pad_size))
g = F.pad(g, (0, pad_size))
- tot_heads = num_heads + pad_size
+ total_sequence_length = sequence_length + pad_size
scale = 1 / (query.shape[-1] ** 0.5)
query = query * scale
@@ -491,7 +491,7 @@ def torch_chunk_gated_delta_rule(
value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
last_recurrent_state = (
- torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
if initial_state is None
else initial_state.to(value)
)
@@ -499,7 +499,7 @@ def torch_chunk_gated_delta_rule(
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
# for each chunk
- for i in range(0, tot_heads // chunk_size):
+ for i in range(0, total_sequence_length // chunk_size):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
@@ -514,7 +514,7 @@ def torch_chunk_gated_delta_rule(
if not output_final_state:
last_recurrent_state = None
core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
- core_attn_out = core_attn_out[:, :, :num_heads]
+ core_attn_out = core_attn_out[:, :, :sequence_length]
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
return core_attn_out, last_recurrent_state
@@ -530,19 +530,19 @@ def torch_recurrent_gated_delta_rule(
x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
]
- batch_size, sequence_length, num_heads, k_head_dim = key.shape
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
v_head_dim = value.shape[-1]
scale = 1 / (query.shape[-1] ** 0.5)
query = query * scale
- core_attn_out = torch.zeros(batch_size, sequence_length, num_heads, v_head_dim).to(value)
+ core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)
last_recurrent_state = (
- torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
if initial_state is None
else initial_state.to(value)
)
- for i in range(num_heads):
+ for i in range(sequence_length):
q_t = query[:, :, i]
k_t = key[:, :, i]
v_t = value[:, :, i]
@@ -970,6 +970,9 @@ def _init_weights(self, module):
if isinstance(module, Qwen3NextGatedDeltaNet):
module.dt_bias.data.fill_(1.0)
module.A_log.data.uniform_(0, 16).log_()
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif isinstance(module, Qwen3NextRMSNorm):
+ module.weight.data.zero_()
class Qwen3NextModel(Qwen3NextPreTrainedModel):
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index e141e229eedf..6d4b6a5e04a3 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -293,15 +293,15 @@ def torch_chunk_gated_delta_rule(
x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
]
- batch_size, sequence_length, num_heads, k_head_dim = key.shape
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
v_head_dim = value.shape[-1]
- pad_size = (chunk_size - num_heads % chunk_size) % chunk_size
+ pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
query = F.pad(query, (0, 0, 0, pad_size))
key = F.pad(key, (0, 0, 0, pad_size))
value = F.pad(value, (0, 0, 0, pad_size))
beta = F.pad(beta, (0, pad_size))
g = F.pad(g, (0, pad_size))
- tot_heads = num_heads + pad_size
+ total_sequence_length = sequence_length + pad_size
scale = 1 / (query.shape[-1] ** 0.5)
query = query * scale
@@ -326,7 +326,7 @@ def torch_chunk_gated_delta_rule(
value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
last_recurrent_state = (
- torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
if initial_state is None
else initial_state.to(value)
)
@@ -334,7 +334,7 @@ def torch_chunk_gated_delta_rule(
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
# for each chunk
- for i in range(0, tot_heads // chunk_size):
+ for i in range(0, total_sequence_length // chunk_size):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
@@ -349,7 +349,7 @@ def torch_chunk_gated_delta_rule(
if not output_final_state:
last_recurrent_state = None
core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
- core_attn_out = core_attn_out[:, :, :num_heads]
+ core_attn_out = core_attn_out[:, :, :sequence_length]
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
return core_attn_out, last_recurrent_state
@@ -365,19 +365,19 @@ def torch_recurrent_gated_delta_rule(
x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
]
- batch_size, sequence_length, num_heads, k_head_dim = key.shape
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
v_head_dim = value.shape[-1]
scale = 1 / (query.shape[-1] ** 0.5)
query = query * scale
- core_attn_out = torch.zeros(batch_size, sequence_length, num_heads, v_head_dim).to(value)
+ core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)
last_recurrent_state = (
- torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
if initial_state is None
else initial_state.to(value)
)
- for i in range(num_heads):
+ for i in range(sequence_length):
q_t = query[:, :, i]
k_t = key[:, :, i]
v_t = value[:, :, i]
@@ -709,6 +709,9 @@ def _init_weights(self, module):
if isinstance(module, Qwen3NextGatedDeltaNet):
module.dt_bias.data.fill_(1.0)
module.A_log.data.uniform_(0, 16).log_()
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif isinstance(module, Qwen3NextRMSNorm):
+ module.weight.data.zero_()
class Qwen3NextModel(Qwen3NextPreTrainedModel):
diff --git a/src/transformers/models/qwen3_omni_moe/__init__.py b/src/transformers/models/qwen3_omni_moe/__init__.py
new file mode 100644
index 000000000000..bd9da3809533
--- /dev/null
+++ b/src/transformers/models/qwen3_omni_moe/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_qwen3_omni_moe import *
+ from .modeling_qwen3_omni_moe import *
+ from .processing_qwen3_omni_moe import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
new file mode 100644
index 000000000000..b530630813da
--- /dev/null
+++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
@@ -0,0 +1,1247 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen3_omni_moe.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group 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 ...configuration_utils import PretrainedConfig, layer_type_validation
+from ...modeling_rope_utils import rope_config_validation
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Qwen3OmniMoeAudioEncoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeAudioEncoder`]. It is used to instantiate a
+ Qwen2.5-Omni-Thinker audio encoder according to the specified arguments, defining the model architecture. Instantiating a
+ configuration with the defaults will yield a similar configuration to that of the audio encoder of the Qwen2-Audio
+ architecture.
+
+ e.g. [Qwen/Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ num_mel_bins (`int`, *optional*, defaults to 128):
+ Number of mel features used per input features. Should correspond to the value used in the
+ `Qwen3OmniMoeProcessor` class.
+ encoder_layers (`int`, *optional*, defaults to 32):
+ Number of encoder layers.
+ encoder_attention_heads (`int`, *optional*, defaults to 20):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ encoder_ffn_dim (`int`, *optional*, defaults to 5120):
+ Dimensionality of the "intermediate" (often named feed-forward) layer in encoder.
+ d_model (`int`, *optional*, defaults to 1280):
+ Dimensionality of the layers.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ activation_function (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ activation_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for activations inside the fully connected layer.
+ scale_embedding (`bool`, *optional*, defaults to `False`):
+ Scale embeddings by diving by sqrt(d_model).
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ max_source_positions (`int`, *optional*, defaults to 1500):
+ The maximum sequence length of log-mel filter-bank features that this model might ever be used with.
+ n_window (`int`, *optional*, defaults to 100):
+ The chunk for conv and flash attn in AudioEncoder.
+ output_dim (`int`, *optional*, defaults to 3584):
+ The output dimension of AudioEncoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeAudioEncoderConfig, Qwen3OmniMoeAudioEncoder
+
+ >>> # Initializing a Qwen3OmniMoeAudioEncoderConfig
+ >>> configuration = Qwen3OmniMoeAudioEncoderConfig()
+
+ >>> # Initializing a Qwen3OmniMoeAudioEncoder (with random weights)
+ >>> model = Qwen3OmniMoeAudioEncoder(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_audio_encoder"
+
+ def __init__(
+ self,
+ num_mel_bins=128,
+ encoder_layers=32,
+ encoder_attention_heads=20,
+ encoder_ffn_dim=5120,
+ d_model=1280,
+ dropout=0,
+ attention_dropout=0,
+ activation_function="gelu",
+ activation_dropout=0,
+ scale_embedding=False,
+ initializer_range=0.02,
+ max_source_positions=1500,
+ n_window=100,
+ output_dim=3584,
+ n_window_infer=400,
+ conv_chunksize=500,
+ downsample_hidden_size=480,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.num_mel_bins = num_mel_bins
+ self.d_model = d_model
+ self.encoder_layers = encoder_layers
+ self.encoder_attention_heads = encoder_attention_heads
+ self.encoder_ffn_dim = encoder_ffn_dim
+ self.dropout = dropout
+ self.attention_dropout = attention_dropout
+ self.activation_function = activation_function
+ self.activation_dropout = activation_dropout
+ self.num_hidden_layers = encoder_layers
+ self.initializer_range = initializer_range
+ self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
+ self.max_source_positions = max_source_positions
+ self.n_window = n_window
+ self.output_dim = output_dim
+ self.n_window_infer = n_window_infer
+ self.conv_chunksize = conv_chunksize
+ self.downsample_hidden_size = downsample_hidden_size
+
+
+class Qwen3OmniMoeVisionEncoderConfig(PretrainedConfig):
+ model_type = "qwen3_omni_moe_vision_encoder"
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ depth=27,
+ hidden_size=1152,
+ hidden_act="gelu_pytorch_tanh",
+ intermediate_size=4304,
+ num_heads=16,
+ in_channels=3,
+ patch_size=16,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+ out_hidden_size=3584,
+ num_position_embeddings=2304,
+ deepstack_visual_indexes=[8, 16, 24],
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.depth = depth
+ self.hidden_size = hidden_size
+ self.hidden_act = hidden_act
+ self.intermediate_size = intermediate_size
+ self.num_heads = num_heads
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.spatial_merge_size = spatial_merge_size
+ self.temporal_patch_size = temporal_patch_size
+ self.out_hidden_size = out_hidden_size
+ self.num_position_embeddings = num_position_embeddings
+ self.initializer_range = initializer_range
+ self.deepstack_visual_indexes = deepstack_visual_indexes
+
+
+class Qwen3OmniMoeTextConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeTextModel`]. It is used to instantiate a
+ Qwen3OmniMoeText model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of [Qwen/Qwen3-15B-A2B](https://huggingface.co/Qwen/Qwen3-15B-A2B).
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 151936):
+ Vocabulary size of the Qwen3OmniMoeText model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`Qwen3OmniMoeTextModel`]
+ hidden_size (`int`, *optional*, defaults to 2048):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 6144):
+ Dimension of the MLP representations.
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of hidden layers in the Transformer encoder.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ num_key_value_heads (`int`, *optional*, defaults to 4):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 32768):
+ The maximum sequence length that this model might ever be used with.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
+ relevant if `config.is_decoder=True`.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether the model's input and output word embeddings should be tied.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+ accordingly.
+ Expected contents:
+ `rope_type` (`str`):
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+ 'llama3'], with 'default' being the original RoPE implementation.
+ `factor` (`float`, *optional*):
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+ original maximum pre-trained length.
+ `original_max_position_embeddings` (`int`, *optional*):
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+ pretraining.
+ `attention_factor` (`float`, *optional*):
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
+ `factor` field to infer the suggested value.
+ `beta_fast` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+ ramp function. If unspecified, it defaults to 32.
+ `beta_slow` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+ ramp function. If unspecified, it defaults to 1.
+ `short_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `long_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `low_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+ `high_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ use_sliding_window (`bool`, *optional*, defaults to `False`):
+ Whether to use sliding window attention.
+ sliding_window (`int`, *optional*, defaults to 4096):
+ Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ decoder_sparse_step (`int`, *optional*, defaults to 1):
+ The frequency of the MoE layer.
+ moe_intermediate_size (`int`, *optional*, defaults to 768):
+ Intermediate size of the routed expert.
+ num_experts_per_tok (`int`, *optional*, defaults to 8):
+ Number of selected experts.
+ num_experts (`int`, *optional*, defaults to 128):
+ Number of routed experts.
+ norm_topk_prob (`bool`, *optional*, defaults to `False`):
+ Whether to normalize the topk probabilities.
+ output_router_logits (`bool`, *optional*, defaults to `False`):
+ Whether or not the router logits should be returned by the model. Enabling this will also
+ allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+ The aux loss factor for the total loss.
+ mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
+ Indicate which layers use Qwen3OmniMoeTextMLP rather than Qwen3OmniMoeTextSparseMoeBlock
+ The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
+ If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeTextModel, Qwen3OmniMoeTextConfig
+
+ >>> # Initializing a Qwen3OmniMoeText style configuration
+ >>> configuration = Qwen3OmniMoeTextConfig()
+
+ >>> # Initializing a model from the Qwen3-15B-A2B" style configuration
+ >>> model = Qwen3OmniMoeTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_text"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ # Default tensor parallel plan for base model `Qwen3OmniMoeText`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.*.gate_proj": "colwise",
+ "layers.*.mlp.experts.*.up_proj": "colwise",
+ "layers.*.mlp.experts.*.down_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=3584,
+ hidden_size=2048,
+ intermediate_size=18944,
+ num_hidden_layers=28,
+ num_attention_heads=28,
+ num_key_value_heads=4,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=1000000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ attention_dropout=0,
+ decoder_sparse_step=1,
+ moe_intermediate_size=768,
+ num_experts_per_tok=8,
+ num_experts=128,
+ norm_topk_prob=True,
+ output_router_logits=False,
+ router_aux_loss_coef=0.001,
+ mlp_only_layers=None,
+ **kwargs,
+ ):
+ super().__init__(
+ tie_word_embeddings=tie_word_embeddings,
+ **kwargs,
+ )
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.sliding_window = sliding_window
+
+ self.num_key_value_heads = num_key_value_heads
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Validate the correctness of rotary position embeddings parameters
+ # BC: if there is a 'type' field, move it to 'rope_type'.
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+ rope_config_validation(self)
+
+ # MoE arguments
+ self.decoder_sparse_step = decoder_sparse_step
+ self.moe_intermediate_size = moe_intermediate_size
+ self.num_experts_per_tok = num_experts_per_tok
+ self.num_experts = num_experts
+ self.norm_topk_prob = norm_topk_prob
+ self.output_router_logits = output_router_logits
+ self.router_aux_loss_coef = router_aux_loss_coef
+ self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+
+class Qwen3OmniMoeThinkerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeThinker`]. It is used to instantiate a
+ Qwen3-Omni-Thinker model according to the specified arguments, defining the model architecture. Instantiating a
+ configuration with the defaults will yield a similar configuration to that of the thinker component of the Qwen3-Omni
+ architecture.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ audio_config (`dict`, *optional*):
+ The config dictionary of the audio backbone.
+ vision_config (`dict`, *optional*):
+ The config dictionary of the vision backbone.
+ text_config (`dict`, *optional*):
+ The config dictionary of the text backbone.
+ audio_token_id (`int`, *optional*, defaults to 151646):
+ The audio token id to encode the audio prompt.
+ image_token_id (`int`, *optional*, defaults to 151655):
+ The image token id to encode the image prompt.
+ video_token_id (`int`, *optional*, defaults to 151656):
+ The video token id to encode the video prompt.
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ The increment of position id per second.
+ audio_start_token_id (`int`, *optional*, defaults to 151647):
+ The audio start token id to encode the audio prompt.
+ user_token_id (`int`, *optional*, defaults to 872):
+ The user token id to encode the user token.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeThinkerModel, Qwen3OmniMoeThinkerConfig
+
+ >>> # Initializing a default Qwen3OmniMoeThinkerConfig
+ >>> configuration = Qwen3OmniMoeThinkerConfig()
+
+ >>> # Initializing a model (with random weights) from the default configuration
+ >>> model = Qwen3OmniMoeThinkerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_thinker"
+ # Override parent's attribute_map as we use audio_token_id directly, not audio_token_index
+ attribute_map = {}
+ sub_configs = {
+ "audio_config": Qwen3OmniMoeAudioEncoderConfig,
+ "vision_config": Qwen3OmniMoeVisionEncoderConfig,
+ "text_config": Qwen3OmniMoeTextConfig,
+ }
+
+ def __init__(
+ self,
+ audio_config=None,
+ vision_config=None,
+ text_config=None,
+ audio_token_id=151646,
+ image_token_id=151655,
+ video_token_id=151656,
+ position_id_per_seconds=25,
+ audio_start_token_id=151647,
+ user_token_id=872,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.user_token_id = user_token_id
+ self.position_id_per_seconds = position_id_per_seconds
+ self.audio_start_token_id = audio_start_token_id
+ self.initializer_range = initializer_range
+
+ if isinstance(vision_config, dict):
+ vision_config = Qwen3OmniMoeVisionEncoderConfig(**vision_config)
+ elif vision_config is None:
+ vision_config = Qwen3OmniMoeVisionEncoderConfig()
+ self.vision_config = vision_config
+
+ if isinstance(audio_config, dict):
+ audio_config = Qwen3OmniMoeAudioEncoderConfig(**audio_config)
+ elif audio_config is None:
+ audio_config = Qwen3OmniMoeAudioEncoderConfig()
+ self.audio_config = audio_config
+
+ if isinstance(text_config, dict):
+ text_config = Qwen3OmniMoeTextConfig(**text_config)
+ elif text_config is None:
+ text_config = Qwen3OmniMoeTextConfig()
+ self.text_config = text_config
+ self.audio_token_id = audio_token_id
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
+
+
+class Qwen3OmniMoeTalkerCodePredictorConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeTalkerCodePredictorModel`]. It is used to instantiate a
+ Qwen3OmniMoeTalkerCodePredictor model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of
+ Qwen3OmniMoeTalkerCodePredictor-8B [Qwen/Qwen3OmniMoeTalkerCodePredictor-8B](https://huggingface.co/Qwen/Qwen3OmniMoeTalkerCodePredictor-8B).
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 151936):
+ Vocabulary size of the Qwen3OmniMoeTalkerCodePredictor model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`Qwen3OmniMoeTalkerCodePredictorModel`]
+ hidden_size (`int`, *optional*, defaults to 4096):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 22016):
+ Dimension of the MLP representations.
+ num_hidden_layers (`int`, *optional*, defaults to 32):
+ Number of hidden layers in the Transformer encoder.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ num_key_value_heads (`int`, *optional*, defaults to 32):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+ head_dim (`int`, *optional*, defaults to 128):
+ The attention head dimension.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 32768):
+ The maximum sequence length that this model might ever be used with.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
+ relevant if `config.is_decoder=True`.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether the model's input and output word embeddings should be tied.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+ accordingly.
+ Expected contents:
+ `rope_type` (`str`):
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+ 'llama3'], with 'default' being the original RoPE implementation.
+ `factor` (`float`, *optional*):
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+ original maximum pre-trained length.
+ `original_max_position_embeddings` (`int`, *optional*):
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+ pretraining.
+ `attention_factor` (`float`, *optional*):
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
+ `factor` field to infer the suggested value.
+ `beta_fast` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+ ramp function. If unspecified, it defaults to 32.
+ `beta_slow` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+ ramp function. If unspecified, it defaults to 1.
+ `short_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `long_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `low_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+ `high_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ use_sliding_window (`bool`, *optional*, defaults to `False`):
+ Whether to use sliding window attention.
+ sliding_window (`int`, *optional*, defaults to 4096):
+ Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+ max_window_layers (`int`, *optional*, defaults to 28):
+ The number of layers using full attention. The first `max_window_layers` layers will use full attention, while any
+ additional layer afterwards will use SWA (Sliding Window Attention).
+ layer_types (`list`, *optional*):
+ Attention pattern for each layer.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeTalkerCodePredictorModel, Qwen3OmniMoeTalkerCodePredictorConfig
+
+ >>> # Initializing a Qwen3OmniMoeTalkerCodePredictor style configuration
+ >>> configuration = Qwen3OmniMoeTalkerCodePredictorConfig()
+
+ >>> # Initializing a model from the Qwen3OmniMoeTalkerCodePredictor-8B style configuration
+ >>> model = Qwen3OmniMoeTalkerCodePredictorModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_talker_code_predictor"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerCodePredictor`
+ 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"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=2048,
+ hidden_size=1024,
+ intermediate_size=3072,
+ num_hidden_layers=5,
+ num_attention_heads=16,
+ num_key_value_heads=8,
+ head_dim=128,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=0.000001,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=10000,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ layer_types=None,
+ attention_dropout=0,
+ num_code_groups=32,
+ **kwargs,
+ ):
+ super().__init__(
+ tie_word_embeddings=tie_word_embeddings,
+ **kwargs,
+ )
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.sliding_window = sliding_window
+
+ # for backward compatibility
+ if num_key_value_heads is None:
+ num_key_value_heads = num_attention_heads
+
+ self.num_key_value_heads = num_key_value_heads
+ self.head_dim = head_dim
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Validate the correctness of rotary position embeddings parameters
+ # BC: if there is a 'type' field, move it to 'rope_type'.
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+ rope_config_validation(self)
+
+ self.layer_types = layer_types
+ if self.layer_types is None:
+ self.layer_types = [
+ "sliding_attention"
+ if self.sliding_window is not None and i >= self.max_window_layers
+ else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+ layer_type_validation(self.layer_types, self.num_hidden_layers)
+ self.num_code_groups = num_code_groups
+
+
+class Qwen3OmniMoeTalkerTextConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeTalkerTextModel`]. It is used to instantiate a
+ Qwen3OmniMoeTalkerText model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of [Qwen/Qwen3-15B-A2B](https://huggingface.co/Qwen/Qwen3-15B-A2B).
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 151936):
+ Vocabulary size of the Qwen3OmniMoeTalkerText model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`Qwen3OmniMoeTalkerTextModel`]
+ hidden_size (`int`, *optional*, defaults to 2048):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 6144):
+ Dimension of the MLP representations.
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of hidden layers in the Transformer encoder.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ num_key_value_heads (`int`, *optional*, defaults to 4):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 32768):
+ The maximum sequence length that this model might ever be used with.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
+ relevant if `config.is_decoder=True`.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether the model's input and output word embeddings should be tied.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+ accordingly.
+ Expected contents:
+ `rope_type` (`str`):
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+ 'llama3'], with 'default' being the original RoPE implementation.
+ `factor` (`float`, *optional*):
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+ original maximum pre-trained length.
+ `original_max_position_embeddings` (`int`, *optional*):
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+ pretraining.
+ `attention_factor` (`float`, *optional*):
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
+ `factor` field to infer the suggested value.
+ `beta_fast` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+ ramp function. If unspecified, it defaults to 32.
+ `beta_slow` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+ ramp function. If unspecified, it defaults to 1.
+ `short_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `long_factor` (`list[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `low_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+ `high_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ use_sliding_window (`bool`, *optional*, defaults to `False`):
+ Whether to use sliding window attention.
+ sliding_window (`int`, *optional*, defaults to 4096):
+ Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ decoder_sparse_step (`int`, *optional*, defaults to 1):
+ The frequency of the MoE layer.
+ moe_intermediate_size (`int`, *optional*, defaults to 768):
+ Intermediate size of the routed expert.
+ num_experts_per_tok (`int`, *optional*, defaults to 8):
+ Number of selected experts.
+ num_experts (`int`, *optional*, defaults to 128):
+ Number of routed experts.
+ norm_topk_prob (`bool`, *optional*, defaults to `False`):
+ Whether to normalize the topk probabilities.
+ output_router_logits (`bool`, *optional*, defaults to `False`):
+ Whether or not the router logits should be returned by the model. Enabling this will also
+ allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+ The aux loss factor for the total loss.
+ mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
+ Indicate which layers use Qwen3OmniMoeTalkerTextMLP rather than Qwen3OmniMoeTalkerTextSparseMoeBlock
+ The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
+ If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeTalkerTextModel, Qwen3OmniMoeTalkerTextConfig
+
+ >>> # Initializing a Qwen3OmniMoeTalkerText style configuration
+ >>> configuration = Qwen3OmniMoeTalkerTextConfig()
+
+ >>> # Initializing a model from the Qwen3-15B-A2B" style configuration
+ >>> model = Qwen3OmniMoeTalkerTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_talker_text"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerText`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.*.gate_proj": "colwise",
+ "layers.*.mlp.experts.*.up_proj": "colwise",
+ "layers.*.mlp.experts.*.down_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=3072,
+ hidden_size=1024,
+ intermediate_size=2048,
+ num_hidden_layers=20,
+ num_attention_heads=16,
+ num_key_value_heads=2,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=0.000001,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=10000,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ attention_dropout=0,
+ decoder_sparse_step=1,
+ moe_intermediate_size=384,
+ num_experts_per_tok=8,
+ num_experts=128,
+ norm_topk_prob=False,
+ output_router_logits=False,
+ router_aux_loss_coef=0.001,
+ mlp_only_layers=None,
+ **kwargs,
+ ):
+ super().__init__(
+ tie_word_embeddings=tie_word_embeddings,
+ **kwargs,
+ )
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.sliding_window = sliding_window
+
+ self.num_key_value_heads = num_key_value_heads
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Validate the correctness of rotary position embeddings parameters
+ # BC: if there is a 'type' field, move it to 'rope_type'.
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+ rope_config_validation(self)
+
+ # MoE arguments
+ self.decoder_sparse_step = decoder_sparse_step
+ self.moe_intermediate_size = moe_intermediate_size
+ self.num_experts_per_tok = num_experts_per_tok
+ self.num_experts = num_experts
+ self.norm_topk_prob = norm_topk_prob
+ self.output_router_logits = output_router_logits
+ self.router_aux_loss_coef = router_aux_loss_coef
+ self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+
+class Qwen3OmniMoeTalkerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeTalker`]. It is used to instantiate a
+ Qwen3-Omni multi-modal talker model capable of handling text, audio, and vision modalities in a unified architecture.
+ The model integrates a text decoder with a code predictor for autoregressive generation of both semantic and acoustic
+ tokens, enabling speech and multimodal content generation. This configuration wraps sub-configurations for the text and
+ code predictor components, allowing modular setup and initialization.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ code_predictor_config (`dict`, *optional*):
+ A dictionary of configuration parameters used to initialize a [`Qwen3OmniMoeTalkerCodePredictorConfig`].
+ If not provided, defaults will be used.
+ text_config (`dict`, *optional*):
+ A dictionary of configuration parameters used to initialize a [`Qwen3OmniMoeTalkerTextConfig`].
+ If not provided, defaults will be used.
+ num_code_groups (`int`, *optional*, defaults to 32):
+ Number of codebook groups used in the predicted acoustic token sequence, corresponding to multi-codebook VQ representation.
+ thinker_hidden_size (`int`, *optional*, defaults to 2048):
+ Hidden dimension size of the thinker module used for intermediate reasoning or latent planning before audio generation.
+ codec_eos_token_id (`int`, *optional*, defaults to 4198):
+ Token ID representing the end-of-speech token in the codec-generated sequence.
+ accept_hidden_layer (`int`, *optional*, defaults to 18):
+ Index of the hidden layer whose output is used for accepting or refining generated tokens during think-and-speak process.
+ codec_nothink_id (`int`, *optional*, defaults to 4203):
+ Token ID indicating no thinking step is required during generation.
+ codec_think_bos_id (`int`, *optional*, defaults to 4204):
+ Token ID marking the beginning of a thinking sequence.
+ codec_think_eos_id (`int`, *optional*, defaults to 4205):
+ Token ID marking the end of a thinking sequence.
+ codec_pad_id (`int`, *optional*, defaults to 4196):
+ Padding token ID used in codec input sequences.
+ codec_bos_id (`int`, *optional*, defaults to 4197):
+ Beginning-of-speech token ID in codec sequences.
+ audio_token_id (`int`, *optional*, defaults to 151646):
+ Special token ID used to indicate the position of audio tokens in the input sequence.
+ image_token_id (`int`, *optional*, defaults to 151655):
+ Special token ID used to represent image inputs in the multimodal context.
+ video_token_id (`int`, *optional*, defaults to 151656):
+ Special token ID used to represent video inputs.
+ vision_start_token_id (`int`, *optional*, defaults to 151652):
+ Token ID indicating the start of a visual input sequence (e.g., image or video embeddings).
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ Number of position IDs allocated per second of audio content, used for temporal alignment in generation.
+ audio_start_token_id (`int`, *optional*, defaults to 151669):
+ Token ID that indicates the start of an audio generation segment in the output.
+ speaker_id (`dict`, *optional*):
+ Speaker name to speaker id dict.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeTalkerConfig, Qwen3OmniMoeTalker
+
+ >>> # Initialize a Qwen3OmniMoeTalkerConfig with default sub-configurations
+ >>> config = Qwen3OmniMoeTalkerConfig(
+ ... num_code_groups=32,
+ ... thinker_hidden_size=2048,
+ ... )
+
+ >>> # Initialize the full Qwen3-Omni Talker model
+ >>> model = Qwen3OmniMoeTalker(config)
+
+ >>> # Access the model configuration
+ >>> config = model.config
+ >>> print(config.text_config) # Access text decoder configuration
+ >>> print(config.code_predictor_config) # Access code predictor configuration
+ ```"""
+
+ sub_configs = {
+ "code_predictor_config": Qwen3OmniMoeTalkerCodePredictorConfig,
+ "text_config": Qwen3OmniMoeTalkerTextConfig,
+ }
+
+ def __init__(
+ self,
+ code_predictor_config=None,
+ text_config=None,
+ num_code_groups=32,
+ thinker_hidden_size=2048,
+ codec_eos_token_id=4198,
+ accept_hidden_layer=18,
+ codec_nothink_id=4203,
+ codec_think_bos_id=4204,
+ codec_think_eos_id=4205,
+ codec_pad_id=4196,
+ codec_bos_id=4197,
+ audio_token_id=151646,
+ image_token_id=151655,
+ video_token_id=151656,
+ vision_start_token_id=151652,
+ position_id_per_seconds=25,
+ audio_start_token_id=151669,
+ speaker_id=None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ if code_predictor_config is None:
+ code_predictor_config = {}
+ self.code_predictor_config = Qwen3OmniMoeTalkerCodePredictorConfig()
+ logger.info("code_predictor_config is None. Initializing code_predictor_config model with default values")
+ elif isinstance(code_predictor_config, Qwen3OmniMoeTalkerCodePredictorConfig):
+ self.code_predictor_config = code_predictor_config
+ else:
+ self.code_predictor_config = Qwen3OmniMoeTalkerCodePredictorConfig(**code_predictor_config)
+
+ if text_config is None:
+ text_config = {}
+ self.text_config = Qwen3OmniMoeTalkerTextConfig()
+ logger.info("talker text_config is None. Initializing talker text model with default values")
+ elif isinstance(text_config, Qwen3OmniMoeTalkerTextConfig):
+ self.text_config = text_config
+ else:
+ self.text_config = Qwen3OmniMoeTalkerTextConfig(**text_config)
+ self.num_code_groups = num_code_groups
+ self.thinker_hidden_size = thinker_hidden_size
+ self.codec_eos_token_id = codec_eos_token_id
+ self.accept_hidden_layer = accept_hidden_layer
+ self.codec_nothink_id = codec_nothink_id
+ self.codec_think_bos_id = codec_think_bos_id
+ self.codec_think_eos_id = codec_think_eos_id
+ self.codec_pad_id = codec_pad_id
+ self.codec_bos_id = codec_bos_id
+ self.audio_token_id = audio_token_id
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
+ self.position_id_per_seconds = position_id_per_seconds
+ self.audio_start_token_id = audio_start_token_id
+ self.vision_start_token_id = vision_start_token_id
+ self.speaker_id = speaker_id
+
+
+class Qwen3OmniMoeCode2WavConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeCode2WavConfig`]. It is used to instantiate a
+ Qwen3-Omni code-to-waveform decoder, responsible for converting discrete audio codes into high-fidelity waveforms.
+ The configuration defines the architecture of the decoder, including parameters for vector quantization, autoregressive modeling,
+ and upsampling layers.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ codebook_size (`int`, *optional*, defaults to 2048):
+ Number of entries in each residual codebook used for acoustic token quantization.
+ hidden_size (`int`, *optional*, defaults to 1024):
+ Dimensionality of the hidden states and embeddings in the autoregressive transformer decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 8000):
+ Maximum sequence length that the autoregressive decoder can handle. Determines positional embedding size.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period for rotary position embeddings (RoPE) applied to attention layers.
+ num_attention_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads for each attention layer in the decoder.
+ num_key_value_heads (`int`, *optional*, defaults to 16):
+ Number of key and value attention heads used in grouped-query attention (if applicable).
+ attention_bias (`bool`, *optional*, defaults to `False`):
+ Whether to use bias in the attention projection layers.
+ sliding_window (`int`, *optional*, defaults to 72):
+ Window size for local attention mechanism, limiting attention context to improve efficiency.
+ intermediate_size (`int`, *optional*, defaults to 3072):
+ Dimensionality of the feed-forward (intermediate) layer in each transformer block.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function used in the feed-forward layers. Supports `"silu"`, `"relu"`, `"gelu"`, etc.
+ layer_scale_initial_scale (`float`, *optional*, defaults to 0.01):
+ Initial value for LayerScale applied in transformer blocks, helping stabilize training.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-5):
+ Epsilon value for RMS normalization layers to prevent division by zero.
+ num_hidden_layers (`int`, *optional*, defaults to 8):
+ Number of transformer blocks in the autoregressive decoder.
+ num_quantizers (`int`, *optional*, defaults to 16):
+ Number of residual vector quantizers used in the vocoder for fine-grained audio reconstruction.
+ upsample_rates (`Tuple[int]`, *optional*, defaults to `(8, 5, 4, 3)`):
+ Rate at which features are upsampled in the final waveform synthesis stage.
+ upsampling_ratios (`Tuple[int]`, *optional*, defaults to `(2, 2)`):
+ Ratios used in transposed convolutional layers to progressively upsample feature maps to waveform.
+ decoder_dim (`int`, *optional*, defaults to 1536):
+ Final dimensionality of the decoder's output before waveform generation.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ Dropout probability applied to attention weights in the decoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeCode2WavConfig, Qwen3OmniMoeCode2WavModel
+
+ >>> # Initializing a default Qwen3OmniMoeCode2WavConfig
+ >>> config = Qwen3OmniMoeCode2WavConfig()
+
+ >>> # Initializing the Code2Wav model with the configuration
+ >>> model = Qwen3OmniMoeCode2WavModel(config)
+
+ >>> # Accessing configuration
+ >>> config = model.config
+ ```"""
+
+ def __init__(
+ self,
+ codebook_size=2048,
+ hidden_size=1024,
+ max_position_embeddings=8000,
+ rope_theta=10000,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ attention_bias=False,
+ sliding_window=72,
+ intermediate_size=3072,
+ hidden_act="silu",
+ layer_scale_initial_scale=0.01,
+ rms_norm_eps=1e-5,
+ num_hidden_layers=8,
+ num_quantizers=16,
+ upsample_rates=(8, 5, 4, 3),
+ upsampling_ratios=(2, 2),
+ decoder_dim=1536,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.codebook_size = codebook_size
+ self.hidden_size = hidden_size
+ self.max_position_embeddings = max_position_embeddings
+ self.rope_theta = rope_theta
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.attention_bias = attention_bias
+ self.sliding_window = sliding_window
+ self.intermediate_size = intermediate_size
+ self.hidden_act = hidden_act
+ self.layer_scale_initial_scale = layer_scale_initial_scale
+ self.rms_norm_eps = rms_norm_eps
+ self.num_hidden_layers = num_hidden_layers
+ self.num_quantizers = num_quantizers
+ self.upsample_rates = upsample_rates
+ self.upsampling_ratios = upsampling_ratios
+ self.decoder_dim = decoder_dim
+ self.attention_dropout = attention_dropout
+
+ @property
+ def layer_types(self):
+ """
+ All layer in code2wav should be sliding attention
+ """
+ return ["sliding_attention"] * self.num_hidden_layers
+
+
+class Qwen3OmniMoeConfig(PretrainedConfig):
+ """
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeForConditionalGeneration`]. It is used to instantiate a Qwen3Omni
+ model according to the specified sub-models configurations, defining the model architecture.
+
+ Instantiating a configuration with the defaults will yield a similar configuration to that of the
+ [Qwen/Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ thinker_config (`dict`, *optional*): Configuration of the underlying thinker sub-model.
+ talker_config (`dict`, *optional*): Configuration of the underlying talker sub-model.
+ code2wav_config (`dict`, *optional*): Configuration of the underlying code2wav sub-model.
+ enable_audio_output (`bool`, *optional*, defaults to `True`): Whether enable audio output and load talker and code2wav module.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... Qwen3OmniMoeThinkerConfig,
+ ... Qwen3OmniMoeTalkerConfig,
+ ... Qwen3OmniMoeCode2WavConfig,
+ ... Qwen3OmniMoeForConditionalGeneration,
+ ... Qwen3OmniMoeConfig,
+ ... )
+
+ >>> # Initializing a Qwen3OmniMoe style configuration
+ >>> configuration = Qwen3OmniMoeConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = Qwen3OmniMoeForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe"
+ sub_configs = {
+ "thinker_config": Qwen3OmniMoeThinkerConfig,
+ "talker_config": Qwen3OmniMoeTalkerConfig,
+ "code2wav_config": Qwen3OmniMoeCode2WavConfig,
+ }
+
+ def __init__(
+ self,
+ thinker_config=None,
+ talker_config=None,
+ code2wav_config=None,
+ enable_audio_output=True,
+ im_start_token_id=151644,
+ im_end_token_id=151645,
+ tts_pad_token_id=151671,
+ tts_bos_token_id=151672,
+ tts_eos_token_id=151673,
+ system_token_id=8948,
+ user_token_id=872,
+ assistant_token_id=77091,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ if thinker_config is None:
+ thinker_config = {}
+ logger.info("thinker_config is None. Initializing thinker model with default values")
+
+ if talker_config is None:
+ talker_config = {}
+ logger.info("talker_config is None. Initializing talker model with default values")
+
+ if code2wav_config is None:
+ code2wav_config = {}
+ logger.info("code2wav_config is None. Initializing code2wav model with default values")
+
+ self.thinker_config = Qwen3OmniMoeThinkerConfig(**thinker_config)
+ self.talker_config = Qwen3OmniMoeTalkerConfig(**talker_config)
+ self.code2wav_config = Qwen3OmniMoeCode2WavConfig(**code2wav_config)
+ self.enable_audio_output = enable_audio_output
+ self.im_start_token_id = im_start_token_id
+ self.im_end_token_id = im_end_token_id
+ self.tts_pad_token_id = tts_pad_token_id
+ self.tts_bos_token_id = tts_bos_token_id
+ self.tts_eos_token_id = tts_eos_token_id
+ self.system_token_id = system_token_id
+ self.user_token_id = user_token_id
+ self.assistant_token_id = assistant_token_id
+
+ def get_text_config(self, decoder=False) -> "PretrainedConfig":
+ """
+ Returns the config that is meant to be used with text IO. On most models, it is the original config instance
+ itself. On specific composite models, it is under a set of valid names.
+
+ Args:
+ decoder (`Optional[bool]`, *optional*, defaults to `False`):
+ If set to `True`, then only search for decoder config names.
+ """
+ # Overridden for deeply nested config like Qwen2-Omni. We don't have any omni model
+ # except for Qwen yet. This has to be generalized if more deeply nested configs are
+ # added. NOTE: currently method used only by vLLM
+ return self.thinker_config.get_text_config()
+
+
+__all__ = ["Qwen3OmniMoeConfig", "Qwen3OmniMoeThinkerConfig", "Qwen3OmniMoeTalkerConfig"]
diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py
new file mode 100644
index 000000000000..408b76031bcc
--- /dev/null
+++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py
@@ -0,0 +1,4066 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen3_omni_moe.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import Callable, Optional, Union
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import Parameter
+from torch.nn import functional as F
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ 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 auto_docstring, can_return_tuple
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import OutputRecorder, TransformersKwargs, check_model_inputs
+from .configuration_qwen3_omni_moe import (
+ Qwen3OmniMoeAudioEncoderConfig,
+ Qwen3OmniMoeCode2WavConfig,
+ Qwen3OmniMoeConfig,
+ Qwen3OmniMoeTalkerCodePredictorConfig,
+ Qwen3OmniMoeTalkerConfig,
+ Qwen3OmniMoeTalkerTextConfig,
+ Qwen3OmniMoeTextConfig,
+ Qwen3OmniMoeThinkerConfig,
+ Qwen3OmniMoeVisionEncoderConfig,
+)
+
+
+@auto_docstring
+class Qwen3OmniMoePreTrainedModel(PreTrainedModel):
+ config: Qwen3OmniMoeConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3OmniMoeDecoderLayer", "Qwen3OmniMoeVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = False
+ _supports_attention_backend = True
+
+
+def _get_feat_extract_output_lengths(input_lengths):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+
+ input_lengths_leave = input_lengths % 100
+ feat_lengths = (input_lengths_leave - 1) // 2 + 1
+ output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
+ return output_lengths
+
+
+class Qwen3OmniMoePreTrainedModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel):
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ self,
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ min_dtype: float,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ min_dtype (`float`):
+ The minimum value representable with the dtype `dtype`.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ causal_mask = torch.full(
+ (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+ )
+ if sequence_length != 1:
+ causal_mask = torch.triu(causal_mask, diagonal=1)
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+ causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[torch.Tensor],
+ grid_hs: list[torch.Tensor],
+ grid_ws: list[torch.Tensor],
+ ):
+ llm_pos_ids_list = []
+ llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+ llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float()
+ t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float()
+ _llm_pos_ids = torch.stack([t_index, h_index, w_index])
+ llm_pos_ids_list.append(_llm_pos_ids + start_idx)
+ llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+ return llm_pos_ids
+
+ def get_chunked_index(
+ self, token_indices: torch.Tensor, tokens_per_chunk: int, remove_index: int
+ ) -> list[tuple[int, int]]:
+ """
+ Splits token index list into chunks based on token value ranges.
+
+ Given a list of token indices, returns a list of (start, end) index tuples representing
+ slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
+
+ For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
+ - the first chunk contains token values < 1000,
+ - the second chunk contains values >= 1000 and < 2000, and so on.
+
+ Parameters:
+ token_indices (`torch.Tensor` of shape `(seq_len, )`): A monotonically increasing list of
+ token index values.
+ t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
+ remove_index (`int`) An index id to subtract from `token_indices` before chunking
+
+ Returns:
+ `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
+ and end (exclusive) indices of a chunk in `token_indices`.
+ """
+
+ def _iter():
+ i, start_idx = 0, 0 # skip bos token
+ current_chunk = 1
+ while i < len(token_indices): # skip eos token
+ if token_indices[i] - remove_index >= current_chunk * tokens_per_chunk:
+ yield (start_idx, i)
+ start_idx = i
+ current_chunk += 1
+ i += 1
+ yield (start_idx, len(token_indices))
+
+ return list(_iter())
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: Optional[torch.LongTensor] = None,
+ second_per_grids: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_seqlens (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ second_per_grids (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ audio_token_id = self.config.audio_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ audio_start_token_id = self.config.audio_start_token_id
+ position_id_per_seconds = self.config.position_id_per_seconds
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is not None:
+ attention_mask = attention_mask == 1
+ position_ids = torch.zeros(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=torch.float,
+ device=input_ids.device,
+ )
+ image_idx, video_idx, audio_idx = 0, 0, 0
+ for i, input_ids in enumerate(total_input_ids):
+ if attention_mask is not None:
+ input_ids = input_ids[attention_mask[i]]
+ image_nums, video_nums, audio_nums = 0, 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ audio_nums = torch.sum(input_ids == audio_start_token_id)
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (
+ (vision_tokens == audio_start_token_id).sum()
+ if use_audio_in_video
+ else (vision_tokens == video_token_id).sum()
+ )
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums
+ multimodal_nums = (
+ image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums
+ )
+ for _ in range(multimodal_nums):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ if (image_token_id in input_tokens or video_token_id in input_tokens) and (
+ remain_videos > 0 or remain_images > 0
+ ):
+ ed_vision_start = input_tokens.index(vision_start_token_id, st)
+ else:
+ ed_vision_start = len(input_tokens) + 1
+ if audio_token_id in input_tokens and remain_audios > 0:
+ ed_audio_start = input_tokens.index(audio_start_token_id, st)
+ else:
+ ed_audio_start = len(input_tokens) + 1
+ min_ed = min(ed_vision_start, ed_audio_start)
+
+ text_len = min_ed - st
+ if text_len != 0:
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+ st_idx += text_len
+ # Audio in Video
+ if min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start:
+ bos_len, eos_len = 2, 2
+ else:
+ bos_len, eos_len = 1, 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+ st_idx += bos_len
+ # Audio Only
+ if min_ed == ed_audio_start:
+ audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
+ llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + audio_len + eos_len)
+ audio_idx += 1
+ remain_audios -= 1
+
+ # Image Only
+ elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id:
+ grid_t = image_grid_thw[image_idx][0]
+ grid_hs = image_grid_thw[:, 1]
+ grid_ws = image_grid_thw[:, 2]
+ t_index = (torch.arange(grid_t) * 1 * position_id_per_seconds).float()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + image_len + eos_len)
+ image_idx += 1
+ remain_images -= 1
+
+ # Video Only
+ elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == video_token_id:
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).float()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + video_len + eos_len)
+ video_idx += 1
+ remain_videos -= 1
+
+ # Audio in Video
+ elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start:
+ audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
+ audio_llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).float()
+ video_llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_data_index, audio_data_index = 0, 0
+ while (
+ video_data_index < video_llm_pos_ids.shape[-1]
+ and audio_data_index < audio_llm_pos_ids.shape[-1]
+ ):
+ if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]:
+ llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1])
+ video_data_index += 1
+ else:
+ llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1])
+ audio_data_index += 1
+ if video_data_index < video_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(
+ video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]]
+ )
+ if audio_data_index < audio_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(
+ audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]]
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+
+ st += int(text_len + bos_len + audio_len + video_len + eos_len)
+
+ audio_idx += 1
+ video_idx += 1
+ remain_videos -= 1
+ remain_audios -= 1
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat([item.float() for item in llm_pos_ids_list], dim=1).reshape(3, -1)
+
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+
+ return position_ids, mrope_position_deltas
+ else:
+ position_ids = attention_mask.float().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
+
+ return position_ids, mrope_position_deltas
+
+
+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: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class Qwen3OmniMoeAudioAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.num_heads = config.encoder_attention_heads
+ self.dropout = config.attention_dropout
+ self.head_dim = self.embed_dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.config = config
+
+ if (self.head_dim * self.num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {self.num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = 0.0
+ self.is_decoder = False
+ self.is_causal = False
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+ """Input shape: Batch x Time x Channel"""
+
+ seq_length, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ key_states = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ value_states = self.v_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output
+
+
+class Qwen3OmniMoeAudioEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.self_attn = Qwen3OmniMoeAudioAttention(config)
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+ `(encoder_attention_heads,)`.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states = self.self_attn(
+ hidden_states=hidden_states,
+ cu_seqlens=cu_seqlens,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16:
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ outputs = (hidden_states,)
+
+ return outputs
+
+
+class SinusoidsPositionEmbedding(nn.Module):
+ def __init__(self, length, channels, max_timescale=10000):
+ super().__init__()
+ if channels % 2 != 0:
+ raise ValueError("SinusoidsPositionEmbedding needs even channels input")
+ log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float())
+ scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+ self.register_buffer(
+ "positional_embedding",
+ torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1),
+ persistent=False,
+ )
+
+ def forward(self, seqlen: int):
+ return self.positional_embedding[:seqlen, :]
+
+
+@auto_docstring(
+ custom_intro="""
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`Qwen3OmniMoeAudioEncoderLayer`].
+ """
+)
+class Qwen3OmniMoeAudioEncoder(Qwen3OmniMoePreTrainedModel):
+ config: Qwen3OmniMoeAudioEncoderConfig
+ main_input_name = "input_features"
+ _no_split_modules = ["Qwen3OmniMoeAudioEncoderLayer"]
+ _supports_sdpa = True
+
+ def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+
+ embed_dim = config.d_model
+ self.num_mel_bins = config.num_mel_bins
+ self.max_source_positions = config.max_source_positions
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+ self.n_window = config.n_window
+ self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
+ self.layers = nn.ModuleList([Qwen3OmniMoeAudioEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.ln_post = nn.LayerNorm(config.d_model)
+ self.gradient_checkpointing = False
+ self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv2d2 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv2d3 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv_out = nn.Linear(
+ config.downsample_hidden_size * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2),
+ config.d_model,
+ bias=False,
+ )
+ self.proj1 = nn.Linear(config.d_model, config.d_model)
+ self.act = ACT2FN[config.activation_function]
+ self.proj2 = nn.Linear(config.d_model, config.output_dim)
+ self.n_window_infer = self.config.n_window_infer
+ self.conv_chunksize = self.config.conv_chunksize
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def _freeze_parameters(self):
+ for param in self.parameters():
+ param.requires_grad = False
+ self._requires_grad = False
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
+ # Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
+ # NOTE: the created attention masl only approximates the ragged FA2 attention by
+ # allowing bidirectional attention within `cu_seqlens` blocks, and not attending between
+ # blocks. Though it will not be a 100% match for FA2's `varlen` path
+ if self.config._attn_implementation == "flash_attention_2":
+ return None
+
+ seq_length = inputs_tensor.shape[0]
+ attention_mask = torch.full(
+ [1, 1, seq_length, seq_length],
+ torch.finfo(inputs_tensor.dtype).min,
+ device=inputs_tensor.device,
+ dtype=inputs_tensor.dtype,
+ )
+ for i in range(1, len(cu_seqlens)):
+ attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
+ return attention_mask
+
+ @auto_docstring
+ def forward(
+ self,
+ input_features,
+ feature_lens=None,
+ aftercnn_lens=None,
+ ):
+ r"""
+ feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
+ mel length
+ aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`):
+ mel length after cnn
+ """
+ aftercnn_lens = _get_feat_extract_output_lengths(feature_lens)
+ chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
+
+ chunk_lengths = torch.tensor(
+ [self.n_window * 2] * chunk_num.sum(),
+ dtype=torch.long,
+ device=feature_lens.device,
+ )
+ tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+ chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
+ chunk_lengths[chunk_lengths == 0] = self.n_window * 2
+
+ chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
+ padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
+ feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths)
+ padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
+ [torch.ones(length, dtype=torch.bool, device=padded_feature.device) for length in feature_lens_after_cnn],
+ batch_first=True,
+ )
+ padded_feature = padded_feature.unsqueeze(1)
+ # Split to chunk to avoid OOM during convolution
+ padded_embeds = []
+ for chunk in padded_feature.split(self.conv_chunksize, dim=0):
+ padded_embed = F.gelu(self.conv2d1(chunk))
+ padded_embed = F.gelu(self.conv2d2(padded_embed))
+ padded_embed = F.gelu(self.conv2d3(padded_embed))
+ padded_embeds.append(padded_embed)
+ padded_embed = torch.cat(padded_embeds, dim=0)
+ b, c, f, t = padded_embed.size()
+ padded_embed = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f))
+
+ positional_embedding = (
+ self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
+ .unsqueeze(0)
+ .to(padded_embed.dtype)
+ )
+ padded_embed = padded_embed + positional_embedding
+ hidden_states = padded_embed[padded_mask_after_cnn]
+ cu_chunk_lens = [0]
+ window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2))
+ for cnn_len in aftercnn_lens:
+ cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn)
+ remainder = cnn_len % window_aftercnn
+ if remainder != 0:
+ cu_chunk_lens += [remainder]
+ cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32)
+
+ for encoder_layer in self.layers:
+ layer_outputs = encoder_layer(
+ hidden_states,
+ cu_seqlens,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.ln_post(hidden_states)
+ hidden_states = self.proj1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.proj2(hidden_states)
+ return BaseModelOutput(last_hidden_state=hidden_states)
+
+ def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"):
+ """
+ Pads a sequence of tensors to their maximum length on indicated `padding_side`.
+ Then prepares a mask so that pad tokens are not attended to.
+ """
+ max_len = tensor_len.max()
+ dim = tensor_list[0].shape[0]
+ padded_tensor = torch.full(
+ size=(len(tensor_list), dim, max_len),
+ fill_value=padding_value,
+ dtype=self.dtype,
+ device=tensor_list[0].device,
+ )
+
+ batch_mask = torch.zeros(
+ (len(tensor_len), max_len),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(tensor_len):
+ batch_mask[i, :length] = 1
+ padded_tensor[i, :, :length] = tensor_list[i]
+
+ feature_lens_after_cnn = (tensor_len - 1) // 2 + 1
+ max_len_after_cnn = feature_lens_after_cnn.max()
+ batch_mask_after_cnn = torch.zeros(
+ (len(tensor_len), max_len_after_cnn),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(feature_lens_after_cnn):
+ batch_mask_after_cnn[i, :length] = 1
+ return (
+ padded_tensor,
+ batch_mask.unsqueeze(1),
+ batch_mask_after_cnn.bool(),
+ )
+
+ # Ignore copy
+ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+ input_lengths = (input_lengths - 1) // 2 + 1
+ output_lengths = (input_lengths - 2) // 2 + 1
+ return input_lengths, output_lengths
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+class Qwen3OmniMoeVisionAttention(nn.Module):
+ def __init__(self, config: Qwen3OmniMoeVisionEncoderConfig) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ if self.config._attn_implementation == "flash_attention_2":
+ # Flash Attention 2: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen3OmniMoeVisionPatchMerger(nn.Module):
+ def __init__(self, config: Qwen3OmniMoeVisionEncoderConfig, use_postshuffle_norm=False) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.ln_q = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+ self.mlp = nn.ModuleList(
+ [
+ nn.Linear(self.hidden_size, self.hidden_size),
+ nn.GELU(),
+ nn.Linear(self.hidden_size, config.out_hidden_size),
+ ]
+ )
+
+ def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ hidden = self.ln_q(hidden.view(-1, self.hidden_size) if self.use_postshuffle_norm else hidden).view(
+ -1, self.hidden_size
+ )
+ for layer in self.mlp:
+ hidden = layer(hidden)
+ return hidden
+
+
+class Qwen3OmniMoeVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3OmniMoeVisionPatchEmbed(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+
+ kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+ self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen3OmniMoeVisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ freqs = torch.outer(seq, self.inv_freq)
+ return freqs
+
+
+class Qwen3OmniMoeVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.attn = Qwen3OmniMoeVisionAttention(config=config)
+ self.mlp = Qwen3OmniMoeVisionMLP(config=config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen3OmniMoeVisionEncoder(Qwen3OmniMoePreTrainedModel):
+ config: Qwen3OmniMoeVisionEncoderConfig
+ _no_split_modules = ["Qwen3OmniMoeVisionBlock"]
+
+ def __init__(self, config, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.merger_list = nn.ModuleList(
+ [
+ Qwen3OmniMoeVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=True,
+ )
+ for _ in range(len(config.deepstack_visual_indexes))
+ ]
+ )
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen3OmniMoeVisionPatchEmbed(
+ config=config,
+ )
+
+ self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
+ self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen3OmniMoeVisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList([Qwen3OmniMoeVisionBlock(config) for _ in range(config.depth)])
+ self.merger = Qwen3OmniMoeVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=False,
+ )
+
+ self.deepstack_visual_indexes = config.deepstack_visual_indexes
+
+ self.gradient_checkpointing = False
+
+ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+ merge_size = self.spatial_merge_size
+
+ max_hw = int(grid_thw[:, 1:].max().item())
+ freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2)
+ device = freq_table.device
+
+ total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
+ pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)
+
+ offset = 0
+ for num_frames, height, width in grid_thw:
+ merged_h, merged_w = height // merge_size, width // merge_size
+
+ block_rows = torch.arange(merged_h, device=device) # block row indices
+ block_cols = torch.arange(merged_w, device=device) # block col indices
+ intra_row = torch.arange(merge_size, device=device) # intra-block row offsets
+ intra_col = torch.arange(merge_size, device=device) # intra-block col offsets
+
+ # Compute full-resolution positions
+ row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
+ col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]
+
+ row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+ col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+
+ coords = torch.stack((row_idx, col_idx), dim=-1)
+
+ if num_frames > 1:
+ coords = coords.repeat(num_frames, 1)
+
+ num_tokens = coords.shape[0]
+ pos_ids[offset : offset + num_tokens] = coords
+ offset += num_tokens
+
+ embeddings = freq_table[pos_ids] # lookup rotary embeddings
+ embeddings = embeddings.flatten(1)
+ return embeddings
+
+ def fast_pos_embed_interpolate(self, grid_thw):
+ grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
+
+ idx_list = [[] for _ in range(4)]
+ weight_list = [[] for _ in range(4)]
+
+ for t, h, w in zip(grid_ts, grid_hs, grid_ws):
+ h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
+ w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
+
+ h_idxs_floor = h_idxs.int()
+ w_idxs_floor = w_idxs.int()
+ h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+ w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+
+ dh = h_idxs - h_idxs_floor
+ dw = w_idxs - w_idxs_floor
+
+ base_h = h_idxs_floor * self.num_grid_per_side
+ base_h_ceil = h_idxs_ceil * self.num_grid_per_side
+
+ indices = [
+ (base_h[None].T + w_idxs_floor[None]).flatten(),
+ (base_h[None].T + w_idxs_ceil[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
+ ]
+
+ weights = [
+ ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
+ ((1 - dh)[None].T * dw[None]).flatten(),
+ (dh[None].T * (1 - dw)[None]).flatten(),
+ (dh[None].T * dw[None]).flatten(),
+ ]
+
+ for i in range(4):
+ idx_list[i].extend(indices[i].tolist())
+ weight_list[i].extend(weights[i].tolist())
+
+ idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
+ weight_tensor = torch.tensor(
+ weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
+ )
+ pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
+ patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
+
+ patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])
+
+ patch_pos_embeds_permute = []
+ merge_size = self.config.spatial_merge_size
+ for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
+ pos_embed = pos_embed.repeat(t, 1)
+ pos_embed = (
+ pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
+ .permute(0, 1, 3, 2, 4, 5)
+ .flatten(0, 4)
+ )
+ patch_pos_embeds_permute.append(pos_embed)
+ patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
+ return patch_pos_embeds
+
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+
+ pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+ hidden_states = hidden_states + pos_embeds
+
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ deepstack_feature_lists = []
+ for layer_num, blk in enumerate(self.blocks):
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ if layer_num in self.deepstack_visual_indexes:
+ deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
+ hidden_states
+ )
+ deepstack_feature_lists.append(deepstack_feature)
+
+ hidden_states = self.merger(hidden_states)
+
+ return hidden_states, deepstack_feature_lists
+
+ @property
+ def deepstack_merger_list(self):
+ return self.merger_list
+
+
+class Qwen3OmniMoeThinkerTextRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen3OmniMoeTextConfig, device=None):
+ super().__init__()
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", "default")
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
+
+ def apply_interleaved_mrope(self, freqs, mrope_section):
+ """Apply interleaved MRoPE to 3D rotary embeddings.
+ Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
+ interleaved [THTHWHTHW...TT], preserving frequency continuity.
+ args:
+ x: (3, bs, seq_len, head_dim // 2)
+ mrope_section: (3,)
+ returns:
+ x_t: (bs, seq_len, head_dim // 2)
+ """
+ freqs_t = freqs[0] # just overwrite the first dimension T
+ for dim, offset in enumerate((1, 2), start=1): # H, W
+ length = mrope_section[dim] * 3
+ idx = slice(offset, length, 3)
+ freqs_t[..., idx] = freqs[dim, ..., idx]
+ return freqs_t
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, Qwen3OmniMoeThinker has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Qwen3OmniMoeThinkerTextMLP(nn.Module):
+ def __init__(self, config, intermediate_size=None):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = intermediate_size if intermediate_size is not None else 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 Qwen3OmniMoeThinkerTextSparseMoeBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.num_experts = config.num_experts
+ self.top_k = config.num_experts_per_tok
+ self.norm_topk_prob = config.norm_topk_prob
+
+ # gating
+ self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+ self.experts = nn.ModuleList(
+ [
+ Qwen3OmniMoeThinkerTextMLP(config, intermediate_size=config.moe_intermediate_size)
+ for _ in range(self.num_experts)
+ ]
+ )
+
+ 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)
+ # router_logits: (batch * sequence_length, n_experts)
+ router_logits = self.gate(hidden_states)
+
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+ if self.norm_topk_prob: # only diff with mixtral sparse moe block!
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+ # we cast back to the input dtype
+ routing_weights = routing_weights.to(hidden_states.dtype)
+
+ final_hidden_states = torch.zeros(
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+ )
+
+ # One hot encode the selected experts to create an expert mask
+ # this will be used to easily index which expert is going to be sollicitated
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+ # Loop over all available experts in the model and perform the computation on each expert
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+ for expert_idx in expert_hit:
+ expert_layer = self.experts[expert_idx]
+ idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+
+ # Index the correct hidden states and compute the expert hidden state for
+ # the current expert. We need to make sure to multiply the output hidden
+ # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+ current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+ current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+ # However `index_add_` only support torch tensors for indexing so we'll use
+ # the `top_x` tensor here.
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+ final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+ return final_hidden_states, router_logits
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3OmniMoeThinkerTextRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ Qwen3OmniMoeThinkerTextRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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.
+ position_ids (`torch.Tensor`, *optional*):
+ Deprecated and unused.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class Qwen3OmniMoeThinkerTextAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config, layer_idx):
+ 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 = Qwen3OmniMoeThinkerTextRMSNorm(
+ self.head_dim, eps=config.rms_norm_eps
+ ) # unlike olmo, only on the head dim!
+ self.k_norm = Qwen3OmniMoeThinkerTextRMSNorm(
+ self.head_dim, eps=config.rms_norm_eps
+ ) # thus post q_norm does not need reshape
+ self.sliding_window = None
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ sliding_window=self.sliding_window, # 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 Qwen3OmniMoeThinkerTextDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Qwen3OmniMoeThinkerTextAttention(config, layer_idx)
+
+ if (layer_idx not in config.mlp_only_layers) and (
+ config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+ ):
+ self.mlp = Qwen3OmniMoeThinkerTextSparseMoeBlock(config)
+ else:
+ self.mlp = Qwen3OmniMoeThinkerTextMLP(config, intermediate_size=config.intermediate_size)
+
+ self.input_layernorm = Qwen3OmniMoeThinkerTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3OmniMoeThinkerTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.FloatTensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+ `(batch, sequence_length)` where padding elements are indicated by 0.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ output_router_logits (`bool`, *optional*):
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss,
+ and should not be returned during inference.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **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)
+ # For the MoE layers, we need to unpack
+ if isinstance(hidden_states, tuple):
+ hidden_states, _ = hidden_states
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class Qwen3OmniMoeThinkerTextPreTrainedModel(PreTrainedModel):
+ config = Qwen3OmniMoeTextConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3OmniMoeThinkerTextDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = False # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "router_logits": OutputRecorder(Qwen3OmniMoeThinkerTextSparseMoeBlock, index=1),
+ "hidden_states": Qwen3OmniMoeThinkerTextDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ }
+ config_class = Qwen3OmniMoeTextConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3OmniMoeTextRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ Qwen3OmniMoeTextRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+@auto_docstring(
+ custom_intro=(
+ "Text part of Qwen3OmniMoeThinker, "
+ "not a pure text-only model, as DeepStack integrates visual features into the early hidden states."
+ )
+)
+class Qwen3OmniMoeThinkerTextModel(Qwen3OmniMoePreTrainedModel):
+ config: Qwen3OmniMoeTextConfig
+ _no_split_modules = ["Qwen3OmniMoeThinkerTextDecoderLayer"]
+ config_class = Qwen3OmniMoeTextConfig
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeThinkerTextDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeThinkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTextConfig):
+ 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(
+ [Qwen3OmniMoeThinkerTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3OmniMoeThinkerTextRotaryEmbedding(config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ # args for deepstack
+ visual_pos_masks: Optional[torch.Tensor] = None,
+ deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Union[tuple, BaseModelOutputWithPast]:
+ r"""
+ visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+ The mask of the visual positions.
+ deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+ The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+ The feature is extracted from the different visual encoder layers, and fed to the decoder
+ hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ text_position_ids = position_ids[0]
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=text_position_ids,
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # decoder layers
+ for layer_idx, decoder_layer in enumerate(self.layers):
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ # add visual features to the hidden states of first several layers
+ if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+ hidden_states = self._deepstack_process(
+ hidden_states,
+ visual_pos_masks,
+ deepstack_visual_embeds[layer_idx],
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def _deepstack_process(self, hidden_states, visual_pos_masks, visual_embeds):
+ visual_pos_masks = visual_pos_masks[..., 0]
+ visual_pos_masks = visual_pos_masks.to(hidden_states.device)
+ visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
+ local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
+ hidden_states[visual_pos_masks, :] = local_this
+ return hidden_states
+
+
+@dataclass
+class Qwen3OmniMoeThinkerCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
+ r"""
+ Args:
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ rope_deltas: Optional[torch.LongTensor] = None
+
+
+def load_balancing_loss_func(
+ gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
+ num_experts: Optional[int] = None,
+ top_k=2,
+ attention_mask: Optional[torch.Tensor] = None,
+) -> Union[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(
+ custom_intro="""
+ The Qwen2.5OmniThinker model which consists of a audio backbone and a language model.
+ """
+)
+class Qwen3OmniMoeThinkerForConditionalGeneration(
+ Qwen3OmniMoePreTrainedModelForConditionalGeneration, GenerationMixin
+):
+ config: Qwen3OmniMoeThinkerConfig
+ base_model_prefix = "thinker"
+ _tied_weights_keys = ["model.embed_tokens.weight", "lm_head.weight"]
+ _no_split_modules = [
+ "Qwen3OmniMoeAudioEncoderLayer",
+ "Qwen3OmniMoeThinkerTextDecoderLayer",
+ ]
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeThinkerTextDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeThinkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.audio_tower = Qwen3OmniMoeAudioEncoder._from_config(config.audio_config)
+ self.visual = Qwen3OmniMoeVisionEncoder._from_config(config.vision_config)
+ self.vocab_size = config.text_config.vocab_size
+ self.model = Qwen3OmniMoeThinkerTextModel._from_config(config.text_config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+ self.spatial_merge_size = config.vision_config.spatial_merge_size
+ self.rope_deltas = None
+ self.num_experts = config.text_config.num_experts
+ self.num_experts_per_tok = config.text_config.num_experts_per_tok
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def get_video_features(
+ self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
+ ):
+ """
+ Encodes videos into continuous embeddings that can be forwarded to the language model.
+
+ Args:
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+ return video_embeds
+
+ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+ """
+ Encodes images into continuous embeddings that can be forwarded to the language model.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+ return image_embeds
+
+ def get_audio_features(
+ self,
+ input_features: torch.FloatTensor,
+ feature_attention_mask: Optional[torch.LongTensor] = None,
+ audio_feature_lengths: Optional[torch.LongTensor] = None,
+ ):
+ """
+ Encodes audios into continuous embeddings that can be forwarded to the language model.
+
+ Args:
+ input_features (`torch.FloatTensor`):
+ The tensors corresponding to the input audios.
+ feature_attention_mask (`torch.LongTensor`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ """
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+ else:
+ audio_feature_lengths = None
+
+ feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ audio_outputs = self.audio_tower(
+ input_features,
+ feature_lens=feature_lens,
+ )
+ audio_features = audio_outputs.last_hidden_state
+
+ return audio_features
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: Optional[torch.FloatTensor] = None,
+ video_features: Optional[torch.FloatTensor] = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ special_audio_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ ).all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+ special_audio_mask = input_ids == self.config.audio_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel():
+ raise ValueError(
+ f"Videos features and image tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
+ )
+
+ special_audio_mask = special_audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ return special_image_mask, special_video_mask, special_audio_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ input_features=None,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ attention_mask=None,
+ feature_attention_mask=None,
+ audio_feature_lengths=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ rope_deltas=None,
+ labels=None,
+ use_cache=None,
+ output_router_logits: Optional[bool] = None,
+ use_audio_in_video=None,
+ cache_position=None,
+ video_second_per_grid=None,
+ **kwargs,
+ ) -> Union[tuple, Qwen3OmniMoeThinkerCausalLMOutputWithPast]:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ 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]`.
+ use_audio_in_video (`bool`, *optional*):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+
+ Example:
+
+ ```python
+ >>> from io import BytesIO
+ >>> from urllib.request import urlopen
+ >>> import librosa
+ >>> from qwen_vl_utils import process_vision_info
+ >>> from transformers import Qwen3OmniMoeProcessor, Qwen3OmniMoeThinkerForConditionalGeneration
+
+ >>> thinker = Qwen3OmniMoeThinkerForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+ >>> processor = Qwen3OmniMoeProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+
+ >>> conversations = [
+ >>> {'role': 'system', 'content': 'You are a helpful voice chat bot, and please respond to me in a casual conversation manner using random voice.'},
+ >>> {"role": "user", "content": [
+ >>> {"type": "image", "image_url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ >>> {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3"},
+ >>> ]},
+ >>> ]
+
+ >>> text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+ >>> audios = [ librosa.load(BytesIO(urlopen( conversations[1]['content'][1]['audio_url'] ).read()), sr=self.processor.feature_extractor.sampling_rate) ]
+ >>> images, videos = process_vision_info(conversations)
+ >>> inputs = processor(text=text, audios=audios, images=images, videos=videos, return_tensors="pt", padding=True)
+
+ >>> # Generate
+ >>> inputs['use_audio_in_video'] = `True` or `False`
+ >>> generation = thinker.generate(**inputs, max_new_tokens=2048)
+ >>> generate_ids = generation[:, inputs.input_ids.size(1):]
+
+ >>> response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ ```"""
+ output_router_logits = (
+ output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits
+ )
+
+ if inputs_embeds is None:
+ # 1. Extract the input embeddings
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ visual_embeds_multiscale = None
+ visual_pos_masks = None
+ # 2. Merge text , audios , image and video
+ if input_features is not None:
+ audio_features = self.get_audio_features(
+ input_features,
+ feature_attention_mask=feature_attention_mask,
+ audio_feature_lengths=audio_feature_lengths,
+ )
+ audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
+
+ if pixel_values is not None:
+ image_embeds, image_embeds_multiscale = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ visual_pos_masks = image_mask
+ visual_embeds_multiscale = image_embeds_multiscale
+
+ if pixel_values_videos is not None:
+ video_embeds, video_embeds_multiscale = self.get_video_features(pixel_values_videos, video_grid_thw)
+
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if visual_embeds_multiscale is None:
+ visual_embeds_multiscale = video_embeds_multiscale
+ visual_pos_masks = video_mask
+ else:
+ visual_pos_masks = video_mask | image_mask
+ visual_embeds_multiscale_joint = ()
+ image_mask_joint = image_mask[visual_pos_masks]
+ video_mask_joint = video_mask[visual_pos_masks]
+ for img_embed, vid_embed in zip(visual_embeds_multiscale, video_embeds_multiscale):
+ embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1])
+ embed_joint[image_mask_joint, :] = img_embed
+ embed_joint[video_mask_joint, :] = vid_embed
+ visual_embeds_multiscale_joint = visual_embeds_multiscale_joint + (embed_joint,)
+ visual_embeds_multiscale = visual_embeds_multiscale_joint
+
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ else:
+ audio_feature_lengths = None
+
+ if attention_mask is not None and position_ids is None:
+ if (
+ cache_position is None
+ or (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ ):
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ 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,
+ cache_position=cache_position,
+ deepstack_visual_embeds=visual_embeds_multiscale,
+ visual_pos_masks=visual_pos_masks,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.get_text_config().vocab_size
+ )
+
+ 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 Qwen3OmniMoeThinkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ aux_loss=aux_loss,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ cache_position=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ input_features=None,
+ feature_attention_mask=None,
+ use_audio_in_video=False,
+ video_second_per_grid=None,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ input_features=input_features,
+ feature_attention_mask=feature_attention_mask,
+ use_audio_in_video=use_audio_in_video,
+ video_second_per_grid=video_second_per_grid,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+
+ if cache_position[0] != 0:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+ model_inputs["input_features"] = None
+
+ return model_inputs
+
+
+class Qwen3OmniMoeTalkerResizeMLP(nn.Module):
+ def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+ super().__init__()
+ self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.text_config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+@dataclass
+class Qwen3OmniMoeTalkerCodePredictorOutputWithPast(CausalLMOutputWithPast):
+ r"""
+ generation_steps (`int`, *optional*)
+ Current generation step of code predictor model.
+ """
+
+ generation_steps: Optional[int] = None
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3OmniMoeRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Qwen3OmniMoeRMSNorm 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 Qwen3OmniMoeTalkerCodePredictorAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Qwen3OmniMoeConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.q_norm = Qwen3OmniMoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
+ self.k_norm = Qwen3OmniMoeRMSNorm(
+ self.head_dim, eps=config.rms_norm_eps
+ ) # thus post q_norm does not need reshape
+ self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ sliding_window=self.sliding_window, # 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 Qwen3OmniMoeMLP(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 Qwen3OmniMoeTalkerCodePredictorDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Qwen3OmniMoeTalkerCodePredictorAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Qwen3OmniMoeMLP(config)
+ self.input_layernorm = Qwen3OmniMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3OmniMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.attention_type = config.layer_types[layer_idx]
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **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,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+class Qwen3OmniMoeRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen3OmniMoeConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @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 torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class Qwen3OmniMoeTalkerCodePredictorModel(Qwen3OmniMoePreTrainedModel):
+ config_class = Qwen3OmniMoeTalkerCodePredictorConfig
+ base_model_prefix = "talker.code_predictor.model"
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeTalkerCodePredictorAttention,
+ "hidden_states": Qwen3OmniMoeTalkerCodePredictorDecoderLayer,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerCodePredictorConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self.layers = nn.ModuleList(
+ [
+ Qwen3OmniMoeTalkerCodePredictorDecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ]
+ )
+ self.norm = Qwen3OmniMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+ self.codec_embedding = nn.ModuleList(
+ [nn.Embedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if input_ids is not None:
+ raise ValueError("`input_ids` is expected to be `None`")
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ # It may already have been prepared by e.g. `generate`
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
+ # Prepare mask arguments
+ mask_kwargs = {
+ "config": self.config,
+ "input_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "cache_position": cache_position,
+ "past_key_values": past_key_values,
+ "position_ids": position_ids,
+ }
+ # Create the masks
+ causal_mask_mapping = {
+ "full_attention": create_causal_mask(**mask_kwargs),
+ }
+
+ 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,
+ attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+ def get_input_embeddings(self):
+ return self.codec_embedding
+
+
+@auto_docstring
+class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+ _tp_plan = {"lm_head": "colwise_rep"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+ config_class = Qwen3OmniMoeTalkerCodePredictorConfig
+ base_model_prefix = "talker.code_predictor"
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeTalkerCodePredictorAttention,
+ "hidden_states": Qwen3OmniMoeTalkerCodePredictorDecoderLayer,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerCodePredictorConfig):
+ super().__init__(config)
+ self.model = Qwen3OmniMoeTalkerCodePredictorModel._from_config(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.ModuleList(
+ [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ labels=None,
+ use_cache=None,
+ cache_position=None,
+ generation_steps=None,
+ **kwargs,
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Args:
+ generation_steps (`int`):
+ generation step of code predictor, 0..num_code_groups-1
+ """
+
+ # Prefill stage
+ if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
+ generation_steps = inputs_embeds.shape[1] - 2 # hidden & layer 0
+ # Generation stage
+ else:
+ inputs_embeds = self.model.get_input_embeddings()[generation_steps - 1](input_ids)
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=None,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ logits = self.lm_head[generation_steps](hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return Qwen3OmniMoeTalkerCodePredictorOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ generation_steps=generation_steps + 1,
+ )
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+ model_kwargs["generation_steps"] = outputs.generation_steps
+ return model_kwargs
+
+
+@dataclass
+class Qwen3OmniMoeTalkerOutputWithPast(MoeCausalLMOutputWithPast):
+ r"""
+ Args:
+ generation_step (`int`, *optional*):
+ Current generation step, used to track which `trailing_text_hidden` should be used.
+ """
+
+ generation_step: Optional[int] = None
+
+
+class Qwen3OmniMoeTalkerRotaryEmbedding(Qwen3OmniMoeThinkerTextRotaryEmbedding):
+ pass
+
+
+class Qwen3OmniMoeTalkerTextMLP(nn.Module):
+ def __init__(self, config, intermediate_size=None):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = intermediate_size if intermediate_size is not None else 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 Qwen3OmniMoeTalkerTextSparseMoeBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.num_experts = config.num_experts
+ self.top_k = config.num_experts_per_tok
+ self.norm_topk_prob = config.norm_topk_prob
+
+ # gating
+ self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+ self.experts = nn.ModuleList(
+ [
+ Qwen3OmniMoeTalkerTextMLP(config, intermediate_size=config.moe_intermediate_size)
+ for _ in range(self.num_experts)
+ ]
+ )
+
+ self.shared_expert = Qwen3OmniMoeTalkerTextMLP(
+ config, intermediate_size=config.shared_expert_intermediate_size
+ )
+ self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
+
+ 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)
+ # router_logits: (batch * sequence_length, n_experts)
+ router_logits = self.gate(hidden_states)
+
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+ if self.norm_topk_prob:
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+ # we cast back to the input dtype
+ routing_weights = routing_weights.to(hidden_states.dtype)
+
+ final_hidden_states = torch.zeros(
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+ )
+
+ # One hot encode the selected experts to create an expert mask
+ # this will be used to easily index which expert is going to be sollicitated
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+ # Loop over all available experts in the model and perform the computation on each expert
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+ for expert_idx in expert_hit:
+ expert_layer = self.experts[expert_idx]
+ idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+
+ # Index the correct hidden states and compute the expert hidden state for
+ # the current expert. We need to make sure to multiply the output hidden
+ # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+ current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+ current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+ # However `index_add_` only support torch tensors for indexing so we'll use
+ # the `top_x` tensor here.
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+
+ shared_expert_output = self.shared_expert(hidden_states)
+ shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
+
+ final_hidden_states = final_hidden_states + shared_expert_output
+
+ final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+ return final_hidden_states, router_logits
+
+
+class Qwen3OmniMoeTalkerDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Qwen3OmniMoeThinkerTextAttention(config, layer_idx)
+
+ if (layer_idx not in config.mlp_only_layers) and (
+ config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+ ):
+ self.mlp = Qwen3OmniMoeThinkerTextSparseMoeBlock(config)
+ else:
+ self.mlp = Qwen3OmniMoeThinkerTextMLP(config, intermediate_size=config.intermediate_size)
+
+ self.input_layernorm = Qwen3OmniMoeThinkerTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3OmniMoeThinkerTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.mlp = Qwen3OmniMoeTalkerTextSparseMoeBlock(config)
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.FloatTensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+ `(batch, sequence_length)` where padding elements are indicated by 0.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ output_router_logits (`bool`, *optional*):
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss,
+ and should not be returned during inference.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ **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)
+ # For the MoE layers, we need to unpack
+ if isinstance(hidden_states, tuple):
+ hidden_states, _ = hidden_states
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring(
+ custom_intro=(
+ "Text part of Qwen3OmniMoe, "
+ "not a pure text-only model, as DeepStack integrates visual features into the early hidden states."
+ )
+)
+class Qwen3OmniMoeTalkerModel(Qwen3OmniMoePreTrainedModel):
+ config: Qwen3OmniMoeTextConfig
+ _no_split_modules = ["Qwen3OmniMoeTalkerDecoderLayer"]
+ config_class = Qwen3OmniMoeTalkerTextConfig
+ base_model_prefix = "talker.model"
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeTalkerDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeTalkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerTextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self.layers = nn.ModuleList(
+ [Qwen3OmniMoeTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3OmniMoeTalkerRotaryEmbedding(config)
+ self.gradient_checkpointing = False
+ self.codec_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ # args for deepstack
+ visual_pos_masks: Optional[torch.Tensor] = None,
+ deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Union[tuple, BaseModelOutputWithPast]:
+ r"""
+ visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+ The mask of the visual positions.
+ deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+ The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+ The feature is extracted from the different visual encoder layers, and fed to the decoder
+ hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ text_position_ids = position_ids[0]
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=text_position_ids,
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # decoder layers
+ for layer_idx, decoder_layer in enumerate(self.layers):
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ # add visual features to the hidden states of first several layers
+ if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+ hidden_states = self._deepstack_process(
+ hidden_states,
+ visual_pos_masks,
+ deepstack_visual_embeds[layer_idx],
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def _deepstack_process(
+ self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
+ ):
+ visual_pos_masks = visual_pos_masks.to(hidden_states.device)
+ visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
+ local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
+ hidden_states[visual_pos_masks, :] = local_this
+ return hidden_states
+
+ def get_input_embeddings(self):
+ return self.codec_embedding
+
+
+@auto_docstring
+class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+ _tp_plan = {"lm_head": "colwise_rep"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+ config_class = Qwen3OmniMoeTalkerConfig
+ base_model_prefix = "talker"
+ _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"]
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeTalkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+ super().__init__(config)
+ self.model = Qwen3OmniMoeTalkerModel._from_config(config.text_config)
+ self.vocab_size = config.text_config.vocab_size
+ self.router_aux_loss_coef = config.text_config.router_aux_loss_coef
+ self.num_experts = config.text_config.num_experts
+ self.num_experts_per_tok = config.text_config.num_experts_per_tok
+ self.text_projection = Qwen3OmniMoeTalkerResizeMLP(config)
+ self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(config)
+ self.codec_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.code_predictor = Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration._from_config(
+ config=config.code_predictor_config
+ )
+ self.rope_deltas = None
+ self.spatial_merge_size = self.config.spatial_merge_size
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ use_audio_in_video=None,
+ audio_feature_lengths=None,
+ video_second_per_grid=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ labels=None,
+ use_cache=None,
+ output_router_logits=None,
+ cache_position=None,
+ residual_codes=None,
+ trailing_text_hidden=None,
+ tts_pad_embed=None,
+ generation_step=None,
+ talker_input_ids=None,
+ **kwargs,
+ ) -> MoeCausalLMOutputWithPast:
+ r"""
+ Args:
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ residual_codes (`torch.Tensor`):
+ The predicted residual codes of previous step.
+ trailing_text_hidden (`torch.Tensor`):
+ Text hidden states from thinker after the first token.
+ tts_pad_embed (`torch.Tensor`):
+ Embedding tensor of `tts_pad_token_id`.
+ generation_step (`int`):
+ Generation step since prefill, used to sync with `trailing_text_hidden`.
+ talker_input_ids (`torch.Tensor`):
+ Input ids from thinker, used to compute 3d RoPE.
+ """
+ # Prefill
+ if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
+ generation_step = -1
+ residual_codes = None
+ if attention_mask is not None:
+ if (
+ cache_position is None
+ or (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ ):
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ talker_input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs: MoeModelOutputWithPast = self.model(
+ input_ids=None,
+ 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,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ logits = self.codec_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ aux_loss = None
+ if output_router_logits:
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.num_experts,
+ self.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
+
+ return Qwen3OmniMoeTalkerOutputWithPast(
+ loss=loss,
+ logits=logits,
+ aux_loss=aux_loss,
+ past_key_values=outputs.past_key_values,
+ hidden_states=(
+ outputs.hidden_states,
+ residual_codes,
+ ), # TODO: hack here to take residual codes out, need refactor.
+ generation_step=generation_step + 1,
+ )
+
+ # Should inherit from PretrainedModel, but cannot inherit multiple classes in modular
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: Optional[torch.LongTensor] = None,
+ second_per_grids: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return Qwen3OmniMoePreTrainedModelForConditionalGeneration.get_rope_index(
+ self,
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_seqlens,
+ second_per_grids,
+ )
+
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[torch.Tensor],
+ grid_hs: list[torch.Tensor],
+ grid_ws: list[torch.Tensor],
+ ):
+ return Qwen3OmniMoePreTrainedModelForConditionalGeneration.get_llm_pos_ids_for_vision(
+ self, start_idx, vision_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+ model_kwargs["hidden_states"] = outputs.hidden_states
+ model_kwargs["generation_step"] = outputs.generation_step
+ return model_kwargs
+
+ def prepare_inputs_for_generation(
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
+ ):
+ hidden_states = kwargs.pop("hidden_states", None)
+ inputs = super().prepare_inputs_for_generation(
+ input_ids, past_key_values, attention_mask, inputs_embeds, cache_position, **kwargs
+ )
+ # Decode stage
+ # TODO(raushan, gante): Refactor this part to a utility function
+ if cache_position[0] != 0:
+ input_ids = input_ids[:, -1:]
+ generation_step = kwargs.get("generation_step")
+ trailing_text_hidden = kwargs.get("trailing_text_hidden")
+ tts_pad_embed = kwargs.get("tts_pad_embed")
+ last_id_hidden = self.get_input_embeddings()(input_ids)
+
+ past_hidden = hidden_states[0][-1][:, -1:].to(last_id_hidden.device) # hidden, last layer, last token
+ predictor_result = self.code_predictor.generate(
+ inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
+ max_new_tokens=self.config.num_code_groups - 1,
+ do_sample=True,
+ top_k=50,
+ top_p=0.8,
+ output_hidden_states=True,
+ return_dict_in_generate=True,
+ )
+ residual_codes = torch.cat((input_ids, predictor_result.sequences.to(input_ids.device)), dim=-1)
+
+ mid_residual_hiddens = [hid[0].to(last_id_hidden.device) for hid in predictor_result.hidden_states[1:]]
+ last_residual_hidden = self.code_predictor.get_input_embeddings()[-1](
+ predictor_result.sequences[..., -1:]
+ ).to(last_id_hidden.device)
+ codec_hiddens = torch.cat(
+ [last_id_hidden] + mid_residual_hiddens + [last_residual_hidden],
+ dim=1,
+ )
+ inputs_embeds = codec_hiddens.sum(1, keepdim=True)
+
+ if generation_step < trailing_text_hidden.shape[1]:
+ inputs_embeds = inputs_embeds + trailing_text_hidden[:, generation_step].unsqueeze(1).to(
+ inputs_embeds.device
+ )
+ else:
+ inputs_embeds = inputs_embeds + tts_pad_embed.to(inputs_embeds.device)
+ inputs["inputs_embeds"] = inputs_embeds
+ inputs["residual_codes"] = residual_codes
+ return inputs
+
+
+class Qwen3OmniMoeCausalConvNet(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ dilation=1,
+ stride=1,
+ groups=1,
+ ):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ )
+ self.stride = stride
+ self.kernel_size = (kernel_size - 1) * dilation + 1
+ self.dilation = dilation
+ self.padding = self.kernel_size - self.stride
+
+ def _get_extra_padding_for_conv1d(self, hidden_state: torch.Tensor) -> int:
+ length = hidden_state.shape[-1]
+ n_frames = (length - self.kernel_size + self.padding) / self.stride + 1
+ ideal_length = (math.ceil(n_frames) - 1) * self.stride + (self.kernel_size - self.padding)
+ return ideal_length - length
+
+ def forward(self, hidden_state):
+ extra_padding = self._get_extra_padding_for_conv1d(hidden_state)
+ hidden_state = F.pad(hidden_state, (self.padding, extra_padding), mode="constant", value=0)
+ return self.conv(hidden_state).contiguous()
+
+
+class Qwen3OmniMoeCausalTransConvNet(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+ super().__init__()
+ self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride)
+
+ pad = kernel_size - stride
+ self.left_pad = math.ceil(pad)
+ self.right_pad = pad = self.left_pad
+
+ def forward(self, hidden_state):
+ hidden_state = self.conv(hidden_state)
+ hidden_state = hidden_state[..., self.left_pad : hidden_state.shape[-1] - self.right_pad]
+ return hidden_state.contiguous()
+
+
+class Qwen3OmniMoeConvNeXtBlock(nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.dwconv = Qwen3OmniMoeCausalConvNet(
+ dim,
+ dim,
+ kernel_size=7,
+ groups=dim,
+ dilation=1,
+ )
+ self.norm = nn.LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, 4 * dim)
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
+
+ def forward(self, hidden_states):
+ input = hidden_states
+
+ hidden_states = self.dwconv(hidden_states)
+ hidden_states = hidden_states.permute(0, 2, 1)
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.pwconv1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.pwconv2(hidden_states)
+
+ hidden_states = self.gamma * hidden_states
+
+ hidden_states = hidden_states.permute(0, 2, 1)
+
+ hidden_states = input + hidden_states
+
+ return hidden_states
+
+
+class Qwen3OmniMoeCode2WavRotatoryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen3OmniMoeConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @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 torch.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 Qwen3OmniMoeCode2WavAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ 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 = nn.Identity()
+ self.k_norm = nn.Identity()
+ self.sliding_window = config.sliding_window
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ sliding_window=self.sliding_window, # 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 Qwen3OmniMoeCode2WavMlp(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
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3OmniMoeCode2WavRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Qwen3OmniMoeCode2WavRMSNorm 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 Qwen3OmniMoeCode2WavLayerScale(nn.Module):
+ """Layer scale from [Touvron et al 2021] (https://huggingface.co/papers/2103.17239).
+ This rescales diagonally the residual outputs close to 0, with a learnt scale.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ channels = config.hidden_size
+ initial_scale = config.layer_scale_initial_scale
+ self.scale = nn.Parameter(torch.full((channels,), initial_scale, requires_grad=True))
+
+ def forward(self, x: torch.Tensor):
+ return self.scale * x
+
+
+class Qwen3OmniMoeCode2WavTransformerLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Qwen3OmniMoeCode2WavAttention(config, layer_idx)
+ self.mlp = Qwen3OmniMoeCode2WavMlp(config)
+ self.input_layernorm = Qwen3OmniMoeCode2WavRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3OmniMoeCode2WavRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.self_attn_layer_scale = Qwen3OmniMoeCode2WavLayerScale(config)
+ self.mlp_layer_scale = Qwen3OmniMoeCode2WavLayerScale(config)
+ self.attention_type = "sliding_attention"
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*):
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+ query_sequence_length, key_sequence_length)` if default attention is used.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ hidden_states = residual + self.self_attn_layer_scale(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + self.mlp_layer_scale(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class Qwen3OmniMoeCode2WavTransformerModel(Qwen3OmniMoePreTrainedModel):
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeCode2WavTransformerLayer,
+ "attentions": Qwen3OmniMoeCode2WavAttention,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig):
+ super().__init__(config)
+ self.layers = nn.ModuleList(
+ [Qwen3OmniMoeCode2WavTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3OmniMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+ self.window_size = config.sliding_window
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ use_cache=None,
+ cache_position=None,
+ **kwargs,
+ ) -> BaseModelOutputWithPast:
+ if input_ids is not None:
+ raise ValueError("input_ids is not expected")
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ # It may already have been prepared by e.g. `generate`
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
+ # Prepare mask arguments
+ mask_kwargs = {
+ "config": self.config,
+ "input_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "cache_position": cache_position,
+ "past_key_values": past_key_values,
+ "position_ids": position_ids,
+ }
+ # Create the masks
+ causal_mask_mapping = {
+ "full_attention": create_causal_mask(**mask_kwargs),
+ }
+ # The sliding window alternating layers are not always activated depending on the config
+ if self.has_sliding_layers:
+ causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+ hidden_states = inputs_embeds
+
+ # 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,
+ attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+class SnakeBeta(nn.Module):
+ """
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://huggingface.co/papers/2006.08195
+ """
+
+ def __init__(self, in_features, alpha=1.0):
+ super().__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ self.beta = Parameter(torch.zeros(in_features) * alpha)
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, hidden_states):
+ """
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta ∶= x + 1/b * sin^2 (xa)
+ """
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
+ torch.sin(hidden_states * alpha), 2
+ )
+
+ return hidden_states
+
+
+class Qwen3OmniMoeCode2WavDecoderResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+
+ self.act1 = SnakeBeta(dim)
+ self.conv1 = Qwen3OmniMoeCausalConvNet(dim, dim, kernel_size=7, dilation=dilation)
+ self.act2 = SnakeBeta(dim)
+ self.conv2 = Qwen3OmniMoeCausalConvNet(dim, dim, kernel_size=1)
+
+ def forward(self, hidden_state):
+ residual = hidden_state
+
+ hidden_state = self.act1(hidden_state)
+ hidden_state = self.conv1(hidden_state)
+ hidden_state = self.act2(hidden_state)
+ hidden_state = self.conv2(hidden_state)
+ return hidden_state + residual
+
+
+class Qwen3OmniMoeCode2WavDecoderBlock(Qwen3OmniMoePreTrainedModel):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ super().__init__(config)
+ in_dim = config.decoder_dim // 2**layer_idx
+ out_dim = config.decoder_dim // 2 ** (layer_idx + 1)
+ upsample_rate = config.upsample_rates[layer_idx]
+
+ block = [
+ SnakeBeta(in_dim),
+ Qwen3OmniMoeCausalTransConvNet(in_dim, out_dim, 2 * upsample_rate, upsample_rate),
+ ]
+
+ for dilation in (1, 3, 9):
+ block.append(Qwen3OmniMoeCode2WavDecoderResidualUnit(out_dim, dilation))
+
+ self.block = nn.ModuleList(block)
+
+ def forward(self, hidden):
+ for block in self.block:
+ hidden = block(hidden)
+ return hidden
+
+
+class Qwen3OmniMoeCode2Wav(Qwen3OmniMoePreTrainedModel):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig):
+ super().__init__(config)
+ self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios)
+ self.pre_transformer = Qwen3OmniMoeCode2WavTransformerModel._from_config(config)
+ self.code_embedding = nn.Embedding(config.codebook_size * config.num_quantizers, config.hidden_size)
+ self.register_buffer(
+ "code_offset", torch.arange(config.num_quantizers).view(1, -1, 1) * config.codebook_size, persistent=False
+ )
+
+ upsample = []
+ for factor in config.upsampling_ratios:
+ upsample.append(
+ nn.ModuleList(
+ [
+ Qwen3OmniMoeCausalTransConvNet(config.hidden_size, config.hidden_size, factor, factor),
+ Qwen3OmniMoeConvNeXtBlock(config.hidden_size),
+ ]
+ )
+ )
+ self.upsample = nn.ModuleList(upsample)
+
+ decoder = [Qwen3OmniMoeCausalConvNet(config.hidden_size, config.decoder_dim, 7)]
+ for i in range(len(config.upsample_rates)):
+ decoder.append(Qwen3OmniMoeCode2WavDecoderBlock(config, i))
+ output_dim = config.decoder_dim // 2 ** len(config.upsample_rates)
+ decoder += [
+ SnakeBeta(output_dim),
+ Qwen3OmniMoeCausalConvNet(output_dim, 1, 7),
+ ]
+ self.decoder = nn.ModuleList(decoder)
+
+ self.post_init()
+
+ def forward(self, codes):
+ if codes.shape[1] != self.config.num_quantizers:
+ raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
+ hidden = self.code_embedding(codes + self.code_offset).mean(1)
+ hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
+ hidden = hidden.permute(0, 2, 1)
+ for blocks in self.upsample:
+ for block in blocks:
+ hidden = block(hidden)
+ wav = hidden
+ for block in self.decoder:
+ wav = block(wav)
+ return wav.clamp(min=-1, max=1)
+
+ def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
+ wavs = []
+ start_index = 0
+ while start_index < codes.shape[-1]:
+ end_index = min(start_index + chunk_size, codes.shape[-1])
+ context_size = left_context_size if start_index - left_context_size > 0 else start_index
+ codes_chunk = codes[..., start_index - context_size : end_index]
+ wav_chunk = self(codes_chunk)
+ wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+ start_index = end_index
+ return torch.cat(wavs, dim=-1)
+
+
+class Qwen3OmniMoeForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin):
+ config_class = Qwen3OmniMoeConfig
+
+ def __init__(self, config: Qwen3OmniMoeConfig):
+ super().__init__(config)
+
+ self.thinker = Qwen3OmniMoeThinkerForConditionalGeneration._from_config(config.thinker_config)
+ self.has_talker = config.enable_audio_output
+ if self.has_talker:
+ self.enable_talker()
+ self.post_init()
+
+ def enable_talker(self):
+ self.talker = Qwen3OmniMoeTalkerForConditionalGeneration._from_config(self.config.talker_config)
+ self.code2wav = Qwen3OmniMoeCode2Wav._from_config(self.config.code2wav_config)
+
+ def disable_talker(self):
+ if hasattr(self, "talker"):
+ del self.talker
+ if hasattr(self, "code2wav"):
+ del self.code2wav
+ self.has_talker = False
+
+ def _get_talker_user_parts(
+ self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+ ):
+ user_talker_part = torch.empty(
+ (1, segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size),
+ device=self.talker.device,
+ dtype=self.talker.dtype,
+ )
+
+ user_mm_mask = multimodal_mask[:, im_start_index:segment_end_index]
+
+ # Multimodal data exists
+ if user_mm_mask.any():
+ user_thinker_hidden_mm = thinker_hidden[:, im_start_index:segment_end_index][user_mm_mask]
+ mm_hidden = self.talker.hidden_projection(user_thinker_hidden_mm).to(self.talker.device)
+ user_talker_part[user_mm_mask] = mm_hidden
+ user_thinker_embed = thinker_embed[:, im_start_index:segment_end_index][~user_mm_mask]
+ user_text_hidden = self.talker.text_projection(user_thinker_embed).to(self.talker.device)
+ user_talker_part[~user_mm_mask] = user_text_hidden
+ return user_talker_part
+
+ def _get_talker_assistant_parts(
+ self, im_start_index, segment_end_index, speaker_id, thinker_embed, tts_pad_embed, tts_bos_embed, tts_eos_embed
+ ):
+ assistant_hidden = self.talker.text_projection(thinker_embed[:, im_start_index:segment_end_index]).to(
+ self.talker.device
+ ) # [1 t d]
+ assistant_text_hidden = torch.cat(
+ (
+ assistant_hidden[:, :3],
+ tts_pad_embed.expand(-1, 4, -1),
+ tts_bos_embed,
+ assistant_hidden[:, 3:4], # First text
+ ),
+ dim=1,
+ )
+ codec_special_tokens = torch.tensor(
+ [
+ [
+ self.config.talker_config.codec_nothink_id,
+ self.config.talker_config.codec_think_bos_id,
+ self.config.talker_config.codec_think_eos_id,
+ speaker_id,
+ self.config.talker_config.codec_pad_id,
+ self.config.talker_config.codec_bos_id,
+ ]
+ ],
+ device=self.talker.device,
+ dtype=torch.long,
+ )
+ assistant_codec_hidden = torch.cat(
+ (
+ torch.zeros(
+ (1, 3, self.config.talker_config.text_config.hidden_size),
+ device=self.talker.device,
+ dtype=self.talker.dtype,
+ ),
+ self.talker.get_input_embeddings()(codec_special_tokens).to(self.talker.device),
+ ),
+ dim=1,
+ )
+ trailing_text_hidden = torch.cat(
+ (
+ assistant_hidden[:, 4:],
+ tts_eos_embed,
+ ),
+ dim=1,
+ )
+
+ input_embeds = assistant_text_hidden + assistant_codec_hidden
+ input_ids = torch.full(
+ (1, assistant_text_hidden.shape[1]),
+ fill_value=self.config.tts_pad_token_id,
+ dtype=torch.long,
+ device=assistant_text_hidden.device,
+ )
+ return input_embeds, input_ids, trailing_text_hidden
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ speaker: str = "Ethan",
+ use_audio_in_video: bool = False,
+ return_audio: Optional[bool] = None,
+ thinker_max_new_tokens: int = 1024,
+ thinker_eos_token_id: int = 151645,
+ talker_max_new_tokens: int = 4096,
+ talker_do_sample: bool = True,
+ talker_top_k: int = 50,
+ talker_top_p: float = 1.0,
+ talker_temperature: float = 0.9,
+ talker_repetition_penalty: float = 1.05,
+ **kwargs,
+ ):
+ if return_audio and not self.has_talker:
+ raise ValueError(
+ "Cannot use talker when talker module not initialized. Use `enable_talker` method or set enable_talker in config to enable talker."
+ )
+ if return_audio is None:
+ return_audio = self.has_talker
+
+ shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+ thinker_kwargs = {
+ "max_new_tokens": thinker_max_new_tokens,
+ "eos_token_id": thinker_eos_token_id,
+ }
+
+ talker_kwargs = {}
+ token2wav_kwargs = {}
+ if return_audio:
+ speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
+ if speaker_id is None:
+ raise NotImplementedError(f"Speaker {speaker} not implemented")
+ if input_ids.shape[0] != 1:
+ raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
+ talker_supppressed_tokens = [
+ i
+ for i in range(
+ self.config.talker_config.text_config.vocab_size - 1024,
+ self.config.talker_config.text_config.vocab_size,
+ )
+ if i != self.config.talker_config.codec_eos_token_id
+ ] # Suppress additional special tokens, should not be predicted
+ talker_kwargs = {
+ "max_new_tokens": talker_max_new_tokens,
+ "do_sample": talker_do_sample,
+ "top_k": talker_top_k,
+ "top_p": talker_top_p,
+ "temperature": talker_temperature,
+ "eos_token_id": self.config.talker_config.codec_eos_token_id,
+ "repetition_penalty": talker_repetition_penalty,
+ "suppress_tokens": talker_supppressed_tokens,
+ "output_hidden_states": True,
+ "return_dict_in_generate": True,
+ }
+ token2wav_kwargs = {}
+
+ for key, value in kwargs.items():
+ if key.startswith("thinker_"):
+ thinker_kwargs[key[len("thinker_") :]] = value
+ elif key.startswith("talker_"):
+ talker_kwargs[key[len("talker_") :]] = value
+ elif key.startswith("token2wav_"):
+ token2wav_kwargs[key[len("token2wav_") :]] = value
+ # Process special input values
+ elif key == "feature_attention_mask":
+ thinker_kwargs[key] = value
+ talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+ elif key in ("input_features", "attention_mask"):
+ thinker_kwargs[key] = value
+ # Put other key to shared kwargs
+ else:
+ shared_kwargs[key] = value
+
+ # Merge kwargs
+ for key, value in shared_kwargs.items():
+ if key not in thinker_kwargs:
+ thinker_kwargs[key] = value
+ if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
+ talker_kwargs[key] = value
+ if key not in token2wav_kwargs:
+ token2wav_kwargs[key] = value
+
+ # 1. Generate from thinker module
+ generate_audio = return_audio and self.has_talker
+ if generate_audio:
+ thinker_kwargs["output_hidden_states"] = True
+ thinker_kwargs["return_dict_in_generate"] = True
+
+ thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+
+ if not generate_audio:
+ return thinker_result, None
+
+ # 2. Prepare talker input
+ thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
+ self.talker.device
+ ) # [1 t d]
+ thinker_hidden = torch.cat(
+ [
+ hidden_states[self.config.talker_config.accept_hidden_layer]
+ for hidden_states in thinker_result.hidden_states
+ ],
+ dim=1,
+ ).to(self.talker.device) # [1 t d]
+ im_start_indexes = torch.cat(
+ (
+ torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
+ torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+ ),
+ dim=-1,
+ ).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+ multimodal_mask = (
+ (thinker_result.sequences == self.config.thinker_config.audio_token_id) |
+ (thinker_result.sequences == self.config.thinker_config.image_token_id) |
+ (thinker_result.sequences == self.config.thinker_config.video_token_id)
+ ).to(self.talker.device) # [1 t] # fmt: skip
+
+ talker_special_tokens = torch.tensor(
+ [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
+ device=self.thinker.device,
+ dtype=input_ids.dtype,
+ )
+ tts_bos_embed, tts_eos_embed, tts_pad_embed = (
+ self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
+ .to(self.talker.device)
+ .chunk(3, dim=1)
+ ) # 3 * [1 1 d]
+
+ talker_input_embeds = [] # [1 t d]
+ talker_input_ids = []
+ # For every chatml parts
+ for i in range(len(im_start_indexes) - 1):
+ im_start_index = im_start_indexes[i]
+ segment_end_index = im_start_indexes[i + 1]
+ role_token = input_ids[0][im_start_index + 1]
+ # Talker should ignore thinker system prompt
+ if role_token == self.config.system_token_id:
+ continue
+ # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
+ elif role_token == self.config.user_token_id:
+ talker_user_part = self._get_talker_user_parts(
+ im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+ )
+ talker_input_embeds.append(talker_user_part)
+ talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
+ # Take assistant output (for now)
+ elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
+ talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
+ im_start_index,
+ segment_end_index,
+ speaker_id,
+ thinker_embed,
+ tts_pad_embed,
+ tts_bos_embed,
+ tts_eos_embed,
+ )
+ talker_input_embeds.append(talker_assistant_embeds)
+ talker_input_ids.append(talker_assistant_ids)
+ # History assistant output (ignore for now)
+ elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
+ continue
+ else:
+ raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
+ talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
+ talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
+ talker_result = self.talker.generate(
+ inputs_embeds=talker_input_embed,
+ trailing_text_hidden=trailing_text_hidden,
+ tts_pad_embed=tts_pad_embed,
+ talker_input_ids=talker_input_id, # Not use input_ids to prevent repetation penalty out of bound
+ **talker_kwargs,
+ )
+ talker_codes = (
+ torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
+ .transpose(1, 2)
+ .to(self.code2wav.device)
+ )
+ talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25)
+
+ return thinker_result, talker_wavs.float()
+
+
+__all__ = [
+ "Qwen3OmniMoeForConditionalGeneration",
+ "Qwen3OmniMoeThinkerTextModel",
+ "Qwen3OmniMoeThinkerForConditionalGeneration",
+ "Qwen3OmniMoeTalkerForConditionalGeneration",
+ "Qwen3OmniMoePreTrainedModel",
+ "Qwen3OmniMoePreTrainedModelForConditionalGeneration",
+ "Qwen3OmniMoeTalkerModel",
+ "Qwen3OmniMoeThinkerTextPreTrainedModel",
+ "Qwen3OmniMoeCode2Wav",
+ "Qwen3OmniMoeCode2WavDecoderBlock",
+ "Qwen3OmniMoeCode2WavTransformerModel",
+ "Qwen3OmniMoeTalkerCodePredictorModel",
+ "Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration",
+]
diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
new file mode 100644
index 000000000000..14a8c3ac1248
--- /dev/null
+++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
@@ -0,0 +1,2787 @@
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group 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 Qwen3Omni model (Audio, Image, Video)."""
+
+import math
+import re
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ...activations import ACT2FN
+from ...audio_utils import AudioInput
+from ...cache_utils import Cache, DynamicCache
+from ...configuration_utils import PretrainedConfig
+from ...feature_extraction_utils import BatchFeature
+from ...generation import GenerationMixin
+from ...image_utils import ImageInput
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+)
+from ...processing_utils import ProcessorMixin, Unpack
+from ...tokenization_utils_base import TextInput
+from ...utils import auto_docstring, can_return_tuple, logging
+from ...utils.generic import OutputRecorder, TransformersKwargs, check_model_inputs
+from ...video_utils import VideoInput, make_batched_videos
+from ..mimi.modeling_mimi import MimiLayerScale
+from ..qwen2_5_omni.configuration_qwen2_5_omni import (
+ Qwen2_5OmniAudioEncoderConfig,
+ Qwen2_5OmniThinkerConfig,
+)
+from ..qwen2_5_omni.modeling_qwen2_5_omni import (
+ Qwen2_5OmniAudioAttention,
+ Qwen2_5OmniAudioEncoder,
+ Qwen2_5OmniPreTrainedModel,
+ Qwen2_5OmniPreTrainedModelForConditionalGeneration,
+ Qwen2_5OmniThinkerForConditionalGeneration,
+ SnakeBeta,
+)
+from ..qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessor, Qwen2_5OmniProcessorKwargs
+from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeSparseMoeBlock
+from ..qwen3.configuration_qwen3 import Qwen3Config
+from ..qwen3.modeling_qwen3 import (
+ Qwen3Attention,
+ Qwen3DecoderLayer,
+ Qwen3ForCausalLM,
+ Qwen3MLP,
+ Qwen3Model,
+ Qwen3RMSNorm,
+ Qwen3RotaryEmbedding,
+)
+from ..qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+from ..qwen3_moe.modeling_qwen3_moe import (
+ Qwen3MoeAttention,
+ Qwen3MoeDecoderLayer,
+ Qwen3MoeForCausalLM,
+ Qwen3MoeMLP,
+ Qwen3MoePreTrainedModel,
+ Qwen3MoeSparseMoeBlock,
+ load_balancing_loss_func,
+)
+from ..qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeVisionConfig
+from ..qwen3_vl_moe.modeling_qwen3_vl_moe import (
+ Qwen3VLMoeTextModel,
+ Qwen3VLMoeTextRotaryEmbedding,
+ Qwen3VLMoeVisionAttention,
+ Qwen3VLMoeVisionModel,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_feat_extract_output_lengths(input_lengths):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+
+ input_lengths_leave = input_lengths % 100
+ feat_lengths = (input_lengths_leave - 1) // 2 + 1
+ output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
+ return output_lengths
+
+
+class Qwen3OmniMoeAudioEncoderConfig(Qwen2_5OmniAudioEncoderConfig):
+ def __init__(
+ self,
+ num_mel_bins=128,
+ encoder_layers=32,
+ encoder_attention_heads=20,
+ encoder_ffn_dim=5120,
+ d_model=1280,
+ dropout=0,
+ attention_dropout=0,
+ activation_function="gelu",
+ activation_dropout=0,
+ scale_embedding=False,
+ initializer_range=0.02,
+ max_source_positions=1500,
+ n_window=100,
+ output_dim=3584,
+ n_window_infer=400,
+ conv_chunksize=500,
+ downsample_hidden_size=480,
+ **kwargs,
+ ):
+ super().__init__(
+ num_mel_bins,
+ encoder_layers,
+ encoder_attention_heads,
+ encoder_ffn_dim,
+ d_model,
+ dropout,
+ attention_dropout,
+ activation_function,
+ activation_dropout,
+ scale_embedding,
+ initializer_range,
+ max_source_positions,
+ n_window,
+ output_dim,
+ **kwargs,
+ )
+ self.n_window_infer = n_window_infer
+ self.conv_chunksize = conv_chunksize
+ self.downsample_hidden_size = downsample_hidden_size
+
+
+class Qwen3OmniMoeVisionEncoderConfig(Qwen3VLMoeVisionConfig):
+ pass
+
+
+class Qwen3OmniMoeTextConfig(Qwen3MoeConfig):
+ def __init__(
+ self,
+ vocab_size=3584,
+ hidden_size=2048,
+ intermediate_size=18944,
+ num_hidden_layers=28,
+ num_attention_heads=28,
+ num_key_value_heads=4,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=1000000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ attention_dropout=0,
+ decoder_sparse_step=1,
+ moe_intermediate_size=768,
+ num_experts_per_tok=8,
+ num_experts=128,
+ norm_topk_prob=True,
+ output_router_logits=False,
+ router_aux_loss_coef=0.001,
+ mlp_only_layers=None,
+ **kwargs,
+ ):
+ super().__init__(
+ vocab_size,
+ hidden_size,
+ intermediate_size,
+ num_hidden_layers,
+ num_attention_heads,
+ num_key_value_heads,
+ hidden_act,
+ max_position_embeddings,
+ initializer_range,
+ rms_norm_eps,
+ use_cache,
+ tie_word_embeddings,
+ rope_theta,
+ rope_scaling,
+ attention_bias,
+ False,
+ sliding_window,
+ attention_dropout,
+ decoder_sparse_step,
+ moe_intermediate_size,
+ num_experts_per_tok,
+ num_experts,
+ norm_topk_prob,
+ output_router_logits,
+ router_aux_loss_coef,
+ mlp_only_layers,
+ **kwargs,
+ )
+ del self.use_sliding_window
+ self.sliding_window = sliding_window
+
+
+class Qwen3OmniMoeThinkerConfig(Qwen2_5OmniThinkerConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeThinker`]. It is used to instantiate a
+ Qwen3-Omni-Thinker model according to the specified arguments, defining the model architecture. Instantiating a
+ configuration with the defaults will yield a similar configuration to that of the thinker component of the Qwen3-Omni
+ architecture.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ audio_config (`dict`, *optional*):
+ The config dictionary of the audio backbone.
+ vision_config (`dict`, *optional*):
+ The config dictionary of the vision backbone.
+ text_config (`dict`, *optional*):
+ The config dictionary of the text backbone.
+ audio_token_id (`int`, *optional*, defaults to 151646):
+ The audio token id to encode the audio prompt.
+ image_token_id (`int`, *optional*, defaults to 151655):
+ The image token id to encode the image prompt.
+ video_token_id (`int`, *optional*, defaults to 151656):
+ The video token id to encode the video prompt.
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ The increment of position id per second.
+ audio_start_token_id (`int`, *optional*, defaults to 151647):
+ The audio start token id to encode the audio prompt.
+ user_token_id (`int`, *optional*, defaults to 872):
+ The user token id to encode the user token.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeThinkerModel, Qwen3OmniMoeThinkerConfig
+
+ >>> # Initializing a default Qwen3OmniMoeThinkerConfig
+ >>> configuration = Qwen3OmniMoeThinkerConfig()
+
+ >>> # Initializing a model (with random weights) from the default configuration
+ >>> model = Qwen3OmniMoeThinkerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe_thinker"
+ # Override parent's attribute_map as we use audio_token_id directly, not audio_token_index
+ attribute_map = {}
+
+ def __init__(
+ self,
+ audio_config=None,
+ vision_config=None,
+ text_config=None,
+ audio_token_id=151646,
+ image_token_id=151655,
+ video_token_id=151656,
+ position_id_per_seconds=25,
+ audio_start_token_id=151647,
+ user_token_id=872,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(
+ audio_config,
+ vision_config,
+ text_config,
+ None,
+ None,
+ None,
+ position_id_per_seconds,
+ None,
+ audio_start_token_id,
+ None,
+ user_token_id,
+ initializer_range,
+ **kwargs,
+ )
+ del self.seconds_per_chunk
+ del self.audio_token_index
+ del self.image_token_index
+ del self.video_token_index
+ del self.audio_end_token_id
+ self.audio_token_id = audio_token_id
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
+
+
+class Qwen3OmniMoeTalkerCodePredictorConfig(Qwen3Config):
+ def __init__(
+ self,
+ vocab_size=2048,
+ hidden_size=1024,
+ intermediate_size=3072,
+ num_hidden_layers=5,
+ num_attention_heads=16,
+ num_key_value_heads=8,
+ head_dim=128,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=0.000001,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=10000,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ layer_types=None,
+ attention_dropout=0,
+ num_code_groups=32,
+ **kwargs,
+ ):
+ super().__init__(
+ vocab_size,
+ hidden_size,
+ intermediate_size,
+ num_hidden_layers,
+ num_attention_heads,
+ num_key_value_heads,
+ head_dim,
+ hidden_act,
+ max_position_embeddings,
+ initializer_range,
+ rms_norm_eps,
+ use_cache,
+ tie_word_embeddings,
+ rope_theta,
+ rope_scaling,
+ attention_bias,
+ False,
+ sliding_window,
+ None,
+ layer_types,
+ attention_dropout,
+ **kwargs,
+ )
+ del self.use_sliding_window
+ del self.max_window_layers
+ self.sliding_window = sliding_window
+ self.num_code_groups = num_code_groups
+
+
+class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig):
+ def __init__(
+ self,
+ vocab_size=3072,
+ hidden_size=1024,
+ intermediate_size=2048,
+ num_hidden_layers=20,
+ num_attention_heads=16,
+ num_key_value_heads=2,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=0.000001,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=10000,
+ rope_scaling=None,
+ attention_bias=False,
+ sliding_window=None,
+ attention_dropout=0,
+ decoder_sparse_step=1,
+ moe_intermediate_size=384,
+ num_experts_per_tok=8,
+ num_experts=128,
+ norm_topk_prob=False,
+ output_router_logits=False,
+ router_aux_loss_coef=0.001,
+ mlp_only_layers=None,
+ **kwargs,
+ ):
+ super().__init__(
+ vocab_size,
+ hidden_size,
+ intermediate_size,
+ num_hidden_layers,
+ num_attention_heads,
+ num_key_value_heads,
+ hidden_act,
+ max_position_embeddings,
+ initializer_range,
+ rms_norm_eps,
+ use_cache,
+ tie_word_embeddings,
+ rope_theta,
+ rope_scaling,
+ attention_bias,
+ False,
+ sliding_window,
+ attention_dropout,
+ decoder_sparse_step,
+ moe_intermediate_size,
+ num_experts_per_tok,
+ num_experts,
+ norm_topk_prob,
+ output_router_logits,
+ router_aux_loss_coef,
+ mlp_only_layers,
+ **kwargs,
+ )
+ del self.use_sliding_window
+ self.sliding_window = sliding_window
+
+
+class Qwen3OmniMoeTalkerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeTalker`]. It is used to instantiate a
+ Qwen3-Omni multi-modal talker model capable of handling text, audio, and vision modalities in a unified architecture.
+ The model integrates a text decoder with a code predictor for autoregressive generation of both semantic and acoustic
+ tokens, enabling speech and multimodal content generation. This configuration wraps sub-configurations for the text and
+ code predictor components, allowing modular setup and initialization.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ code_predictor_config (`dict`, *optional*):
+ A dictionary of configuration parameters used to initialize a [`Qwen3OmniMoeTalkerCodePredictorConfig`].
+ If not provided, defaults will be used.
+ text_config (`dict`, *optional*):
+ A dictionary of configuration parameters used to initialize a [`Qwen3OmniMoeTalkerTextConfig`].
+ If not provided, defaults will be used.
+ num_code_groups (`int`, *optional*, defaults to 32):
+ Number of codebook groups used in the predicted acoustic token sequence, corresponding to multi-codebook VQ representation.
+ thinker_hidden_size (`int`, *optional*, defaults to 2048):
+ Hidden dimension size of the thinker module used for intermediate reasoning or latent planning before audio generation.
+ codec_eos_token_id (`int`, *optional*, defaults to 4198):
+ Token ID representing the end-of-speech token in the codec-generated sequence.
+ accept_hidden_layer (`int`, *optional*, defaults to 18):
+ Index of the hidden layer whose output is used for accepting or refining generated tokens during think-and-speak process.
+ codec_nothink_id (`int`, *optional*, defaults to 4203):
+ Token ID indicating no thinking step is required during generation.
+ codec_think_bos_id (`int`, *optional*, defaults to 4204):
+ Token ID marking the beginning of a thinking sequence.
+ codec_think_eos_id (`int`, *optional*, defaults to 4205):
+ Token ID marking the end of a thinking sequence.
+ codec_pad_id (`int`, *optional*, defaults to 4196):
+ Padding token ID used in codec input sequences.
+ codec_bos_id (`int`, *optional*, defaults to 4197):
+ Beginning-of-speech token ID in codec sequences.
+ audio_token_id (`int`, *optional*, defaults to 151646):
+ Special token ID used to indicate the position of audio tokens in the input sequence.
+ image_token_id (`int`, *optional*, defaults to 151655):
+ Special token ID used to represent image inputs in the multimodal context.
+ video_token_id (`int`, *optional*, defaults to 151656):
+ Special token ID used to represent video inputs.
+ vision_start_token_id (`int`, *optional*, defaults to 151652):
+ Token ID indicating the start of a visual input sequence (e.g., image or video embeddings).
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ Number of position IDs allocated per second of audio content, used for temporal alignment in generation.
+ audio_start_token_id (`int`, *optional*, defaults to 151669):
+ Token ID that indicates the start of an audio generation segment in the output.
+ speaker_id (`dict`, *optional*):
+ Speaker name to speaker id dict.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeTalkerConfig, Qwen3OmniMoeTalker
+
+ >>> # Initialize a Qwen3OmniMoeTalkerConfig with default sub-configurations
+ >>> config = Qwen3OmniMoeTalkerConfig(
+ ... num_code_groups=32,
+ ... thinker_hidden_size=2048,
+ ... )
+
+ >>> # Initialize the full Qwen3-Omni Talker model
+ >>> model = Qwen3OmniMoeTalker(config)
+
+ >>> # Access the model configuration
+ >>> config = model.config
+ >>> print(config.text_config) # Access text decoder configuration
+ >>> print(config.code_predictor_config) # Access code predictor configuration
+ ```"""
+
+ sub_configs = {
+ "code_predictor_config": Qwen3OmniMoeTalkerCodePredictorConfig,
+ "text_config": Qwen3OmniMoeTalkerTextConfig,
+ }
+
+ def __init__(
+ self,
+ code_predictor_config=None,
+ text_config=None,
+ num_code_groups=32,
+ thinker_hidden_size=2048,
+ codec_eos_token_id=4198,
+ accept_hidden_layer=18,
+ codec_nothink_id=4203,
+ codec_think_bos_id=4204,
+ codec_think_eos_id=4205,
+ codec_pad_id=4196,
+ codec_bos_id=4197,
+ audio_token_id=151646,
+ image_token_id=151655,
+ video_token_id=151656,
+ vision_start_token_id=151652,
+ position_id_per_seconds=25,
+ audio_start_token_id=151669,
+ speaker_id=None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ if code_predictor_config is None:
+ code_predictor_config = {}
+ self.code_predictor_config = Qwen3OmniMoeTalkerCodePredictorConfig()
+ logger.info("code_predictor_config is None. Initializing code_predictor_config model with default values")
+ elif isinstance(code_predictor_config, Qwen3OmniMoeTalkerCodePredictorConfig):
+ self.code_predictor_config = code_predictor_config
+ else:
+ self.code_predictor_config = Qwen3OmniMoeTalkerCodePredictorConfig(**code_predictor_config)
+
+ if text_config is None:
+ text_config = {}
+ self.text_config = Qwen3OmniMoeTalkerTextConfig()
+ logger.info("talker text_config is None. Initializing talker text model with default values")
+ elif isinstance(text_config, Qwen3OmniMoeTalkerTextConfig):
+ self.text_config = text_config
+ else:
+ self.text_config = Qwen3OmniMoeTalkerTextConfig(**text_config)
+ self.num_code_groups = num_code_groups
+ self.thinker_hidden_size = thinker_hidden_size
+ self.codec_eos_token_id = codec_eos_token_id
+ self.accept_hidden_layer = accept_hidden_layer
+ self.codec_nothink_id = codec_nothink_id
+ self.codec_think_bos_id = codec_think_bos_id
+ self.codec_think_eos_id = codec_think_eos_id
+ self.codec_pad_id = codec_pad_id
+ self.codec_bos_id = codec_bos_id
+ self.audio_token_id = audio_token_id
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
+ self.position_id_per_seconds = position_id_per_seconds
+ self.audio_start_token_id = audio_start_token_id
+ self.vision_start_token_id = vision_start_token_id
+ self.speaker_id = speaker_id
+
+
+class Qwen3OmniMoeCode2WavConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeCode2WavConfig`]. It is used to instantiate a
+ Qwen3-Omni code-to-waveform decoder, responsible for converting discrete audio codes into high-fidelity waveforms.
+ The configuration defines the architecture of the decoder, including parameters for vector quantization, autoregressive modeling,
+ and upsampling layers.
+
+ e.g. [Qwen/Qwen3-Omni-7B](https://huggingface.co/Qwen/Qwen3-Omni-7B)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ codebook_size (`int`, *optional*, defaults to 2048):
+ Number of entries in each residual codebook used for acoustic token quantization.
+ hidden_size (`int`, *optional*, defaults to 1024):
+ Dimensionality of the hidden states and embeddings in the autoregressive transformer decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 8000):
+ Maximum sequence length that the autoregressive decoder can handle. Determines positional embedding size.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period for rotary position embeddings (RoPE) applied to attention layers.
+ num_attention_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads for each attention layer in the decoder.
+ num_key_value_heads (`int`, *optional*, defaults to 16):
+ Number of key and value attention heads used in grouped-query attention (if applicable).
+ attention_bias (`bool`, *optional*, defaults to `False`):
+ Whether to use bias in the attention projection layers.
+ sliding_window (`int`, *optional*, defaults to 72):
+ Window size for local attention mechanism, limiting attention context to improve efficiency.
+ intermediate_size (`int`, *optional*, defaults to 3072):
+ Dimensionality of the feed-forward (intermediate) layer in each transformer block.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function used in the feed-forward layers. Supports `"silu"`, `"relu"`, `"gelu"`, etc.
+ layer_scale_initial_scale (`float`, *optional*, defaults to 0.01):
+ Initial value for LayerScale applied in transformer blocks, helping stabilize training.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-5):
+ Epsilon value for RMS normalization layers to prevent division by zero.
+ num_hidden_layers (`int`, *optional*, defaults to 8):
+ Number of transformer blocks in the autoregressive decoder.
+ num_quantizers (`int`, *optional*, defaults to 16):
+ Number of residual vector quantizers used in the vocoder for fine-grained audio reconstruction.
+ upsample_rates (`Tuple[int]`, *optional*, defaults to `(8, 5, 4, 3)`):
+ Rate at which features are upsampled in the final waveform synthesis stage.
+ upsampling_ratios (`Tuple[int]`, *optional*, defaults to `(2, 2)`):
+ Ratios used in transposed convolutional layers to progressively upsample feature maps to waveform.
+ decoder_dim (`int`, *optional*, defaults to 1536):
+ Final dimensionality of the decoder's output before waveform generation.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ Dropout probability applied to attention weights in the decoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3OmniMoeCode2WavConfig, Qwen3OmniMoeCode2WavModel
+
+ >>> # Initializing a default Qwen3OmniMoeCode2WavConfig
+ >>> config = Qwen3OmniMoeCode2WavConfig()
+
+ >>> # Initializing the Code2Wav model with the configuration
+ >>> model = Qwen3OmniMoeCode2WavModel(config)
+
+ >>> # Accessing configuration
+ >>> config = model.config
+ ```"""
+
+ def __init__(
+ self,
+ codebook_size=2048,
+ hidden_size=1024,
+ max_position_embeddings=8000,
+ rope_theta=10000,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ attention_bias=False,
+ sliding_window=72,
+ intermediate_size=3072,
+ hidden_act="silu",
+ layer_scale_initial_scale=0.01,
+ rms_norm_eps=1e-5,
+ num_hidden_layers=8,
+ num_quantizers=16,
+ upsample_rates=(8, 5, 4, 3),
+ upsampling_ratios=(2, 2),
+ decoder_dim=1536,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.codebook_size = codebook_size
+ self.hidden_size = hidden_size
+ self.max_position_embeddings = max_position_embeddings
+ self.rope_theta = rope_theta
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.attention_bias = attention_bias
+ self.sliding_window = sliding_window
+ self.intermediate_size = intermediate_size
+ self.hidden_act = hidden_act
+ self.layer_scale_initial_scale = layer_scale_initial_scale
+ self.rms_norm_eps = rms_norm_eps
+ self.num_hidden_layers = num_hidden_layers
+ self.num_quantizers = num_quantizers
+ self.upsample_rates = upsample_rates
+ self.upsampling_ratios = upsampling_ratios
+ self.decoder_dim = decoder_dim
+ self.attention_dropout = attention_dropout
+
+ @property
+ def layer_types(self):
+ """
+ All layer in code2wav should be sliding attention
+ """
+ return ["sliding_attention"] * self.num_hidden_layers
+
+
+class Qwen3OmniMoeConfig(PretrainedConfig):
+ """
+ This is the configuration class to store the configuration of a [`Qwen3OmniMoeForConditionalGeneration`]. It is used to instantiate a Qwen3Omni
+ model according to the specified sub-models configurations, defining the model architecture.
+
+ Instantiating a configuration with the defaults will yield a similar configuration to that of the
+ [Qwen/Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ thinker_config (`dict`, *optional*): Configuration of the underlying thinker sub-model.
+ talker_config (`dict`, *optional*): Configuration of the underlying talker sub-model.
+ code2wav_config (`dict`, *optional*): Configuration of the underlying code2wav sub-model.
+ enable_audio_output (`bool`, *optional*, defaults to `True`): Whether enable audio output and load talker and code2wav module.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... Qwen3OmniMoeThinkerConfig,
+ ... Qwen3OmniMoeTalkerConfig,
+ ... Qwen3OmniMoeCode2WavConfig,
+ ... Qwen3OmniMoeForConditionalGeneration,
+ ... Qwen3OmniMoeConfig,
+ ... )
+
+ >>> # Initializing a Qwen3OmniMoe style configuration
+ >>> configuration = Qwen3OmniMoeConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = Qwen3OmniMoeForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_omni_moe"
+ sub_configs = {
+ "thinker_config": Qwen3OmniMoeThinkerConfig,
+ "talker_config": Qwen3OmniMoeTalkerConfig,
+ "code2wav_config": Qwen3OmniMoeCode2WavConfig,
+ }
+
+ def __init__(
+ self,
+ thinker_config=None,
+ talker_config=None,
+ code2wav_config=None,
+ enable_audio_output=True,
+ im_start_token_id=151644,
+ im_end_token_id=151645,
+ tts_pad_token_id=151671,
+ tts_bos_token_id=151672,
+ tts_eos_token_id=151673,
+ system_token_id=8948,
+ user_token_id=872,
+ assistant_token_id=77091,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ if thinker_config is None:
+ thinker_config = {}
+ logger.info("thinker_config is None. Initializing thinker model with default values")
+
+ if talker_config is None:
+ talker_config = {}
+ logger.info("talker_config is None. Initializing talker model with default values")
+
+ if code2wav_config is None:
+ code2wav_config = {}
+ logger.info("code2wav_config is None. Initializing code2wav model with default values")
+
+ self.thinker_config = Qwen3OmniMoeThinkerConfig(**thinker_config)
+ self.talker_config = Qwen3OmniMoeTalkerConfig(**talker_config)
+ self.code2wav_config = Qwen3OmniMoeCode2WavConfig(**code2wav_config)
+ self.enable_audio_output = enable_audio_output
+ self.im_start_token_id = im_start_token_id
+ self.im_end_token_id = im_end_token_id
+ self.tts_pad_token_id = tts_pad_token_id
+ self.tts_bos_token_id = tts_bos_token_id
+ self.tts_eos_token_id = tts_eos_token_id
+ self.system_token_id = system_token_id
+ self.user_token_id = user_token_id
+ self.assistant_token_id = assistant_token_id
+
+ def get_text_config(self, decoder=False) -> "PretrainedConfig":
+ """
+ Returns the config that is meant to be used with text IO. On most models, it is the original config instance
+ itself. On specific composite models, it is under a set of valid names.
+
+ Args:
+ decoder (`Optional[bool]`, *optional*, defaults to `False`):
+ If set to `True`, then only search for decoder config names.
+ """
+ # Overridden for deeply nested config like Qwen2-Omni. We don't have any omni model
+ # except for Qwen yet. This has to be generalized if more deeply nested configs are
+ # added. NOTE: currently method used only by vLLM
+ return self.thinker_config.get_text_config()
+
+
+class Qwen3OmniMoePreTrainedModel(Qwen2_5OmniPreTrainedModel):
+ pass
+
+
+class Qwen3OmniMoePreTrainedModelForConditionalGeneration(Qwen2_5OmniPreTrainedModelForConditionalGeneration):
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[torch.Tensor],
+ grid_hs: list[torch.Tensor],
+ grid_ws: list[torch.Tensor],
+ ):
+ llm_pos_ids_list = []
+ llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+ llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float()
+ t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float()
+ _llm_pos_ids = torch.stack([t_index, h_index, w_index])
+ llm_pos_ids_list.append(_llm_pos_ids + start_idx)
+ llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+ return llm_pos_ids
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: Optional[torch.LongTensor] = None,
+ second_per_grids: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_seqlens (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ second_per_grids (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ audio_token_id = self.config.audio_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ audio_start_token_id = self.config.audio_start_token_id
+ position_id_per_seconds = self.config.position_id_per_seconds
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is not None:
+ attention_mask = attention_mask == 1
+ position_ids = torch.zeros(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=torch.float,
+ device=input_ids.device,
+ )
+ image_idx, video_idx, audio_idx = 0, 0, 0
+ for i, input_ids in enumerate(total_input_ids):
+ if attention_mask is not None:
+ input_ids = input_ids[attention_mask[i]]
+ image_nums, video_nums, audio_nums = 0, 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ audio_nums = torch.sum(input_ids == audio_start_token_id)
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (
+ (vision_tokens == audio_start_token_id).sum()
+ if use_audio_in_video
+ else (vision_tokens == video_token_id).sum()
+ )
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums
+ multimodal_nums = (
+ image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums
+ )
+ for _ in range(multimodal_nums):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ if (image_token_id in input_tokens or video_token_id in input_tokens) and (
+ remain_videos > 0 or remain_images > 0
+ ):
+ ed_vision_start = input_tokens.index(vision_start_token_id, st)
+ else:
+ ed_vision_start = len(input_tokens) + 1
+ if audio_token_id in input_tokens and remain_audios > 0:
+ ed_audio_start = input_tokens.index(audio_start_token_id, st)
+ else:
+ ed_audio_start = len(input_tokens) + 1
+ min_ed = min(ed_vision_start, ed_audio_start)
+
+ text_len = min_ed - st
+ if text_len != 0:
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+ st_idx += text_len
+ # Audio in Video
+ if min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start:
+ bos_len, eos_len = 2, 2
+ else:
+ bos_len, eos_len = 1, 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+ st_idx += bos_len
+ # Audio Only
+ if min_ed == ed_audio_start:
+ audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
+ llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + audio_len + eos_len)
+ audio_idx += 1
+ remain_audios -= 1
+
+ # Image Only
+ elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id:
+ grid_t = image_grid_thw[image_idx][0]
+ grid_hs = image_grid_thw[:, 1]
+ grid_ws = image_grid_thw[:, 2]
+ t_index = (torch.arange(grid_t) * 1 * position_id_per_seconds).float()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + image_len + eos_len)
+ image_idx += 1
+ remain_images -= 1
+
+ # Video Only
+ elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == video_token_id:
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).float()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st += int(text_len + bos_len + video_len + eos_len)
+ video_idx += 1
+ remain_videos -= 1
+
+ # Audio in Video
+ elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start:
+ audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
+ audio_llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).float()
+ video_llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_data_index, audio_data_index = 0, 0
+ while (
+ video_data_index < video_llm_pos_ids.shape[-1]
+ and audio_data_index < audio_llm_pos_ids.shape[-1]
+ ):
+ if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]:
+ llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1])
+ video_data_index += 1
+ else:
+ llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1])
+ audio_data_index += 1
+ if video_data_index < video_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(
+ video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]]
+ )
+ if audio_data_index < audio_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(
+ audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]]
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+
+ st += int(text_len + bos_len + audio_len + video_len + eos_len)
+
+ audio_idx += 1
+ video_idx += 1
+ remain_videos -= 1
+ remain_audios -= 1
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat([item.float() for item in llm_pos_ids_list], dim=1).reshape(3, -1)
+
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+
+ return position_ids, mrope_position_deltas
+ else:
+ position_ids = attention_mask.float().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
+
+ return position_ids, mrope_position_deltas
+
+
+class Qwen3OmniMoeAudioAttention(Qwen2_5OmniAudioAttention):
+ def __init__(self, config):
+ super().__init__(config)
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+
+class Qwen3OmniMoeAudioEncoder(Qwen2_5OmniAudioEncoder):
+ def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
+ super().__init__(config)
+ del self.proj
+ del self.avg_pooler
+ del self.audio_bos_eos_token
+ del self.conv1
+ del self.conv2
+ self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv2d2 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv2d3 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
+ self.conv_out = nn.Linear(
+ config.downsample_hidden_size * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2),
+ config.d_model,
+ bias=False,
+ )
+ self.proj1 = nn.Linear(config.d_model, config.d_model)
+ self.act = ACT2FN[config.activation_function]
+ self.proj2 = nn.Linear(config.d_model, config.output_dim)
+ self.n_window_infer = self.config.n_window_infer
+ self.conv_chunksize = self.config.conv_chunksize
+
+ def forward(
+ self,
+ input_features,
+ feature_lens=None,
+ aftercnn_lens=None,
+ ):
+ aftercnn_lens = _get_feat_extract_output_lengths(feature_lens)
+ chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
+
+ chunk_lengths = torch.tensor(
+ [self.n_window * 2] * chunk_num.sum(),
+ dtype=torch.long,
+ device=feature_lens.device,
+ )
+ tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+ chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
+ chunk_lengths[chunk_lengths == 0] = self.n_window * 2
+
+ chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
+ padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
+ feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths)
+ padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
+ [torch.ones(length, dtype=torch.bool, device=padded_feature.device) for length in feature_lens_after_cnn],
+ batch_first=True,
+ )
+ padded_feature = padded_feature.unsqueeze(1)
+ # Split to chunk to avoid OOM during convolution
+ padded_embeds = []
+ for chunk in padded_feature.split(self.conv_chunksize, dim=0):
+ padded_embed = F.gelu(self.conv2d1(chunk))
+ padded_embed = F.gelu(self.conv2d2(padded_embed))
+ padded_embed = F.gelu(self.conv2d3(padded_embed))
+ padded_embeds.append(padded_embed)
+ padded_embed = torch.cat(padded_embeds, dim=0)
+ b, c, f, t = padded_embed.size()
+ padded_embed = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f))
+
+ positional_embedding = (
+ self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
+ .unsqueeze(0)
+ .to(padded_embed.dtype)
+ )
+ padded_embed = padded_embed + positional_embedding
+ hidden_states = padded_embed[padded_mask_after_cnn]
+ cu_chunk_lens = [0]
+ window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2))
+ for cnn_len in aftercnn_lens:
+ cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn)
+ remainder = cnn_len % window_aftercnn
+ if remainder != 0:
+ cu_chunk_lens += [remainder]
+ cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32)
+
+ for encoder_layer in self.layers:
+ layer_outputs = encoder_layer(
+ hidden_states,
+ cu_seqlens,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.ln_post(hidden_states)
+ hidden_states = self.proj1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.proj2(hidden_states)
+ return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class Qwen3OmniMoeVisionAttention(Qwen3VLMoeVisionAttention):
+ def __init__(self, config: Qwen3OmniMoeVisionEncoderConfig):
+ super().__init__(config)
+
+
+class Qwen3OmniMoeVisionPatchMerger(nn.Module):
+ def __init__(self, config: Qwen3OmniMoeVisionEncoderConfig, use_postshuffle_norm=False) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.ln_q = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+ self.mlp = nn.ModuleList(
+ [
+ nn.Linear(self.hidden_size, self.hidden_size),
+ nn.GELU(),
+ nn.Linear(self.hidden_size, config.out_hidden_size),
+ ]
+ )
+
+ def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ hidden = self.ln_q(hidden.view(-1, self.hidden_size) if self.use_postshuffle_norm else hidden).view(
+ -1, self.hidden_size
+ )
+ for layer in self.mlp:
+ hidden = layer(hidden)
+ return hidden
+
+
+class Qwen3OmniMoeVisionEncoder(Qwen3VLMoeVisionModel):
+ config: Qwen3OmniMoeVisionEncoderConfig
+ _no_split_modules = ["Qwen3OmniMoeVisionBlock"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ self.merger_list = nn.ModuleList(
+ [
+ Qwen3OmniMoeVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=True,
+ )
+ for _ in range(len(config.deepstack_visual_indexes))
+ ]
+ )
+ super().__init__(config, *inputs, **kwargs)
+ del self.deepstack_merger_list
+
+ @property
+ def deepstack_merger_list(self):
+ return self.merger_list
+
+
+class Qwen3OmniMoeThinkerTextRotaryEmbedding(Qwen3VLMoeTextRotaryEmbedding):
+ pass
+
+
+class Qwen3OmniMoeThinkerTextSparseMoeBlock(Qwen3MoeSparseMoeBlock):
+ pass
+
+
+class Qwen3OmniMoeThinkerTextAttention(Qwen3MoeAttention):
+ def __init__(self, config, layer_idx):
+ super().__init__(config, layer_idx)
+ self.sliding_window = None
+
+
+class Qwen3OmniMoeThinkerTextDecoderLayer(Qwen3MoeDecoderLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__(config, layer_idx)
+ self.self_attn = Qwen3OmniMoeThinkerTextAttention(config, layer_idx)
+
+
+class Qwen3OmniMoeThinkerTextPreTrainedModel(Qwen3MoePreTrainedModel):
+ config_class = Qwen3OmniMoeTextConfig
+ config = Qwen3OmniMoeTextConfig
+
+
+class Qwen3OmniMoeThinkerTextModel(Qwen3VLMoeTextModel):
+ config_class = Qwen3OmniMoeTextConfig
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeThinkerTextDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeThinkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTextConfig):
+ super().__init__(config)
+ self.layers = nn.ModuleList(
+ [Qwen3OmniMoeThinkerTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = Qwen3OmniMoeThinkerTextRotaryEmbedding(config)
+
+ def _deepstack_process(self, hidden_states, visual_pos_masks, visual_embeds):
+ visual_pos_masks = visual_pos_masks[..., 0]
+ return super()._deepstack_process(hidden_states, visual_pos_masks, visual_embeds)
+
+
+@dataclass
+class Qwen3OmniMoeThinkerCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
+ r"""
+ Args:
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ rope_deltas: Optional[torch.LongTensor] = None
+
+
+class Qwen3OmniMoeThinkerForConditionalGeneration(Qwen2_5OmniThinkerForConditionalGeneration):
+ _no_split_modules = [
+ "Qwen3OmniMoeAudioEncoderLayer",
+ "Qwen3OmniMoeThinkerTextDecoderLayer",
+ ]
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeThinkerTextDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeThinkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_experts = config.text_config.num_experts
+ self.num_experts_per_tok = config.text_config.num_experts_per_tok
+
+ def get_audio_features(
+ self,
+ input_features: torch.FloatTensor,
+ feature_attention_mask: Optional[torch.LongTensor] = None,
+ audio_feature_lengths: Optional[torch.LongTensor] = None,
+ ):
+ """
+ Encodes audios into continuous embeddings that can be forwarded to the language model.
+
+ Args:
+ input_features (`torch.FloatTensor`):
+ The tensors corresponding to the input audios.
+ feature_attention_mask (`torch.LongTensor`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ """
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+ else:
+ audio_feature_lengths = None
+
+ feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ audio_outputs = self.audio_tower(
+ input_features,
+ feature_lens=feature_lens,
+ )
+ audio_features = audio_outputs.last_hidden_state
+
+ return audio_features
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ input_features=None,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ attention_mask=None,
+ feature_attention_mask=None,
+ audio_feature_lengths=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ rope_deltas=None,
+ labels=None,
+ use_cache=None,
+ output_router_logits: Optional[bool] = None,
+ use_audio_in_video=None,
+ cache_position=None,
+ video_second_per_grid=None,
+ **kwargs,
+ ) -> Union[tuple, Qwen3OmniMoeThinkerCausalLMOutputWithPast]:
+ output_router_logits = (
+ output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits
+ )
+
+ if inputs_embeds is None:
+ # 1. Extract the input embeddings
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ visual_embeds_multiscale = None
+ visual_pos_masks = None
+ # 2. Merge text , audios , image and video
+ if input_features is not None:
+ audio_features = self.get_audio_features(
+ input_features,
+ feature_attention_mask=feature_attention_mask,
+ audio_feature_lengths=audio_feature_lengths,
+ )
+ audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
+
+ if pixel_values is not None:
+ image_embeds, image_embeds_multiscale = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ visual_pos_masks = image_mask
+ visual_embeds_multiscale = image_embeds_multiscale
+
+ if pixel_values_videos is not None:
+ video_embeds, video_embeds_multiscale = self.get_video_features(pixel_values_videos, video_grid_thw)
+
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if visual_embeds_multiscale is None:
+ visual_embeds_multiscale = video_embeds_multiscale
+ visual_pos_masks = video_mask
+ else:
+ visual_pos_masks = video_mask | image_mask
+ visual_embeds_multiscale_joint = ()
+ image_mask_joint = image_mask[visual_pos_masks]
+ video_mask_joint = video_mask[visual_pos_masks]
+ for img_embed, vid_embed in zip(visual_embeds_multiscale, video_embeds_multiscale):
+ embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1])
+ embed_joint[image_mask_joint, :] = img_embed
+ embed_joint[video_mask_joint, :] = vid_embed
+ visual_embeds_multiscale_joint = visual_embeds_multiscale_joint + (embed_joint,)
+ visual_embeds_multiscale = visual_embeds_multiscale_joint
+
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ else:
+ audio_feature_lengths = None
+
+ if attention_mask is not None and position_ids is None:
+ if (
+ cache_position is None
+ or (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ ):
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ 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,
+ cache_position=cache_position,
+ deepstack_visual_embeds=visual_embeds_multiscale,
+ visual_pos_masks=visual_pos_masks,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.get_text_config().vocab_size
+ )
+
+ 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 Qwen3OmniMoeThinkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ aux_loss=aux_loss,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=self.rope_deltas,
+ )
+
+
+class Qwen3OmniMoeTalkerResizeMLP(nn.Module):
+ def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+ super().__init__()
+ self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.text_config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+@dataclass
+class Qwen3OmniMoeTalkerCodePredictorOutputWithPast(CausalLMOutputWithPast):
+ r"""
+ generation_steps (`int`, *optional*)
+ Current generation step of code predictor model.
+ """
+
+ generation_steps: Optional[int] = None
+
+
+class Qwen3OmniMoeTalkerCodePredictorAttention(Qwen3Attention):
+ pass
+
+
+class Qwen3OmniMoeTalkerCodePredictorDecoderLayer(Qwen3DecoderLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__(config, layer_idx)
+ self.self_attn = Qwen3OmniMoeTalkerCodePredictorAttention(config=config, layer_idx=layer_idx)
+
+
+class Qwen3OmniMoeTalkerCodePredictorModel(Qwen3Model):
+ config_class = Qwen3OmniMoeTalkerCodePredictorConfig
+ base_model_prefix = "talker.code_predictor.model"
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeTalkerCodePredictorAttention,
+ "hidden_states": Qwen3OmniMoeTalkerCodePredictorDecoderLayer,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerCodePredictorConfig):
+ super().__init__(config)
+ del self.embed_tokens
+ self.layers = nn.ModuleList(
+ [
+ Qwen3OmniMoeTalkerCodePredictorDecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ]
+ )
+ self.codec_embedding = nn.ModuleList(
+ [nn.Embedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)]
+ )
+
+ def get_input_embeddings(self):
+ return self.codec_embedding
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if input_ids is not None:
+ raise ValueError("`input_ids` is expected to be `None`")
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ # It may already have been prepared by e.g. `generate`
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
+ # Prepare mask arguments
+ mask_kwargs = {
+ "config": self.config,
+ "input_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "cache_position": cache_position,
+ "past_key_values": past_key_values,
+ "position_ids": position_ids,
+ }
+ # Create the masks
+ causal_mask_mapping = {
+ "full_attention": create_causal_mask(**mask_kwargs),
+ }
+
+ 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,
+ attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3ForCausalLM):
+ config_class = Qwen3OmniMoeTalkerCodePredictorConfig
+ base_model_prefix = "talker.code_predictor"
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeTalkerCodePredictorAttention,
+ "hidden_states": Qwen3OmniMoeTalkerCodePredictorDecoderLayer,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerCodePredictorConfig):
+ super().__init__(config)
+ self.model = Qwen3OmniMoeTalkerCodePredictorModel._from_config(config)
+ self.lm_head = nn.ModuleList(
+ [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
+ )
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ labels=None,
+ use_cache=None,
+ cache_position=None,
+ generation_steps=None,
+ **kwargs,
+ ):
+ r"""
+ Args:
+ generation_steps (`int`):
+ generation step of code predictor, 0..num_code_groups-1
+ """
+
+ # Prefill stage
+ if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
+ generation_steps = inputs_embeds.shape[1] - 2 # hidden & layer 0
+ # Generation stage
+ else:
+ inputs_embeds = self.model.get_input_embeddings()[generation_steps - 1](input_ids)
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=None,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ logits = self.lm_head[generation_steps](hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return Qwen3OmniMoeTalkerCodePredictorOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ generation_steps=generation_steps + 1,
+ )
+
+ def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+ model_kwargs["generation_steps"] = outputs.generation_steps
+ return model_kwargs
+
+
+@dataclass
+class Qwen3OmniMoeTalkerOutputWithPast(MoeCausalLMOutputWithPast):
+ r"""
+ Args:
+ generation_step (`int`, *optional*):
+ Current generation step, used to track which `trailing_text_hidden` should be used.
+ """
+
+ generation_step: Optional[int] = None
+
+
+class Qwen3OmniMoeTalkerRotaryEmbedding(Qwen3OmniMoeThinkerTextRotaryEmbedding):
+ pass
+
+
+class Qwen3OmniMoeTalkerTextMLP(Qwen3MoeMLP):
+ pass
+
+
+class Qwen3OmniMoeTalkerTextSparseMoeBlock(Qwen2MoeSparseMoeBlock):
+ pass
+
+
+class Qwen3OmniMoeTalkerDecoderLayer(Qwen3MoeDecoderLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__(config, layer_idx)
+ self.self_attn = Qwen3OmniMoeThinkerTextAttention(config, layer_idx)
+ self.mlp = Qwen3OmniMoeTalkerTextSparseMoeBlock(config)
+
+
+class Qwen3OmniMoeTalkerModel(Qwen3VLMoeTextModel):
+ config_class = Qwen3OmniMoeTalkerTextConfig
+ base_model_prefix = "talker.model"
+ _no_split_modules = ["Qwen3OmniMoeTalkerDecoderLayer"]
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeTalkerDecoderLayer,
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeTalkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerTextConfig):
+ super().__init__(config)
+ del self.embed_tokens
+ self.codec_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.layers = nn.ModuleList(
+ [Qwen3OmniMoeTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.rotary_emb = Qwen3OmniMoeTalkerRotaryEmbedding(config)
+
+ def get_input_embeddings(self):
+ return self.codec_embedding
+
+
+class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM):
+ config_class = Qwen3OmniMoeTalkerConfig
+ base_model_prefix = "talker"
+ _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"]
+ _can_record_outputs = {
+ "attentions": Qwen3OmniMoeThinkerTextAttention,
+ "router_logits": OutputRecorder(Qwen3OmniMoeTalkerTextSparseMoeBlock, index=1),
+ }
+
+ def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+ super().__init__(config)
+ del self.lm_head
+ self.model = Qwen3OmniMoeTalkerModel._from_config(config.text_config)
+ self.text_projection = Qwen3OmniMoeTalkerResizeMLP(config)
+ self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(config)
+ self.codec_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.code_predictor = Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration._from_config(
+ config=config.code_predictor_config
+ )
+ self.rope_deltas = None
+ self.spatial_merge_size = self.config.spatial_merge_size
+ self.vocab_size = config.text_config.vocab_size
+ self.router_aux_loss_coef = config.text_config.router_aux_loss_coef
+ self.num_experts = config.text_config.num_experts
+ self.num_experts_per_tok = config.text_config.num_experts_per_tok
+
+ # Should inherit from PretrainedModel, but cannot inherit multiple classes in modular
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: Optional[torch.LongTensor] = None,
+ second_per_grids: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return Qwen3OmniMoePreTrainedModelForConditionalGeneration.get_rope_index(
+ self,
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_seqlens,
+ second_per_grids,
+ )
+
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[torch.Tensor],
+ grid_hs: list[torch.Tensor],
+ grid_ws: list[torch.Tensor],
+ ):
+ return Qwen3OmniMoePreTrainedModelForConditionalGeneration.get_llm_pos_ids_for_vision(
+ self, start_idx, vision_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ use_audio_in_video=None,
+ audio_feature_lengths=None,
+ video_second_per_grid=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ labels=None,
+ use_cache=None,
+ output_router_logits=None,
+ cache_position=None,
+ residual_codes=None,
+ trailing_text_hidden=None,
+ tts_pad_embed=None,
+ generation_step=None,
+ talker_input_ids=None,
+ **kwargs,
+ ):
+ r"""
+ Args:
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ residual_codes (`torch.Tensor`):
+ The predicted residual codes of previous step.
+ trailing_text_hidden (`torch.Tensor`):
+ Text hidden states from thinker after the first token.
+ tts_pad_embed (`torch.Tensor`):
+ Embedding tensor of `tts_pad_token_id`.
+ generation_step (`int`):
+ Generation step since prefill, used to sync with `trailing_text_hidden`.
+ talker_input_ids (`torch.Tensor`):
+ Input ids from thinker, used to compute 3d RoPE.
+ """
+ # Prefill
+ if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
+ generation_step = -1
+ residual_codes = None
+ if attention_mask is not None:
+ if (
+ cache_position is None
+ or (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ ):
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ talker_input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs: MoeModelOutputWithPast = self.model(
+ input_ids=None,
+ 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,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ logits = self.codec_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ aux_loss = None
+ if output_router_logits:
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.num_experts,
+ self.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
+
+ return Qwen3OmniMoeTalkerOutputWithPast(
+ loss=loss,
+ logits=logits,
+ aux_loss=aux_loss,
+ past_key_values=outputs.past_key_values,
+ hidden_states=(
+ outputs.hidden_states,
+ residual_codes,
+ ), # TODO: hack here to take residual codes out, need refactor.
+ generation_step=generation_step + 1,
+ )
+
+ def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+ model_kwargs["hidden_states"] = outputs.hidden_states
+ model_kwargs["generation_step"] = outputs.generation_step
+ return model_kwargs
+
+ def prepare_inputs_for_generation(
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
+ ):
+ hidden_states = kwargs.pop("hidden_states", None)
+ inputs = super().prepare_inputs_for_generation(
+ input_ids, past_key_values, attention_mask, inputs_embeds, cache_position, **kwargs
+ )
+ # Decode stage
+ # TODO(raushan, gante): Refactor this part to a utility function
+ if cache_position[0] != 0:
+ input_ids = input_ids[:, -1:]
+ generation_step = kwargs.get("generation_step")
+ trailing_text_hidden = kwargs.get("trailing_text_hidden")
+ tts_pad_embed = kwargs.get("tts_pad_embed")
+ last_id_hidden = self.get_input_embeddings()(input_ids)
+
+ past_hidden = hidden_states[0][-1][:, -1:].to(last_id_hidden.device) # hidden, last layer, last token
+ predictor_result = self.code_predictor.generate(
+ inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
+ max_new_tokens=self.config.num_code_groups - 1,
+ do_sample=True,
+ top_k=50,
+ top_p=0.8,
+ output_hidden_states=True,
+ return_dict_in_generate=True,
+ )
+ residual_codes = torch.cat((input_ids, predictor_result.sequences.to(input_ids.device)), dim=-1)
+
+ mid_residual_hiddens = [hid[0].to(last_id_hidden.device) for hid in predictor_result.hidden_states[1:]]
+ last_residual_hidden = self.code_predictor.get_input_embeddings()[-1](
+ predictor_result.sequences[..., -1:]
+ ).to(last_id_hidden.device)
+ codec_hiddens = torch.cat(
+ [last_id_hidden] + mid_residual_hiddens + [last_residual_hidden],
+ dim=1,
+ )
+ inputs_embeds = codec_hiddens.sum(1, keepdim=True)
+
+ if generation_step < trailing_text_hidden.shape[1]:
+ inputs_embeds = inputs_embeds + trailing_text_hidden[:, generation_step].unsqueeze(1).to(
+ inputs_embeds.device
+ )
+ else:
+ inputs_embeds = inputs_embeds + tts_pad_embed.to(inputs_embeds.device)
+ inputs["inputs_embeds"] = inputs_embeds
+ inputs["residual_codes"] = residual_codes
+ return inputs
+
+
+class Qwen3OmniMoeCausalConvNet(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ dilation=1,
+ stride=1,
+ groups=1,
+ ):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ )
+ self.stride = stride
+ self.kernel_size = (kernel_size - 1) * dilation + 1
+ self.dilation = dilation
+ self.padding = self.kernel_size - self.stride
+
+ def _get_extra_padding_for_conv1d(self, hidden_state: torch.Tensor) -> int:
+ length = hidden_state.shape[-1]
+ n_frames = (length - self.kernel_size + self.padding) / self.stride + 1
+ ideal_length = (math.ceil(n_frames) - 1) * self.stride + (self.kernel_size - self.padding)
+ return ideal_length - length
+
+ def forward(self, hidden_state):
+ extra_padding = self._get_extra_padding_for_conv1d(hidden_state)
+ hidden_state = F.pad(hidden_state, (self.padding, extra_padding), mode="constant", value=0)
+ return self.conv(hidden_state).contiguous()
+
+
+class Qwen3OmniMoeCausalTransConvNet(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+ super().__init__()
+ self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride)
+
+ pad = kernel_size - stride
+ self.left_pad = math.ceil(pad)
+ self.right_pad = pad = self.left_pad
+
+ def forward(self, hidden_state):
+ hidden_state = self.conv(hidden_state)
+ hidden_state = hidden_state[..., self.left_pad : hidden_state.shape[-1] - self.right_pad]
+ return hidden_state.contiguous()
+
+
+class Qwen3OmniMoeConvNeXtBlock(nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.dwconv = Qwen3OmniMoeCausalConvNet(
+ dim,
+ dim,
+ kernel_size=7,
+ groups=dim,
+ dilation=1,
+ )
+ self.norm = nn.LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, 4 * dim)
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
+
+ def forward(self, hidden_states):
+ input = hidden_states
+
+ hidden_states = self.dwconv(hidden_states)
+ hidden_states = hidden_states.permute(0, 2, 1)
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.pwconv1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.pwconv2(hidden_states)
+
+ hidden_states = self.gamma * hidden_states
+
+ hidden_states = hidden_states.permute(0, 2, 1)
+
+ hidden_states = input + hidden_states
+
+ return hidden_states
+
+
+class Qwen3OmniMoeCode2WavRotatoryEmbedding(Qwen3RotaryEmbedding):
+ pass
+
+
+class Qwen3OmniMoeCode2WavAttention(Qwen3Attention):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ super().__init__(config, layer_idx)
+ self.q_norm = nn.Identity()
+ self.k_norm = nn.Identity()
+ self.sliding_window = config.sliding_window
+
+
+class Qwen3OmniMoeCode2WavMlp(Qwen3MLP):
+ pass
+
+
+class Qwen3OmniMoeCode2WavRMSNorm(Qwen3RMSNorm):
+ pass
+
+
+class Qwen3OmniMoeCode2WavLayerScale(MimiLayerScale):
+ pass
+
+
+class Qwen3OmniMoeCode2WavTransformerLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Qwen3OmniMoeCode2WavAttention(config, layer_idx)
+ self.mlp = Qwen3OmniMoeCode2WavMlp(config)
+ self.input_layernorm = Qwen3OmniMoeCode2WavRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3OmniMoeCode2WavRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.self_attn_layer_scale = Qwen3OmniMoeCode2WavLayerScale(config)
+ self.mlp_layer_scale = Qwen3OmniMoeCode2WavLayerScale(config)
+ self.attention_type = "sliding_attention"
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*):
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+ query_sequence_length, key_sequence_length)` if default attention is used.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ hidden_states = residual + self.self_attn_layer_scale(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + self.mlp_layer_scale(hidden_states)
+
+ return hidden_states
+
+
+class Qwen3OmniMoeCode2WavTransformerModel(Qwen3Model):
+ _can_record_outputs = {
+ "hidden_states": Qwen3OmniMoeCode2WavTransformerLayer,
+ "attentions": Qwen3OmniMoeCode2WavAttention,
+ }
+
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig):
+ super().__init__(config)
+ del self.vocab_size
+ del self.padding_idx
+ del self.embed_tokens
+ self.window_size = config.sliding_window
+ self.layers = nn.ModuleList(
+ [Qwen3OmniMoeCode2WavTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ use_cache=None,
+ cache_position=None,
+ **kwargs,
+ ):
+ if input_ids is not None:
+ raise ValueError("input_ids is not expected")
+ return super().forward(
+ input_ids,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ inputs_embeds,
+ use_cache,
+ cache_position,
+ **kwargs,
+ )
+
+
+class SnakeBeta(SnakeBeta):
+ pass
+
+
+class Qwen3OmniMoeCode2WavDecoderResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+
+ self.act1 = SnakeBeta(dim)
+ self.conv1 = Qwen3OmniMoeCausalConvNet(dim, dim, kernel_size=7, dilation=dilation)
+ self.act2 = SnakeBeta(dim)
+ self.conv2 = Qwen3OmniMoeCausalConvNet(dim, dim, kernel_size=1)
+
+ def forward(self, hidden_state):
+ residual = hidden_state
+
+ hidden_state = self.act1(hidden_state)
+ hidden_state = self.conv1(hidden_state)
+ hidden_state = self.act2(hidden_state)
+ hidden_state = self.conv2(hidden_state)
+ return hidden_state + residual
+
+
+class Qwen3OmniMoeCode2WavDecoderBlock(Qwen3OmniMoePreTrainedModel):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig, layer_idx):
+ super().__init__(config)
+ in_dim = config.decoder_dim // 2**layer_idx
+ out_dim = config.decoder_dim // 2 ** (layer_idx + 1)
+ upsample_rate = config.upsample_rates[layer_idx]
+
+ block = [
+ SnakeBeta(in_dim),
+ Qwen3OmniMoeCausalTransConvNet(in_dim, out_dim, 2 * upsample_rate, upsample_rate),
+ ]
+
+ for dilation in (1, 3, 9):
+ block.append(Qwen3OmniMoeCode2WavDecoderResidualUnit(out_dim, dilation))
+
+ self.block = nn.ModuleList(block)
+
+ def forward(self, hidden):
+ for block in self.block:
+ hidden = block(hidden)
+ return hidden
+
+
+class Qwen3OmniMoeCode2Wav(Qwen3OmniMoePreTrainedModel):
+ def __init__(self, config: Qwen3OmniMoeCode2WavConfig):
+ super().__init__(config)
+ self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios)
+ self.pre_transformer = Qwen3OmniMoeCode2WavTransformerModel._from_config(config)
+ self.code_embedding = nn.Embedding(config.codebook_size * config.num_quantizers, config.hidden_size)
+ self.register_buffer(
+ "code_offset", torch.arange(config.num_quantizers).view(1, -1, 1) * config.codebook_size, persistent=False
+ )
+
+ upsample = []
+ for factor in config.upsampling_ratios:
+ upsample.append(
+ nn.ModuleList(
+ [
+ Qwen3OmniMoeCausalTransConvNet(config.hidden_size, config.hidden_size, factor, factor),
+ Qwen3OmniMoeConvNeXtBlock(config.hidden_size),
+ ]
+ )
+ )
+ self.upsample = nn.ModuleList(upsample)
+
+ decoder = [Qwen3OmniMoeCausalConvNet(config.hidden_size, config.decoder_dim, 7)]
+ for i in range(len(config.upsample_rates)):
+ decoder.append(Qwen3OmniMoeCode2WavDecoderBlock(config, i))
+ output_dim = config.decoder_dim // 2 ** len(config.upsample_rates)
+ decoder += [
+ SnakeBeta(output_dim),
+ Qwen3OmniMoeCausalConvNet(output_dim, 1, 7),
+ ]
+ self.decoder = nn.ModuleList(decoder)
+
+ self.post_init()
+
+ def forward(self, codes):
+ if codes.shape[1] != self.config.num_quantizers:
+ raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
+ hidden = self.code_embedding(codes + self.code_offset).mean(1)
+ hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
+ hidden = hidden.permute(0, 2, 1)
+ for blocks in self.upsample:
+ for block in blocks:
+ hidden = block(hidden)
+ wav = hidden
+ for block in self.decoder:
+ wav = block(wav)
+ return wav.clamp(min=-1, max=1)
+
+ def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
+ wavs = []
+ start_index = 0
+ while start_index < codes.shape[-1]:
+ end_index = min(start_index + chunk_size, codes.shape[-1])
+ context_size = left_context_size if start_index - left_context_size > 0 else start_index
+ codes_chunk = codes[..., start_index - context_size : end_index]
+ wav_chunk = self(codes_chunk)
+ wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+ start_index = end_index
+ return torch.cat(wavs, dim=-1)
+
+
+class Qwen3OmniMoeForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin):
+ config_class = Qwen3OmniMoeConfig
+
+ def __init__(self, config: Qwen3OmniMoeConfig):
+ super().__init__(config)
+
+ self.thinker = Qwen3OmniMoeThinkerForConditionalGeneration._from_config(config.thinker_config)
+ self.has_talker = config.enable_audio_output
+ if self.has_talker:
+ self.enable_talker()
+ self.post_init()
+
+ def enable_talker(self):
+ self.talker = Qwen3OmniMoeTalkerForConditionalGeneration._from_config(self.config.talker_config)
+ self.code2wav = Qwen3OmniMoeCode2Wav._from_config(self.config.code2wav_config)
+
+ def disable_talker(self):
+ if hasattr(self, "talker"):
+ del self.talker
+ if hasattr(self, "code2wav"):
+ del self.code2wav
+ self.has_talker = False
+
+ def _get_talker_user_parts(
+ self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+ ):
+ user_talker_part = torch.empty(
+ (1, segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size),
+ device=self.talker.device,
+ dtype=self.talker.dtype,
+ )
+
+ user_mm_mask = multimodal_mask[:, im_start_index:segment_end_index]
+
+ # Multimodal data exists
+ if user_mm_mask.any():
+ user_thinker_hidden_mm = thinker_hidden[:, im_start_index:segment_end_index][user_mm_mask]
+ mm_hidden = self.talker.hidden_projection(user_thinker_hidden_mm).to(self.talker.device)
+ user_talker_part[user_mm_mask] = mm_hidden
+ user_thinker_embed = thinker_embed[:, im_start_index:segment_end_index][~user_mm_mask]
+ user_text_hidden = self.talker.text_projection(user_thinker_embed).to(self.talker.device)
+ user_talker_part[~user_mm_mask] = user_text_hidden
+ return user_talker_part
+
+ def _get_talker_assistant_parts(
+ self, im_start_index, segment_end_index, speaker_id, thinker_embed, tts_pad_embed, tts_bos_embed, tts_eos_embed
+ ):
+ assistant_hidden = self.talker.text_projection(thinker_embed[:, im_start_index:segment_end_index]).to(
+ self.talker.device
+ ) # [1 t d]
+ assistant_text_hidden = torch.cat(
+ (
+ assistant_hidden[:, :3],
+ tts_pad_embed.expand(-1, 4, -1),
+ tts_bos_embed,
+ assistant_hidden[:, 3:4], # First text
+ ),
+ dim=1,
+ )
+ codec_special_tokens = torch.tensor(
+ [
+ [
+ self.config.talker_config.codec_nothink_id,
+ self.config.talker_config.codec_think_bos_id,
+ self.config.talker_config.codec_think_eos_id,
+ speaker_id,
+ self.config.talker_config.codec_pad_id,
+ self.config.talker_config.codec_bos_id,
+ ]
+ ],
+ device=self.talker.device,
+ dtype=torch.long,
+ )
+ assistant_codec_hidden = torch.cat(
+ (
+ torch.zeros(
+ (1, 3, self.config.talker_config.text_config.hidden_size),
+ device=self.talker.device,
+ dtype=self.talker.dtype,
+ ),
+ self.talker.get_input_embeddings()(codec_special_tokens).to(self.talker.device),
+ ),
+ dim=1,
+ )
+ trailing_text_hidden = torch.cat(
+ (
+ assistant_hidden[:, 4:],
+ tts_eos_embed,
+ ),
+ dim=1,
+ )
+
+ input_embeds = assistant_text_hidden + assistant_codec_hidden
+ input_ids = torch.full(
+ (1, assistant_text_hidden.shape[1]),
+ fill_value=self.config.tts_pad_token_id,
+ dtype=torch.long,
+ device=assistant_text_hidden.device,
+ )
+ return input_embeds, input_ids, trailing_text_hidden
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ speaker: str = "Ethan",
+ use_audio_in_video: bool = False,
+ return_audio: Optional[bool] = None,
+ thinker_max_new_tokens: int = 1024,
+ thinker_eos_token_id: int = 151645,
+ talker_max_new_tokens: int = 4096,
+ talker_do_sample: bool = True,
+ talker_top_k: int = 50,
+ talker_top_p: float = 1.0,
+ talker_temperature: float = 0.9,
+ talker_repetition_penalty: float = 1.05,
+ **kwargs,
+ ):
+ if return_audio and not self.has_talker:
+ raise ValueError(
+ "Cannot use talker when talker module not initialized. Use `enable_talker` method or set enable_talker in config to enable talker."
+ )
+ if return_audio is None:
+ return_audio = self.has_talker
+
+ shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+ thinker_kwargs = {
+ "max_new_tokens": thinker_max_new_tokens,
+ "eos_token_id": thinker_eos_token_id,
+ }
+
+ talker_kwargs = {}
+ token2wav_kwargs = {}
+ if return_audio:
+ speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
+ if speaker_id is None:
+ raise NotImplementedError(f"Speaker {speaker} not implemented")
+ if input_ids.shape[0] != 1:
+ raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
+ talker_supppressed_tokens = [
+ i
+ for i in range(
+ self.config.talker_config.text_config.vocab_size - 1024,
+ self.config.talker_config.text_config.vocab_size,
+ )
+ if i != self.config.talker_config.codec_eos_token_id
+ ] # Suppress additional special tokens, should not be predicted
+ talker_kwargs = {
+ "max_new_tokens": talker_max_new_tokens,
+ "do_sample": talker_do_sample,
+ "top_k": talker_top_k,
+ "top_p": talker_top_p,
+ "temperature": talker_temperature,
+ "eos_token_id": self.config.talker_config.codec_eos_token_id,
+ "repetition_penalty": talker_repetition_penalty,
+ "suppress_tokens": talker_supppressed_tokens,
+ "output_hidden_states": True,
+ "return_dict_in_generate": True,
+ }
+ token2wav_kwargs = {}
+
+ for key, value in kwargs.items():
+ if key.startswith("thinker_"):
+ thinker_kwargs[key[len("thinker_") :]] = value
+ elif key.startswith("talker_"):
+ talker_kwargs[key[len("talker_") :]] = value
+ elif key.startswith("token2wav_"):
+ token2wav_kwargs[key[len("token2wav_") :]] = value
+ # Process special input values
+ elif key == "feature_attention_mask":
+ thinker_kwargs[key] = value
+ talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+ elif key in ("input_features", "attention_mask"):
+ thinker_kwargs[key] = value
+ # Put other key to shared kwargs
+ else:
+ shared_kwargs[key] = value
+
+ # Merge kwargs
+ for key, value in shared_kwargs.items():
+ if key not in thinker_kwargs:
+ thinker_kwargs[key] = value
+ if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
+ talker_kwargs[key] = value
+ if key not in token2wav_kwargs:
+ token2wav_kwargs[key] = value
+
+ # 1. Generate from thinker module
+ generate_audio = return_audio and self.has_talker
+ if generate_audio:
+ thinker_kwargs["output_hidden_states"] = True
+ thinker_kwargs["return_dict_in_generate"] = True
+
+ thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+
+ if not generate_audio:
+ return thinker_result, None
+
+ # 2. Prepare talker input
+ thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
+ self.talker.device
+ ) # [1 t d]
+ thinker_hidden = torch.cat(
+ [
+ hidden_states[self.config.talker_config.accept_hidden_layer]
+ for hidden_states in thinker_result.hidden_states
+ ],
+ dim=1,
+ ).to(self.talker.device) # [1 t d]
+ im_start_indexes = torch.cat(
+ (
+ torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
+ torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+ ),
+ dim=-1,
+ ).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+ multimodal_mask = (
+ (thinker_result.sequences == self.config.thinker_config.audio_token_id) |
+ (thinker_result.sequences == self.config.thinker_config.image_token_id) |
+ (thinker_result.sequences == self.config.thinker_config.video_token_id)
+ ).to(self.talker.device) # [1 t] # fmt: skip
+
+ talker_special_tokens = torch.tensor(
+ [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
+ device=self.thinker.device,
+ dtype=input_ids.dtype,
+ )
+ tts_bos_embed, tts_eos_embed, tts_pad_embed = (
+ self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
+ .to(self.talker.device)
+ .chunk(3, dim=1)
+ ) # 3 * [1 1 d]
+
+ talker_input_embeds = [] # [1 t d]
+ talker_input_ids = []
+ # For every chatml parts
+ for i in range(len(im_start_indexes) - 1):
+ im_start_index = im_start_indexes[i]
+ segment_end_index = im_start_indexes[i + 1]
+ role_token = input_ids[0][im_start_index + 1]
+ # Talker should ignore thinker system prompt
+ if role_token == self.config.system_token_id:
+ continue
+ # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
+ elif role_token == self.config.user_token_id:
+ talker_user_part = self._get_talker_user_parts(
+ im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+ )
+ talker_input_embeds.append(talker_user_part)
+ talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
+ # Take assistant output (for now)
+ elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
+ talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
+ im_start_index,
+ segment_end_index,
+ speaker_id,
+ thinker_embed,
+ tts_pad_embed,
+ tts_bos_embed,
+ tts_eos_embed,
+ )
+ talker_input_embeds.append(talker_assistant_embeds)
+ talker_input_ids.append(talker_assistant_ids)
+ # History assistant output (ignore for now)
+ elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
+ continue
+ else:
+ raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
+ talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
+ talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
+ talker_result = self.talker.generate(
+ inputs_embeds=talker_input_embed,
+ trailing_text_hidden=trailing_text_hidden,
+ tts_pad_embed=tts_pad_embed,
+ talker_input_ids=talker_input_id, # Not use input_ids to prevent repetation penalty out of bound
+ **talker_kwargs,
+ )
+ talker_codes = (
+ torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
+ .transpose(1, 2)
+ .to(self.code2wav.device)
+ )
+ talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25)
+
+ return thinker_result, talker_wavs.float()
+
+
+class Qwen3OmniMoeProcessorKwargs(Qwen2_5OmniProcessorKwargs):
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ "padding_side": "left",
+ },
+ "videos_kwargs": {
+ "seconds_per_chunk": 2.0,
+ "position_id_per_seconds": 13.0,
+ "use_audio_in_video": False,
+ "size": {
+ "shortest_edge": 128 * 32 * 32,
+ "longest_edge": 768 * 32 * 32,
+ },
+ },
+ "audio_kwargs": {
+ "sampling_rate": 16000,
+ "padding": True,
+ "return_attention_mask": True,
+ },
+ }
+
+
+class Qwen3OmniMoeProcessor(Qwen2_5OmniProcessor, ProcessorMixin):
+ def replace_multimodal_special_tokens(
+ self,
+ text,
+ audio_lengths,
+ image_grid_thw,
+ video_grid_thw,
+ video_second_per_grid,
+ use_audio_in_video,
+ position_id_per_seconds,
+ seconds_per_chunk,
+ ):
+ # Extend mm token length
+ merge_length_image = self.image_processor.merge_size**2
+ merge_length_video = self.video_processor.merge_size**2
+
+ processed_text = []
+ for sample in text:
+ positions = []
+ special_tokens = [re.escape(tok) for tok in [self.audio_token, self.image_token, self.video_token]]
+ pattern = "|".join(special_tokens)
+ positions = sorted([(match.start(), match.group()) for match in re.finditer(pattern, sample)])
+ positions.sort(key=lambda x: x[0])
+
+ for _, special_token in positions:
+ if special_token == self.audio_token:
+ sample = sample.replace(self.audio_token, "<|audio_placeholder|>" * next(audio_lengths), 1)
+ elif special_token == self.image_token:
+ image_seq_length = next(image_grid_thw).prod() // merge_length_image
+ sample = sample.replace(self.image_token, "<|image_placeholder|>" * image_seq_length, 1)
+ elif special_token == self.video_token:
+ if not use_audio_in_video:
+ video_seq_length = next(video_grid_thw).prod() // merge_length_video
+ sample = sample.replace(self.video_token, "<|video_placeholder|>" * video_seq_length, 1)
+ else:
+ audio_token_indices = np.arange(next(audio_lengths))
+ curr_video_grid_thw = next(video_grid_thw)
+ height = curr_video_grid_thw[1] // self.video_processor.merge_size
+ width = curr_video_grid_thw[2] // self.video_processor.merge_size
+ video_token_indices = np.arange(curr_video_grid_thw[0]).reshape(-1, 1, 1)
+ video_token_indices = np.broadcast_to(
+ video_token_indices, (video_token_indices.shape[0], height, width)
+ ).reshape(-1)
+ video_token_indices = (
+ video_token_indices * next(video_second_per_grid) * position_id_per_seconds
+ )
+
+ video_data_index, audio_data_index = 0, 0
+ placeholder_string = self.vision_bos_token + self.audio_bos_token
+ while video_data_index < len(video_token_indices) and audio_data_index < len(
+ audio_token_indices
+ ):
+ if video_token_indices[video_data_index] <= audio_token_indices[audio_data_index]:
+ placeholder_string += "<|video_placeholder|>"
+ video_data_index += 1
+ else:
+ placeholder_string += "<|audio_placeholder|>"
+ audio_data_index += 1
+ if video_data_index < len(video_token_indices):
+ placeholder_string += "<|video_placeholder|>" * (
+ len(video_token_indices) - video_data_index
+ )
+ if audio_data_index < len(audio_token_indices):
+ placeholder_string += "<|audio_placeholder|>" * (
+ len(audio_token_indices) - audio_data_index
+ )
+ placeholder_string += self.audio_eos_token + self.vision_eos_token
+ sample = sample.replace(
+ self.vision_bos_token + self.video_token + self.vision_eos_token,
+ placeholder_string,
+ 1,
+ )
+
+ sample = sample.replace("<|audio_placeholder|>", self.audio_token)
+ sample = sample.replace("<|image_placeholder|>", self.image_token)
+ sample = sample.replace("<|video_placeholder|>", self.video_token)
+ processed_text.append(sample)
+ return processed_text
+
+ def __call__(
+ self,
+ text: TextInput = None,
+ images: ImageInput = None,
+ videos: VideoInput = None,
+ audio: AudioInput = None,
+ **kwargs,
+ ):
+ """
+ Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
+ and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+ the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to
+ WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. To prepare the vision inputs,
+ this method forwards the `vision_infos` and `kwargs` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`]
+ if `vision_infos` is not `None`. Please refer to the doctsring
+ of the above two methods for more information.
+
+ Args:
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. Both channels-first and channels-last formats are supported.
+ videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+ tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
+ audio (`np.ndarray`, `List[np.ndarray]`):
+ The audio or batch of audio to be prepared. Each audio can be a NumPy array.
+ """
+
+ if text is None:
+ raise ValueError("You need to specify either a `text` input to process.")
+
+ output_kwargs = self._merge_kwargs(
+ Qwen3OmniMoeProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ seconds_per_chunk = output_kwargs["videos_kwargs"].pop("seconds_per_chunk")
+ position_id_per_seconds = output_kwargs["videos_kwargs"].pop("position_id_per_seconds")
+ use_audio_in_video = output_kwargs["videos_kwargs"].pop("use_audio_in_video")
+ fps = output_kwargs["videos_kwargs"].get("fps", 1.0)
+
+ if audio is not None:
+ output_kwargs["audio_kwargs"]["padding"] = True # Setting to True to avoid default truncation
+ audio_inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
+ audio_inputs["feature_attention_mask"] = audio_inputs.pop(
+ "attention_mask"
+ ) # rename feature_attention_mask to prevent conflicts later on
+ audio_inputs["input_features"] = audio_inputs.pop(
+ "input_features"
+ ) # rename input_features to prevent conflicts later on
+ audio_lengths = iter(_get_feat_extract_output_lengths(audio_inputs["feature_attention_mask"].sum(-1)))
+ else:
+ audio_inputs = {}
+ audio_lengths = iter([])
+
+ if images is not None:
+ images_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = iter(images_inputs["image_grid_thw"])
+ else:
+ images_inputs = {}
+ image_grid_thw = iter([])
+
+ if videos is not None:
+ videos = make_batched_videos(videos)
+ videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+ fps = [fps] * len(videos)
+ videos_inputs["video_second_per_grid"] = [
+ self.video_processor.temporal_patch_size / fps[i] for i in range(len(fps))
+ ]
+ video_grid_thw = iter(videos_inputs["video_grid_thw"])
+ video_second_per_grid = iter(videos_inputs["video_second_per_grid"])
+ else:
+ videos_inputs = {}
+ video_grid_thw = iter([])
+ video_second_per_grid = iter([])
+
+ if not isinstance(text, list):
+ text = [text]
+
+ text = self.replace_multimodal_special_tokens(
+ text,
+ audio_lengths,
+ image_grid_thw,
+ video_grid_thw,
+ video_second_per_grid=video_second_per_grid,
+ use_audio_in_video=use_audio_in_video,
+ position_id_per_seconds=position_id_per_seconds,
+ seconds_per_chunk=seconds_per_chunk,
+ )
+
+ texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ return BatchFeature(
+ data={**texts_inputs, **images_inputs, **videos_inputs, **audio_inputs},
+ tensor_type=kwargs.get("return_tensors"),
+ )
+
+ def apply_chat_template(self, conversations, chat_template=None, **kwargs):
+ return ProcessorMixin.apply_chat_template(self, conversations, chat_template, **kwargs)
+
+
+__all__ = [
+ "Qwen3OmniMoeConfig",
+ "Qwen3OmniMoeThinkerConfig",
+ "Qwen3OmniMoeTalkerConfig",
+ "Qwen3OmniMoeForConditionalGeneration",
+ "Qwen3OmniMoeThinkerTextModel",
+ "Qwen3OmniMoeThinkerForConditionalGeneration",
+ "Qwen3OmniMoeTalkerForConditionalGeneration",
+ "Qwen3OmniMoePreTrainedModel",
+ "Qwen3OmniMoePreTrainedModelForConditionalGeneration",
+ "Qwen3OmniMoeTalkerModel",
+ "Qwen3OmniMoeThinkerTextPreTrainedModel",
+ "Qwen3OmniMoeProcessor",
+ "Qwen3OmniMoeCode2Wav",
+ "Qwen3OmniMoeCode2WavDecoderBlock",
+ "Qwen3OmniMoeCode2WavTransformerModel",
+ "Qwen3OmniMoeTalkerCodePredictorModel",
+ "Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration",
+]
diff --git a/src/transformers/models/qwen3_omni_moe/processing_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/processing_qwen3_omni_moe.py
new file mode 100644
index 000000000000..86041fc3de16
--- /dev/null
+++ b/src/transformers/models/qwen3_omni_moe/processing_qwen3_omni_moe.py
@@ -0,0 +1,360 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen3_omni_moe.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group 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 re
+from typing import Optional, Union
+
+import numpy as np
+
+from ...audio_utils import AudioInput
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, VideosKwargs
+from ...tokenization_utils_base import TextInput
+from ...video_utils import VideoInput, make_batched_videos
+
+
+class Qwen3OmniMoeVideosKwargs(VideosKwargs):
+ fps: Optional[list[Union[int, float]]]
+ use_audio_in_video: Optional[bool]
+ seconds_per_chunk: Optional[float]
+ position_id_per_seconds: Optional[int]
+ min_pixels: Optional[int]
+ max_pixels: Optional[int]
+ patch_size: Optional[int]
+ temporal_patch_size: Optional[int]
+ merge_size: Optional[int]
+
+
+class Qwen3OmniMoeImagesKwargs(ImagesKwargs):
+ min_pixels: Optional[int]
+ max_pixels: Optional[int]
+ patch_size: Optional[int]
+ temporal_patch_size: Optional[int]
+ merge_size: Optional[int]
+
+
+class Qwen3OmniMoeProcessorKwargs(ProcessingKwargs, total=False):
+ videos_kwargs: Qwen3OmniMoeVideosKwargs
+ images_kwargs: Qwen3OmniMoeImagesKwargs
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ "padding_side": "left",
+ },
+ "videos_kwargs": {
+ "seconds_per_chunk": 2.0,
+ "position_id_per_seconds": 13.0,
+ "use_audio_in_video": False,
+ "size": {
+ "shortest_edge": 128 * 32 * 32,
+ "longest_edge": 768 * 32 * 32,
+ },
+ },
+ "audio_kwargs": {
+ "sampling_rate": 16000,
+ "padding": True,
+ "return_attention_mask": True,
+ },
+ }
+
+
+def _get_feat_extract_output_lengths(input_lengths):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+
+ input_lengths_leave = input_lengths % 100
+ feat_lengths = (input_lengths_leave - 1) // 2 + 1
+ output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
+ return output_lengths
+
+
+class Qwen3OmniMoeProcessor(ProcessorMixin):
+ r"""
+ Constructs a Qwen2.5Omni processor.
+ [`Qwen3OmniMoeProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`], [`WhisperFeatureExtractor`], and [`Qwen2TokenizerFast`]. See the
+ [`~Qwen3OmniMoeProcessor.__call__`] and [`~Qwen3OmniMoeProcessor.decode`] for more information.
+
+ Args:
+ image_processor ([`Qwen2VLImageProcessor`], *optional*):
+ The image processor.
+ video_processor ([`Qwen2VLVideoProcessor`], *optional*):
+ The video processor.
+ feature_extractor ([`WhisperFeatureExtractor`], *optional*):
+ The audio feature extractor.
+ tokenizer ([`Qwen2TokenizerFast`], *optional*):
+ The text tokenizer.
+ chat_template (`Optional[str]`, *optional*):
+ The Jinja template to use for formatting the conversation. If not provided, the default chat template is used.
+ """
+
+ attributes = ["image_processor", "video_processor", "feature_extractor", "tokenizer"]
+ image_processor_class = "AutoImageProcessor"
+ video_processor_class = "AutoVideoProcessor"
+ feature_extractor_class = "WhisperFeatureExtractor"
+ tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
+
+ def __init__(
+ self, image_processor=None, video_processor=None, feature_extractor=None, tokenizer=None, chat_template=None
+ ):
+ super().__init__(image_processor, video_processor, feature_extractor, tokenizer, chat_template=chat_template)
+ self.image_token = self.tokenizer.image_token
+ self.audio_token = self.tokenizer.audio_token
+ self.video_token = self.tokenizer.video_token
+ self.vision_bos_token = self.tokenizer.vision_bos_token
+ self.vision_eos_token = self.tokenizer.vision_eos_token
+ self.audio_bos_token = self.tokenizer.audio_bos_token
+ self.audio_eos_token = self.tokenizer.audio_eos_token
+
+ def __call__(
+ self,
+ text: TextInput = None,
+ images: ImageInput = None,
+ videos: VideoInput = None,
+ audio: AudioInput = None,
+ **kwargs,
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
+ and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+ the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to
+ WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. To prepare the vision inputs,
+ this method forwards the `vision_infos` and `kwargs` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`]
+ if `vision_infos` is not `None`. Please refer to the doctsring
+ of the above two methods for more information.
+
+ Args:
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. Both channels-first and channels-last formats are supported.
+ videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+ tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
+ audio (`np.ndarray`, `List[np.ndarray]`):
+ The audio or batch of audio to be prepared. Each audio can be a NumPy array.
+ """
+
+ if text is None:
+ raise ValueError("You need to specify either a `text` input to process.")
+
+ output_kwargs = self._merge_kwargs(
+ Qwen3OmniMoeProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ seconds_per_chunk = output_kwargs["videos_kwargs"].pop("seconds_per_chunk")
+ position_id_per_seconds = output_kwargs["videos_kwargs"].pop("position_id_per_seconds")
+ use_audio_in_video = output_kwargs["videos_kwargs"].pop("use_audio_in_video")
+ fps = output_kwargs["videos_kwargs"].get("fps", 1.0)
+
+ if audio is not None:
+ output_kwargs["audio_kwargs"]["padding"] = True # Setting to True to avoid default truncation
+ audio_inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
+ audio_inputs["feature_attention_mask"] = audio_inputs.pop(
+ "attention_mask"
+ ) # rename feature_attention_mask to prevent conflicts later on
+ audio_inputs["input_features"] = audio_inputs.pop(
+ "input_features"
+ ) # rename input_features to prevent conflicts later on
+ audio_lengths = iter(_get_feat_extract_output_lengths(audio_inputs["feature_attention_mask"].sum(-1)))
+ else:
+ audio_inputs = {}
+ audio_lengths = iter([])
+
+ if images is not None:
+ images_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = iter(images_inputs["image_grid_thw"])
+ else:
+ images_inputs = {}
+ image_grid_thw = iter([])
+
+ if videos is not None:
+ videos = make_batched_videos(videos)
+ videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+ fps = [fps] * len(videos)
+ videos_inputs["video_second_per_grid"] = [
+ self.video_processor.temporal_patch_size / fps[i] for i in range(len(fps))
+ ]
+ video_grid_thw = iter(videos_inputs["video_grid_thw"])
+ video_second_per_grid = iter(videos_inputs["video_second_per_grid"])
+ else:
+ videos_inputs = {}
+ video_grid_thw = iter([])
+ video_second_per_grid = iter([])
+
+ if not isinstance(text, list):
+ text = [text]
+
+ text = self.replace_multimodal_special_tokens(
+ text,
+ audio_lengths,
+ image_grid_thw,
+ video_grid_thw,
+ video_second_per_grid=video_second_per_grid,
+ use_audio_in_video=use_audio_in_video,
+ position_id_per_seconds=position_id_per_seconds,
+ seconds_per_chunk=seconds_per_chunk,
+ )
+
+ texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ return BatchFeature(
+ data={**texts_inputs, **images_inputs, **videos_inputs, **audio_inputs},
+ tensor_type=kwargs.get("return_tensors"),
+ )
+
+ def replace_multimodal_special_tokens(
+ self,
+ text,
+ audio_lengths,
+ image_grid_thw,
+ video_grid_thw,
+ video_second_per_grid,
+ use_audio_in_video,
+ position_id_per_seconds,
+ seconds_per_chunk,
+ ):
+ # Extend mm token length
+ merge_length_image = self.image_processor.merge_size**2
+ merge_length_video = self.video_processor.merge_size**2
+
+ processed_text = []
+ for sample in text:
+ positions = []
+ special_tokens = [re.escape(tok) for tok in [self.audio_token, self.image_token, self.video_token]]
+ pattern = "|".join(special_tokens)
+ positions = sorted([(match.start(), match.group()) for match in re.finditer(pattern, sample)])
+ positions.sort(key=lambda x: x[0])
+
+ for _, special_token in positions:
+ if special_token == self.audio_token:
+ sample = sample.replace(self.audio_token, "<|audio_placeholder|>" * next(audio_lengths), 1)
+ elif special_token == self.image_token:
+ image_seq_length = next(image_grid_thw).prod() // merge_length_image
+ sample = sample.replace(self.image_token, "<|image_placeholder|>" * image_seq_length, 1)
+ elif special_token == self.video_token:
+ if not use_audio_in_video:
+ video_seq_length = next(video_grid_thw).prod() // merge_length_video
+ sample = sample.replace(self.video_token, "<|video_placeholder|>" * video_seq_length, 1)
+ else:
+ audio_token_indices = np.arange(next(audio_lengths))
+ curr_video_grid_thw = next(video_grid_thw)
+ height = curr_video_grid_thw[1] // self.video_processor.merge_size
+ width = curr_video_grid_thw[2] // self.video_processor.merge_size
+ video_token_indices = np.arange(curr_video_grid_thw[0]).reshape(-1, 1, 1)
+ video_token_indices = np.broadcast_to(
+ video_token_indices, (video_token_indices.shape[0], height, width)
+ ).reshape(-1)
+ video_token_indices = (
+ video_token_indices * next(video_second_per_grid) * position_id_per_seconds
+ )
+
+ video_data_index, audio_data_index = 0, 0
+ placeholder_string = self.vision_bos_token + self.audio_bos_token
+ while video_data_index < len(video_token_indices) and audio_data_index < len(
+ audio_token_indices
+ ):
+ if video_token_indices[video_data_index] <= audio_token_indices[audio_data_index]:
+ placeholder_string += "<|video_placeholder|>"
+ video_data_index += 1
+ else:
+ placeholder_string += "<|audio_placeholder|>"
+ audio_data_index += 1
+ if video_data_index < len(video_token_indices):
+ placeholder_string += "<|video_placeholder|>" * (
+ len(video_token_indices) - video_data_index
+ )
+ if audio_data_index < len(audio_token_indices):
+ placeholder_string += "<|audio_placeholder|>" * (
+ len(audio_token_indices) - audio_data_index
+ )
+ placeholder_string += self.audio_eos_token + self.vision_eos_token
+ sample = sample.replace(
+ self.vision_bos_token + self.video_token + self.vision_eos_token,
+ placeholder_string,
+ 1,
+ )
+
+ sample = sample.replace("<|audio_placeholder|>", self.audio_token)
+ sample = sample.replace("<|image_placeholder|>", self.image_token)
+ sample = sample.replace("<|video_placeholder|>", self.video_token)
+ processed_text.append(sample)
+ return processed_text
+
+ def get_chunked_index(self, token_indices: np.ndarray, tokens_per_chunk: int) -> list[tuple[int, int]]:
+ """
+ Splits token index list into chunks based on token value ranges.
+
+ Given a list of token indices, returns a list of (start, end) index tuples representing
+ slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
+
+ For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
+ - the first chunk contains token values < 1000,
+ - the second chunk contains values >= 1000 and < 2000, and so on.
+
+ Parameters:
+ token_indices (`np.ndarray`): A monotonically increasing list of token index values.
+ t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
+
+ Returns:
+ `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
+ and end (exclusive) indices of a chunk in `token_indices`.
+ """
+
+ def _iter():
+ i, start_idx = 0, 0 # skip bos token
+ current_chunk = 1
+ while i < len(token_indices): # skip eos token
+ if token_indices[i] >= current_chunk * tokens_per_chunk:
+ yield (start_idx, i)
+ start_idx = i
+ current_chunk += 1
+ i += 1
+ yield (start_idx, len(token_indices))
+
+ return list(_iter())
+
+ def apply_chat_template(self, conversations, chat_template=None, **kwargs):
+ return super().apply_chat_template(conversations, chat_template, **kwargs)
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ feature_extractor_input_names = self.feature_extractor.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+ return list(
+ dict.fromkeys(
+ tokenizer_input_names
+ + feature_extractor_input_names
+ + image_processor_input_names
+ + ["feature_attention_mask"]
+ + ["video_second_per_grid"]
+ )
+ )
+
+
+__all__ = ["Qwen3OmniMoeProcessor"]
diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
index d3bc3b6b044f..a6eec74f8009 100644
--- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
+++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
@@ -37,7 +37,7 @@
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, is_torchdynamo_compiling
+from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling
from ...utils.deprecation import deprecate_kwarg
from ...utils.generic import check_model_inputs
from .configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig
@@ -1104,7 +1104,7 @@ def get_placeholder_mask(
return special_image_mask, special_video_mask
@auto_docstring
- @can_return_tuple
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1235,8 +1235,6 @@ def forward(
return Qwen3VLModelOutputWithPast(
last_hidden_state=outputs.last_hidden_state,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=self.rope_deltas,
)
@@ -1313,8 +1311,7 @@ def language_model(self):
def visual(self):
return self.model.visual
- @can_return_tuple
- @auto_docstring
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1372,8 +1369,6 @@ def forward(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=outputs.rope_deltas,
)
diff --git a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py
index 7a2fa852739e..5d76287d0b88 100644
--- a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py
+++ b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py
@@ -33,7 +33,7 @@
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
from ...processing_utils import ProcessingKwargs, Unpack, VideosKwargs
from ...tokenization_utils_base import PreTokenizedInput, TextInput
-from ...utils import auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import auto_docstring, is_torchdynamo_compiling, logging
from ...utils.generic import check_model_inputs
from ...video_utils import VideoInput
from ..qwen2_5_vl.modeling_qwen2_5_vl import (
@@ -1006,7 +1006,7 @@ def get_video_features(
return self.get_image_features(pixel_values_videos, video_grid_thw)
@auto_docstring
- @can_return_tuple
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1137,8 +1137,6 @@ def forward(
return Qwen3VLModelOutputWithPast(
last_hidden_state=outputs.last_hidden_state,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=self.rope_deltas,
)
@@ -1151,6 +1149,7 @@ class Qwen3VLForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
config: Qwen3VLConfig
_checkpoint_conversion_mapping = {}
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1208,8 +1207,6 @@ def forward(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=outputs.rope_deltas,
)
@@ -1398,7 +1395,7 @@ def __call__(
index = 0
for i in range(len(text)):
while self.video_token in text[i]:
- metadata = video_metadata[i]
+ metadata = video_metadata[index]
if metadata.fps is None:
logger.warning_once(
"Qwen3VL requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
diff --git a/src/transformers/models/qwen3_vl/processing_qwen3_vl.py b/src/transformers/models/qwen3_vl/processing_qwen3_vl.py
index cac82e738f39..17bdd975eef3 100644
--- a/src/transformers/models/qwen3_vl/processing_qwen3_vl.py
+++ b/src/transformers/models/qwen3_vl/processing_qwen3_vl.py
@@ -198,7 +198,7 @@ def __call__(
index = 0
for i in range(len(text)):
while self.video_token in text[i]:
- metadata = video_metadata[i]
+ metadata = video_metadata[index]
if metadata.fps is None:
logger.warning_once(
"Qwen3VL requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
index c4a31e8f9f92..25358aa79bff 100644
--- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
@@ -80,6 +80,8 @@ class Qwen3VLMoeTextConfig(PretrainedConfig):
Number of routed experts.
norm_topk_prob (`bool`, *optional*, defaults to `True`):
Whether to normalize the topk probabilities.
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+ The aux loss factor for the total loss.
mlp_only_layers (`List[int]`, *optional*, defaults to `[]`):
Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock
The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
@@ -178,6 +180,7 @@ def __init__(
num_experts_per_tok=4,
num_experts=60,
norm_topk_prob=True,
+ router_aux_loss_coef=0.001,
mlp_only_layers=None,
rope_scaling=None,
head_dim=None,
@@ -213,6 +216,7 @@ def __init__(
self.num_experts_per_tok = num_experts_per_tok
self.num_experts = num_experts
self.norm_topk_prob = norm_topk_prob
+ self.router_aux_loss_coef = router_aux_loss_coef
self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py
index 08c647ea50ac..88e3f6e19f0e 100644
--- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py
@@ -37,7 +37,7 @@
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, is_torchdynamo_compiling
+from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling
from ...utils.deprecation import deprecate_kwarg
from ...utils.generic import OutputRecorder, check_model_inputs
from .configuration_qwen3_vl_moe import Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig
@@ -64,25 +64,6 @@ def extra_repr(self):
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-class Qwen3VLMoeTextRouter(nn.Linear):
- def __init__(self, config):
- super().__init__(config.hidden_size, config.num_experts, bias=False)
- self.hidden_size = config.hidden_size
- self.top_k = config.num_experts_per_tok
- # since all the models use norm_topk_prob, we don't need to have a extra check for it
- # self.norm_topk_prob = config.norm_topk_prob
-
- def forward(self, hidden_states):
- hidden_states = hidden_states.reshape(-1, self.hidden_size)
- router_logits = super().forward(hidden_states)
- routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
- routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
- routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
- routing_weights = routing_weights.to(hidden_states.dtype)
- router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)
- return router_weights, router_logits, router_indices
-
-
class Qwen3VLMoeTextExperts(nn.Module):
def __init__(self, config):
super().__init__()
@@ -150,11 +131,23 @@ def __init__(self, config):
super().__init__()
self.hidden_size = config.hidden_size
self.num_experts = config.num_experts
- self.gate = Qwen3VLMoeTextRouter(config)
+ self.top_k = config.num_experts_per_tok
+ self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
self.experts = Qwen3VLMoeTextExperts(config)
+ # since all the models use norm_topk_prob, we don't need to have a extra check for it
+ # self.norm_topk_prob = config.norm_topk_prob
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- router_weights, router_logits, router_indices = self.gate(hidden_states)
+ batch_size = hidden_states.shape[0]
+ hidden_states = hidden_states.reshape(-1, self.hidden_size)
+ router_logits = self.gate(hidden_states)
+ routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
+ routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
+ routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+ routing_weights = routing_weights.to(hidden_states.dtype)
+ router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)
+ hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)
routed_out = self.experts(hidden_states, router_weights, router_indices)
return routed_out, router_logits
@@ -1002,6 +995,36 @@ def _deepstack_process(
return hidden_states
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen3VLMoe causal language model (or autoregressive) outputs.
+ """
+)
+class Qwen3VLMoeCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+ aux_loss: Optional[torch.FloatTensor] = None
+
+
@dataclass
@auto_docstring(
custom_intro="""
@@ -1247,7 +1270,7 @@ def get_placeholder_mask(
return special_image_mask, special_video_mask
@auto_docstring
- @can_return_tuple
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1378,39 +1401,90 @@ def forward(
return Qwen3VLMoeModelOutputWithPast(
last_hidden_state=outputs.last_hidden_state,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=self.rope_deltas,
)
-@dataclass
-@auto_docstring(
- custom_intro="""
- Base class for Qwen3VLMoe causal language model (or autoregressive) outputs.
- """
-)
-class Qwen3VLMoeCausalLMOutputWithPast(ModelOutput):
+def load_balancing_loss_func(
+ gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
+ num_experts: Optional[int] = None,
+ top_k=2,
+ attention_mask: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, int]:
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).
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
- Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
- `past_key_values` input) to speed up sequential decoding.
- rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
- The rope index difference between sequence length and multimodal rope.
+ 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
- loss: Optional[torch.FloatTensor] = None
- logits: Optional[torch.FloatTensor] = None
- past_key_values: Optional[Cache] = None
- hidden_states: Optional[tuple[torch.FloatTensor]] = None
- attentions: Optional[tuple[torch.FloatTensor]] = None
- rope_deltas: Optional[torch.LongTensor] = None
+ if isinstance(gate_logits, tuple):
+ compute_device = gate_logits[0].device
+ concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+ if attention_mask is None:
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
+ else:
+ batch_size, sequence_length = attention_mask.shape
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+ expert_attention_mask = (
+ attention_mask[None, :, :, None, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+ .reshape(-1, top_k, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+ expert_attention_mask, dim=0
+ )
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+ router_per_expert_attention_mask = (
+ attention_mask[None, :, :, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+ .reshape(-1, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+ router_per_expert_attention_mask, dim=0
+ )
+
+ overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+ return overall_loss * num_experts
class Qwen3VLMoeForConditionalGeneration(Qwen3VLMoePreTrainedModel, GenerationMixin):
@@ -1456,8 +1530,7 @@ def language_model(self):
def visual(self):
return self.model.visual
- @can_return_tuple
- @auto_docstring
+ @check_model_inputs
def forward(
self,
input_ids: torch.LongTensor = None,
@@ -1485,8 +1558,46 @@ def forward(
The temporal, height and width of feature shape of each video in LLM.
Example:
- TODO: Add example
- """
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+
+ >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+ },
+ {"type": "text", "text": "Describe this image in short."},
+ ],
+ }
+ ]
+
+ >>> # Preparation for inference
+ >>> inputs = processor.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ return_dict=True,
+ return_tensors="pt"
+ )
+ >>> inputs = inputs.to(model.device)
+
+ >>> # Generate
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
+ >>> generated_ids_trimmed = [
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+ ]
+ >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
+ ```"""
+
outputs = self.model(
input_ids=input_ids,
pixel_values=pixel_values,
@@ -1511,12 +1622,24 @@ def forward(
if labels is not None:
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+ aux_loss = None
+ if kwargs.get("output_router_logits", False):
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.config.text_config.num_experts,
+ self.config.text_config.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
+ loss.device
+ ) # make sure to reside in the same device
+
return Qwen3VLMoeCausalLMOutputWithPast(
loss=loss,
+ aux_loss=aux_loss,
logits=logits,
past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
rope_deltas=outputs.rope_deltas,
)
diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
index 456d7c60aa89..72d3452bdc50 100644
--- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
@@ -14,21 +14,27 @@
# limitations under the License.
"""PyTorch Qwen3-VL-MOE model."""
+from typing import Optional, Union
+
import torch
import torch.nn as nn
from ...activations import ACT2FN
+from ...cache_utils import Cache
from ...configuration_utils import PretrainedConfig
from ...modeling_rope_utils import rope_config_validation
from ...modeling_utils import PreTrainedModel
-from ...utils import logging
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, logging
from ..qwen3_moe.modeling_qwen3_moe import (
Qwen3MoeDecoderLayer,
Qwen3MoePreTrainedModel,
Qwen3MoeRMSNorm,
+ load_balancing_loss_func,
)
from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
from ..qwen3_vl.modeling_qwen3_vl import (
+ Qwen3VLCausalLMOutputWithPast,
Qwen3VLForConditionalGeneration,
Qwen3VLModel,
Qwen3VLTextAttention,
@@ -98,6 +104,8 @@ class Qwen3VLMoeTextConfig(PretrainedConfig):
Number of routed experts.
norm_topk_prob (`bool`, *optional*, defaults to `True`):
Whether to normalize the topk probabilities.
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+ The aux loss factor for the total loss.
mlp_only_layers (`List[int]`, *optional*, defaults to `[]`):
Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock
The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
@@ -196,6 +204,7 @@ def __init__(
num_experts_per_tok=4,
num_experts=60,
norm_topk_prob=True,
+ router_aux_loss_coef=0.001,
mlp_only_layers=None,
rope_scaling=None,
head_dim=None,
@@ -231,6 +240,7 @@ def __init__(
self.num_experts_per_tok = num_experts_per_tok
self.num_experts = num_experts
self.norm_topk_prob = norm_topk_prob
+ self.router_aux_loss_coef = router_aux_loss_coef
self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
@@ -288,25 +298,6 @@ class Qwen3VLMoeTextRMSNorm(Qwen3MoeRMSNorm):
pass
-class Qwen3VLMoeTextRouter(nn.Linear):
- def __init__(self, config):
- super().__init__(config.hidden_size, config.num_experts, bias=False)
- self.hidden_size = config.hidden_size
- self.top_k = config.num_experts_per_tok
- # since all the models use norm_topk_prob, we don't need to have a extra check for it
- # self.norm_topk_prob = config.norm_topk_prob
-
- def forward(self, hidden_states):
- hidden_states = hidden_states.reshape(-1, self.hidden_size)
- router_logits = super().forward(hidden_states)
- routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
- routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
- routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
- routing_weights = routing_weights.to(hidden_states.dtype)
- router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)
- return router_weights, router_logits, router_indices
-
-
class Qwen3VLMoeTextExperts(nn.Module):
def __init__(self, config):
super().__init__()
@@ -374,11 +365,23 @@ def __init__(self, config):
super().__init__()
self.hidden_size = config.hidden_size
self.num_experts = config.num_experts
- self.gate = Qwen3VLMoeTextRouter(config)
+ self.top_k = config.num_experts_per_tok
+ self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
self.experts = Qwen3VLMoeTextExperts(config)
+ # since all the models use norm_topk_prob, we don't need to have a extra check for it
+ # self.norm_topk_prob = config.norm_topk_prob
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- router_weights, router_logits, router_indices = self.gate(hidden_states)
+ batch_size = hidden_states.shape[0]
+ hidden_states = hidden_states.reshape(-1, self.hidden_size)
+ router_logits = self.gate(hidden_states)
+ routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
+ routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
+ routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+ routing_weights = routing_weights.to(hidden_states.dtype)
+ router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)
+ hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)
routed_out = self.experts(hidden_states, router_weights, router_indices)
return routed_out, router_logits
@@ -415,12 +418,126 @@ class Qwen3VLMoeTextModel(Qwen3VLTextModel):
pass
+class Qwen3VLMoeCausalLMOutputWithPast(Qwen3VLCausalLMOutputWithPast):
+ aux_loss: Optional[torch.FloatTensor] = None
+
+
class Qwen3VLMoeModel(Qwen3VLModel):
pass
class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
- pass
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+
+ Example:
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+
+ >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+ },
+ {"type": "text", "text": "Describe this image in short."},
+ ],
+ }
+ ]
+
+ >>> # Preparation for inference
+ >>> inputs = processor.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ return_dict=True,
+ return_tensors="pt"
+ )
+ >>> inputs = inputs.to(model.device)
+
+ >>> # Generate
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
+ >>> generated_ids_trimmed = [
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+ ]
+ >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
+ ```"""
+
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ aux_loss = None
+ if kwargs.get("output_router_logits", False):
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.config.text_config.num_experts,
+ self.config.text_config.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
+ loss.device
+ ) # make sure to reside in the same device
+
+ return Qwen3VLMoeCausalLMOutputWithPast(
+ loss=loss,
+ aux_loss=aux_loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=outputs.rope_deltas,
+ )
__all__ = [
diff --git a/src/transformers/models/rag/modeling_rag.py b/src/transformers/models/rag/modeling_rag.py
index 13389107a2cb..25c2d66dd701 100644
--- a/src/transformers/models/rag/modeling_rag.py
+++ b/src/transformers/models/rag/modeling_rag.py
@@ -1088,9 +1088,7 @@ def _mask_pads(ll, smooth_obj):
@staticmethod
def _cat_and_pad(tensors, pad_token_id):
- output = (
- tensors[0].new(sum([t.shape[0] for t in tensors]), max([t.shape[1] for t in tensors])).fill_(pad_token_id)
- )
+ output = tensors[0].new(sum(t.shape[0] for t in tensors), max(t.shape[1] for t in tensors)).fill_(pad_token_id)
ind = 0
for t in tensors:
output[ind : ind + t.shape[0], : t.shape[1]] = t
diff --git a/src/transformers/models/rag/retrieval_rag.py b/src/transformers/models/rag/retrieval_rag.py
index e397d111a0a4..1c4548cba6f1 100644
--- a/src/transformers/models/rag/retrieval_rag.py
+++ b/src/transformers/models/rag/retrieval_rag.py
@@ -509,10 +509,7 @@ def postprocess_docs(self, docs, input_strings, prefix, n_docs, return_tensors=N
def cat_input_and_doc(doc_title, doc_text, input_string, prefix):
# TODO(Patrick): if we train more RAG models, I want to put the input first to take advantage of effortless truncation
# TODO(piktus): better handling of truncation
- if doc_title.startswith('"'):
- doc_title = doc_title[1:]
- if doc_title.endswith('"'):
- doc_title = doc_title[:-1]
+ doc_title = doc_title.removeprefix('"').removesuffix('"')
if prefix is None:
prefix = ""
out = (prefix + doc_title + self.config.title_sep + doc_text + self.config.doc_sep + input_string).replace(
diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
index d7d1ce33e8f0..88364515459a 100644
--- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
+++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
@@ -556,8 +556,9 @@ def _init_weights(self, module):
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
elif isinstance(module, RecurrentGemmaRMSNorm):
- module.weight.data.fill_(1.0)
+ module.weight.data.zero_()
def _setup_cache(self, config, batch, device, dtype):
layers = getattr(self, "model", self).layers
diff --git a/src/transformers/models/reformer/modeling_reformer.py b/src/transformers/models/reformer/modeling_reformer.py
index 990f21359bc0..031754e9faa0 100755
--- a/src/transformers/models/reformer/modeling_reformer.py
+++ b/src/transformers/models/reformer/modeling_reformer.py
@@ -2067,10 +2067,10 @@ def forward(
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
elif input_ids is not None:
self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size() # noqa: F841
+ input_shape = input_ids.size()
device = input_ids.device
elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1] # noqa: F841
+ input_shape = inputs_embeds.size()[:-1]
device = inputs_embeds.device
else:
raise ValueError("You have to specify either input_ids or inputs_embeds")
diff --git a/src/transformers/models/regnet/convert_regnet_seer_10b_to_pytorch.py b/src/transformers/models/regnet/convert_regnet_seer_10b_to_pytorch.py
index ed4bc48035d0..a27296dae8e4 100644
--- a/src/transformers/models/regnet/convert_regnet_seer_10b_to_pytorch.py
+++ b/src/transformers/models/regnet/convert_regnet_seer_10b_to_pytorch.py
@@ -164,11 +164,9 @@ def convert_weights_and_push(save_directory: Path, model_name: Optional[str] = N
num_labels = 1000
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text())
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
ImageNetPreTrainedConfig = partial(RegNetConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/regnet/convert_regnet_to_pytorch.py b/src/transformers/models/regnet/convert_regnet_to_pytorch.py
index 9d6659d7685d..d74e6ad263f0 100644
--- a/src/transformers/models/regnet/convert_regnet_to_pytorch.py
+++ b/src/transformers/models/regnet/convert_regnet_to_pytorch.py
@@ -224,11 +224,9 @@ def convert_weights_and_push(save_directory: Path, model_name: Optional[str] = N
expected_shape = (1, num_labels)
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text())
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
ImageNetPreTrainedConfig = partial(RegNetConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/resnet/convert_resnet_to_pytorch.py b/src/transformers/models/resnet/convert_resnet_to_pytorch.py
index 11b09c372c31..1e02a3e8b6c0 100644
--- a/src/transformers/models/resnet/convert_resnet_to_pytorch.py
+++ b/src/transformers/models/resnet/convert_resnet_to_pytorch.py
@@ -128,11 +128,9 @@ def convert_weights_and_push(save_directory: Path, model_name: Optional[str] = N
expected_shape = (1, num_labels)
repo_id = "huggingface/label-files"
- num_labels = num_labels
id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
- id2label = id2label
label2id = {v: k for k, v in id2label.items()}
ImageNetPreTrainedConfig = partial(ResNetConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
diff --git a/src/transformers/models/rt_detr/image_processing_rt_detr_fast.py b/src/transformers/models/rt_detr/image_processing_rt_detr_fast.py
index 68c5497b0205..9aae271deacc 100644
--- a/src/transformers/models/rt_detr/image_processing_rt_detr_fast.py
+++ b/src/transformers/models/rt_detr/image_processing_rt_detr_fast.py
@@ -8,6 +8,7 @@
from typing import Any, Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -31,17 +32,11 @@
validate_annotations,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, requires_backends
+from ...utils import TensorType, auto_docstring, requires_backends
from ...utils.import_utils import requires
from .image_processing_rt_detr import get_size_with_aspect_ratio
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class RTDetrFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
r"""
format (`str`, *optional*, defaults to `AnnotationFormat.COCO_DETECTION`):
@@ -247,13 +242,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/rt_detr/modular_rt_detr.py b/src/transformers/models/rt_detr/modular_rt_detr.py
index 760e4a6675cf..61bd055144f0 100644
--- a/src/transformers/models/rt_detr/modular_rt_detr.py
+++ b/src/transformers/models/rt_detr/modular_rt_detr.py
@@ -2,6 +2,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from transformers.models.detr.image_processing_detr_fast import DetrFastImageProcessorKwargs, DetrImageProcessorFast
@@ -22,18 +23,11 @@
from ...processing_utils import Unpack
from ...utils import (
TensorType,
- is_torchvision_v2_available,
logging,
requires_backends,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION,)
diff --git a/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py b/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py
index 87d35db22363..33044a4d1271 100644
--- a/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py
+++ b/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py
@@ -36,7 +36,7 @@
"14B": 40,
}
-HIDEN_SIZE_MAPPING = {
+HIDDEN_SIZE_MAPPING = {
"169M": 768,
"430M": 1024,
"1B5": 2048,
@@ -106,7 +106,7 @@ def convert_rmkv_checkpoint_to_hf_format(
config = RwkvConfig(
vocab_size=vocab_size,
num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size],
- hidden_size=HIDEN_SIZE_MAPPING[size],
+ hidden_size=HIDDEN_SIZE_MAPPING[size],
)
config.save_pretrained(output_dir)
diff --git a/src/transformers/models/rwkv/modeling_rwkv.py b/src/transformers/models/rwkv/modeling_rwkv.py
index 816b22f1b2dd..6c1edc74508c 100644
--- a/src/transformers/models/rwkv/modeling_rwkv.py
+++ b/src/transformers/models/rwkv/modeling_rwkv.py
@@ -532,7 +532,7 @@ def set_input_embeddings(self, new_embeddings):
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.LongTensor] = None, # noqa
+ attention_mask: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
state: Optional[list[torch.FloatTensor]] = None,
use_cache: Optional[bool] = None,
@@ -730,7 +730,7 @@ def prepare_inputs_for_generation(self, input_ids, state=None, inputs_embeds=Non
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.LongTensor] = None, # noqa
+ attention_mask: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
state: Optional[list[torch.FloatTensor]] = None,
labels: Optional[torch.LongTensor] = None,
diff --git a/src/transformers/models/sam/image_processing_sam.py b/src/transformers/models/sam/image_processing_sam.py
index c9b54f561fb6..6acb775b43db 100644
--- a/src/transformers/models/sam/image_processing_sam.py
+++ b/src/transformers/models/sam/image_processing_sam.py
@@ -757,7 +757,7 @@ def generate_crop_boxes(
Generates a list of crop boxes of different sizes. Each layer has (2**i)**2 boxes for the ith layer.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
Input original image
target_size (`int`):
Target size of the resized image
@@ -831,7 +831,7 @@ def filter_masks(
List of IoU scores.
original_size (`tuple[int,int]`):
Size of the original image.
- cropped_box_image (`np.array`):
+ cropped_box_image (`np.ndarray`):
The cropped image.
pred_iou_thresh (`float`, *optional*, defaults to 0.88):
The threshold for the iou scores.
@@ -891,7 +891,7 @@ def _filter_masks_pt(
List of IoU scores.
original_size (`tuple[int,int]`):
Size of the original image.
- cropped_box_image (`np.array`):
+ cropped_box_image (`np.ndarray`):
The cropped image.
pred_iou_thresh (`float`, *optional*, defaults to 0.88):
The threshold for the iou scores.
diff --git a/src/transformers/models/sam/image_processing_sam_fast.py b/src/transformers/models/sam/image_processing_sam_fast.py
index ba75e73c8680..5cfd5314899c 100644
--- a/src/transformers/models/sam/image_processing_sam_fast.py
+++ b/src/transformers/models/sam/image_processing_sam_fast.py
@@ -23,6 +23,7 @@
import torch
from torch.nn import functional as F
from torchvision.ops.boxes import batched_nms
+from torchvision.transforms.v2 import functional as F_t
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -39,13 +40,7 @@
pil_torch_interpolation_mapping,
)
from ...processing_utils import Unpack
-from ...utils import auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F_t
-else:
- from torchvision.transforms import functional as F_t
+from ...utils import auto_docstring
class SamFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
@@ -223,9 +218,7 @@ def _preprocess_image_like_inputs(
{
"do_normalize": False,
"do_rescale": False,
- "interpolation": F_t.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F_t.InterpolationMode.NEAREST,
+ "interpolation": F_t.InterpolationMode.NEAREST_EXACT,
"size": segmentation_maps_kwargs.pop("mask_size"),
"pad_size": segmentation_maps_kwargs.pop("mask_pad_size"),
}
diff --git a/src/transformers/models/sam/processing_sam.py b/src/transformers/models/sam/processing_sam.py
index 603adde95040..b3b3728fe273 100644
--- a/src/transformers/models/sam/processing_sam.py
+++ b/src/transformers/models/sam/processing_sam.py
@@ -205,7 +205,7 @@ def _pad_points_and_labels(self, input_points, input_labels, point_pad_value):
r"""
The method pads the 2D points and labels to the maximum number of points in the batch.
"""
- expected_nb_points = max([point.shape[0] for point in input_points])
+ expected_nb_points = max(point.shape[0] for point in input_points)
processed_input_points = []
for i, point in enumerate(input_points):
if point.shape[0] != expected_nb_points:
diff --git a/src/transformers/models/sam2/configuration_sam2.py b/src/transformers/models/sam2/configuration_sam2.py
index 39fbc9dfc2f5..e14583181d38 100644
--- a/src/transformers/models/sam2/configuration_sam2.py
+++ b/src/transformers/models/sam2/configuration_sam2.py
@@ -214,7 +214,7 @@ def __init__(
backbone_config["model_type"] = backbone_config.get("model_type", "sam2_hiera_det_model")
backbone_config = CONFIG_MAPPING[backbone_config["model_type"]](**backbone_config)
elif isinstance(backbone_config, Sam2HieraDetConfig):
- backbone_config = backbone_config
+ pass
elif backbone_config is None:
backbone_config = Sam2HieraDetConfig()
@@ -379,8 +379,6 @@ class Sam2Config(PretrainedConfig):
Dictionary of configuration options used to initialize [`Sam2MaskDecoderConfig`].
initializer_range (`float`, *optional*, defaults to 0.02):
Standard deviation for parameter initialization.
- kwargs (*optional*):
- Dictionary of keyword arguments.
Example:
@@ -434,8 +432,6 @@ def __init__(
if isinstance(vision_config, dict):
vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
- elif isinstance(vision_config, PretrainedConfig):
- vision_config = vision_config
if isinstance(prompt_encoder_config, Sam2PromptEncoderConfig):
prompt_encoder_config = prompt_encoder_config.to_dict()
if isinstance(mask_decoder_config, Sam2MaskDecoderConfig):
diff --git a/src/transformers/models/sam2/image_processing_sam2_fast.py b/src/transformers/models/sam2/image_processing_sam2_fast.py
index a55188f4e786..a773e8ad54d7 100644
--- a/src/transformers/models/sam2/image_processing_sam2_fast.py
+++ b/src/transformers/models/sam2/image_processing_sam2_fast.py
@@ -40,10 +40,7 @@
pil_torch_interpolation_mapping,
)
from ...processing_utils import Unpack
-from ...utils import (
- TensorType,
- auto_docstring,
-)
+from ...utils import TensorType, auto_docstring
class Sam2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
diff --git a/src/transformers/models/sam2/modeling_sam2.py b/src/transformers/models/sam2/modeling_sam2.py
index ef16466d344c..fe42cc39cacf 100644
--- a/src/transformers/models/sam2/modeling_sam2.py
+++ b/src/transformers/models/sam2/modeling_sam2.py
@@ -37,10 +37,7 @@
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,
- auto_docstring,
-)
+from ...utils import ModelOutput, auto_docstring
from ...utils.generic import TransformersKwargs, check_model_inputs
from ..auto import AutoModel
from .configuration_sam2 import (
diff --git a/src/transformers/models/sam2_video/configuration_sam2_video.py b/src/transformers/models/sam2_video/configuration_sam2_video.py
index a47858c6340e..2712165b44c5 100644
--- a/src/transformers/models/sam2_video/configuration_sam2_video.py
+++ b/src/transformers/models/sam2_video/configuration_sam2_video.py
@@ -335,8 +335,6 @@ def __init__(
if isinstance(vision_config, dict):
vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
- elif isinstance(vision_config, PretrainedConfig):
- vision_config = vision_config
if isinstance(prompt_encoder_config, Sam2VideoPromptEncoderConfig):
prompt_encoder_config = prompt_encoder_config.to_dict()
if isinstance(mask_decoder_config, Sam2VideoMaskDecoderConfig):
diff --git a/src/transformers/models/sam2_video/convert_sam2_video_to_hf.py b/src/transformers/models/sam2_video/convert_sam2_video_to_hf.py
index 322aa5507978..cc2ee0c7c612 100644
--- a/src/transformers/models/sam2_video/convert_sam2_video_to_hf.py
+++ b/src/transformers/models/sam2_video/convert_sam2_video_to_hf.py
@@ -190,7 +190,7 @@ def replace_keys(state_dict, config):
if re.match(output_vision_encoder_neck_pattern, key):
key = key.replace(".conv.", ".")
- # memory_encoder.out_proj.weight -> memory_encoder.projection.weight
+ # memory_encoder.o_proj.weight -> memory_encoder.projection.weight
if re.match(output_memory_encoder_projection_pattern, key):
key = key.replace(".o_proj.", ".projection.")
diff --git a/src/transformers/models/sam2_video/modeling_sam2_video.py b/src/transformers/models/sam2_video/modeling_sam2_video.py
index f4c1261d6779..79d5b015f889 100644
--- a/src/transformers/models/sam2_video/modeling_sam2_video.py
+++ b/src/transformers/models/sam2_video/modeling_sam2_video.py
@@ -39,10 +39,7 @@
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,
- auto_docstring,
-)
+from ...utils import ModelOutput, auto_docstring
from ...utils.generic import OutputRecorder, TransformersKwargs
from ..auto import AutoModel
from .configuration_sam2_video import Sam2VideoConfig, Sam2VideoMaskDecoderConfig, Sam2VideoPromptEncoderConfig
@@ -134,8 +131,10 @@ def __init__(
dtype: Union[torch.dtype, str] = "float32",
max_vision_features_cache_size: int = 1,
):
- # store as a list to avoid double memory allocation with torch.cat when adding new frames
- self.processed_frames = list(video.to(video_storage_device, dtype=dtype)) if video is not None else None
+ # store as a dictionary to avoid double memory allocation with torch.cat when adding new frames
+ self.processed_frames = (
+ dict(enumerate(video.to(video_storage_device, dtype=dtype))) if video is not None else None
+ )
self.video_height = video_height
self.video_width = video_width
@@ -293,18 +292,21 @@ def get_output(
return value
# Video frame management
- def add_new_frame(self, pixel_values: torch.Tensor) -> int:
+ def add_new_frame(self, pixel_values: torch.Tensor, frame_idx: Optional[int] = None) -> int:
"""Add new frame with automatic device placement."""
pixel_values = pixel_values.to(self.video_storage_device, dtype=self.dtype, non_blocking=True)
if pixel_values.dim() == 4:
pixel_values = pixel_values.squeeze(0)
+ if frame_idx is None:
+ frame_idx = len(self.processed_frames) if self.processed_frames is not None else 0
+
if self.processed_frames is None:
- self.processed_frames = [pixel_values]
+ self.processed_frames = {frame_idx: pixel_values}
else:
- self.processed_frames.append(pixel_values)
+ self.processed_frames[frame_idx] = pixel_values
- return self.num_frames - 1
+ return frame_idx
def get_frame(self, frame_idx: int) -> torch.Tensor:
"""Get frame from video."""
@@ -1714,7 +1716,7 @@ def forward(
Whether to propagate in reverse.
"""
if frame is not None:
- frame_idx = inference_session.add_new_frame(frame)
+ frame_idx = inference_session.add_new_frame(frame, frame_idx)
if frame is not None and inference_session.get_obj_num() == 0:
raise ValueError("No objects are provided for tracking; please add inputs first.")
@@ -2097,6 +2099,195 @@ def _use_mask_as_output(
image_embeddings=high_res_features + [backbone_features],
)
+ def _gather_memory_frame_outputs(
+ self,
+ inference_session: Sam2VideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ track_in_reverse_time: bool = False,
+ ) -> list[tuple[int, dict]]:
+ """
+ Get memory frames from conditioning and non-conditioning outputs.
+
+ Returns:
+ List of (relative_temporal_offset, output_data) tuples.
+ """
+ temporal_positions_and_previous_outputs = []
+
+ # Add conditioning frame outputs (no limit here, as is the case in the original checkpoints)
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ if not conditioning_outputs:
+ raise ValueError(
+ "maskmem_features in conditioning outputs cannot be empty when not is_initial_conditioning_frame"
+ )
+
+ # Store (temporal_position, output_data) tuples
+ temporal_positions_and_previous_outputs = [(0, out) for out in conditioning_outputs.values()]
+
+ # Add non-conditioning memory frames (up to self.num_maskmem - 1)
+ # These are typically frames tracked by the model without direct user input.
+ # Frames are selected with a stride, prioritizing the most recent ones. Here we only support stride = 1 for simplicity.
+ for relative_temporal_offset in range(self.num_maskmem - 1, 0, -1):
+ # relative_temporal_offset: how many frames before (or after if reversing) the current frame
+ if not track_in_reverse_time:
+ previous_frame_idx = frame_idx - relative_temporal_offset
+ else:
+ previous_frame_idx = frame_idx + relative_temporal_offset
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ output_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ previous_frame_idx, None
+ )
+
+ temporal_positions_and_previous_outputs.append((relative_temporal_offset, output_data))
+
+ return temporal_positions_and_previous_outputs
+
+ def _build_memory_attention_inputs(
+ self,
+ temporal_positions_and_previous_outputs: list[tuple[int, dict]],
+ device: torch.device,
+ ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+ """
+ Concatenate memory features and positional embeddings from previous frames.
+
+ Returns:
+ Tuple of (memories_to_concatenate, memory_positional_embeddings_to_concatenate).
+ """
+ memories_to_concatenate = []
+ memory_positional_embeddings_to_concatenate = []
+
+ for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
+ if prev_output_data is None:
+ continue # Skip if no output data for this temporal position (e.g., padding frames)
+
+ # Load memory features (potentially from CPU to GPU)
+ # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
+ memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
+ memories_to_concatenate.append(memory_features)
+
+ # Spatial positional encoding (potentially from CPU to GPU)
+ spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
+
+ # Add temporal positional encoding
+ # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
+ combined_memory_pos_embed = (
+ spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
+ )
+ memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
+
+ return memories_to_concatenate, memory_positional_embeddings_to_concatenate
+
+ def _get_object_pointers(
+ self,
+ inference_session: Sam2VideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ num_total_frames: int,
+ device: torch.device,
+ track_in_reverse_time: bool = False,
+ streaming: bool = False,
+ ) -> tuple[list[int], list[torch.Tensor], int]:
+ """
+ Get object pointers and their positional embeddings from past frames.
+
+ Returns:
+ Tuple of (temporal_offsets, pointer_tokens, max_object_pointers_to_use).
+ """
+ temporal_position_sign_multiplier = -1 if track_in_reverse_time else 1
+
+ # Determine max object pointers to use
+ if streaming:
+ max_object_pointers_to_use = self.config.max_object_pointers_in_encoder
+ else:
+ max_object_pointers_to_use = min(num_total_frames, self.config.max_object_pointers_in_encoder)
+
+ temporal_offsets: list[int] = []
+ pointer_tokens: list[torch.Tensor] = []
+
+ # Add object pointers from selected conditioning frames
+ # Optionally, only include pointers from past frames during evaluation
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ eligible_conditioning_outputs = conditioning_outputs
+ if not self.training:
+ eligible_conditioning_outputs = {
+ temporal_idx: out
+ for temporal_idx, out in conditioning_outputs.items()
+ if (temporal_idx >= frame_idx if track_in_reverse_time else temporal_idx <= frame_idx)
+ }
+
+ for temporal_idx, out_data in eligible_conditioning_outputs.items():
+ temporal_difference = (frame_idx - temporal_idx) * temporal_position_sign_multiplier
+ temporal_offsets.append(temporal_difference)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ # Add object pointers from non-conditioning frames (up to max_object_pointers_to_use - 1)
+ for t_diff_offset in range(1, max_object_pointers_to_use):
+ ref_frame_idx = frame_idx + t_diff_offset if track_in_reverse_time else frame_idx - t_diff_offset
+ if ref_frame_idx < 0 or (
+ not streaming and num_total_frames is not None and ref_frame_idx >= num_total_frames
+ ):
+ break # Stop if frame index is out of bounds
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ out_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ ref_frame_idx, None
+ )
+ if out_data is not None:
+ temporal_offsets.append(t_diff_offset)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ return temporal_offsets, pointer_tokens, max_object_pointers_to_use
+
+ def _process_object_pointers(
+ self,
+ temporal_offsets: list[int],
+ pointer_tokens: list[torch.Tensor],
+ max_object_pointers_to_use: int,
+ batch_size: int,
+ num_channels: int,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Process object pointers and compute their positional embeddings.
+
+ Returns:
+ Tuple of (object_pointers, object_pointers_pos_embed).
+ """
+ if not pointer_tokens:
+ return None, None
+
+ # Stack object pointers: List of (Batch, Channels) -> (SeqLen_ptr, Batch, Channels)
+ object_pointers = torch.stack(pointer_tokens, dim=0)
+
+ if self.config.enable_temporal_pos_encoding_for_object_pointers:
+ max_temporal_diff = float(max_object_pointers_to_use - 1)
+ # Determine dimensionality for temporal positional encoding of pointers
+ pointer_tpos_dim = num_channels
+
+ # Normalize temporal differences before sine PE calculation
+ normalized_temporal_diffs = (
+ torch.tensor(temporal_offsets, device=device, dtype=torch.float32) / max_temporal_diff
+ )
+ sine_pe = get_1d_sine_pe(normalized_temporal_diffs, dim=pointer_tpos_dim).to(object_pointers.dtype)
+ projected_sine_pe = self.temporal_positional_encoding_projection_layer(sine_pe)
+ object_pointers_pos_embed = projected_sine_pe.unsqueeze(1).expand(-1, batch_size, self.mem_dim)
+ else:
+ object_pointers_pos_embed = object_pointers.new_zeros(
+ len(temporal_offsets), batch_size, self.mem_dim, dtype=object_pointers.dtype
+ )
+
+ if self.mem_dim < num_channels:
+ # If memory dimension is smaller, reshape/split pointers and repeat positional encoding
+ num_splits = num_channels // self.mem_dim
+ object_pointers = object_pointers.reshape(-1, batch_size, num_splits, self.mem_dim)
+ object_pointers = object_pointers.permute(0, 2, 1, 3).flatten(
+ 0, 1
+ ) # (SeqLen_ptr*num_splits, Batch, MemDim)
+ object_pointers_pos_embed = object_pointers_pos_embed.repeat_interleave(num_splits, dim=0)
+
+ return object_pointers, object_pointers_pos_embed
+
def _prepare_memory_conditioned_features(
self,
inference_session: Sam2VideoInferenceSession,
@@ -2157,135 +2348,9 @@ def _prepare_memory_conditioned_features(
)
return current_feature_map
- num_object_pointer_tokens = 0
- temporal_position_sign_multiplier = -1 if track_in_reverse_time else 1
-
- # Step 1: Condition the visual features of the current frame on previous memories
- if not is_initial_conditioning_frame:
- # Retrieve memories encoded from previous frames
- memories_to_concatenate = []
- memory_positional_embeddings_to_concatenate = []
-
- # Ensure there are conditioning frame outputs to process
- conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
- if not conditioning_outputs:
- raise ValueError(
- "maskmem_features in conditioning outputs cannot be empty when not is_initial_conditioning_frame"
- )
-
- # Select a maximum number of temporally closest conditioning frames for cross-attention (no limit here, as is the case in the original checkpoints)
- # Store (temporal_position, output_data) tuples
- temporal_positions_and_previous_outputs = [(0, out) for out in conditioning_outputs.values()]
-
- # Add non-conditioning memory frames (up to self.num_maskmem - 1)
- # These are typically frames tracked by the model without direct user input.
- # Frames are selected with a stride, prioritizing the most recent ones. Here we only support stride = 1 for simplicity.
- for relative_temporal_offset in range(self.num_maskmem - 1, 0, -1):
- # relative_temporal_offset: how many frames before (or after if reversing) the current frame
- if not track_in_reverse_time:
- previous_frame_idx = frame_idx - relative_temporal_offset
- else:
- previous_frame_idx = frame_idx + relative_temporal_offset
-
- # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
- output_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
- previous_frame_idx, None
- )
-
- temporal_positions_and_previous_outputs.append((relative_temporal_offset, output_data))
-
- for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
- if prev_output_data is None:
- continue # Skip if no output data for this temporal position (e.g., padding frames)
-
- # Load memory features (potentially from CPU to GPU)
- # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
- memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
- memories_to_concatenate.append(memory_features)
-
- # Spatial positional encoding (potentially from CPU to GPU)
- spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
-
- # Add temporal positional encoding
- # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
- combined_memory_pos_embed = (
- spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
- )
- memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
-
- # Construct the list of past object pointers to be used in attention
- if streaming:
- max_object_pointers_to_use = self.config.max_object_pointers_in_encoder
- else:
- max_object_pointers_to_use = min(num_total_frames, self.config.max_object_pointers_in_encoder)
- temporal_diff_and_pointers = []
-
- # Add object pointers from selected conditioning frames
- # Optionally, only include pointers from past frames during evaluation
- eligible_conditioning_outputs = conditioning_outputs
- if not self.training:
- eligible_conditioning_outputs = {
- temporal_idx: out
- for temporal_idx, out in conditioning_outputs.items()
- if (temporal_idx >= frame_idx if track_in_reverse_time else temporal_idx <= frame_idx)
- }
-
- for temporal_idx, out_data in eligible_conditioning_outputs.items():
- temporal_difference = (frame_idx - temporal_idx) * temporal_position_sign_multiplier
- temporal_diff_and_pointers.append((temporal_difference, out_data["object_pointer"]))
-
- # Add object pointers from non-conditioning frames (up to max_object_pointers_to_use - 1)
- for t_diff_offset in range(1, max_object_pointers_to_use):
- ref_frame_idx = frame_idx + t_diff_offset if track_in_reverse_time else frame_idx - t_diff_offset
- if ref_frame_idx < 0 or (
- not streaming and num_total_frames is not None and ref_frame_idx >= num_total_frames
- ):
- break # Stop if frame index is out of bounds
-
- # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
- out_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
- ref_frame_idx, None
- )
- if out_data is not None:
- temporal_diff_and_pointers.append((t_diff_offset, out_data["object_pointer"]))
-
- if temporal_diff_and_pointers:
- temporal_differences, object_pointers_list = zip(*temporal_diff_and_pointers)
- # Stack object pointers: List of (Batch, Channels) -> (SeqLen_ptr, Batch, Channels)
- object_pointers = torch.stack(object_pointers_list, dim=0)
-
- if self.config.enable_temporal_pos_encoding_for_object_pointers:
- max_temporal_diff = float(max_object_pointers_to_use - 1)
- # Determine dimensionality for temporal positional encoding of pointers
- pointer_tpos_dim = num_channels
-
- # Normalize temporal differences before sine PE calculation
- normalized_temporal_diffs = (
- torch.tensor(temporal_differences, device=device, dtype=torch.float32) / max_temporal_diff
- )
- sine_pe = get_1d_sine_pe(normalized_temporal_diffs, dim=pointer_tpos_dim).to(object_pointers.dtype)
- projected_sine_pe = self.temporal_positional_encoding_projection_layer(sine_pe)
- object_pointers_pos_embed = projected_sine_pe.unsqueeze(1).expand(-1, batch_size, self.mem_dim)
- else:
- object_pointers_pos_embed = object_pointers.new_zeros(
- len(temporal_differences), batch_size, self.mem_dim, dtype=object_pointers.dtype
- )
-
- if self.mem_dim < num_channels:
- # If memory dimension is smaller, reshape/split pointers and repeat positional encoding
- num_splits = num_channels // self.mem_dim
- object_pointers = object_pointers.reshape(-1, batch_size, num_splits, self.mem_dim)
- object_pointers = object_pointers.permute(0, 2, 1, 3).flatten(
- 0, 1
- ) # (SeqLen_ptr*num_splits, Batch, MemDim)
- object_pointers_pos_embed = object_pointers_pos_embed.repeat_interleave(num_splits, dim=0)
-
- memories_to_concatenate.append(object_pointers)
- memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
- num_object_pointer_tokens = object_pointers.shape[0]
- else:
+ # Step 1: Handle initial conditioning frames
+ if is_initial_conditioning_frame:
# For initial conditioning frames, no prior memory is used directly in this block.
- # The model might handle this with a special token or mechanism.
# If configured, directly add a learnable "no memory" embedding.
# current_vision_features has shape (SeqLen, Batch, Channels)
conditioned_feature_map_flat = current_vision_features + self.no_memory_embedding
@@ -2295,11 +2360,36 @@ def _prepare_memory_conditioned_features(
)
return conditioned_feature_map
- # Step 2: Concatenate all retrieved memories and their positional embeddings.
+ # Step 2: Get memory frames and concatenate their features
+ temporal_positions_and_previous_outputs = self._gather_memory_frame_outputs(
+ inference_session, obj_idx, frame_idx, track_in_reverse_time
+ )
+
+ memories_to_concatenate, memory_positional_embeddings_to_concatenate = self._build_memory_attention_inputs(
+ temporal_positions_and_previous_outputs, device
+ )
+
+ # Step 3: Get and process object pointers
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use = self._get_object_pointers(
+ inference_session, obj_idx, frame_idx, num_total_frames, device, track_in_reverse_time, streaming
+ )
+
+ num_object_pointer_tokens = 0
+ if pointer_tokens:
+ object_pointers, object_pointers_pos_embed = self._process_object_pointers(
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use, batch_size, num_channels, device
+ )
+
+ if object_pointers is not None:
+ memories_to_concatenate.append(object_pointers)
+ memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
+ num_object_pointer_tokens = object_pointers.shape[0]
+
+ # Step 4: Concatenate all retrieved memories and their positional embeddings
combined_memory = torch.cat(memories_to_concatenate, dim=0)
combined_memory_positional_embeddings = torch.cat(memory_positional_embeddings_to_concatenate, dim=0)
- # Step 3: Forward through the memory attention mechanism.
+ # Step 5: Forward through the memory attention mechanism
conditioned_feature_map_flat = self.memory_attention(
current_vision_features=current_vision_features,
current_vision_position_embeddings=current_vision_positional_embeddings,
diff --git a/src/transformers/models/sam2_video/modular_sam2_video.py b/src/transformers/models/sam2_video/modular_sam2_video.py
index c0c9b3e1ef7a..b95a9f778251 100644
--- a/src/transformers/models/sam2_video/modular_sam2_video.py
+++ b/src/transformers/models/sam2_video/modular_sam2_video.py
@@ -36,8 +36,6 @@
from ...utils import (
ModelOutput,
auto_docstring,
- is_torchvision_available,
- is_torchvision_v2_available,
logging,
)
from ...utils.generic import OutputRecorder, TransformersKwargs
@@ -59,12 +57,6 @@
from ..sam2.processing_sam2 import Sam2Processor
-if is_torchvision_available() and is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
@@ -264,8 +256,6 @@ def __init__(
if isinstance(vision_config, dict):
vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
- elif isinstance(vision_config, PretrainedConfig):
- vision_config = vision_config
if isinstance(prompt_encoder_config, Sam2VideoPromptEncoderConfig):
prompt_encoder_config = prompt_encoder_config.to_dict()
if isinstance(mask_decoder_config, Sam2VideoMaskDecoderConfig):
@@ -405,8 +395,10 @@ def __init__(
dtype: Union[torch.dtype, str] = "float32",
max_vision_features_cache_size: int = 1,
):
- # store as a list to avoid double memory allocation with torch.cat when adding new frames
- self.processed_frames = list(video.to(video_storage_device, dtype=dtype)) if video is not None else None
+ # store as a dictionary to avoid double memory allocation with torch.cat when adding new frames
+ self.processed_frames = (
+ dict(enumerate(video.to(video_storage_device, dtype=dtype))) if video is not None else None
+ )
self.video_height = video_height
self.video_width = video_width
@@ -564,18 +556,21 @@ def get_output(
return value
# Video frame management
- def add_new_frame(self, pixel_values: torch.Tensor) -> int:
+ def add_new_frame(self, pixel_values: torch.Tensor, frame_idx: Optional[int] = None) -> int:
"""Add new frame with automatic device placement."""
pixel_values = pixel_values.to(self.video_storage_device, dtype=self.dtype, non_blocking=True)
if pixel_values.dim() == 4:
pixel_values = pixel_values.squeeze(0)
+ if frame_idx is None:
+ frame_idx = len(self.processed_frames) if self.processed_frames is not None else 0
+
if self.processed_frames is None:
- self.processed_frames = [pixel_values]
+ self.processed_frames = {frame_idx: pixel_values}
else:
- self.processed_frames.append(pixel_values)
+ self.processed_frames[frame_idx] = pixel_values
- return self.num_frames - 1
+ return frame_idx
def get_frame(self, frame_idx: int) -> torch.Tensor:
"""Get frame from video."""
@@ -1801,6 +1796,195 @@ def _use_mask_as_output(
image_embeddings=high_res_features + [backbone_features],
)
+ def _gather_memory_frame_outputs(
+ self,
+ inference_session: Sam2VideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ track_in_reverse_time: bool = False,
+ ) -> list[tuple[int, dict]]:
+ """
+ Get memory frames from conditioning and non-conditioning outputs.
+
+ Returns:
+ List of (relative_temporal_offset, output_data) tuples.
+ """
+ temporal_positions_and_previous_outputs = []
+
+ # Add conditioning frame outputs (no limit here, as is the case in the original checkpoints)
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ if not conditioning_outputs:
+ raise ValueError(
+ "maskmem_features in conditioning outputs cannot be empty when not is_initial_conditioning_frame"
+ )
+
+ # Store (temporal_position, output_data) tuples
+ temporal_positions_and_previous_outputs = [(0, out) for out in conditioning_outputs.values()]
+
+ # Add non-conditioning memory frames (up to self.num_maskmem - 1)
+ # These are typically frames tracked by the model without direct user input.
+ # Frames are selected with a stride, prioritizing the most recent ones. Here we only support stride = 1 for simplicity.
+ for relative_temporal_offset in range(self.num_maskmem - 1, 0, -1):
+ # relative_temporal_offset: how many frames before (or after if reversing) the current frame
+ if not track_in_reverse_time:
+ previous_frame_idx = frame_idx - relative_temporal_offset
+ else:
+ previous_frame_idx = frame_idx + relative_temporal_offset
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ output_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ previous_frame_idx, None
+ )
+
+ temporal_positions_and_previous_outputs.append((relative_temporal_offset, output_data))
+
+ return temporal_positions_and_previous_outputs
+
+ def _build_memory_attention_inputs(
+ self,
+ temporal_positions_and_previous_outputs: list[tuple[int, dict]],
+ device: torch.device,
+ ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+ """
+ Concatenate memory features and positional embeddings from previous frames.
+
+ Returns:
+ Tuple of (memories_to_concatenate, memory_positional_embeddings_to_concatenate).
+ """
+ memories_to_concatenate = []
+ memory_positional_embeddings_to_concatenate = []
+
+ for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
+ if prev_output_data is None:
+ continue # Skip if no output data for this temporal position (e.g., padding frames)
+
+ # Load memory features (potentially from CPU to GPU)
+ # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
+ memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
+ memories_to_concatenate.append(memory_features)
+
+ # Spatial positional encoding (potentially from CPU to GPU)
+ spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
+
+ # Add temporal positional encoding
+ # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
+ combined_memory_pos_embed = (
+ spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
+ )
+ memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
+
+ return memories_to_concatenate, memory_positional_embeddings_to_concatenate
+
+ def _get_object_pointers(
+ self,
+ inference_session: Sam2VideoInferenceSession,
+ obj_idx: int,
+ frame_idx: int,
+ num_total_frames: int,
+ device: torch.device,
+ track_in_reverse_time: bool = False,
+ streaming: bool = False,
+ ) -> tuple[list[int], list[torch.Tensor], int]:
+ """
+ Get object pointers and their positional embeddings from past frames.
+
+ Returns:
+ Tuple of (temporal_offsets, pointer_tokens, max_object_pointers_to_use).
+ """
+ temporal_position_sign_multiplier = -1 if track_in_reverse_time else 1
+
+ # Determine max object pointers to use
+ if streaming:
+ max_object_pointers_to_use = self.config.max_object_pointers_in_encoder
+ else:
+ max_object_pointers_to_use = min(num_total_frames, self.config.max_object_pointers_in_encoder)
+
+ temporal_offsets: list[int] = []
+ pointer_tokens: list[torch.Tensor] = []
+
+ # Add object pointers from selected conditioning frames
+ # Optionally, only include pointers from past frames during evaluation
+ conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
+ eligible_conditioning_outputs = conditioning_outputs
+ if not self.training:
+ eligible_conditioning_outputs = {
+ temporal_idx: out
+ for temporal_idx, out in conditioning_outputs.items()
+ if (temporal_idx >= frame_idx if track_in_reverse_time else temporal_idx <= frame_idx)
+ }
+
+ for temporal_idx, out_data in eligible_conditioning_outputs.items():
+ temporal_difference = (frame_idx - temporal_idx) * temporal_position_sign_multiplier
+ temporal_offsets.append(temporal_difference)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ # Add object pointers from non-conditioning frames (up to max_object_pointers_to_use - 1)
+ for t_diff_offset in range(1, max_object_pointers_to_use):
+ ref_frame_idx = frame_idx + t_diff_offset if track_in_reverse_time else frame_idx - t_diff_offset
+ if ref_frame_idx < 0 or (
+ not streaming and num_total_frames is not None and ref_frame_idx >= num_total_frames
+ ):
+ break # Stop if frame index is out of bounds
+
+ # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
+ out_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
+ ref_frame_idx, None
+ )
+ if out_data is not None:
+ temporal_offsets.append(t_diff_offset)
+ pointer_tokens.append(out_data["object_pointer"].to(device))
+
+ return temporal_offsets, pointer_tokens, max_object_pointers_to_use
+
+ def _process_object_pointers(
+ self,
+ temporal_offsets: list[int],
+ pointer_tokens: list[torch.Tensor],
+ max_object_pointers_to_use: int,
+ batch_size: int,
+ num_channels: int,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Process object pointers and compute their positional embeddings.
+
+ Returns:
+ Tuple of (object_pointers, object_pointers_pos_embed).
+ """
+ if not pointer_tokens:
+ return None, None
+
+ # Stack object pointers: List of (Batch, Channels) -> (SeqLen_ptr, Batch, Channels)
+ object_pointers = torch.stack(pointer_tokens, dim=0)
+
+ if self.config.enable_temporal_pos_encoding_for_object_pointers:
+ max_temporal_diff = float(max_object_pointers_to_use - 1)
+ # Determine dimensionality for temporal positional encoding of pointers
+ pointer_tpos_dim = num_channels
+
+ # Normalize temporal differences before sine PE calculation
+ normalized_temporal_diffs = (
+ torch.tensor(temporal_offsets, device=device, dtype=torch.float32) / max_temporal_diff
+ )
+ sine_pe = get_1d_sine_pe(normalized_temporal_diffs, dim=pointer_tpos_dim).to(object_pointers.dtype)
+ projected_sine_pe = self.temporal_positional_encoding_projection_layer(sine_pe)
+ object_pointers_pos_embed = projected_sine_pe.unsqueeze(1).expand(-1, batch_size, self.mem_dim)
+ else:
+ object_pointers_pos_embed = object_pointers.new_zeros(
+ len(temporal_offsets), batch_size, self.mem_dim, dtype=object_pointers.dtype
+ )
+
+ if self.mem_dim < num_channels:
+ # If memory dimension is smaller, reshape/split pointers and repeat positional encoding
+ num_splits = num_channels // self.mem_dim
+ object_pointers = object_pointers.reshape(-1, batch_size, num_splits, self.mem_dim)
+ object_pointers = object_pointers.permute(0, 2, 1, 3).flatten(
+ 0, 1
+ ) # (SeqLen_ptr*num_splits, Batch, MemDim)
+ object_pointers_pos_embed = object_pointers_pos_embed.repeat_interleave(num_splits, dim=0)
+
+ return object_pointers, object_pointers_pos_embed
+
def _prepare_memory_conditioned_features(
self,
inference_session: Sam2VideoInferenceSession,
@@ -1861,135 +2045,9 @@ def _prepare_memory_conditioned_features(
)
return current_feature_map
- num_object_pointer_tokens = 0
- temporal_position_sign_multiplier = -1 if track_in_reverse_time else 1
-
- # Step 1: Condition the visual features of the current frame on previous memories
- if not is_initial_conditioning_frame:
- # Retrieve memories encoded from previous frames
- memories_to_concatenate = []
- memory_positional_embeddings_to_concatenate = []
-
- # Ensure there are conditioning frame outputs to process
- conditioning_outputs = inference_session.output_dict_per_obj[obj_idx]["cond_frame_outputs"]
- if not conditioning_outputs:
- raise ValueError(
- "maskmem_features in conditioning outputs cannot be empty when not is_initial_conditioning_frame"
- )
-
- # Select a maximum number of temporally closest conditioning frames for cross-attention (no limit here, as is the case in the original checkpoints)
- # Store (temporal_position, output_data) tuples
- temporal_positions_and_previous_outputs = [(0, out) for out in conditioning_outputs.values()]
-
- # Add non-conditioning memory frames (up to self.num_maskmem - 1)
- # These are typically frames tracked by the model without direct user input.
- # Frames are selected with a stride, prioritizing the most recent ones. Here we only support stride = 1 for simplicity.
- for relative_temporal_offset in range(self.num_maskmem - 1, 0, -1):
- # relative_temporal_offset: how many frames before (or after if reversing) the current frame
- if not track_in_reverse_time:
- previous_frame_idx = frame_idx - relative_temporal_offset
- else:
- previous_frame_idx = frame_idx + relative_temporal_offset
-
- # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
- output_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
- previous_frame_idx, None
- )
-
- temporal_positions_and_previous_outputs.append((relative_temporal_offset, output_data))
-
- for relative_temporal_offset, prev_output_data in temporal_positions_and_previous_outputs:
- if prev_output_data is None:
- continue # Skip if no output data for this temporal position (e.g., padding frames)
-
- # Load memory features (potentially from CPU to GPU)
- # Features are flattened: (Batch, Channels, H, W) -> (H*W, Batch, Channels)
- memory_features = prev_output_data["maskmem_features"].to(device, non_blocking=True)
- memories_to_concatenate.append(memory_features)
-
- # Spatial positional encoding (potentially from CPU to GPU)
- spatial_memory_pos_embed = prev_output_data["maskmem_pos_enc"].to(device, non_blocking=True)
-
- # Add temporal positional encoding
- # self.memory_temporal_positional_encoding shape: (NumMaskMem, 1, 1, MemDim)
- combined_memory_pos_embed = (
- spatial_memory_pos_embed + self.memory_temporal_positional_encoding[relative_temporal_offset - 1]
- )
- memory_positional_embeddings_to_concatenate.append(combined_memory_pos_embed)
-
- # Construct the list of past object pointers to be used in attention
- if streaming:
- max_object_pointers_to_use = self.config.max_object_pointers_in_encoder
- else:
- max_object_pointers_to_use = min(num_total_frames, self.config.max_object_pointers_in_encoder)
- temporal_diff_and_pointers = []
-
- # Add object pointers from selected conditioning frames
- # Optionally, only include pointers from past frames during evaluation
- eligible_conditioning_outputs = conditioning_outputs
- if not self.training:
- eligible_conditioning_outputs = {
- temporal_idx: out
- for temporal_idx, out in conditioning_outputs.items()
- if (temporal_idx >= frame_idx if track_in_reverse_time else temporal_idx <= frame_idx)
- }
-
- for temporal_idx, out_data in eligible_conditioning_outputs.items():
- temporal_difference = (frame_idx - temporal_idx) * temporal_position_sign_multiplier
- temporal_diff_and_pointers.append((temporal_difference, out_data["object_pointer"]))
-
- # Add object pointers from non-conditioning frames (up to max_object_pointers_to_use - 1)
- for t_diff_offset in range(1, max_object_pointers_to_use):
- ref_frame_idx = frame_idx + t_diff_offset if track_in_reverse_time else frame_idx - t_diff_offset
- if ref_frame_idx < 0 or (
- not streaming and num_total_frames is not None and ref_frame_idx >= num_total_frames
- ):
- break # Stop if frame index is out of bounds
-
- # check if the output is already stored without using get_output to avoid unnecessary memory transfers between CPU and GPU
- out_data = inference_session.output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].get(
- ref_frame_idx, None
- )
- if out_data is not None:
- temporal_diff_and_pointers.append((t_diff_offset, out_data["object_pointer"]))
-
- if temporal_diff_and_pointers:
- temporal_differences, object_pointers_list = zip(*temporal_diff_and_pointers)
- # Stack object pointers: List of (Batch, Channels) -> (SeqLen_ptr, Batch, Channels)
- object_pointers = torch.stack(object_pointers_list, dim=0)
-
- if self.config.enable_temporal_pos_encoding_for_object_pointers:
- max_temporal_diff = float(max_object_pointers_to_use - 1)
- # Determine dimensionality for temporal positional encoding of pointers
- pointer_tpos_dim = num_channels
-
- # Normalize temporal differences before sine PE calculation
- normalized_temporal_diffs = (
- torch.tensor(temporal_differences, device=device, dtype=torch.float32) / max_temporal_diff
- )
- sine_pe = get_1d_sine_pe(normalized_temporal_diffs, dim=pointer_tpos_dim).to(object_pointers.dtype)
- projected_sine_pe = self.temporal_positional_encoding_projection_layer(sine_pe)
- object_pointers_pos_embed = projected_sine_pe.unsqueeze(1).expand(-1, batch_size, self.mem_dim)
- else:
- object_pointers_pos_embed = object_pointers.new_zeros(
- len(temporal_differences), batch_size, self.mem_dim, dtype=object_pointers.dtype
- )
-
- if self.mem_dim < num_channels:
- # If memory dimension is smaller, reshape/split pointers and repeat positional encoding
- num_splits = num_channels // self.mem_dim
- object_pointers = object_pointers.reshape(-1, batch_size, num_splits, self.mem_dim)
- object_pointers = object_pointers.permute(0, 2, 1, 3).flatten(
- 0, 1
- ) # (SeqLen_ptr*num_splits, Batch, MemDim)
- object_pointers_pos_embed = object_pointers_pos_embed.repeat_interleave(num_splits, dim=0)
-
- memories_to_concatenate.append(object_pointers)
- memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
- num_object_pointer_tokens = object_pointers.shape[0]
- else:
+ # Step 1: Handle initial conditioning frames
+ if is_initial_conditioning_frame:
# For initial conditioning frames, no prior memory is used directly in this block.
- # The model might handle this with a special token or mechanism.
# If configured, directly add a learnable "no memory" embedding.
# current_vision_features has shape (SeqLen, Batch, Channels)
conditioned_feature_map_flat = current_vision_features + self.no_memory_embedding
@@ -1999,11 +2057,36 @@ def _prepare_memory_conditioned_features(
)
return conditioned_feature_map
- # Step 2: Concatenate all retrieved memories and their positional embeddings.
+ # Step 2: Get memory frames and concatenate their features
+ temporal_positions_and_previous_outputs = self._gather_memory_frame_outputs(
+ inference_session, obj_idx, frame_idx, track_in_reverse_time
+ )
+
+ memories_to_concatenate, memory_positional_embeddings_to_concatenate = self._build_memory_attention_inputs(
+ temporal_positions_and_previous_outputs, device
+ )
+
+ # Step 3: Get and process object pointers
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use = self._get_object_pointers(
+ inference_session, obj_idx, frame_idx, num_total_frames, device, track_in_reverse_time, streaming
+ )
+
+ num_object_pointer_tokens = 0
+ if pointer_tokens:
+ object_pointers, object_pointers_pos_embed = self._process_object_pointers(
+ temporal_offsets, pointer_tokens, max_object_pointers_to_use, batch_size, num_channels, device
+ )
+
+ if object_pointers is not None:
+ memories_to_concatenate.append(object_pointers)
+ memory_positional_embeddings_to_concatenate.append(object_pointers_pos_embed)
+ num_object_pointer_tokens = object_pointers.shape[0]
+
+ # Step 4: Concatenate all retrieved memories and their positional embeddings
combined_memory = torch.cat(memories_to_concatenate, dim=0)
combined_memory_positional_embeddings = torch.cat(memory_positional_embeddings_to_concatenate, dim=0)
- # Step 3: Forward through the memory attention mechanism.
+ # Step 5: Forward through the memory attention mechanism
conditioned_feature_map_flat = self.memory_attention(
current_vision_features=current_vision_features,
current_vision_position_embeddings=current_vision_positional_embeddings,
@@ -2213,7 +2296,7 @@ def forward(
Whether to propagate in reverse.
"""
if frame is not None:
- frame_idx = inference_session.add_new_frame(frame)
+ frame_idx = inference_session.add_new_frame(frame, frame_idx)
if frame is not None and inference_session.get_obj_num() == 0:
raise ValueError("No objects are provided for tracking; please add inputs first.")
diff --git a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py
index 5078d437e978..fbe1d2fb714d 100755
--- a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py
+++ b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py
@@ -648,8 +648,6 @@ def forward(
output_attentions: bool = False,
conv_attention_mask: Optional[torch.Tensor] = None,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
@@ -2187,7 +2185,7 @@ def __init__(self, config):
kernel_size=kernel_size,
padding=(kernel_size - 1) // 2,
)
- self.activation_fuction = nn.ReLU()
+ self.activation_function = nn.ReLU()
self.ln1 = nn.LayerNorm(embed_dim)
self.dropout_module = nn.Dropout(p=var_pred_dropout)
self.conv2 = nn.Conv1d(
@@ -2202,10 +2200,10 @@ def __init__(self, config):
def forward(self, hidden_states: Tensor) -> Tensor:
# Input: B x T x C; Output: B x T
hidden_states = self.conv1(hidden_states.transpose(1, 2))
- hidden_states = self.activation_fuction(hidden_states).transpose(1, 2)
+ hidden_states = self.activation_function(hidden_states).transpose(1, 2)
hidden_states = self.dropout_module(self.ln1(hidden_states))
hidden_states = self.conv2(hidden_states.transpose(1, 2))
- hidden_states = self.activation_fuction(hidden_states).transpose(1, 2)
+ hidden_states = self.activation_function(hidden_states).transpose(1, 2)
hidden_states = self.dropout_module(self.ln2(hidden_states))
return self.proj(hidden_states).squeeze(dim=2)
diff --git a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py
index 7aa15cb84ddd..bf883374b68f 100644
--- a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py
+++ b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py
@@ -493,8 +493,6 @@ def forward(
output_attentions: bool = False,
conv_attention_mask: Optional[torch.Tensor] = None,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
@@ -557,7 +555,6 @@ def _apply_chunk_attention(self, attention_mask, hidden_states):
if self.config.speech_encoder_left_chunk_num >= 0:
start_indices = (chunk_indices - self.config.speech_encoder_left_chunk_num).clamp_(min=0)
start_indices = start_indices * self.config.speech_encoder_chunk_size
- start_indices = start_indices
start_indices = start_indices.unsqueeze(1).expand(-1, sequence_len)
end_indices = ((chunk_indices + 1) * self.config.speech_encoder_chunk_size).clamp_(max=sequence_len)
@@ -2380,7 +2377,7 @@ def __init__(self, embed_dim, hidden_dim, kernel_size, var_pred_dropout):
kernel_size=kernel_size,
padding="same",
)
- self.activation_fuction = nn.ReLU()
+ self.activation_function = nn.ReLU()
self.ln1 = nn.LayerNorm(hidden_dim)
self.dropout_module = nn.Dropout(p=var_pred_dropout)
self.conv2 = nn.Conv1d(
@@ -2397,12 +2394,12 @@ def forward(self, hidden_states: Tensor, padding_mask: Optional[Tensor] = None)
if padding_mask is not None:
hidden_states = hidden_states.masked_fill(~padding_mask.bool().unsqueeze(-1), 0.0)
hidden_states = self.conv1(hidden_states.transpose(1, 2))
- hidden_states = self.activation_fuction(hidden_states).transpose(1, 2)
+ hidden_states = self.activation_function(hidden_states).transpose(1, 2)
hidden_states = self.dropout_module(self.ln1(hidden_states))
if padding_mask is not None:
hidden_states = hidden_states.masked_fill(~padding_mask.bool().unsqueeze(-1), 0.0)
hidden_states = self.conv2(hidden_states.transpose(1, 2))
- hidden_states = self.activation_fuction(hidden_states).transpose(1, 2)
+ hidden_states = self.activation_function(hidden_states).transpose(1, 2)
hidden_states = self.dropout_module(self.ln2(hidden_states))
return self.proj(hidden_states).squeeze(dim=2)
diff --git a/src/transformers/models/segformer/image_processing_segformer_fast.py b/src/transformers/models/segformer/image_processing_segformer_fast.py
index da4bef3e9ee8..36befda0eaa3 100644
--- a/src/transformers/models/segformer/image_processing_segformer_fast.py
+++ b/src/transformers/models/segformer/image_processing_segformer_fast.py
@@ -22,6 +22,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -40,13 +41,7 @@
is_torch_tensor,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import TensorType, auto_docstring
class SegformerFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
@@ -135,9 +130,7 @@ def _preprocess_image_like_inputs(
"do_normalize": False,
"do_rescale": False,
# Nearest interpolation is used for segmentation maps instead of BILINEAR.
- "interpolation": F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ "interpolation": F.InterpolationMode.NEAREST_EXACT,
}
)
processed_segmentation_maps = self._preprocess(
diff --git a/src/transformers/models/segformer/modular_segformer.py b/src/transformers/models/segformer/modular_segformer.py
index 341e6949d8b7..831d046fd9a7 100644
--- a/src/transformers/models/segformer/modular_segformer.py
+++ b/src/transformers/models/segformer/modular_segformer.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from transformers.models.beit.image_processing_beit_fast import BeitFastImageProcessorKwargs, BeitImageProcessorFast
@@ -36,16 +37,9 @@
from ...processing_utils import Unpack
from ...utils import (
TensorType,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class SegformerFastImageProcessorKwargs(BeitFastImageProcessorKwargs):
pass
@@ -96,9 +90,7 @@ def _preprocess_image_like_inputs(
"do_normalize": False,
"do_rescale": False,
# Nearest interpolation is used for segmentation maps instead of BILINEAR.
- "interpolation": F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST,
+ "interpolation": F.InterpolationMode.NEAREST_EXACT,
}
)
processed_segmentation_maps = self._preprocess(
diff --git a/src/transformers/models/sew_d/modeling_sew_d.py b/src/transformers/models/sew_d/modeling_sew_d.py
index f8b71241c79e..0458e8764cb4 100644
--- a/src/transformers/models/sew_d/modeling_sew_d.py
+++ b/src/transformers/models/sew_d/modeling_sew_d.py
@@ -509,7 +509,7 @@ def forward(ctx, input, mask, dim):
@staticmethod
def backward(ctx, grad_output):
(output,) = ctx.saved_tensors
- inputGrad = softmax_backward_data(ctx, grad_output, output, ctx.dim, output)
+ inputGrad = softmax_backward_data(ctx, grad_output, output)
return inputGrad, None, None
@staticmethod
@@ -755,7 +755,6 @@ def forward(
if rel_att is not None:
attention_scores = attention_scores + rel_att
- attention_scores = attention_scores
attention_scores = attention_scores.view(
-1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
)
diff --git a/src/transformers/models/siglip2/image_processing_siglip2_fast.py b/src/transformers/models/siglip2/image_processing_siglip2_fast.py
index 64dcfa1ad566..45261fab2cd0 100644
--- a/src/transformers/models/siglip2/image_processing_siglip2_fast.py
+++ b/src/transformers/models/siglip2/image_processing_siglip2_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -32,17 +33,11 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from .image_processing_siglip2 import get_image_size_for_max_num_patches
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/smolvlm/processing_smolvlm.py b/src/transformers/models/smolvlm/processing_smolvlm.py
index 97f0eaa9e7b2..42dcecce6a3b 100644
--- a/src/transformers/models/smolvlm/processing_smolvlm.py
+++ b/src/transformers/models/smolvlm/processing_smolvlm.py
@@ -316,7 +316,7 @@ def __call__(
text = [text]
elif not isinstance(text, list) and not isinstance(text[0], str):
raise ValueError("Invalid input text. Please provide a string, or a list of strings")
- n_images_in_text = sum([sample.count(self.image_token) for sample in text])
+ n_images_in_text = sum(sample.count(self.image_token) for sample in text)
if n_images_in_text > 0 and (images is None and videos is None):
raise ValueError(f"We detected {n_images_in_text} tokens in the text but no images/videos were passed")
diff --git a/src/transformers/models/smolvlm/video_processing_smolvlm.py b/src/transformers/models/smolvlm/video_processing_smolvlm.py
index 7e8e544b8fc7..522a344b09b5 100644
--- a/src/transformers/models/smolvlm/video_processing_smolvlm.py
+++ b/src/transformers/models/smolvlm/video_processing_smolvlm.py
@@ -17,21 +17,16 @@
import numpy as np
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict
from ...processing_utils import Unpack, VideosKwargs
-from ...utils import TensorType, is_torchvision_v2_available, logging
+from ...utils import TensorType, logging
from ...video_processing_utils import BaseVideoProcessor
from ...video_utils import VideoMetadata, group_videos_by_shape, reorder_videos
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
DEFAULT_SYSTEM_MESSAGE = "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."
diff --git a/src/transformers/models/superpoint/image_processing_superpoint_fast.py b/src/transformers/models/superpoint/image_processing_superpoint_fast.py
index a752e08ac5f0..54f95fa75af6 100644
--- a/src/transformers/models/superpoint/image_processing_superpoint_fast.py
+++ b/src/transformers/models/superpoint/image_processing_superpoint_fast.py
@@ -33,17 +33,13 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
if TYPE_CHECKING:
from .modeling_superpoint import SuperPointKeypointDescriptionOutput
-if is_torchvision_v2_available():
- import torchvision.transforms.v2.functional as F
-else:
- import torchvision.transforms.functional as F
+import torchvision.transforms.v2.functional as F
def is_grayscale(
diff --git a/src/transformers/models/swin2sr/image_processing_swin2sr_fast.py b/src/transformers/models/swin2sr/image_processing_swin2sr_fast.py
index c10bd5081754..82c9d733d367 100644
--- a/src/transformers/models/swin2sr/image_processing_swin2sr_fast.py
+++ b/src/transformers/models/swin2sr/image_processing_swin2sr_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, ChannelDimension, get_image_size
from ...image_processing_utils_fast import (
@@ -30,17 +31,11 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
)
from ...utils.deprecation import deprecate_kwarg
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/swin2sr/modeling_swin2sr.py b/src/transformers/models/swin2sr/modeling_swin2sr.py
index e010a1d8a01e..ff8e4cb6d6ae 100644
--- a/src/transformers/models/swin2sr/modeling_swin2sr.py
+++ b/src/transformers/models/swin2sr/modeling_swin2sr.py
@@ -862,7 +862,7 @@ def __init__(self, scale, num_features):
self.scale = scale
if (scale & (scale - 1)) == 0:
# scale = 2^n
- for i in range(int(math.log(scale, 2))):
+ for i in range(int(math.log2(scale))):
self.add_module(f"convolution_{i}", nn.Conv2d(num_features, 4 * num_features, 3, 1, 1))
self.add_module(f"pixelshuffle_{i}", nn.PixelShuffle(2))
elif scale == 3:
@@ -873,7 +873,7 @@ def __init__(self, scale, num_features):
def forward(self, hidden_state):
if (self.scale & (self.scale - 1)) == 0:
- for i in range(int(math.log(self.scale, 2))):
+ for i in range(int(math.log2(self.scale))):
hidden_state = self.__getattr__(f"convolution_{i}")(hidden_state)
hidden_state = self.__getattr__(f"pixelshuffle_{i}")(hidden_state)
diff --git a/src/transformers/models/switch_transformers/convert_switch_transformers_original_flax_checkpoint_to_pytorch.py b/src/transformers/models/switch_transformers/convert_switch_transformers_original_flax_checkpoint_to_pytorch.py
index 71d304ea96c6..1570877607e2 100644
--- a/src/transformers/models/switch_transformers/convert_switch_transformers_original_flax_checkpoint_to_pytorch.py
+++ b/src/transformers/models/switch_transformers/convert_switch_transformers_original_flax_checkpoint_to_pytorch.py
@@ -29,6 +29,146 @@
logging.set_verbosity_info()
+def load_flax_weights_in_pytorch_model(pt_model, flax_state):
+ """Load flax checkpoints in a PyTorch model"""
+
+ try:
+ import torch
+ except (ImportError, ModuleNotFoundError):
+ logger.error(
+ "Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"
+ " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/index.html#installation for installation"
+ " instructions."
+ )
+ raise
+
+ # check if we have bf16 weights
+ is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values()
+ if any(is_type_bf16):
+ # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
+ # and bf16 is not fully supported in PT yet.
+ logger.warning(
+ "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` "
+ "before loading those in PyTorch model."
+ )
+ flax_state = jax.tree_util.tree_map(
+ lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state
+ )
+
+ flax_state_dict = flatten_dict(flax_state)
+ pt_model_dict = pt_model.state_dict()
+
+ load_model_with_head_into_base_model = (pt_model.base_model_prefix in flax_state) and (
+ pt_model.base_model_prefix not in {k.split(".")[0] for k in pt_model_dict}
+ )
+ load_base_model_into_model_with_head = (pt_model.base_model_prefix not in flax_state) and (
+ pt_model.base_model_prefix in {k.split(".")[0] for k in pt_model_dict}
+ )
+
+ # keep track of unexpected & missing keys
+ unexpected_keys = []
+ missing_keys = set(pt_model_dict.keys())
+
+ for flax_key_tuple, flax_tensor in flax_state_dict.items():
+ has_base_model_prefix = flax_key_tuple[0] == pt_model.base_model_prefix
+ require_base_model_prefix = ".".join((pt_model.base_model_prefix,) + flax_key_tuple) in pt_model_dict
+
+ # adapt flax_key to prepare for loading from/to base model only
+ if load_model_with_head_into_base_model and has_base_model_prefix:
+ flax_key_tuple = flax_key_tuple[1:]
+ elif load_base_model_into_model_with_head and require_base_model_prefix:
+ flax_key_tuple = (pt_model.base_model_prefix,) + flax_key_tuple
+
+ # rename flax weights to PyTorch format
+ if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(flax_key_tuple) not in pt_model_dict:
+ # conv layer
+ flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
+ flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1))
+ elif flax_key_tuple[-1] == "kernel" and ".".join(flax_key_tuple) not in pt_model_dict:
+ # linear layer
+ flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
+ flax_tensor = flax_tensor.T
+ elif flax_key_tuple[-1] in ["scale", "embedding"]:
+ flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
+
+ # adding batch stats from flax batch norm to pt
+ elif "mean" in flax_key_tuple[-1]:
+ flax_key_tuple = flax_key_tuple[:-1] + ("running_mean",)
+ elif "var" in flax_key_tuple[-1]:
+ flax_key_tuple = flax_key_tuple[:-1] + ("running_var",)
+
+ if "batch_stats" in flax_state:
+ flax_key = ".".join(flax_key_tuple[1:]) # Remove the params/batch_stats header
+ else:
+ flax_key = ".".join(flax_key_tuple)
+
+ # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
+ special_pt_names = {}
+ # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
+ for key in pt_model_dict:
+ key_components = key.split(".")
+ name = None
+ if key_components[-3::2] == ["parametrizations", "original0"]:
+ name = key_components[-2] + "_g"
+ elif key_components[-3::2] == ["parametrizations", "original1"]:
+ name = key_components[-2] + "_v"
+ if name is not None:
+ key_components = key_components[:-3] + [name]
+ key_to_check = ".".join(key_components)
+ special_pt_names[key_to_check] = key
+
+ if flax_key in special_pt_names:
+ flax_key = special_pt_names[flax_key]
+
+ if flax_key in pt_model_dict:
+ if flax_tensor.shape != pt_model_dict[flax_key].shape:
+ raise ValueError(
+ f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected "
+ f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."
+ )
+ else:
+ # add weight to pytorch dict
+ flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor
+ pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
+ # remove from missing keys
+ missing_keys.remove(flax_key)
+ else:
+ # weight is not expected by PyTorch model
+ unexpected_keys.append(flax_key)
+
+ pt_model.load_state_dict(pt_model_dict)
+
+ # re-transform missing_keys to list
+ missing_keys = list(missing_keys)
+
+ if len(unexpected_keys) > 0:
+ logger.warning(
+ "Some weights of the Flax model were not used when initializing the PyTorch model"
+ f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"
+ f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"
+ " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"
+ f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"
+ " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
+ " FlaxBertForSequenceClassification model)."
+ )
+ else:
+ logger.warning(f"All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n")
+ if len(missing_keys) > 0:
+ logger.warning(
+ f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"
+ f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"
+ " use it for predictions and inference."
+ )
+ else:
+ logger.warning(
+ f"All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"
+ "If your task is similar to the task the model of the checkpoint was trained on, "
+ f"you can already use {pt_model.__class__.__name__} for predictions without further training."
+ )
+
+ return pt_model
+
+
# should not include what is already done by the `from_pt` argument
MOE_LAYER_NAME_MAPPING = {
"/attention/": "/0/SelfAttention/",
diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py
index f3c6e3fb1a2a..f5bdfe6e18e0 100644
--- a/src/transformers/models/t5/modeling_t5.py
+++ b/src/transformers/models/t5/modeling_t5.py
@@ -265,7 +265,7 @@ def forward(self, hidden_states):
try:
from apex.normalization import FusedRMSNorm
- T5LayerNorm = FusedRMSNorm # noqa
+ T5LayerNorm = FusedRMSNorm
logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of T5LayerNorm")
except ImportError:
diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py
index 217a24df0417..2085cc8aa517 100644
--- a/src/transformers/models/t5gemma/configuration_t5gemma.py
+++ b/src/transformers/models/t5gemma/configuration_t5gemma.py
@@ -323,9 +323,5 @@ def __setattr__(self, key, value):
setattr(self.decoder, key, value)
super().__setattr__(key, value)
- def get_text_config(self, *args, **kwargs):
- # Always return self, regardless of the decoder option.
- return self
-
__all__ = ["T5GemmaConfig", "T5GemmaModuleConfig"]
diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py
index ba023447c2bc..b6be86e9cdd7 100644
--- a/src/transformers/models/t5gemma/modeling_t5gemma.py
+++ b/src/transformers/models/t5gemma/modeling_t5gemma.py
@@ -611,6 +611,9 @@ def _init_weights(self, module):
if not self.config.tie_word_embeddings:
scale = module.out_proj.weight.shape[0] ** -0.5
module.out_proj.weight.data.normal_(mean=0.0, std=std * scale)
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
def _shift_right(self, input_ids):
"""
diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py
index 4ac42d99239c..a7b11f9a4f3d 100644
--- a/src/transformers/models/t5gemma/modular_t5gemma.py
+++ b/src/transformers/models/t5gemma/modular_t5gemma.py
@@ -206,10 +206,6 @@ def __setattr__(self, key, value):
setattr(self.decoder, key, value)
super().__setattr__(key, value)
- def get_text_config(self, *args, **kwargs):
- # Always return self, regardless of the decoder option.
- return self
-
class T5GemmaRMSNorm(Gemma2RMSNorm):
pass
@@ -491,6 +487,9 @@ def _init_weights(self, module):
if not self.config.tie_word_embeddings:
scale = module.out_proj.weight.shape[0] ** -0.5
module.out_proj.weight.data.normal_(mean=0.0, std=std * scale)
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
def _shift_right(self, input_ids):
"""
@@ -1234,7 +1233,7 @@ def forward(
"T5GemmaForConditionalGeneration",
"T5GemmaModel",
"T5GemmaEncoderModel",
- "T5GemmaPreTrainedModel", # noqa: F822
+ "T5GemmaPreTrainedModel",
"T5GemmaForSequenceClassification",
"T5GemmaForTokenClassification",
]
diff --git a/src/transformers/models/tapas/convert_tapas_original_tf_checkpoint_to_pytorch.py b/src/transformers/models/tapas/convert_tapas_original_tf_checkpoint_to_pytorch.py
index 34bf77cccd6b..2691ffc74829 100644
--- a/src/transformers/models/tapas/convert_tapas_original_tf_checkpoint_to_pytorch.py
+++ b/src/transformers/models/tapas/convert_tapas_original_tf_checkpoint_to_pytorch.py
@@ -31,6 +31,140 @@
logging.set_verbosity_info()
+def load_tf_weights_in_tapas(model, config, tf_checkpoint_path):
+ """
+ Load tf checkpoints in a PyTorch model. This is an adaptation from load_tf_weights_in_bert
+
+ - add cell selection and aggregation heads
+ - take into account additional token type embedding layers
+ """
+ try:
+ import re
+
+ import numpy as np
+ import tensorflow as tf
+ except ImportError:
+ logger.error(
+ "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+ "https://www.tensorflow.org/install/ for installation instructions."
+ )
+ raise
+ tf_path = os.path.abspath(tf_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+ for name, shape in init_vars:
+ logger.info(f"Loading TF weight {name} with shape {shape}")
+ array = tf.train.load_variable(tf_path, name)
+ names.append(name)
+ arrays.append(array)
+
+ for name, array in zip(names, arrays):
+ name = name.split("/")
+ # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculate m and v
+ # which are not required for using pretrained model
+ if any(
+ n
+ in [
+ "adam_v",
+ "adam_m",
+ "AdamWeightDecayOptimizer",
+ "AdamWeightDecayOptimizer_1",
+ "global_step",
+ "seq_relationship",
+ ]
+ for n in name
+ ):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ # in case the model is TapasForSequenceClassification, we skip output_bias and output_weights
+ # since these are not used for classification
+ if isinstance(model, TapasForSequenceClassification):
+ if any(n in ["output_bias", "output_weights"] for n in name):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ # in case the model is TapasModel, we skip output_bias, output_weights, output_bias_cls and output_weights_cls
+ # since this model does not have MLM and NSP heads
+ if isinstance(model, TapasModel):
+ if any(n in ["output_bias", "output_weights", "output_bias_cls", "output_weights_cls"] for n in name):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ # in case the model is TapasForMaskedLM, we skip the pooler
+ if isinstance(model, TapasForMaskedLM):
+ if any(n == "pooler" for n in name):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ # if first scope name starts with "bert", change it to "tapas"
+ if name[0] == "bert":
+ name[0] = "tapas"
+ pointer = model
+ for m_name in name:
+ if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
+ scope_names = re.split(r"_(\d+)", m_name)
+ else:
+ scope_names = [m_name]
+ if scope_names[0] == "kernel" or scope_names[0] == "gamma":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "beta":
+ pointer = getattr(pointer, "bias")
+ # cell selection heads
+ elif scope_names[0] == "output_bias":
+ if not isinstance(model, TapasForMaskedLM):
+ pointer = getattr(pointer, "output_bias")
+ else:
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "output_weights":
+ pointer = getattr(pointer, "output_weights")
+ elif scope_names[0] == "column_output_bias":
+ pointer = getattr(pointer, "column_output_bias")
+ elif scope_names[0] == "column_output_weights":
+ pointer = getattr(pointer, "column_output_weights")
+ # aggregation head
+ elif scope_names[0] == "output_bias_agg":
+ pointer = getattr(pointer, "aggregation_classifier")
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "output_weights_agg":
+ pointer = getattr(pointer, "aggregation_classifier")
+ pointer = getattr(pointer, "weight")
+ # classification head
+ elif scope_names[0] == "output_bias_cls":
+ pointer = getattr(pointer, "classifier")
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "output_weights_cls":
+ pointer = getattr(pointer, "classifier")
+ pointer = getattr(pointer, "weight")
+ else:
+ try:
+ pointer = getattr(pointer, scope_names[0])
+ except AttributeError:
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ if len(scope_names) >= 2:
+ num = int(scope_names[1])
+ pointer = pointer[num]
+ if m_name[-11:] == "_embeddings":
+ pointer = getattr(pointer, "weight")
+ elif m_name[-13:] in [f"_embeddings_{i}" for i in range(7)]:
+ pointer = getattr(pointer, "weight")
+ elif m_name == "kernel":
+ array = np.transpose(array)
+ try:
+ if pointer.shape != array.shape:
+ raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
+ except AssertionError as e:
+ e.args += (pointer.shape, array.shape)
+ raise
+ logger.info(f"Initialize PyTorch weight {name}")
+ # Added a check to see whether the array is a scalar (because bias terms in Tapas checkpoints can be
+ # scalar => should first be converted to numpy arrays)
+ if np.isscalar(array):
+ array = np.array(array)
+ pointer.data = torch.from_numpy(array)
+ return model
+
+
def convert_tf_checkpoint_to_pytorch(
task, reset_position_index_per_cell, tf_checkpoint_path, tapas_config_file, pytorch_dump_path
):
diff --git a/src/transformers/models/textnet/image_processing_textnet_fast.py b/src/transformers/models/textnet/image_processing_textnet_fast.py
index 2f5ef22ef5e3..baa6276736f7 100644
--- a/src/transformers/models/textnet/image_processing_textnet_fast.py
+++ b/src/transformers/models/textnet/image_processing_textnet_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import BaseImageProcessorFast, DefaultFastImageProcessorKwargs
@@ -37,16 +38,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
class TextNetFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
"""
size_divisor (`int`, *optional*, defaults to 32):
diff --git a/src/transformers/models/timm_wrapper/configuration_timm_wrapper.py b/src/transformers/models/timm_wrapper/configuration_timm_wrapper.py
index 5fa115a05431..34e640ade8bf 100644
--- a/src/transformers/models/timm_wrapper/configuration_timm_wrapper.py
+++ b/src/transformers/models/timm_wrapper/configuration_timm_wrapper.py
@@ -41,6 +41,8 @@ class TimmWrapperConfig(PretrainedConfig):
imagenet models is set to `None` due to occlusions in the label descriptions.
Args:
+ architecture (`str`, *optional*, defaults to `"resnet50"`):
+ The timm architecture to load.
initializer_range (`float`, *optional*, defaults to 0.02):
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
do_pooling (`bool`, *optional*, defaults to `True`):
@@ -65,11 +67,13 @@ class TimmWrapperConfig(PretrainedConfig):
def __init__(
self,
+ architecture: str = "resnet50",
initializer_range: float = 0.02,
do_pooling: bool = True,
model_args: Optional[dict[str, Any]] = None,
**kwargs,
):
+ self.architecture = architecture
self.initializer_range = initializer_range
self.do_pooling = do_pooling
self.model_args = model_args # named "model_args" for BC with timm
@@ -117,8 +121,8 @@ def from_dict(cls, config_dict: dict[str, Any], **kwargs):
def to_dict(self) -> dict[str, Any]:
output = super().to_dict()
- output["num_classes"] = self.num_labels
- output["label_names"] = list(self.id2label.values())
+ output.setdefault("num_classes", self.num_labels)
+ output.setdefault("label_names", list(self.id2label.values()))
output.pop("id2label", None)
output.pop("label2id", None)
return output
diff --git a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
index 7839bf7813f2..d388ff05297f 100644
--- a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
+++ b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
@@ -55,6 +55,28 @@ class TimmWrapperModelOutput(ModelOutput):
attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+def _create_timm_model_with_error_handling(config: "TimmWrapperConfig", **model_kwargs):
+ """
+ Creates a timm model and provides a clear error message if the model is not found,
+ suggesting a library update.
+ """
+ try:
+ model = timm.create_model(
+ config.architecture,
+ pretrained=False,
+ **model_kwargs,
+ )
+ return model
+ except RuntimeError as e:
+ if "Unknown model" in str(e):
+ # A good general check for unknown models.
+ raise ImportError(
+ f"The model architecture '{config.architecture}' is not supported in your version of timm ({timm.__version__}). "
+ "Please upgrade timm to a more recent version with `pip install -U timm`."
+ ) from e
+ raise e
+
+
@auto_docstring
class TimmWrapperPreTrainedModel(PreTrainedModel):
main_input_name = "pixel_values"
@@ -138,7 +160,8 @@ def __init__(self, config: TimmWrapperConfig):
super().__init__(config)
# using num_classes=0 to avoid creating classification head
extra_init_kwargs = config.model_args or {}
- self.timm_model = timm.create_model(config.architecture, pretrained=False, num_classes=0, **extra_init_kwargs)
+ self.features_only = extra_init_kwargs.get("features_only", False)
+ self.timm_model = _create_timm_model_with_error_handling(config, num_classes=0, **extra_init_kwargs)
self.post_init()
@auto_docstring
@@ -211,20 +234,25 @@ def forward(
pixel_values = pixel_values.to(self.device, self.dtype)
- if output_hidden_states:
- # to enable hidden states selection
- if isinstance(output_hidden_states, (list, tuple)):
- kwargs["indices"] = output_hidden_states
- last_hidden_state, hidden_states = self.timm_model.forward_intermediates(pixel_values, **kwargs)
- else:
- last_hidden_state = self.timm_model.forward_features(pixel_values, **kwargs)
- hidden_states = None
-
- if do_pooling:
- # classification head is not created, applying pooling only
- pooler_output = self.timm_model.forward_head(last_hidden_state)
- else:
+ if self.features_only:
+ last_hidden_state = self.timm_model.forward(pixel_values, **kwargs)
+ hidden_states = last_hidden_state if output_hidden_states else None
pooler_output = None
+ else:
+ if output_hidden_states:
+ # to enable hidden states selection
+ if isinstance(output_hidden_states, (list, tuple)):
+ kwargs["indices"] = output_hidden_states
+ last_hidden_state, hidden_states = self.timm_model.forward_intermediates(pixel_values, **kwargs)
+ else:
+ last_hidden_state = self.timm_model.forward_features(pixel_values, **kwargs)
+ hidden_states = None
+
+ if do_pooling:
+ # classification head is not created, applying pooling only
+ pooler_output = self.timm_model.forward_head(last_hidden_state)
+ else:
+ pooler_output = None
if not return_dict:
outputs = (last_hidden_state, pooler_output, hidden_states)
@@ -254,8 +282,8 @@ def __init__(self, config: TimmWrapperConfig):
)
extra_init_kwargs = config.model_args or {}
- self.timm_model = timm.create_model(
- config.architecture, pretrained=False, num_classes=config.num_labels, **extra_init_kwargs
+ self.timm_model = _create_timm_model_with_error_handling(
+ config, num_classes=config.num_labels, **extra_init_kwargs
)
self.num_labels = config.num_labels
self.post_init()
diff --git a/src/transformers/models/tvp/image_processing_tvp.py b/src/transformers/models/tvp/image_processing_tvp.py
index d3f698873d55..0d4ee226525d 100644
--- a/src/transformers/models/tvp/image_processing_tvp.py
+++ b/src/transformers/models/tvp/image_processing_tvp.py
@@ -277,7 +277,7 @@ def _preprocess_image(
do_pad: bool = True,
pad_size: Optional[dict[str, int]] = None,
constant_values: Optional[Union[float, Iterable[float]]] = None,
- pad_mode: PaddingMode = None,
+ pad_mode: Optional[PaddingMode] = None,
do_normalize: Optional[bool] = None,
do_flip_channel_order: Optional[bool] = None,
image_mean: Optional[Union[float, list[float]]] = None,
@@ -349,7 +349,7 @@ def preprocess(
do_pad: Optional[bool] = None,
pad_size: Optional[dict[str, int]] = None,
constant_values: Optional[Union[float, Iterable[float]]] = None,
- pad_mode: PaddingMode = None,
+ pad_mode: Optional[PaddingMode] = None,
do_normalize: Optional[bool] = None,
do_flip_channel_order: Optional[bool] = None,
image_mean: Optional[Union[float, list[float]]] = None,
diff --git a/src/transformers/models/tvp/image_processing_tvp_fast.py b/src/transformers/models/tvp/image_processing_tvp_fast.py
index e7fe7e621d8c..5d74e6efb71f 100644
--- a/src/transformers/models/tvp/image_processing_tvp_fast.py
+++ b/src/transformers/models/tvp/image_processing_tvp_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -34,13 +35,7 @@
make_nested_list_of_images,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available
-
-
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
+from ...utils import TensorType, auto_docstring
class TvpFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py
index c70a7a83fa9c..eaad6c5335a4 100644
--- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py
+++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py
@@ -342,6 +342,13 @@ class VaultGemmaPreTrainedModel(PreTrainedModel):
"attentions": VaultGemmaAttention,
}
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ if "RMSNorm" in module.__class__.__name__:
+ module.weight.data.zero_()
+
@auto_docstring
class VaultGemmaModel(VaultGemmaPreTrainedModel):
diff --git a/src/transformers/models/videomae/modeling_videomae.py b/src/transformers/models/videomae/modeling_videomae.py
index 97c227f1d8bf..e69e36daa471 100755
--- a/src/transformers/models/videomae/modeling_videomae.py
+++ b/src/transformers/models/videomae/modeling_videomae.py
@@ -402,6 +402,7 @@ class VideoMAEPreTrainedModel(PreTrainedModel):
base_model_prefix = "videomae"
main_input_name = "pixel_values"
supports_gradient_checkpointing = True
+ _no_split_modules = ["VideoMAEEmbeddings", "VideoMAELayer"]
_supports_sdpa = True
_supports_flash_attn = True
_supports_flex_attn = True
diff --git a/src/transformers/models/vilt/image_processing_vilt_fast.py b/src/transformers/models/vilt/image_processing_vilt_fast.py
index 79e601648c55..6926b655ce45 100644
--- a/src/transformers/models/vilt/image_processing_vilt_fast.py
+++ b/src/transformers/models/vilt/image_processing_vilt_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -30,15 +31,9 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
# Set maximum size based on the typical aspect ratio of the COCO dataset
MAX_LONGER_EDGE = 1333
MAX_SHORTER_EDGE = 800
diff --git a/src/transformers/models/vitmatte/configuration_vitmatte.py b/src/transformers/models/vitmatte/configuration_vitmatte.py
index f63c3e4eb85d..85b0b0f58d89 100644
--- a/src/transformers/models/vitmatte/configuration_vitmatte.py
+++ b/src/transformers/models/vitmatte/configuration_vitmatte.py
@@ -15,6 +15,7 @@
"""VitMatte model configuration"""
import copy
+from typing import Optional
from ...configuration_utils import PretrainedConfig
from ...utils import logging
@@ -80,7 +81,7 @@ class VitMatteConfig(PretrainedConfig):
def __init__(
self,
- backbone_config: PretrainedConfig = None,
+ backbone_config: Optional[PretrainedConfig] = None,
backbone=None,
use_pretrained_backbone=False,
use_timm_backbone=False,
diff --git a/src/transformers/models/vitmatte/image_processing_vitmatte_fast.py b/src/transformers/models/vitmatte/image_processing_vitmatte_fast.py
index ae8797789df8..c5a7256a612b 100644
--- a/src/transformers/models/vitmatte/image_processing_vitmatte_fast.py
+++ b/src/transformers/models/vitmatte/image_processing_vitmatte_fast.py
@@ -17,6 +17,7 @@
from typing import Optional, Union
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature
from ...image_processing_utils_fast import (
@@ -37,16 +38,10 @@
TensorType,
auto_docstring,
filter_out_non_signature_kwargs,
- is_torchvision_v2_available,
logging,
)
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
logger = logging.get_logger(__name__)
@@ -91,14 +86,14 @@ def size_divisibility(self, value):
def _pad_image(
self,
- images: "torch.tensor",
+ images: torch.Tensor,
size_divisibility: int = 32,
- ) -> "torch.tensor":
+ ) -> torch.Tensor:
"""
Pads an image or batched images constantly so that width and height are divisible by size_divisibility
Args:
- image (`torch,tensor`):
+ image (`torch.Tensor`):
Image to pad.
size_divisibility (`int`, *optional*, defaults to 32):
The width and height of the image will be padded to be divisible by this number.
diff --git a/src/transformers/models/vitpose/image_processing_vitpose.py b/src/transformers/models/vitpose/image_processing_vitpose.py
index 5bdefe3064bb..8309ba9d8e9e 100644
--- a/src/transformers/models/vitpose/image_processing_vitpose.py
+++ b/src/transformers/models/vitpose/image_processing_vitpose.py
@@ -380,12 +380,12 @@ def affine_transform(
size: dict[str, int],
data_format: Optional[ChannelDimension] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.array:
+ ) -> np.ndarray:
"""
Apply an affine transformation to an image.
Args:
- image (`np.array`):
+ image (`np.ndarray`):
Image to transform.
center (`tuple[float]`):
Center of the bounding box (x, y).
diff --git a/src/transformers/models/vivit/convert_vivit_flax_to_pytorch.py b/src/transformers/models/vivit/convert_vivit_flax_to_pytorch.py
index bf6aa8e4a36b..517f8fdb1537 100644
--- a/src/transformers/models/vivit/convert_vivit_flax_to_pytorch.py
+++ b/src/transformers/models/vivit/convert_vivit_flax_to_pytorch.py
@@ -129,7 +129,7 @@ def transform_state_encoder_block(state_dict, i):
def get_n_layers(state_dict):
- return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"]])
+ return sum(1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"])
def transform_state(state_dict, classification_head=False):
diff --git a/src/transformers/models/vjepa2/modeling_vjepa2.py b/src/transformers/models/vjepa2/modeling_vjepa2.py
index bde505b4ea54..eedc94b845a4 100644
--- a/src/transformers/models/vjepa2/modeling_vjepa2.py
+++ b/src/transformers/models/vjepa2/modeling_vjepa2.py
@@ -1125,7 +1125,7 @@ def forward(
return encoder_output
def get_vision_features(self, pixel_values_videos) -> torch.Tensor:
- encoder_output = self.forward(pixel_values_videos)
+ encoder_output = self.forward(pixel_values_videos, skip_predictor=True)
return encoder_output.last_hidden_state
diff --git a/src/transformers/models/voxtral/processing_voxtral.py b/src/transformers/models/voxtral/processing_voxtral.py
index 0cf2d121f9da..0fd3515c0af9 100644
--- a/src/transformers/models/voxtral/processing_voxtral.py
+++ b/src/transformers/models/voxtral/processing_voxtral.py
@@ -88,7 +88,7 @@ def __init__(
super().__init__(feature_extractor, tokenizer)
- def _retreive_input_features(self, audio, max_source_positions, **kwargs):
+ def _retrieve_input_features(self, audio, max_source_positions, **kwargs):
"""
Handles specific logic of Voxtral expected input features: audio arrays should be padded to next multiple of 480000 (duration is a multiple of 30s), see VoxtralProcessorKwargs' default audio_kwargs.
Then mel input features are extracted and stacked along batch dimension, splitting into chunks of max_source_positions.
@@ -222,7 +222,7 @@ def apply_chat_template(
data = dict(encoded_instruct_inputs)
if audio is not None:
max_source_positions = audio_kwargs.pop("max_source_positions")
- data["input_features"] = self._retreive_input_features(audio, max_source_positions, **audio_kwargs)
+ data["input_features"] = self._retrieve_input_features(audio, max_source_positions, **audio_kwargs)
return BatchFeature(data=data, tensor_type=return_tensors)
@@ -423,7 +423,7 @@ def apply_transcription_request(
# extract the input features
max_source_positions = audio_kwargs.pop("max_source_positions")
- data["input_features"] = self._retreive_input_features(
+ data["input_features"] = self._retrieve_input_features(
audio_arrays, max_source_positions, **audio_kwargs
)
diff --git a/src/transformers/models/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
index d8c58a333e07..7ff8f2e6e439 100755
--- a/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+++ b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
@@ -21,6 +21,7 @@
import numpy as np
import torch
+from safetensors.torch import load_file as safe_load_file
from torch import nn
from torch.nn import CrossEntropyLoss
@@ -50,7 +51,6 @@
cached_file,
check_torch_load_is_safe,
is_peft_available,
- is_safetensors_available,
is_torch_flex_attn_available,
logging,
)
@@ -60,10 +60,6 @@
WAV2VEC2_ADAPTER_PT_FILE = "adapter.{}.bin"
WAV2VEC2_ADAPTER_SAFE_FILE = "adapter.{}.safetensors"
-if is_safetensors_available():
- from safetensors.torch import load_file as safe_load_file
-
-
if is_torch_flex_attn_available():
from ...integrations.flex_attention import make_flex_block_causal_mask
@@ -1224,7 +1220,7 @@ def load_adapter(self, target_lang: str, force_load=True, **kwargs):
token = kwargs.pop("token", None)
use_auth_token = kwargs.pop("use_auth_token", None)
revision = kwargs.pop("revision", None)
- use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
+ use_safetensors = kwargs.pop("use_safetensors", None)
if use_auth_token is not None:
warnings.warn(
diff --git a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py
index e8f67e2d73cd..3448089c632b 100644
--- a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py
+++ b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py
@@ -428,8 +428,6 @@ def forward(
output_attentions: bool = False,
conv_attention_mask: Optional[torch.Tensor] = None,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
diff --git a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py
index b9b60a6bd3ad..79f70da7cb84 100644
--- a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py
+++ b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py
@@ -326,8 +326,6 @@ def forward(
output_attentions: bool = False,
conv_attention_mask: Optional[torch.Tensor] = None,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
diff --git a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py
index b786e415546e..62357c8e0dcb 100644
--- a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py
+++ b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py
@@ -27,11 +27,7 @@
XVectorOutput,
)
from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- auto_docstring,
- is_peft_available,
-)
+from ...utils import ModelOutput, auto_docstring, is_peft_available
from .configuration_wav2vec2_conformer import Wav2Vec2ConformerConfig
@@ -602,8 +598,6 @@ def forward(
relative_position_embeddings: Optional[torch.Tensor] = None,
output_attentions: bool = False,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
diff --git a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py
index 2c009c004453..bfa6c20737d8 100644
--- a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py
+++ b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py
@@ -410,8 +410,6 @@ def forward(
relative_position_embeddings: Optional[torch.Tensor] = None,
output_attentions: bool = False,
):
- hidden_states = hidden_states
-
# 1. Feed-Forward 1 layer
residual = hidden_states
hidden_states = self.ffn1_layer_norm(hidden_states)
diff --git a/src/transformers/models/x_clip/configuration_x_clip.py b/src/transformers/models/x_clip/configuration_x_clip.py
index 66db819168e5..3fa6bb1544a8 100644
--- a/src/transformers/models/x_clip/configuration_x_clip.py
+++ b/src/transformers/models/x_clip/configuration_x_clip.py
@@ -294,7 +294,7 @@ def __init__(
# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
+ if key in text_config and value != text_config[key] and key != "transformers_version":
# If specified in `text_config_dict`
if key in text_config_dict:
message = (
@@ -326,7 +326,7 @@ def __init__(
# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
for key, value in _vision_config_dict.items():
- if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
# If specified in `vision_config_dict`
if key in vision_config_dict:
message = (
diff --git a/src/transformers/models/xcodec/modeling_xcodec.py b/src/transformers/models/xcodec/modeling_xcodec.py
index 8909162db724..4e1d376a3d08 100644
--- a/src/transformers/models/xcodec/modeling_xcodec.py
+++ b/src/transformers/models/xcodec/modeling_xcodec.py
@@ -332,7 +332,6 @@ def _init_weights(self, module):
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
module.bias.data.zero_()
-
elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
@@ -341,6 +340,23 @@ def _init_weights(self, module):
if module.bias is not None:
k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
nn.init.uniform_(module.bias, a=-k, b=k)
+ elif module.__class__.__name__ == "Snake1d":
+ module.alpha.data.fill_(1.0)
+ elif isinstance(module, nn.ConvTranspose1d):
+ module.reset_parameters()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=0.02)
+ elif isinstance(module, XcodecModel):
+ # The conv1d are not handled correctly, as `self.acoustic_encoder/decoder` are initialized from a PreTrainedModel,
+ # but then only the submodules are used (which are not PreTrainedModels...) -> here we reinit them as in DacModel
+ for submodule in module.acoustic_encoder.modules():
+ if isinstance(submodule, nn.Conv1d):
+ nn.init.trunc_normal_(submodule.weight, std=0.02)
+ nn.init.constant_(submodule.bias, 0)
+ for submodule in module.acoustic_decoder.modules():
+ if isinstance(submodule, nn.Conv1d):
+ nn.init.trunc_normal_(submodule.weight, std=0.02)
+ nn.init.constant_(submodule.bias, 0)
def apply_weight_norm(self):
"""Apply weight norm in the acoustic encoder and decoder because the original checkpoint has weight norm applied."""
@@ -396,6 +412,9 @@ def __init__(self, config):
self.fc2 = nn.Linear(config.hidden_size, config.acoustic_model_config.hidden_size)
self.quantizer = XcodecResidualVectorQuantization(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
@staticmethod
def _adjust_dac_decoder(decoder: nn.Module):
r"""
diff --git a/src/transformers/models/xlstm/modeling_xlstm.py b/src/transformers/models/xlstm/modeling_xlstm.py
index 5bb438efce7e..fd577c0c0bac 100644
--- a/src/transformers/models/xlstm/modeling_xlstm.py
+++ b/src/transformers/models/xlstm/modeling_xlstm.py
@@ -169,7 +169,7 @@ def mlstm_chunkwise_parallel_fw_H(
eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
_device = matQ.device
- nc, chunk_size = num_chunks, chunk_size
+ nc = num_chunks
batch_size, nh, dqk, dhv = matC_states.shape
matC_k_states = matC_states.view(batch_size, nh, nc, dqk // nc, dhv)
vecN_k_states = vecN_states.view(batch_size, nh, nc, dqk // nc)
diff --git a/src/transformers/models/yolos/image_processing_yolos_fast.py b/src/transformers/models/yolos/image_processing_yolos_fast.py
index fda06dfc522a..59bb3868e75e 100644
--- a/src/transformers/models/yolos/image_processing_yolos_fast.py
+++ b/src/transformers/models/yolos/image_processing_yolos_fast.py
@@ -9,6 +9,7 @@
import torch
from torchvision.io import read_image
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import (
@@ -32,16 +33,10 @@
validate_annotations,
)
from ...processing_utils import Unpack
-from ...utils import TensorType, auto_docstring, is_torchvision_v2_available, logging
+from ...utils import TensorType, auto_docstring, logging
from ...utils.import_utils import requires
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
@@ -475,13 +470,7 @@ def resize_annotation(
resample (`InterpolationMode`, defaults to `F.InterpolationMode.NEAREST_EXACT`):
The resampling filter to use when resizing the masks.
"""
- interpolation = (
- interpolation
- if interpolation is not None
- else F.InterpolationMode.NEAREST_EXACT
- if is_torchvision_v2_available()
- else F.InterpolationMode.NEAREST
- )
+ interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
new_annotation = {}
diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py
index a69b7a0a3f86..dc95e1e550fa 100644
--- a/src/transformers/models/zamba/modeling_zamba.py
+++ b/src/transformers/models/zamba/modeling_zamba.py
@@ -355,7 +355,7 @@ def __init__(self, config: ZambaConfig, layer_idx):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
+ "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
" is None. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config"
)
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index 33e7e4b5a351..60e546f32120 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -563,7 +563,7 @@ def __init__(self, config: Zamba2Config, layer_idx: Optional[int] = None):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/zamba2/modular_zamba2.py b/src/transformers/models/zamba2/modular_zamba2.py
index 3cada0c0dd43..d05b23721142 100644
--- a/src/transformers/models/zamba2/modular_zamba2.py
+++ b/src/transformers/models/zamba2/modular_zamba2.py
@@ -346,7 +346,7 @@ def __init__(self, config: Zamba2Config, layer_idx: Optional[int] = None):
if not is_fast_path_available:
logger.warning_once(
- "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
" https://github.com/Dao-AILab/causal-conv1d"
)
diff --git a/src/transformers/models/zoedepth/image_processing_zoedepth_fast.py b/src/transformers/models/zoedepth/image_processing_zoedepth_fast.py
index 7967932729e5..045dbfdacd4d 100644
--- a/src/transformers/models/zoedepth/image_processing_zoedepth_fast.py
+++ b/src/transformers/models/zoedepth/image_processing_zoedepth_fast.py
@@ -21,6 +21,7 @@
import numpy as np
import torch
+from torchvision.transforms.v2 import functional as F
from ...image_processing_utils import (
BatchFeature,
@@ -44,7 +45,6 @@
from ...utils import (
TensorType,
auto_docstring,
- is_torchvision_v2_available,
logging,
requires_backends,
)
@@ -52,12 +52,6 @@
from .modeling_zoedepth import ZoeDepthDepthEstimatorOutput
-if is_torchvision_v2_available():
- from torchvision.transforms.v2 import functional as F
-else:
- from torchvision.transforms import functional as F
-
-
logger = logging.get_logger(__name__)
diff --git a/src/transformers/models/zoedepth/modeling_zoedepth.py b/src/transformers/models/zoedepth/modeling_zoedepth.py
index 7bbad31c2ee0..992ef3eaf631 100644
--- a/src/transformers/models/zoedepth/modeling_zoedepth.py
+++ b/src/transformers/models/zoedepth/modeling_zoedepth.py
@@ -294,7 +294,7 @@ def __init__(self, config: ZoeDepthConfig):
self.config = config
# postprocessing: only required in case of a non-hierarchical backbone (e.g. ViT, BEiT)
- if config.backbone_config is not None and config.backbone_config.model_type in ["swinv2"]:
+ if config.backbone_config is not None and config.backbone_config.model_type == "swinv2":
self.reassemble_stage = None
else:
self.reassemble_stage = ZoeDepthReassembleStage(config)
diff --git a/src/transformers/onnx/config.py b/src/transformers/onnx/config.py
index 2a47127b3855..9392e5bdc986 100644
--- a/src/transformers/onnx/config.py
+++ b/src/transformers/onnx/config.py
@@ -609,7 +609,7 @@ def outputs(self) -> Mapping[str, Mapping[int, str]]:
return common_outputs
@property
- def num_layers(self) -> tuple[int]:
+ def num_layers(self) -> tuple[int, ...]:
try:
num_layers = super().num_layers
num_layers = (num_layers, num_layers)
@@ -625,7 +625,7 @@ def num_layers(self) -> tuple[int]:
return num_layers
@property
- def num_attention_heads(self) -> tuple[int]:
+ def num_attention_heads(self) -> tuple[int, ...]:
try:
num_attention_heads = super().num_attention_heads
num_attention_heads = (num_attention_heads, num_attention_heads)
diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py
index 92da22477f55..32e9c276e98d 100755
--- a/src/transformers/pipelines/__init__.py
+++ b/src/transformers/pipelines/__init__.py
@@ -946,7 +946,6 @@ def pipeline(
# Retrieve the task
if task in custom_tasks:
- normalized_task = task
targeted_task, task_options = clean_custom_task(custom_tasks[task])
if pipeline_class is None:
if not trust_remote_code:
diff --git a/src/transformers/pipelines/audio_utils.py b/src/transformers/pipelines/audio_utils.py
index 63e18c03524b..dad6f9694520 100644
--- a/src/transformers/pipelines/audio_utils.py
+++ b/src/transformers/pipelines/audio_utils.py
@@ -7,7 +7,7 @@
import numpy as np
-def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
+def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.ndarray:
"""
Helper function to read an audio file through ffmpeg.
"""
@@ -173,7 +173,7 @@ def ffmpeg_microphone_live(
Return:
A generator yielding dictionaries of the following form
- `{"sampling_rate": int, "raw": np.array(), "partial" bool}` With optionally a `"stride" (int, int)` key if
+ `{"sampling_rate": int, "raw": np.ndarray, "partial" bool}` With optionally a `"stride" (int, int)` key if
`stride_length_s` is defined.
`stride` and `raw` are all expressed in `samples`, and `partial` is a boolean saying if the current yield item
diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py
index 944c7a90a184..ed8c97251a53 100644
--- a/src/transformers/pipelines/base.py
+++ b/src/transformers/pipelines/base.py
@@ -64,7 +64,7 @@
from ..models.auto.modeling_tf_auto import TFAutoModel
-if is_torch_available():
+if is_torch_available() or TYPE_CHECKING:
import torch
from torch.utils.data import DataLoader, Dataset
@@ -186,7 +186,7 @@ def inner(items):
# input_values, input_pixels, input_ids, ...
padded = {}
for key in keys:
- if key in {"input_ids"}:
+ if key == "input_ids":
# ImageGPT uses a feature extractor
if tokenizer is None and feature_extractor is not None:
_padding_value = f_padding_value
@@ -1122,7 +1122,7 @@ def save_pretrained(
self,
save_directory: Union[str, os.PathLike],
safe_serialization: bool = True,
- **kwargs,
+ **kwargs: Any,
):
"""
Save the pipeline's model and tokenizer.
diff --git a/src/transformers/pipelines/fill_mask.py b/src/transformers/pipelines/fill_mask.py
index cc69cf6d2792..d9c18acb792d 100644
--- a/src/transformers/pipelines/fill_mask.py
+++ b/src/transformers/pipelines/fill_mask.py
@@ -188,7 +188,7 @@ def postprocess(self, model_outputs, top_k=5, target_ids=None):
return result[0]
return result
- def get_target_ids(self, targets, top_k=None):
+ def get_target_ids(self, targets):
if isinstance(targets, str):
targets = [targets]
try:
@@ -238,7 +238,7 @@ def _sanitize_parameters(self, top_k=None, targets=None, tokenizer_kwargs=None):
postprocess_params = {}
if targets is not None:
- target_ids = self.get_target_ids(targets, top_k)
+ target_ids = self.get_target_ids(targets)
postprocess_params["target_ids"] = target_ids
if top_k is not None:
diff --git a/src/transformers/pipelines/mask_generation.py b/src/transformers/pipelines/mask_generation.py
index 3a65fdff617a..59ce0a27a6e3 100644
--- a/src/transformers/pipelines/mask_generation.py
+++ b/src/transformers/pipelines/mask_generation.py
@@ -195,8 +195,8 @@ def preprocess(
points_per_batch=64,
crops_n_layers: int = 0,
crop_overlap_ratio: float = 512 / 1500,
- points_per_crop: Optional[int] = 32,
- crop_n_points_downscale_factor: Optional[int] = 1,
+ points_per_crop: int = 32,
+ crop_n_points_downscale_factor: int = 1,
timeout: Optional[float] = None,
):
image = load_image(image, timeout=timeout)
diff --git a/src/transformers/pipelines/question_answering.py b/src/transformers/pipelines/question_answering.py
index ee86074a4c58..1958fbd1fcc8 100644
--- a/src/transformers/pipelines/question_answering.py
+++ b/src/transformers/pipelines/question_answering.py
@@ -678,7 +678,7 @@ def span_to_answer(self, text: str, start: int, end: int) -> dict[str, Union[str
words = []
token_idx = char_start_idx = char_end_idx = chars_idx = 0
- for i, word in enumerate(text.split(" ")):
+ for word in text.split(" "):
token = self.tokenizer.tokenize(word)
# Append words if they are in the span
diff --git a/src/transformers/pipelines/table_question_answering.py b/src/transformers/pipelines/table_question_answering.py
index da579423d2d4..4eba8ad64cf2 100644
--- a/src/transformers/pipelines/table_question_answering.py
+++ b/src/transformers/pipelines/table_question_answering.py
@@ -131,8 +131,8 @@ class TableQuestionAnsweringPipeline(Pipeline):
max_new_tokens=256,
)
- def __init__(self, args_parser=TableQuestionAnsweringArgumentHandler(), *args, **kwargs):
- super().__init__(*args, **kwargs)
+ def __init__(self, args_parser=TableQuestionAnsweringArgumentHandler(), **kwargs):
+ super().__init__(**kwargs)
self._args_parser = args_parser
if self.framework == "tf":
@@ -381,7 +381,7 @@ def _sanitize_parameters(self, sequential=None, padding=None, truncation=None, *
return preprocess_params, forward_params, {}
- def preprocess(self, pipeline_input, sequential=None, padding=True, truncation=None):
+ def preprocess(self, pipeline_input, padding=True, truncation=None):
if truncation is None:
if self.type == "tapas":
truncation = "drop_rows_to_fit"
diff --git a/src/transformers/pipelines/text_generation.py b/src/transformers/pipelines/text_generation.py
index 7d703ba50117..45ec58b702a2 100644
--- a/src/transformers/pipelines/text_generation.py
+++ b/src/transformers/pipelines/text_generation.py
@@ -186,7 +186,7 @@ def _sanitize_parameters(
generate_kwargs["prefix_length"] = prefix_inputs["input_ids"].shape[-1]
if handle_long_generation is not None:
- if handle_long_generation not in {"hole"}:
+ if handle_long_generation != "hole":
raise ValueError(
f"{handle_long_generation} is not a valid value for `handle_long_generation` parameter expected"
" [None, 'hole']"
@@ -241,7 +241,7 @@ def _parse_and_tokenize(self, *args, **kwargs):
Parse arguments and tokenize
"""
# Parse arguments
- if self.model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
+ if self.model.__class__.__name__ == "TransfoXLLMHeadModel":
kwargs.update({"add_space_before_punct_symbol": True})
return super()._parse_and_tokenize(*args, **kwargs)
diff --git a/src/transformers/pipelines/token_classification.py b/src/transformers/pipelines/token_classification.py
index efa70ca1851f..31ba1c481107 100644
--- a/src/transformers/pipelines/token_classification.py
+++ b/src/transformers/pipelines/token_classification.py
@@ -141,8 +141,8 @@ class TokenClassificationPipeline(ChunkPipeline):
_load_feature_extractor = False
_load_tokenizer = True
- def __init__(self, args_parser=TokenClassificationArgumentHandler(), *args, **kwargs):
- super().__init__(*args, **kwargs)
+ def __init__(self, args_parser=TokenClassificationArgumentHandler(), **kwargs):
+ super().__init__(**kwargs)
self.check_model_type(
TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
@@ -160,7 +160,7 @@ def _sanitize_parameters(
ignore_subwords: Optional[bool] = None,
aggregation_strategy: Optional[AggregationStrategy] = None,
offset_mapping: Optional[list[tuple[int, int]]] = None,
- is_split_into_words: Optional[bool] = False,
+ is_split_into_words: bool = False,
stride: Optional[int] = None,
delimiter: Optional[str] = None,
):
diff --git a/src/transformers/pipelines/zero_shot_classification.py b/src/transformers/pipelines/zero_shot_classification.py
index 20675d4a2928..d7a609bcd167 100644
--- a/src/transformers/pipelines/zero_shot_classification.py
+++ b/src/transformers/pipelines/zero_shot_classification.py
@@ -87,9 +87,9 @@ class ZeroShotClassificationPipeline(ChunkPipeline):
_load_feature_extractor = False
_load_tokenizer = True
- def __init__(self, args_parser=ZeroShotClassificationArgumentHandler(), *args, **kwargs):
+ def __init__(self, args_parser=ZeroShotClassificationArgumentHandler(), **kwargs):
self._args_parser = args_parser
- super().__init__(*args, **kwargs)
+ super().__init__(**kwargs)
if self.entailment_id == -1:
logger.warning(
"Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to "
diff --git a/src/transformers/processing_utils.py b/src/transformers/processing_utils.py
index 86cdb372034c..8b40b6535f1b 100644
--- a/src/transformers/processing_utils.py
+++ b/src/transformers/processing_utils.py
@@ -311,7 +311,7 @@ class AudioKwargs(TypedDict, total=False):
"""
sampling_rate: Optional[int]
- raw_speech: Optional[Union["np.ndarray", list[float], list["np.ndarray"], list[list[float]]]]
+ raw_speech: Optional[Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]]]
padding: Optional[Union[bool, str, PaddingStrategy]]
max_length: Optional[int]
truncation: Optional[bool]
@@ -963,6 +963,7 @@ def get_processor_dict(
local_files_only=local_files_only,
revision=revision,
cache_dir=cache_dir,
+ token=token,
):
additional_chat_template_files[template] = f"{CHAT_TEMPLATE_DIR}/{template}.jinja"
except EntryNotFoundError:
@@ -1267,7 +1268,7 @@ class MyProcessingKwargs(ProcessingKwargs, CommonKwargs, TextKwargs, ImagesKwarg
used_keys = set()
# get defaults from set model processor kwargs if they exist
- for modality in default_kwargs: # noqa: PLC0206
+ for modality in default_kwargs:
default_kwargs[modality] = ModelProcessorKwargs._defaults.get(modality, {}).copy()
# update defaults with arguments from tokenizer init
for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__:
diff --git a/src/transformers/pytorch_utils.py b/src/transformers/pytorch_utils.py
index 87136d079f10..b1f41117d4cf 100644
--- a/src/transformers/pytorch_utils.py
+++ b/src/transformers/pytorch_utils.py
@@ -50,7 +50,7 @@
_torch_distributed_available = torch.distributed.is_available()
-def softmax_backward_data(parent, grad_output, output, dim, self):
+def softmax_backward_data(parent, grad_output, output):
"""
A function that calls the internal `_softmax_backward_data` PyTorch method and that adjusts the arguments according
to the torch version detected.
@@ -58,7 +58,7 @@ def softmax_backward_data(parent, grad_output, output, dim, self):
from torch import _softmax_backward_data
- return _softmax_backward_data(grad_output, output, parent.dim, self.dtype)
+ return _softmax_backward_data(grad_output, output, parent.dim, output.dtype)
def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py
index 323faa9c17e2..b9dd7ae10f9e 100644
--- a/src/transformers/quantizers/base.py
+++ b/src/transformers/quantizers/base.py
@@ -128,27 +128,6 @@ def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> li
"""
return missing_keys
- def update_unexpected_keys(self, model, unexpected_keys: list[str], prefix: str) -> list[str]:
- """
- Override this method if you want to adjust the `unexpected_keys`.
-
- Args:
- unexpected_keys (`list[str]`, *optional*):
- The list of unexpected keys in the checkpoint compared to the state dict of the model
- """
- return unexpected_keys
-
- def update_missing_keys_after_loading(self, model, missing_keys: list[str], prefix: str) -> list[str]:
- """
- Override this method if you want to adjust the `missing_keys` after loading the model params,
- but before the model is post-processed.
-
- Args:
- missing_keys (`list[str]`, *optional*):
- The list of missing keys in the checkpoint compared to the state dict of the model
- """
- return missing_keys
-
def update_expected_keys(self, model, expected_keys: list[str], loaded_keys: list[str]) -> list[str]:
"""
Override this method if you want to adjust the `update_expected_keys`.
@@ -161,6 +140,9 @@ def update_expected_keys(self, model, expected_keys: list[str], loaded_keys: lis
"""
return expected_keys
+ def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
+ return unexpected_keys
+
def get_special_dtypes_update(self, model, dtype: "torch.dtype") -> dict[str, "torch.dtype"]:
"""
returns dtypes for modules that are not quantized - used for the computation of the device_map in case
@@ -182,25 +164,26 @@ def adjust_max_memory(self, max_memory: dict[str, Union[int, str]]) -> dict[str,
"""adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization"""
return max_memory
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def check_quantized_param(self, *args, **kwargs) -> bool:
+ """DEPRECATED -> remove in v5"""
+ logger.warning_once(
+ "`check_quantized_param` is deprecated in favor of `param_needs_quantization`, which is a much "
+ "more self.explanatory name for what the method achieves. It will be removed in v5"
+ )
+ return self.param_needs_quantization(*args, **kwargs)
+
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
"""
- checks if a loaded state_dict component is part of quantized param + some validation; only defined if
- requires_parameters_quantization == True for quantization methods that require to create a new parameters
- for quantization.
+ Check whether a given param needs quantization as defined by `create_quantized_param`.
"""
return False
- def create_quantized_param(self, *args, **kwargs) -> "torch.nn.Parameter":
+ def create_quantized_param(self, *args, **kwargs):
"""
- takes needed components from state_dict and creates quantized param; only applicable if
- requires_parameters_quantization == True
+ Take needed components from state_dict (those from which `param_needs_quantization` is True) and create
+ quantized param.
+ It usually also load the new param directly in the `model`.
+ Note: only applicable if requires_parameters_quantization == True.
"""
if not self.requires_parameters_quantization:
raise AttributeError(
@@ -342,6 +325,10 @@ def get_state_dict_and_metadata(self, model, safe_serialization=False):
"""Get state dict and metadata. Useful when we need to modify a bit the state dict due to quantization"""
return None, {}
+ def update_state_dict_with_metadata(self, state_dict, metadata):
+ """Update state dict with metadata. Default behaviour returns state_dict"""
+ return state_dict
+
@abstractmethod
def _process_model_before_weight_loading(self, model, **kwargs): ...
diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py
index 74879fa17ac4..ba240e572480 100644
--- a/src/transformers/quantizers/quantizer_bnb_4bit.py
+++ b/src/transformers/quantizers/quantizer_bnb_4bit.py
@@ -12,8 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
+from collections import defaultdict
from functools import cached_property
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
from packaging import version
@@ -67,6 +68,15 @@ def __init__(self, quantization_config, **kwargs):
if self.quantization_config.llm_int8_skip_modules is not None:
self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
+ # This describes the additional items that are saved on the state dict (on the params themselves)
+ self.bnb_keys = [
+ f"quant_state.bitsandbytes__{self.quantization_config.bnb_4bit_quant_type}",
+ "absmax",
+ "quant_map",
+ ]
+ if self.quantization_config.bnb_4bit_use_double_quant:
+ self.bnb_keys.extend(["nested_absmax", "nested_quant_map"])
+
def validate_environment(self, *args, **kwargs):
if not is_accelerate_available():
raise ImportError(
@@ -138,26 +148,17 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
"calculation. You may encounter unexpected behavior, or pass your own device map"
)
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
+ return [k for k in unexpected_keys if not any(k.endswith(x) for x in self.bnb_keys)]
+
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
import bitsandbytes as bnb
- module, tensor_name = get_module_from_name(model, param_name)
- if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Params4bit):
- # Add here check for loaded components' dtypes once serialization is implemented
+ # They are on the params themselves, so we cannot easily extract the module from the name
+ if any(param_name.endswith(x) for x in self.bnb_keys):
return True
- elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias":
- # bias could be loaded by regular set_module_tensor_to_device() from accelerate,
- # but it would wrongly use uninitialized weight there.
- return True
- else:
- return False
+ module, name = get_module_from_name(model, param_name)
+ return isinstance(module, bnb.nn.Linear4bit) and name != "bias"
def create_quantized_param(
self,
@@ -165,81 +166,51 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
- """
import bitsandbytes as bnb
+ is_quant_stat = any(param_name.endswith(x) for x in self.bnb_keys)
+ full_name = param_name
+ if is_quant_stat:
+ param_name = (
+ param_name.rsplit(".", 1)[0] if "quant_state." not in param_name else param_name.rsplit(".", 2)[0]
+ )
module, tensor_name = get_module_from_name(model, param_name)
- if tensor_name not in module._parameters:
- raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
-
- old_value = getattr(module, tensor_name)
-
# `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)).
if isinstance(target_device, int) and is_torch_npu_available():
target_device = f"npu:{target_device}"
- if tensor_name == "bias":
- if param_value is None:
- new_value = old_value.to(target_device)
- else:
- new_value = param_value.to(target_device)
-
- new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad)
- module._parameters[tensor_name] = new_value
- return
- if not isinstance(module._parameters[tensor_name], bnb.nn.Params4bit):
- raise ValueError("this function only loads `Linear4bit components`")
- if (
- old_value.device == torch.device("meta")
- and target_device not in ["meta", torch.device("meta")]
- and param_value is None
- ):
- raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.")
-
- # construct `new_value` for the module._parameters[tensor_name]:
+ # construct `new_value` for the module._parameters[tensor_name]
if self.pre_quantized:
- # 4bit loading. Collecting components for restoring quantized weight
- # This can be expanded to make a universal call for any quantized weight loading
-
- if not self.is_serializable:
- raise ValueError(
- "Detected int4 weights but the version of bitsandbytes is not compatible with int4 serialization. "
- "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
- )
-
- if (param_name + ".quant_state.bitsandbytes__fp4" not in state_dict) and (
- param_name + ".quant_state.bitsandbytes__nf4" not in state_dict
- ):
- raise ValueError(
- f"Supplied state dict for {param_name} does not contain `bitsandbytes__*` and possibly other `quantized_stats` components."
+ module_name = param_name.rsplit(".", 1)[0]
+ # Save the states for later quantization when they are all gathered
+ if not hasattr(self, "param_quant_stats"):
+ self.param_quant_stats = defaultdict(dict)
+ self.param_quant_stats[module_name].update({full_name: param_value})
+
+ # We are ready for quantization in this case (note, the +1 is for the weight itself)
+ if len(self.param_quant_stats[module_name]) == len(self.bnb_keys) + 1:
+ param_kwargs = {}
+ if self.is_bnb_supports_quant_storage_module:
+ param_kwargs["module"] = module
+
+ weight = self.param_quant_stats[module_name].pop(f"{module_name}.weight")
+ new_value = bnb.nn.Params4bit.from_prequantized(
+ data=weight,
+ quantized_stats=self.param_quant_stats[module_name],
+ requires_grad=False,
+ device=target_device,
+ **param_kwargs,
)
-
- quantized_stats = {}
- for k, v in state_dict.items():
- if param_name + "." in k:
- quantized_stats[k] = v
- if unexpected_keys is not None and k in unexpected_keys:
- unexpected_keys.remove(k)
-
- param_kwargs = {}
- if self.is_bnb_supports_quant_storage_module:
- param_kwargs["module"] = module
-
- new_value = bnb.nn.Params4bit.from_prequantized(
- data=param_value,
- quantized_stats=quantized_stats,
- requires_grad=False,
- device=target_device,
- **param_kwargs,
- )
+ # Set it
+ module._parameters[tensor_name] = new_value
+ # Delete the states
+ del self.param_quant_stats[module_name]
else:
new_value = param_value.to("cpu")
+ old_value = getattr(module, tensor_name)
# Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
# Since weights are saved in the correct "orientation", we skip transposing when loading.
@@ -247,9 +218,10 @@ def create_quantized_param(
new_value = new_value.T
kwargs = old_value.__dict__
+ kwargs.pop("_is_hf_initialized", None)
new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)
- module._parameters[tensor_name] = new_value
+ module._parameters[tensor_name] = new_value
# Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.adjust_max_memory
def adjust_max_memory(self, max_memory: dict[str, Union[int, str]]) -> dict[str, Union[int, str]]:
@@ -321,7 +293,6 @@ def _process_model_before_weight_loading(
model = replace_with_bnb_linear(
model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
)
- # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here
model.config.quantization_config = self.quantization_config
diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py
index 1d269765f57f..c7e3bb0c7af8 100644
--- a/src/transformers/quantizers/quantizer_bnb_8bit.py
+++ b/src/transformers/quantizers/quantizer_bnb_8bit.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
from packaging import version
@@ -164,27 +164,15 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
logger.info("target_dtype {target_dtype} is replaced by `torch.int8` for 8-bit BnB quantization")
return torch.int8
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
+ bnb_keys = ["SCB", "weight_format"]
+ return [k for k in unexpected_keys if not any(k.endswith(x) for x in bnb_keys)]
+
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
import bitsandbytes as bnb
- module, tensor_name = get_module_from_name(model, param_name)
- if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Int8Params):
- if self.pre_quantized:
- if param_name.replace("weight", "SCB") not in state_dict:
- raise ValueError("Missing quantization component `SCB`")
- if param_value.dtype != torch.int8:
- raise ValueError(
- f"Incompatible dtype `{param_value.dtype}` when loading 8-bit prequantized weight. Expected `torch.int8`."
- )
- return True
- return False
+ module, name = get_module_from_name(model, param_name)
+ return isinstance(module, bnb.nn.Linear8bitLt) and name != "bias"
def create_quantized_param(
self,
@@ -192,62 +180,40 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
- needs aux items from state dicts, if found - removes them from unexpected_keys
- """
import bitsandbytes as bnb
- fp16_statistics_key = param_name.replace("weight", "SCB")
- fp16_weights_format_key = param_name.replace("weight", "weight_format")
-
- fp16_statistics = state_dict.get(fp16_statistics_key)
- fp16_weights_format = state_dict.get(fp16_weights_format_key)
-
module, tensor_name = get_module_from_name(model, param_name)
- if tensor_name not in module._parameters:
- raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
-
- old_value = getattr(module, tensor_name)
- if not isinstance(module._parameters[tensor_name], bnb.nn.Int8Params):
- raise TypeError(f"Parameter `{tensor_name}` should only be a `bnb.nn.Int8Params` instance.")
- if (
- old_value.device == torch.device("meta")
- and target_device not in ["meta", torch.device("meta")]
- and param_value is None
- ):
- raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.")
-
- new_value = param_value.to("cpu")
if self.pre_quantized and not self.is_serializable():
raise ValueError(
"Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
"Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
)
+ # Those 2 can only happen when self.pre_quantized == True
+ if tensor_name == "SCB":
+ setattr(module.weight, "SCB", param_value.to(target_device))
+ return
+ # It's not used, but it's getting serialized for BC reason...
+ elif tensor_name == "weight_format":
+ return
# Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
# Since weights are saved in the correct "orientation", we skip transposing when loading.
- if issubclass(module.source_cls, Conv1D):
- if fp16_statistics is None:
- new_value = new_value.T
+ if issubclass(module.source_cls, Conv1D) and not self.pre_quantized:
+ param_value = param_value.T
+ old_value = getattr(module, tensor_name)
kwargs = old_value.__dict__
- new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(target_device)
-
+ kwargs.pop("_is_hf_initialized", None)
+ # Need to pop SCB and reset it because of bnb internals that modifies its value when switching devices ...
+ SCB = kwargs.pop("SCB", None)
+ new_value = bnb.nn.Int8Params(param_value.to("cpu"), requires_grad=False, **kwargs).to(target_device)
+ if SCB is not None:
+ setattr(new_value, "SCB", SCB)
+ # Set it to the module
module._parameters[tensor_name] = new_value
- if fp16_statistics is not None:
- setattr(module.weight, "SCB", fp16_statistics.to(target_device))
- if unexpected_keys is not None:
- unexpected_keys.remove(fp16_statistics_key)
-
- # We just need to pop the `weight_format` keys from the state dict to remove unneeded
- # messages. The correct format is correctly retrieved during the first forward pass.
- if fp16_weights_format is not None and unexpected_keys is not None:
- unexpected_keys.remove(fp16_weights_format_key)
def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
model.is_loaded_in_8bit = True
@@ -284,7 +250,6 @@ def _process_model_before_weight_loading(
model = replace_with_bnb_linear(
model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
)
- # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here
model.config.quantization_config = self.quantization_config
diff --git a/src/transformers/quantizers/quantizer_eetq.py b/src/transformers/quantizers/quantizer_eetq.py
index 00a8117be9d2..d62c7ff9e88e 100644
--- a/src/transformers/quantizers/quantizer_eetq.py
+++ b/src/transformers/quantizers/quantizer_eetq.py
@@ -11,7 +11,7 @@
# WITHOUT 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, Any, Optional
+from typing import TYPE_CHECKING, Optional
from .base import HfQuantizer
@@ -106,26 +106,15 @@ def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with EETQ.")
return dtype
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from eetq import EetqLinear
module, tensor_name = get_module_from_name(model, param_name)
if isinstance(module, EetqLinear):
if self.pre_quantized or tensor_name == "bias":
- if tensor_name == "weight" and param_value.dtype != torch.int8:
- raise ValueError("Expect quantized weights but got an unquantized weight")
return False
else:
- if tensor_name == "weight_scale":
- raise ValueError("Expect unquantized weights but got a quantized weight_scale")
return True
return False
@@ -135,17 +124,22 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- quantizes weights into qweight and weight_scales
- """
- from eetq import quantize_and_preprocess_weights
+ from eetq import EetqLinear, quantize_and_preprocess_weights
module, tensor_name = get_module_from_name(model, param_name)
new_value, weight_scale = quantize_and_preprocess_weights(param_value)
+ # Samity check
+ if isinstance(module, EetqLinear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.int8:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ else:
+ if tensor_name == "weight_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+
module._buffers[tensor_name] = new_value.to(target_device)
module.register("weight_scales", weight_scale.to(target_device))
diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py
index 0c1047f9503f..22c90aa446dd 100644
--- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py
+++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py
@@ -11,7 +11,7 @@
# WITHOUT 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, Any, Optional
+from typing import TYPE_CHECKING, Optional
from .base import HfQuantizer
@@ -105,33 +105,20 @@ def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
)
return dtype
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
module, tensor_name = get_module_from_name(model, param_name)
if isinstance(module, FbgemmFp8Linear):
if self.pre_quantized or tensor_name == "bias":
- if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
- raise ValueError("Expect quantized weights but got an unquantized weight")
return False
else:
- if tensor_name == "weight_scale":
- raise ValueError("Expect unquantized weights but got a quantized weight_scale")
return True
if isinstance(module, FbgemmFp8Llama4TextExperts):
if self.pre_quantized or tensor_name == "bias":
return False
else:
- if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale":
- raise ValueError("Expect unquantized weights but got a quantized weight_scale")
return True
return False
@@ -141,16 +128,25 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- Quantizes weights into weight and weight_scale
- """
-
- from ..integrations import FbgemmFp8Llama4TextExperts
+ from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
module, tensor_name = get_module_from_name(model, param_name)
+
+ # Sanity checks
+ if isinstance(module, FbgemmFp8Linear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ else:
+ if tensor_name == "weight_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+ if isinstance(module, FbgemmFp8Llama4TextExperts):
+ if not (self.pre_quantized or tensor_name == "bias"):
+ if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+
if isinstance(module, FbgemmFp8Llama4TextExperts):
if tensor_name == "gate_up_proj":
# Process each expert separately
@@ -194,8 +190,6 @@ def create_quantized_param(
module._parameters[tensor_name] = torch.nn.Parameter(new_value.to(target_device))
- if unexpected_keys is not None and param_name in unexpected_keys:
- unexpected_keys.remove(param_name)
del param_name
def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py
index dc30221b590e..4804f0d90469 100644
--- a/src/transformers/quantizers/quantizer_finegrained_fp8.py
+++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging
from .base import HfQuantizer
@@ -87,14 +87,21 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- Quantizes weights to FP8 format using Block-wise quantization
- """
+ from ..integrations.finegrained_fp8 import FP8Linear
from ..modeling_utils import _load_parameter_into_model
+ # Sanity checks
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module, FP8Linear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ else:
+ if tensor_name == "weight_scale_inv":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+
param_value = param_value.to(target_device)
# Get FP8 min/max values
@@ -135,26 +142,14 @@ def create_quantized_param(
_load_parameter_into_model(model, param_name, quantized_param)
_load_parameter_into_model(model, param_name.rsplit(".", 1)[0] + ".weight_scale_inv", scale)
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from ..integrations.finegrained_fp8 import FP8Linear
module, tensor_name = get_module_from_name(model, param_name)
-
if isinstance(module, FP8Linear):
if self.pre_quantized or tensor_name == "bias":
- if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
- raise ValueError("Expect quantized weights but got an unquantized weight")
return False
else:
- if tensor_name == "weight_scale_inv":
- raise ValueError("Expect unquantized weights but got a quantized weight_scale")
return True
return False
diff --git a/src/transformers/quantizers/quantizer_fp_quant.py b/src/transformers/quantizers/quantizer_fp_quant.py
index 4ac029ac7436..a7bc077776fe 100644
--- a/src/transformers/quantizers/quantizer_fp_quant.py
+++ b/src/transformers/quantizers/quantizer_fp_quant.py
@@ -11,7 +11,7 @@
# WITHOUT 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, Any, Optional
+from typing import TYPE_CHECKING, Optional
from .base import HfQuantizer
from .quantizers_utils import get_module_from_name
@@ -37,7 +37,7 @@ class FPQuantHfQuantizer(HfQuantizer):
requires_calibration = False
requires_parameters_quantization = True
- is_qat_trainable = False
+ is_qat_trainable = True
required_packages = ["fp_quant"]
def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
@@ -63,12 +63,16 @@ def validate_environment(self, device_map, **kwargs):
if not is_fp_quant_available():
raise ImportError("Using `fp_quant` quantization requires fp_quant: `pip install fp_quant`")
- if device_map is None:
+ if device_map is None and not self.quantization_config.pseudoquantization:
raise ValueError(
"You are attempting to load a FPQuant model without setting device_map."
" Please set device_map comprised of 'cuda' devices."
)
- elif isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+ elif (
+ isinstance(device_map, dict)
+ and ("cpu" in device_map.values() or "disk" in device_map.values())
+ and not self.quantization_config.pseudoquantization
+ ):
raise ValueError(
"You are attempting to load a FPQuant model with a device_map that contains a CPU or disk device."
" This is not supported. Please remove the CPU or disk device from the device_map."
@@ -89,8 +93,7 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
module, _ = get_module_from_name(model, param_name)
@@ -122,9 +125,6 @@ def create_quantized_param(
# Let pre-forward handle the quantization and set None where necessary
module.pre_forward()
- if unexpected_keys is not None and param_name in unexpected_keys:
- unexpected_keys.remove(param_name)
-
def _process_model_before_weight_loading(
self,
model: "PreTrainedModel",
@@ -158,19 +158,17 @@ def should_exclude(key: str) -> bool:
@property
def is_trainable(self, model: Optional["PreTrainedModel"] = None):
- return False
+ trainable = self.quantization_config.store_master_weights
+ if not trainable:
+ logger.warning(
+ "You are attempting to train a model with FPQuant quantization. This is only supported when `store_master_weights=True`. Please set `store_master_weights=True` to train the model."
+ )
+ return trainable
def is_serializable(self, safe_serialization=None):
return True
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from fp_quant import FPQuantLinear
module, tensor_name = get_module_from_name(model, param_name)
diff --git a/src/transformers/quantizers/quantizer_higgs.py b/src/transformers/quantizers/quantizer_higgs.py
index cca104df7db4..41e2d86cf1ec 100644
--- a/src/transformers/quantizers/quantizer_higgs.py
+++ b/src/transformers/quantizers/quantizer_higgs.py
@@ -11,7 +11,7 @@
# WITHOUT 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, Any, Optional
+from typing import TYPE_CHECKING, Optional
from ..utils.logging import tqdm
from .base import HfQuantizer
@@ -87,14 +87,10 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
from ..integrations import quantize_with_higgs
- """
- Quantizes weights into weight and weight_scale
- """
flute_dict = quantize_with_higgs(
param_value.to(target_device),
self.quantization_config.bits,
@@ -117,9 +113,6 @@ def create_quantized_param(
else:
raise ValueError(f"Unexpected key {key} in module {module}")
- if unexpected_keys is not None and param_name in unexpected_keys:
- unexpected_keys.remove(param_name)
-
def _process_model_before_weight_loading(
self,
model: "PreTrainedModel",
@@ -184,18 +177,11 @@ def is_trainable(self) -> bool:
def is_serializable(self, safe_serialization=None):
return True
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from ..integrations import HiggsLinear
module, tensor_name = get_module_from_name(model, param_name)
- if isinstance(module, HiggsLinear) and tensor_name == "weight" and param_value.dtype != torch.int16:
+ if isinstance(module, HiggsLinear) and tensor_name == "weight":
# Only quantize weights of HiggsLinear modules that are not already quantized
return True
else:
diff --git a/src/transformers/quantizers/quantizer_hqq.py b/src/transformers/quantizers/quantizer_hqq.py
index fa1d276c6a1a..22c55000b2d0 100755
--- a/src/transformers/quantizers/quantizer_hqq.py
+++ b/src/transformers/quantizers/quantizer_hqq.py
@@ -12,10 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import TYPE_CHECKING, Any
+from collections import defaultdict
+from typing import TYPE_CHECKING
from ..integrations import prepare_for_hqq_linear
-from ..utils import is_accelerate_available, is_hqq_available, is_torch_available, logging
+from ..utils import is_hqq_available, is_torch_available, logging
from .base import HfQuantizer
from .quantizers_utils import get_module_from_name
@@ -24,29 +25,28 @@
from ..modeling_utils import PreTrainedModel
-if is_accelerate_available():
- from accelerate.hooks import remove_hook_from_module
-
if is_torch_available():
import torch
-logger = logging.get_logger(__name__)
+if is_hqq_available():
+ from hqq.core.quantize import HQQLinear
+ # This is a compatibility hack. HQQ-quantized linear layers do not have a `weight` attribute,
+ # but some models attempt to access `weight.dtype` during the forward pass. To prevent runtime errors,
+ # we patch HQQLinear with a dummy `weight` property that returns an empty tensor with the correct dtype and device.
+ @property
+ def weight(self):
+ return torch.empty(0, dtype=self.compute_dtype, device=self.device)
-# Finds the parent of a node module named "name"
-def find_parent(model, name):
- module_tree = name.split(".")[:-1]
- parent = model
- for m in module_tree:
- parent = parent._modules[m]
- return parent
+ HQQLinear.weight = weight
+
+logger = logging.get_logger(__name__)
class HqqHfQuantizer(HfQuantizer):
"""
HQQ quantizer base HF class.
nn.Linear modules are first tagged with quant_config in _process_model_before_weight_loading().
- The actual quantization and offloading to the GPU is done in check_quantized_param().
"""
use_keep_in_fp32_modules = False
@@ -55,15 +55,15 @@ class HqqHfQuantizer(HfQuantizer):
required_packages = ["hqq"]
def __init__(self, quantization_config, **kwargs):
- super().__init__(quantization_config, **kwargs)
- self.dtype = None
- self.using_multi_gpu = False
-
- def validate_environment(self, *args, **kwargs):
- if not (is_hqq_available()):
+ if not is_hqq_available():
raise ImportError(
"A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`."
)
+ super().__init__(quantization_config, **kwargs)
+ self.dtype = None
+ self.using_multi_gpu = False
+ # Keys that are serialized specifically by hqq
+ self.hqq_keys = HQQLinear(None, None).state_dict_keys() - {"bias"}
if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
raise ValueError(
@@ -111,75 +111,56 @@ def _find_hqq_quantizable_layers(model, layers):
_find_hqq_quantizable_layers(module, layers)
new_keys = set(expected_keys)
- if is_hqq_available():
- from hqq.core.quantize import HQQLinear
-
- # Name modules
- for name, module in model.named_modules():
- module.name = name
-
- # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params
- _valid_modules = set()
- _find_hqq_quantizable_layers(model, _valid_modules)
-
- # Remove skipped modules
- _skipped_modules = set()
- for _module in _valid_modules:
- for _skip_module in model.config.quantization_config["skip_modules"]:
- if _skip_module in _module:
- _skipped_modules.add(_module)
- _valid_modules -= _skipped_modules
-
- # Append new expected layers based on _ref_keys
- _ref_keys = HQQLinear(
- linear_layer=None,
- quant_config=None,
- compute_dtype=torch.float16,
- device="cpu",
- del_orig=False,
- ).state_dict_keys() - {"bias"}
-
- # Clean-up
- _rm_keys = set()
- for key in new_keys:
- if any(_module in key for _module in _valid_modules):
- _rm_keys.add(key)
- new_keys -= _rm_keys
- # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear
-
- # Re-populate Linear/HQQLinear
- for _module in _valid_modules:
- if _module + ".weight" in loaded_keys:
- new_keys.add(_module + ".weight")
- else:
- new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})
- if _module + ".bias" in loaded_keys:
- new_keys.add(_module + ".bias")
- return list(new_keys)
+ # Name modules
+ for name, module in model.named_modules():
+ module.name = name
+
+ # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params
+ _valid_modules = set()
+ _find_hqq_quantizable_layers(model, _valid_modules)
+
+ # Remove skipped modules
+ _skipped_modules = set()
+ for _module in _valid_modules:
+ for _skip_module in model.config.quantization_config["skip_modules"]:
+ if _skip_module in _module:
+ _skipped_modules.add(_module)
+ _valid_modules -= _skipped_modules
+
+ # Append new expected layers based on _ref_keys
+ _ref_keys = HQQLinear(
+ linear_layer=None,
+ quant_config=None,
+ compute_dtype=torch.float16,
+ device="cpu",
+ del_orig=False,
+ ).state_dict_keys() - {"bias"}
+
+ # Clean-up
+ _rm_keys = set()
+ for key in new_keys:
+ if any(_module in key for _module in _valid_modules):
+ _rm_keys.add(key)
+ new_keys -= _rm_keys
+ # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear
+
+ # Re-populate Linear/HQQLinear
+ for _module in _valid_modules:
+ if _module + ".weight" in loaded_keys:
+ new_keys.add(_module + ".weight")
+ else:
+ new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})
+ if _module + ".bias" in loaded_keys:
+ new_keys.add(_module + ".bias")
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
- if is_hqq_available():
- from hqq.core.quantize import HQQLinear
- module, tensor_name = get_module_from_name(model, param_name)
+ return list(new_keys)
- if self.pre_quantized:
- return (isinstance(module, (torch.nn.Linear, HQQLinear))) and tensor_name != "weight"
- else:
- return (
- isinstance(module, torch.nn.Linear)
- and tensor_name == "weight"
- # bias doesn't need to be quantized, we use this as a workaround to avoid loading bias into HQQLinear assuming it was loaded
- # in the state_dict directly with the weight because hqq overwrote load_state_dict for this layer
- or (isinstance(module, HQQLinear) and tensor_name == "bias")
- )
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
+ module, _ = get_module_from_name(model, param_name)
+ # Since we do not prepare the modules in advance, we need every param of the Linear layer to go through
+ # `create_quantized_param`, even when `self.is_quantized == True`
+ return isinstance(module, torch.nn.Linear)
def create_quantized_param(
self,
@@ -187,48 +168,33 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: list[str],
+ **kwargs,
):
- """
- Each nn.Linear layer is processed here.
- We first check if the corresponding module state_dict contains already HQQ quantized parameters.
- If not, we create a temp linear layer with the module state_dict params and use it for quantization
- """
-
- if is_hqq_available():
- from hqq.core.quantize import HQQLinear
-
- # TODO: This is a compatibility hack. HQQ-quantized linear layers do not have a `weight` attribute,
- # but some models attempt to access `weight.dtype` during the forward pass. To prevent runtime errors,
- # we patch HQQLinear with a dummy `weight` property that returns an empty tensor with the correct dtype and device.
- @property
- def weight(_self: HQQLinear):
- return torch.empty(0, dtype=_self.compute_dtype, device=_self.device)
-
- HQQLinear.weight = weight
-
module, tensor_name = get_module_from_name(model, param_name)
- layer_name = ".".join(param_name.split(".")[:-1])
- parent_module = find_parent(model, layer_name)
- node = layer_name.split(".")[-1]
+ module_name = param_name.rsplit(".", 1)[0]
+ parent_module, node = get_module_from_name(model, module_name)
- if tensor_name == "bias":
- # this should already be set
- return
+ quant_config = model.config.quantization_config["quant_config"]
+ skip_modules = model.config.quantization_config["skip_modules"]
- # set module state_dict
- module_state_dict = {}
- for k, v in state_dict.items():
- if layer_name + "." in k:
- module_state_dict[k.split(".")[-1]] = v
- if unexpected_keys is not None and k in unexpected_keys:
- unexpected_keys.remove(k)
+ # In this case we do not quantize this layer (it's explicitly skipped) -> simply load param
+ if any(skip_module in module.name for skip_module in skip_modules):
+ module.load_state_dict(
+ {tensor_name: param_value.to(device=target_device, dtype=self.dtype)}, strict=False, assign=True
+ )
+ return
+ # We need this hack as the model is not pre-prepared as an empty skeleton on meta device
if self.pre_quantized:
- if isinstance(module, HQQLinear):
- return
- else:
+ # Save them for later
+ if not hasattr(self, "hqq_params"):
+ self.hqq_params = defaultdict(dict)
+ self.hqq_params[module_name].update({tensor_name: param_value})
+ hqq_params = self.hqq_params[module_name]
+
+ # If they are all present and saved, make it a HQQLinear layer! (we cannot do it param after param because
+ # hqq does not support it...)
+ if all(k in hqq_params for k in self.hqq_keys) and ("bias" in hqq_params or module.bias is None):
hqq_layer = HQQLinear(
linear_layer=None,
quant_config=None,
@@ -236,43 +202,32 @@ def weight(_self: HQQLinear):
device=target_device,
del_orig=False,
)
+ hqq_layer.load_state_dict(hqq_params)
- hqq_layer.load_state_dict(module_state_dict)
-
- if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
- hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+ if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
+ hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+ if self.using_multi_gpu:
+ hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
- if self.using_multi_gpu:
- hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
+ setattr(parent_module, node, hqq_layer)
+ del self.hqq_params[module_name], module
+ return
- setattr(parent_module, node, hqq_layer)
+ # Load param in the module (without caring about device or dtype, it will be changed later)
+ module.load_state_dict({tensor_name: param_value}, strict=False, assign=True)
- # cleanup
- del module.__dict__, module
- torch.cuda.empty_cache()
- return
+ # If both the weight and bias have already been loaded, time to quantize!
+ module_is_ready = module.weight.device.type != "meta" and (
+ module.bias is None or module.bias.device.type != "meta"
+ )
- # Step 1: populate module with weight/bias from module state dict
- for key, tensor in module_state_dict.items():
- setattr(module, key, torch.nn.Parameter(tensor))
+ if module_is_ready:
+ module_tag = ".".join(module.name.split(".")[-2:])
+ if "weight_quant_params" in quant_config:
+ module_quant_config = quant_config
+ elif module_tag in quant_config:
+ module_quant_config = quant_config[module_tag]
- # Step 2: Replace module with either HQQLinear or move it to device. We do this via setattr on the parent as doing on it on the module
- # directly doesn't work.
- quant_config = model.config.quantization_config["quant_config"]
- skip_modules = model.config.quantization_config["skip_modules"]
- module_tag = ".".join(module.name.split(".")[-2:])
- module_quant_config = None
- if "weight_quant_params" in quant_config:
- module_quant_config = quant_config
- elif module_tag in quant_config:
- module_quant_config = quant_config[module_tag]
-
- for skip_module in skip_modules:
- if skip_module in module.name:
- module_quant_config = None
- break
-
- if module_quant_config is not None:
hqq_layer = HQQLinear(
module,
quant_config=module_quant_config,
@@ -289,16 +244,7 @@ def weight(_self: HQQLinear):
setattr(parent_module, node, hqq_layer)
- else:
- module = module.to(dtype=self.dtype, device=target_device)
- setattr(parent_module, node, module)
-
- torch.cuda.empty_cache()
-
- # Remove accelerate hook and uses a simpler forward pass. Otherwise, this breaks with multi-gpu
def _patch_layer_for_multigpu(self, hqq_layer):
- hqq_layer = remove_hook_from_module(hqq_layer)
-
def forward_with_device(self, x):
out = torch.matmul(x.to(self.device), self.dequantize().t())
if self.bias is not None:
diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py
index d0d370a11df6..04cf8ec56c96 100644
--- a/src/transformers/quantizers/quantizer_mxfp4.py
+++ b/src/transformers/quantizers/quantizer_mxfp4.py
@@ -11,7 +11,7 @@
# WITHOUT 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, Any, Optional
+from typing import TYPE_CHECKING, Optional
from .base import HfQuantizer
@@ -72,7 +72,7 @@ def validate_environment(self, *args, **kwargs):
if self.quantization_config.dequantize:
return
- if not torch.cuda.is_available():
+ if not (torch.cuda.is_available() or torch.xpu.is_available()):
if self.pre_quantized:
logger.warning_once(
"Using MXFP4 quantized models requires a GPU, we will default to dequantizing the model to bf16"
@@ -85,15 +85,19 @@ def validate_environment(self, *args, **kwargs):
if not is_accelerate_available():
raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`")
- compute_capability = torch.cuda.get_device_capability()
- gpu_is_supported = compute_capability >= (7, 5)
- kernels_available = is_triton_available("3.4.0") and is_kernels_available()
+ if torch.xpu.is_available():
+ gpu_is_supported = True
+ kernels_available = is_triton_available("3.5.0") and is_kernels_available()
+ else:
+ compute_capability = torch.cuda.get_device_capability()
+ gpu_is_supported = compute_capability >= (7, 5)
+ kernels_available = is_triton_available("3.4.0") and is_kernels_available()
if self.pre_quantized:
# On unsupported GPUs or without kernels, we will dequantize the model to bf16
if not gpu_is_supported:
logger.warning_once(
- "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200). "
+ "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
"We will default to dequantizing the model to bf16."
)
self.quantization_config.dequantize = True
@@ -101,18 +105,20 @@ def validate_environment(self, *args, **kwargs):
if not kernels_available:
logger.warning_once(
- "MXFP4 quantization requires triton >= 3.4.0 and kernels installed, we will default to dequantizing the model to bf16"
+ "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU requires Triton >= 3.5.0, we will default to dequantizing the model to bf16"
)
self.quantization_config.dequantize = True
return
elif not gpu_is_supported:
# we can't quantize the model in this case so we raise an error
raise ValueError(
- "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200)"
+ "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
)
elif not kernels_available:
# we can't quantize the model in this case so we raise an error
- raise ValueError("MXFP4 quantization requires triton >= 3.4.0 and kernels installed")
+ raise ValueError(
+ "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU requires Triton >= 3.5.0"
+ )
if not self.pre_quantized:
self._lazy_import_kernels()
@@ -120,8 +126,8 @@ def validate_environment(self, *args, **kwargs):
device_map = kwargs.get("device_map")
if device_map is None:
logger.warning_once(
- "You have loaded an FP4 model on CPU and have a CUDA device available, make sure to set "
- "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. "
+ "You have loaded an FP4 model on CPU and have a CUDA/XPU device available, make sure to set "
+ "your model on a GPU/XPU device in order to run your model. To remove this warning, pass device_map = 'cuda' or device_map = 'xpu'. "
)
elif device_map is not None:
if (
@@ -147,14 +153,7 @@ def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
)
return dtype
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
from ..integrations import Mxfp4GptOssExperts
from ..models.gpt_oss.modeling_gpt_oss import GptOssExperts
@@ -177,8 +176,6 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
**kwargs,
):
from ..integrations import (
@@ -264,6 +261,8 @@ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs
# clean cache due to triton ops
if torch.cuda.is_available():
torch.cuda.empty_cache()
+ elif torch.xpu.is_available():
+ torch.xpu.empty_cache()
def update_expected_keys(self, model: "PreTrainedModel", expected_keys: list[str], checkpoint_keys: list[str]):
# Replace expected_keys for experts' gate_up_proj and down_proj with their _blocks and _scales variants
@@ -379,7 +378,7 @@ def update_param_name(self, param_name: str) -> str:
return param_name.replace("down_proj", "down_proj_blocks")
return param_name
- def get_state_dict_and_metadata(self, model):
+ def get_state_dict_and_metadata(self, model, safe_serialization: bool = False):
from ..integrations import Mxfp4GptOssExperts
state_dict = model.state_dict()
diff --git a/src/transformers/quantizers/quantizer_quanto.py b/src/transformers/quantizers/quantizer_quanto.py
index 43e084891469..451179aaf723 100644
--- a/src/transformers/quantizers/quantizer_quanto.py
+++ b/src/transformers/quantizers/quantizer_quanto.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
from packaging import version
@@ -103,29 +103,10 @@ def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> li
not_missing_keys.append(missing)
return [k for k in missing_keys if k not in not_missing_keys]
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
- """
- Check if a parameter needs to be quantized.
- """
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
if is_optimum_quanto_available():
from optimum.quanto import QModuleMixin
- device_map = kwargs.get("device_map")
- param_device = kwargs.get("param_device")
- # we don't quantize the model if the module is going to be offloaded to the cpu
- if device_map is not None and param_device is not None:
- device_map_values = set(device_map.values())
- if param_device == "cpu" and len(device_map_values) > 1:
- if not (device_map_values == {"cpu"} or device_map_values == {"cpu", "disk"}):
- return False
-
module, tensor_name = get_module_from_name(model, param_name)
# We only quantize the weights and the bias is not quantized.
if isinstance(module, QModuleMixin) and "weight" in tensor_name:
@@ -144,15 +125,11 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- *args,
**kwargs,
):
- """
- Create the quantized parameter by calling .freeze() after setting it to the module.
- """
- from accelerate.utils import set_module_tensor_to_device
+ from ..modeling_utils import _load_parameter_into_model
- set_module_tensor_to_device(model, param_name, target_device, param_value)
+ _load_parameter_into_model(model, param_name, param_value.to(target_device))
module, _ = get_module_from_name(model, param_name)
module.freeze()
module.weight.requires_grad = False
diff --git a/src/transformers/quantizers/quantizer_quark.py b/src/transformers/quantizers/quantizer_quark.py
index 0a14a01ee450..8ed6249bf5b9 100644
--- a/src/transformers/quantizers/quantizer_quark.py
+++ b/src/transformers/quantizers/quantizer_quark.py
@@ -13,23 +13,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
-from ..file_utils import is_torch_available
from .base import HfQuantizer
if TYPE_CHECKING:
from ..modeling_utils import PreTrainedModel
- if is_torch_available():
- import torch
+from ..utils import is_quark_available, logging
-from ..utils import is_accelerate_available, is_quark_available, logging
-
-
-if is_accelerate_available():
- from accelerate.utils import set_module_tensor_to_device
logger = logging.get_logger(__name__)
@@ -82,25 +75,18 @@ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwarg
return model
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
return True
- def create_quantized_param(
- self, model, param, param_name, param_device, state_dict, unexpected_keys
- ) -> "torch.nn.Parameter":
+ def create_quantized_param(self, model, param, param_name, param_device, **kwargs):
+ from ..modeling_utils import _load_parameter_into_model
+
postfix = param_name.split(".")[-1]
if postfix in CHECKPOINT_KEYS:
param_name = param_name.replace(postfix, CHECKPOINT_KEYS[postfix])
- set_module_tensor_to_device(model, param_name, param_device, value=param)
+ _load_parameter_into_model(model, param_name, param.to(param_device))
def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
return model
diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py
index cba023a7d811..6538d0c6122b 100644
--- a/src/transformers/quantizers/quantizer_torchao.py
+++ b/src/transformers/quantizers/quantizer_torchao.py
@@ -14,6 +14,7 @@
import importlib
import re
import types
+from collections import defaultdict
from typing import TYPE_CHECKING, Optional, Union
from packaging import version
@@ -25,16 +26,26 @@
if TYPE_CHECKING:
from ..modeling_utils import PreTrainedModel
-from typing import Any
+from safetensors import safe_open
from ..utils import is_torch_available, is_torchao_available, logging
-from ..utils.quantization_config import TorchAoConfig
if is_torch_available():
import torch
import torch.nn as nn
+if is_torchao_available():
+ import torchao
+
+ if version.parse(importlib.metadata.version("torchao")) >= version.parse("0.14.0"):
+ from torchao.prototype.safetensors.safetensors_support import (
+ flatten_tensor_state_dict,
+ unflatten_tensor_state_dict,
+ )
+ from torchao.prototype.safetensors.safetensors_utils import is_metadata_torchao
+
+
logger = logging.get_logger(__name__)
@@ -53,15 +64,6 @@ def fuzzy_match_size(config_name: str) -> Optional[str]:
return None
-# Finds the parent of a node module named "name"
-def find_parent(model, name):
- module_tree = name.split(".")[:-1]
- parent = model
- for m in module_tree:
- parent = parent._modules[m]
- return parent
-
-
def _quantization_type(weight):
from torchao.dtypes import AffineQuantizedTensor
from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor
@@ -81,6 +83,15 @@ def _linear_extra_repr(self):
return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}"
+if is_torchao_available():
+ SUPPORTED_SAFE_SERIALIZATION_CONFIGS = [
+ torchao.quantization.Float8WeightOnlyConfig,
+ torchao.quantization.Float8DynamicActivationFloat8WeightConfig,
+ ]
+
+ TORCHAO_VERSION = version.parse(importlib.metadata.version("torchao"))
+
+
class TorchAoHfQuantizer(HfQuantizer):
"""
Quantizer for torchao: https://github.com/pytorch/ao/
@@ -93,6 +104,20 @@ class TorchAoHfQuantizer(HfQuantizer):
def __init__(self, quantization_config, **kwargs):
super().__init__(quantization_config, **kwargs)
+ if isinstance(self.quantization_config.quant_type, str):
+ is_int_4 = "int4" in self.quantization_config.quant_type
+ else:
+ config_name = self.quantization_config.quant_type.__class__.__name__
+ is_int_4 = fuzzy_match_size(config_name) == "4"
+
+ # TODO: better way to get the serialized key names? Hard to read from torchao codebase
+ if is_int_4:
+ self.weight_ao_keys = ["qdata", "scale", "zero_point"]
+ else:
+ self.weight_ao_keys = ["qdata", "scale"]
+ # Instead of serializing the simple torch.Tensor like usual, torchao adds a `:_data` suffix so we need this
+ self.full_ao_keys = self.weight_ao_keys + ["_data"]
+
def validate_environment(self, *args, **kwargs):
if not is_torchao_available():
raise ImportError("Loading an torchao quantized model requires torchao library (`pip install torchao`)")
@@ -137,6 +162,21 @@ def update_dtype(self, dtype):
dtype = torch.float32
return dtype
+ def get_state_dict_and_metadata(self, model, safe_serialization: Optional[bool] = False):
+ """
+ If the model is safe serializable, we flatten the state dict of tensor subclasses so that it is compatible with
+ the safetensors format.
+ """
+ if type(self.quantization_config.quant_type) in SUPPORTED_SAFE_SERIALIZATION_CONFIGS and safe_serialization:
+ if TORCHAO_VERSION >= version.parse("0.14.0"):
+ return flatten_tensor_state_dict(model.state_dict())
+ else:
+ raise RuntimeError(
+ f"In order to use safetensors with torchao, please use torchao version >= 0.14.0. Current version: {TORCHAO_VERSION}"
+ )
+ else:
+ return None, {}
+
def adjust_target_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"):
from accelerate.utils import CustomDtype
@@ -194,31 +234,25 @@ def _process_model_before_weight_loading(
]
return
- def check_quantized_param(
- self,
- model: "PreTrainedModel",
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ) -> bool:
+ def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
+ return [k for k in unexpected_keys if not any(k.endswith(x) for x in self.full_ao_keys)]
+
+ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
if self.quantization_config.quant_type == "autoquant":
return False
- param_device = kwargs.pop("param_device", None)
# check if the param_name is not in self.modules_to_not_convert
- if any((key + "." in param_name) or (key == param_name) for key in self.modules_to_not_convert):
- return False
- elif param_device == "cpu" and self.offload:
- # We don't quantize weights that we offload
+ if any(key + "." in param_name or key == param_name for key in self.modules_to_not_convert):
return False
+ elif any(param_name.endswith(f":{x}") for x in self.full_ao_keys):
+ return True
else:
# we only quantize the weight of nn.Linear and nn.Embedding
module, tensor_name = get_module_from_name(model, param_name)
_QUANTIZABLE = [torch.nn.Linear]
if self.quantization_config.include_input_output_embeddings:
_QUANTIZABLE.append(torch.nn.Embedding)
- return isinstance(module, tuple(_QUANTIZABLE)) and (tensor_name == "weight")
+ return isinstance(module, tuple(_QUANTIZABLE)) and tensor_name == "weight"
def create_quantized_param(
self,
@@ -226,30 +260,56 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: list[str],
+ **kwargs,
):
"""
Each nn.Linear layer that needs to be quantized is processed here.
First, we set the value the weight tensor, then we move it to the target device. Finally, we quantize the module.
"""
- if self.quantization_config.quant_type == "autoquant":
- return
-
from torchao.quantization import quantize_
+ full_name = param_name
+ # Those are the pre quantized weights
+ if ":" in param_name:
+ param_name = param_name.rsplit(":", 1)[0]
module, tensor_name = get_module_from_name(model, param_name)
+
if self.pre_quantized:
- module._parameters[tensor_name] = torch.nn.Parameter(
- param_value.to(device=target_device), requires_grad=param_value.requires_grad
- )
+ # If it's a bias, no need to do anything special (except removing the ":_data" part of the key, but was
+ # already done) - if it's unsafe-serialized (i.e. not safetensors), not need for anything either
+ is_unsafe_serialization = ":" not in full_name
+ if tensor_name == "bias" or is_unsafe_serialization:
+ module._parameters[tensor_name] = torch.nn.Parameter(
+ param_value.to(target_device), requires_grad=param_value.requires_grad
+ )
+ return
+ # Sanity check for the new serialization format
+ elif not (TORCHAO_VERSION >= version.parse("0.14.0") and is_metadata_torchao(self.metadata)):
+ raise ValueError("To use `safetensors` serialization, you should have `torchao>=0.14.0` installed")
+
+ # Save the states for later quantization when they are all gathered
+ if not hasattr(self, "ao_params"):
+ self.ao_params = defaultdict(dict)
+ self.ao_params[param_name].update({full_name: param_value})
+
+ # We are ready for quantization in this case (we retrieved all the needed keys)
+ if len(self.ao_params[param_name]) == len(self.weight_ao_keys):
+ new_param = unflatten_tensor_state_dict(self.ao_params[param_name], self.metadata)[param_name]
+ # Set it
+ module._parameters[tensor_name] = torch.nn.Parameter(
+ new_param.to(target_device), requires_grad=new_param.requires_grad
+ )
+
+ # Free memory
+ del self.ao_params[param_name]
+
+ # Add repr to the module
if isinstance(module, nn.Linear):
module.extra_repr = types.MethodType(_linear_extra_repr, module)
else:
- assert isinstance(self.quantization_config, TorchAoConfig)
module._parameters[tensor_name] = torch.nn.Parameter(
param_value, requires_grad=param_value.requires_grad
- ).to(device=target_device)
+ ).to(target_device)
# if we are quantizing tied parameters, to avoid tying the quantized weights
# the correct order to do it is
# 1. load the weight to model
@@ -297,10 +357,17 @@ def _process_model_after_weight_loading(self, model, **kwargs):
def is_serializable(self, safe_serialization=None) -> bool:
if safe_serialization:
- logger.warning(
- "torchao quantized model does not support safe serialization, please set `safe_serialization` to False"
- )
- return False
+ _is_torchao_serializable = type(
+ self.quantization_config.quant_type
+ ) in SUPPORTED_SAFE_SERIALIZATION_CONFIGS and TORCHAO_VERSION >= version.parse("0.14.0")
+ if not _is_torchao_serializable:
+ logger.warning(
+ f"torchao quantized model only supports safe serialization for {SUPPORTED_SAFE_SERIALIZATION_CONFIGS}, \
+ and torchao version >= 0.14.0, please set `safe_serialization` to False for \
+ {type(self.quantization_config.quant_type)} and {TORCHAO_VERSION}."
+ )
+ return _is_torchao_serializable
+
_is_torchao_serializable = version.parse(importlib.metadata.version("huggingface_hub")) >= version.parse(
"0.25.0"
)
@@ -364,3 +431,13 @@ def is_trainable(self) -> bool:
@property
def is_compileable(self) -> bool:
return True
+
+ def set_metadata(self, checkpoint_files: list[str]):
+ if checkpoint_files[0].endswith(".safetensors"):
+ metadata = {}
+ for checkpoint in checkpoint_files:
+ with safe_open(checkpoint, framework="pt") as f:
+ metadata_ = f.metadata() or {}
+ metadata.update(metadata_)
+ # Save it
+ self.metadata = metadata
diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py
index d8ec62124556..499716789b9b 100644
--- a/src/transformers/testing_utils.py
+++ b/src/transformers/testing_utils.py
@@ -208,6 +208,17 @@
# Not critical, only usable on the sandboxed CI instance.
TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL"
+
+# Used in CausalLMModelTester (and related classes/methods) to infer the common model classes from the base model class
+_COMMON_MODEL_NAMES_MAP = {
+ "config_class": "Config",
+ "causal_lm_class": "ForCausalLM",
+ "question_answering_class": "ForQuestionAnswering",
+ "sequence_classification_class": "ForSequenceClassification",
+ "token_classification_class": "ForTokenClassification",
+}
+
+
if is_torch_available():
import torch
@@ -1629,7 +1640,7 @@ def evaluate_side_effect_factory(
# final message
# it can handle a single string or a multiline buffer
def apply_print_resets(buf):
- return re.sub(r"^.*\r", "", buf, 0, re.M)
+ return re.sub(r"^.*\r", "", buf, 0, re.MULTILINE)
def assert_screenout(out, what):
@@ -1638,58 +1649,6 @@ def assert_screenout(out, what):
assert match_str != -1, f"expecting to find {what} in output: f{out_pr}"
-def set_model_tester_for_less_flaky_test(test_case):
- # NOTE: this function edits the config object, which may lead to hard-to-debug side-effects. Use with caution.
- # Do not use in tests/models where objects behave very differently based on the config's hidden layer settings
- # (e.g. KV caches, sliding window attention, ...)
-
- # TODO (if possible): Avoid exceptional cases
- exceptional_classes = [
- "ZambaModelTester",
- "Zamba2ModelTester",
- "RwkvModelTester",
- "AriaVisionText2TextModelTester",
- "GPTNeoModelTester",
- "DPTModelTester",
- "Qwen3NextModelTester",
- ]
- if test_case.model_tester.__class__.__name__ in exceptional_classes:
- return
-
- target_num_hidden_layers = 1
- if hasattr(test_case.model_tester, "out_features") or hasattr(test_case.model_tester, "out_indices"):
- target_num_hidden_layers = None
-
- if hasattr(test_case.model_tester, "num_hidden_layers") and target_num_hidden_layers is not None:
- test_case.model_tester.num_hidden_layers = target_num_hidden_layers
- if (
- hasattr(test_case.model_tester, "vision_config")
- and "num_hidden_layers" in test_case.model_tester.vision_config
- and target_num_hidden_layers is not None
- ):
- test_case.model_tester.vision_config = copy.deepcopy(test_case.model_tester.vision_config)
- if isinstance(test_case.model_tester.vision_config, dict):
- test_case.model_tester.vision_config["num_hidden_layers"] = 1
- else:
- test_case.model_tester.vision_config.num_hidden_layers = 1
- if (
- hasattr(test_case.model_tester, "text_config")
- and "num_hidden_layers" in test_case.model_tester.text_config
- and target_num_hidden_layers is not None
- ):
- test_case.model_tester.text_config = copy.deepcopy(test_case.model_tester.text_config)
- if isinstance(test_case.model_tester.text_config, dict):
- test_case.model_tester.text_config["num_hidden_layers"] = 1
- else:
- test_case.model_tester.text_config.num_hidden_layers = 1
-
- # A few model class specific handling
-
- # For Albert
- if hasattr(test_case.model_tester, "num_hidden_groups"):
- test_case.model_tester.num_hidden_groups = test_case.model_tester.num_hidden_layers
-
-
def set_config_for_less_flaky_test(config):
target_attrs = [
"rms_norm_eps",
@@ -2398,7 +2357,7 @@ def summary_failures_short(tr):
msg = tr._getfailureheadline(rep)
tr.write_sep("_", msg, red=True, bold=True)
# chop off the optional leading extra frames, leaving only the last one
- longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S)
+ longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.MULTILINE | re.DOTALL)
tr._tw.line(longrepr)
# note: not printing out any rep.sections to keep the report short
@@ -2546,7 +2505,7 @@ def pytest_xdist_worker_id():
if `-n 1` or `pytest-xdist` isn't being used.
"""
worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
- worker = re.sub(r"^gw", "", worker, 0, re.M)
+ worker = re.sub(r"^gw", "", worker, 0, re.MULTILINE)
return int(worker)
@@ -2847,8 +2806,6 @@ def wrapper(*args, **kwargs):
else:
test = " ".join(os.environ.get("PYTEST_CURRENT_TEST").split(" ")[:-1])
try:
- import copy
-
env = copy.deepcopy(os.environ)
env["_INSIDE_SUB_PROCESS"] = "1"
# This prevents the entries in `short test summary info` given by the subprocess being truncated. so the
@@ -2864,7 +2821,7 @@ def wrapper(*args, **kwargs):
test = test.split("::")[1:]
command[idx] = "::".join([f"{func.__globals__['__file__']}"] + test)
command = [f"{sys.executable}", "-m", "pytest"] + command
- command = [x for x in command if x not in ["--no-summary"]]
+ command = [x for x in command if x != "--no-summary"]
# Otherwise, simply run the test with no option at all
else:
command = [f"{sys.executable}", "-m", "pytest", f"{test}"]
@@ -3340,7 +3297,9 @@ def unpack_device_properties(
class Expectations(UserDict[PackedDeviceProperties, Any]):
def get_expectation(self) -> Any:
"""
- Find best matching expectation based on environment device properties.
+ Find best matching expectation based on environment device properties. We look at device_type, major and minor
+ versions of the drivers. Expectations are stored as a dictionary with keys of the form
+ (device_type, (major, minor)). If the major and minor versions are not provided, we use None.
"""
return self.find_expectation(get_device_properties())
@@ -3453,15 +3412,27 @@ def _get_test_info():
stack_from_inspect = inspect.stack()
# but visit from the top frame to the most recent frame
+ actual_test_file, _actual_test_class = test_file, test_class
test_frame, test_obj, test_method = None, None, None
for frame in reversed(stack_from_inspect):
- if test_file in str(frame).replace(r"\\", "/"):
- if test_name == frame.frame.f_locals["self"]._testMethodName:
- test_frame = frame
- # The test instance
- test_obj = frame.frame.f_locals["self"]
- test_method = getattr(test_obj, test_name)
- break
+ # if test_file in str(frame).replace(r"\\", "/"):
+ # check frame's function + if it has `self` as locals; double check if self has the (function) name
+ # TODO: Question: How about expanded?
+ if (
+ frame.function == test_name
+ and "self" in frame.frame.f_locals
+ and hasattr(frame.frame.f_locals["self"], test_name)
+ ):
+ # if test_name == frame.frame.f_locals["self"]._testMethodName:
+ test_frame = frame
+ # The test instance
+ test_obj = frame.frame.f_locals["self"]
+ # TODO: Do we get the (relative?) path or it's just a file name?
+ # TODO: Does `test_obj` always have `tearDown` object?
+ actual_test_file = frame.filename
+ # TODO: check `test_method` will work used at the several places!
+ test_method = getattr(test_obj, test_name)
+ break
if test_frame is not None:
line_number = test_frame.lineno
@@ -3475,9 +3446,12 @@ def _get_test_info():
# From the most outer (i.e. python's `runpy.py`) frame to most inner frame (i.e. the frame of this method)
# Between `the test method being called` and `before entering `patched``.
for frame in reversed(stack_from_inspect):
- if test_file in str(frame).replace(r"\\", "/"):
- if "self" in frame.frame.f_locals and test_name == frame.frame.f_locals["self"]._testMethodName:
- to_capture = True
+ if (
+ frame.function == test_name
+ and "self" in frame.frame.f_locals
+ and hasattr(frame.frame.f_locals["self"], test_name)
+ ):
+ to_capture = True
# TODO: check simply with the name is not robust.
elif "patched" == frame.frame.f_code.co_name:
frame_of_patched_obj = frame
@@ -3511,7 +3485,7 @@ def _get_test_info():
# Get the code context in the test function/method.
from _pytest._code.source import Source
- with open(test_file) as fp:
+ with open(actual_test_file) as fp:
s = fp.read()
source = Source(s)
test_code_context = "\n".join(source.getstatement(test_lineno - 1).lines)
@@ -3522,9 +3496,7 @@ def _get_test_info():
source = Source(s)
caller_code_context = "\n".join(source.getstatement(caller_lineno - 1).lines)
- test_info = (
- f"test:\n\n{full_test_name}\n\n{'-' * 80}\n\ntest context: {test_file}:{test_lineno}\n\n{test_code_context}"
- )
+ test_info = f"test:\n\n{full_test_name}\n\n{'-' * 80}\n\ntest context: {actual_test_file}:{test_lineno}\n\n{test_code_context}"
test_info = f"{test_info}\n\n{'-' * 80}\n\ncaller context: {caller_path}:{caller_lineno}\n\n{caller_code_context}"
return (
@@ -3745,6 +3717,17 @@ def patched(*args, **kwargs):
info = _parse_call_info_func(orig_method, args, kwargs, call_argument_expressions, target_args)
info = _prepare_debugging_info(test_info, info)
+ # If the test is running in a CI environment (e.g. not a manual run), let's raise and fail the test, so it
+ # behaves as usual.
+ # On Github Actions or CircleCI, this is set automatically.
+ # When running manually, it's the user to determine if to set it.
+ # This is to avoid the patched function being called `with self.assertRaises(AssertionError):` and fails
+ # because of the missing expected `AssertionError`.
+ # TODO (ydshieh): If there is way to raise only when we are inside such context managers?
+ # TODO (ydshieh): How not to record the failure if it happens inside `self.assertRaises(AssertionError)`?
+ if os.getenv("CI") == "true":
+ raise captured_exception.with_traceback(test_traceback)
+
# Save this, so we can raise at the end of the current test
captured_failure = {
"result": "failed",
@@ -3827,6 +3810,18 @@ def patch_testing_methods_to_collect_info():
_patch_with_call_info(torch.testing, "assert_close", _parse_call_info, target_args=("actual", "expected"))
_patch_with_call_info(unittest.case.TestCase, "assertEqual", _parse_call_info, target_args=("first", "second"))
+ _patch_with_call_info(unittest.case.TestCase, "assertListEqual", _parse_call_info, target_args=("list1", "list2"))
+ _patch_with_call_info(
+ unittest.case.TestCase, "assertTupleEqual", _parse_call_info, target_args=("tuple1", "tuple2")
+ )
+ _patch_with_call_info(unittest.case.TestCase, "assertSetEqual", _parse_call_info, target_args=("set1", "set1"))
+ _patch_with_call_info(unittest.case.TestCase, "assertDictEqual", _parse_call_info, target_args=("d1", "d2"))
+ _patch_with_call_info(unittest.case.TestCase, "assertIn", _parse_call_info, target_args=("member", "container"))
+ _patch_with_call_info(unittest.case.TestCase, "assertNotIn", _parse_call_info, target_args=("member", "container"))
+ _patch_with_call_info(unittest.case.TestCase, "assertLess", _parse_call_info, target_args=("a", "b"))
+ _patch_with_call_info(unittest.case.TestCase, "assertLessEqual", _parse_call_info, target_args=("a", "b"))
+ _patch_with_call_info(unittest.case.TestCase, "assertGreater", _parse_call_info, target_args=("a", "b"))
+ _patch_with_call_info(unittest.case.TestCase, "assertGreaterEqual", _parse_call_info, target_args=("a", "b"))
def torchrun(script: str, nproc_per_node: int, is_torchrun: bool = True, env: Optional[dict] = None):
@@ -4139,7 +4134,7 @@ def use_one_line_repr(obj):
if element_types[0] in [int, float]:
# one-line repr. without width limit
return no_new_line_in_elements
- elif element_types[0] in [str]:
+ elif element_types[0] is str:
if len(obj) == 1:
# one single string element --> one-line repr. without width limit
return no_new_line_in_elements
diff --git a/src/transformers/tokenization_mistral_common.py b/src/transformers/tokenization_mistral_common.py
index a362a7c8b066..bf76c19e6c55 100644
--- a/src/transformers/tokenization_mistral_common.py
+++ b/src/transformers/tokenization_mistral_common.py
@@ -433,7 +433,7 @@ def encode(
def decode(
self,
- token_ids: Union[int, list[int], "np.ndarray", "torch.Tensor"],
+ token_ids: Union[int, list[int], np.ndarray, "torch.Tensor"],
skip_special_tokens: bool = False,
clean_up_tokenization_spaces: Optional[bool] = None,
**kwargs,
@@ -475,7 +475,7 @@ def decode(
def batch_decode(
self,
- sequences: Union[list[int], list[list[int]], "np.ndarray", "torch.Tensor"],
+ sequences: Union[list[int], list[list[int]], np.ndarray, "torch.Tensor"],
skip_special_tokens: bool = False,
clean_up_tokenization_spaces: Optional[bool] = None,
**kwargs,
@@ -1794,7 +1794,7 @@ def from_pretrained(
if "tekken.json" in valid_tokenizer_files:
tokenizer_file = "tekken.json"
else:
- tokenizer_file = sorted(valid_tokenizer_files)[-1]
+ tokenizer_file = max(valid_tokenizer_files)
logger.warning(
f"Multiple tokenizer files found in directory: {pretrained_model_name_or_path}. Using {tokenizer_file}."
)
@@ -1824,7 +1824,7 @@ def save_pretrained(
repo_url: Optional[str] = None,
organization: Optional[str] = None,
**kwargs,
- ) -> tuple[str]:
+ ) -> tuple[str, ...]:
"""
Save the full tokenizer state.
diff --git a/src/transformers/tokenization_utils.py b/src/transformers/tokenization_utils.py
index 08627d62c123..b89e57093152 100644
--- a/src/transformers/tokenization_utils.py
+++ b/src/transformers/tokenization_utils.py
@@ -587,11 +587,11 @@ def _add_tokens(self, new_tokens: Union[list[str], list[AddedToken]], special_to
self._update_total_vocab_size()
return added_tokens
- def _update_trie(self, unique_no_split_tokens: Optional[str] = []):
+ def _update_trie(self, unique_no_split_tokens: Optional[list[str]] = None):
for token in self._added_tokens_decoder.values():
if token.content not in self.tokens_trie._tokens:
self.tokens_trie.add(token.content)
- for token in unique_no_split_tokens:
+ for token in unique_no_split_tokens or []:
if token not in self.tokens_trie._tokens:
self.tokens_trie.add(token)
diff --git a/src/transformers/tokenization_utils_base.py b/src/transformers/tokenization_utils_base.py
index e4df51c7f867..62afc153265d 100644
--- a/src/transformers/tokenization_utils_base.py
+++ b/src/transformers/tokenization_utils_base.py
@@ -146,7 +146,7 @@ def __str__(self):
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"]]
+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"
@@ -1573,14 +1573,12 @@ def apply_chat_template(
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)
+ [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. Please
- see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG)
- for examples of passing documents with chat templates.
+ 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.
@@ -2042,6 +2040,7 @@ def from_pretrained(
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"
@@ -2177,7 +2176,7 @@ def _from_pretrained(
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:
+ with open(template_file, encoding="utf8") 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"]
@@ -2457,7 +2456,7 @@ def save_pretrained(
filename_prefix: Optional[str] = None,
push_to_hub: bool = False,
**kwargs,
- ) -> tuple[str]:
+ ) -> tuple[str, ...]:
"""
Save the full tokenizer state.
@@ -2622,10 +2621,10 @@ def save_pretrained(
def _save_pretrained(
self,
save_directory: Union[str, os.PathLike],
- file_names: tuple[str],
+ file_names: tuple[str, ...],
legacy_format: Optional[bool] = None,
filename_prefix: Optional[str] = None,
- ) -> tuple[str]:
+ ) -> tuple[str, ...]:
"""
Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens.
@@ -2654,7 +2653,7 @@ def _save_pretrained(
return file_names + vocab_files + (added_tokens_file,)
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str, ...]:
"""
Save only the vocabulary of the tokenizer (vocabulary + added tokens).
@@ -2668,7 +2667,7 @@ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] =
An optional prefix to add to the named of the saved files.
Returns:
- `Tuple(str)`: Paths to the files saved.
+ `tuple(str)`: Paths to the files saved.
"""
raise NotImplementedError
@@ -3894,7 +3893,7 @@ def batch_decode(
def decode(
self,
- token_ids: Union[int, list[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
+ token_ids: Union[int, list[int], np.ndarray, "torch.Tensor"],
skip_special_tokens: bool = False,
clean_up_tokenization_spaces: Optional[bool] = None,
**kwargs,
diff --git a/src/transformers/tokenization_utils_fast.py b/src/transformers/tokenization_utils_fast.py
index 22c63f10da0c..fe4873d61b37 100644
--- a/src/transformers/tokenization_utils_fast.py
+++ b/src/transformers/tokenization_utils_fast.py
@@ -695,10 +695,10 @@ def _decode(
def _save_pretrained(
self,
save_directory: Union[str, os.PathLike],
- file_names: tuple[str],
+ file_names: tuple[str, ...],
legacy_format: Optional[bool] = None,
filename_prefix: Optional[str] = None,
- ) -> tuple[str]:
+ ) -> tuple[str, ...]:
"""
Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens as well as in a unique JSON
file containing {config + vocab + added-tokens}.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 49e14ce56574..249b7e553270 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -47,6 +47,7 @@
import huggingface_hub.utils as hf_hub_utils
import numpy as np
+import safetensors.torch
import torch
import torch.distributed as dist
from huggingface_hub import ModelCard, create_repo, upload_folder
@@ -160,7 +161,6 @@
is_liger_kernel_available,
is_lomo_available,
is_peft_available,
- is_safetensors_available,
is_sagemaker_dp_enabled,
is_sagemaker_mp_enabled,
is_schedulefree_available,
@@ -216,14 +216,9 @@
else:
IS_SAGEMAKER_MP_POST_1_10 = False
-
-if is_safetensors_available():
- import safetensors.torch
-
if is_peft_available():
from peft import PeftModel
-
if is_accelerate_available():
from accelerate import Accelerator, skip_first_batches
from accelerate import __version__ as accelerate_version
@@ -241,10 +236,9 @@
DATA_SAMPLERS = [RandomSampler]
if version.parse(accelerate_version) > version.parse("1.3.0"):
from accelerate.utils import TorchTensorParallelPlugin
- if version.parse(accelerate_version) > version.parse("0.23.0"):
- from accelerate.data_loader import SeedableRandomSampler
+ from accelerate.data_loader import SeedableRandomSampler
- DATA_SAMPLERS += [SeedableRandomSampler]
+ DATA_SAMPLERS += [SeedableRandomSampler]
if is_deepspeed_available():
from accelerate.utils import DeepSpeedSchedulerWrapper
@@ -908,6 +902,11 @@ def remove_callback(self, callback):
self.callback_handler.remove_callback(callback)
def _move_model_to_device(self, model, device):
+ if getattr(model, "hf_device_map", None) is not None:
+ logger.warning(
+ "The model is already on multiple devices. Skipping the move to device specified in `args`."
+ )
+ return
model = model.to(device)
# Moving a model to an XLA device disconnects the tied weights, so we have to retie them.
if self.args.parallel_mode == ParallelMode.TPU and hasattr(model, "tie_weights"):
@@ -923,7 +922,7 @@ def _align_special_tokens(self):
uses the new tokens as well.
"""
if isinstance(self.processing_class, ProcessorMixin):
- tokenizer = self.processing_class.tokenizer
+ tokenizer: PreTrainedTokenizerBase = self.processing_class.tokenizer
else:
tokenizer = self.processing_class
model_has_generation_config = (
@@ -2216,7 +2215,7 @@ def train(
resume_from_checkpoint: Optional[Union[str, bool]] = None,
trial: Union["optuna.Trial", dict[str, Any], None] = None,
ignore_keys_for_eval: Optional[list[str]] = None,
- **kwargs,
+ **kwargs: Any,
):
"""
Main training entry point.
@@ -2412,7 +2411,7 @@ def _inner_training_loop(
" (torchrun or torch.distributed.launch (deprecated))."
)
else:
- debug_overflow = DebugUnderflowOverflow(self.model) # noqa
+ DebugUnderflowOverflow(self.model)
delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled
@@ -2484,8 +2483,7 @@ def _inner_training_loop(
model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
self.model, self.optimizer, self.lr_scheduler
)
- elif self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
- # In this case we are in DDP + LOMO, which should be supported
+ else:
self.optimizer = self.accelerator.prepare(self.optimizer)
if self.is_fsdp_enabled:
@@ -2533,7 +2531,6 @@ def _inner_training_loop(
start_time = time.time()
epochs_trained = 0
steps_trained_in_current_epoch = 0
- steps_trained_progress_bar = None
# Check if continuing training from a checkpoint
if resume_from_checkpoint is not None and os.path.isfile(
@@ -2594,18 +2591,18 @@ def _inner_training_loop(
)
self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
- if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
- self._load_rng_state(resume_from_checkpoint)
-
+ step = -1
rng_to_sync = False
- steps_skipped = 0
- if steps_trained_in_current_epoch > 0:
- epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch)
- steps_skipped = steps_trained_in_current_epoch
- steps_trained_in_current_epoch = 0
- rng_to_sync = True
- step = -1
+ # Handle resumption from checkpoint
+ if epoch == epochs_trained and resume_from_checkpoint is not None:
+ if steps_trained_in_current_epoch > 0 and not args.ignore_data_skip:
+ epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch)
+ step = steps_trained_in_current_epoch - 1
+ rng_to_sync = True
+ elif steps_trained_in_current_epoch == 0:
+ self._load_rng_state(resume_from_checkpoint)
+
epoch_iterator = iter(epoch_dataloader)
# We chunkify the epoch iterator into gradient accumulation steps `n` batches
remainder = steps_in_epoch % args.gradient_accumulation_steps
@@ -2658,22 +2655,11 @@ def _inner_training_loop(
input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64)
self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item()
+
if rng_to_sync:
self._load_rng_state(resume_from_checkpoint)
rng_to_sync = False
- # Skip past any already trained steps if resuming training
- if steps_trained_in_current_epoch > 0:
- steps_trained_in_current_epoch -= 1
- if steps_trained_progress_bar is not None:
- steps_trained_progress_bar.update(1)
- if steps_trained_in_current_epoch == 0:
- self._load_rng_state(resume_from_checkpoint)
- continue
- elif steps_trained_progress_bar is not None:
- steps_trained_progress_bar.close()
- steps_trained_progress_bar = None
-
if step % args.gradient_accumulation_steps == 0:
self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
@@ -2765,7 +2751,7 @@ def _inner_training_loop(
model.zero_grad()
self.state.global_step += 1
- self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
+ self.state.epoch = epoch + (step + 1) / steps_in_epoch
self.control = self.callback_handler.on_step_end(args, self.state, self.control)
self._maybe_log_save_evaluate(
tr_loss,
@@ -2823,14 +2809,6 @@ def _inner_training_loop(
logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
- # Wait for everyone to get here so we are sure the model has been saved by process 0.
- if is_torch_xla_available():
- xm.rendezvous("load_best_model_at_end")
- elif args.parallel_mode == ParallelMode.DISTRIBUTED:
- dist.barrier()
- elif is_sagemaker_mp_enabled():
- smp.barrier()
-
self._load_best_model()
# add remaining tr_loss
@@ -3347,6 +3325,15 @@ def _save_checkpoint(self, model, trial):
self.save_model(output_dir, _internal_call=True)
if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH] and self.state.best_global_step:
+ # Wait for everyone to get here so we are sure the model has been saved by process 0
+ # before we check if the best_checkpoint_dir exists
+ if is_torch_xla_available():
+ xm.rendezvous("load_best_model_at_end")
+ elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
+ dist.barrier()
+ elif is_sagemaker_mp_enabled():
+ smp.barrier()
+
best_checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.best_global_step}"
best_checkpoint_dir = os.path.join(run_dir, best_checkpoint_folder)
@@ -3796,7 +3783,7 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
"""
if self.state.epoch is not None:
logs["epoch"] = self.state.epoch
- if self.args.include_num_input_tokens_seen:
+ if self.args.include_num_input_tokens_seen != "no":
logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
if start_time is not None:
logs.update(speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen))
@@ -3985,10 +3972,7 @@ def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True):
arguments, depending on the situation.
"""
if self.use_cpu_amp:
- # TODO Matt: This syntax is deprecated and the preferred version is
- # torch.amp.autocast("cpu", cache_enabled=cache_enabled, dtype=self.amp_dtype)
- # but this is unavailable on Torch 2.1 or earlier. We can change this when we stop supporting 2.1.
- ctx_manager = torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
+ ctx_manager = torch.autocast(device_type="cpu", cache_enabled=cache_enabled, dtype=self.amp_dtype)
else:
ctx_manager = contextlib.nullcontext()
@@ -4129,16 +4113,27 @@ def compute_loss(
if self.args.past_index >= 0:
self._past = outputs[self.args.past_index]
- if labels is not None:
+ # User-defined compute_loss function
+ if self.compute_loss_func is not None:
+ if labels is None:
+ logger.warning(
+ "Trainer: `compute_loss_func` is defined but `labels=None`. "
+ "Your custom loss function will still be called with labels=None. "
+ )
+ loss = self.compute_loss_func(
+ outputs,
+ labels,
+ num_items_in_batch=num_items_in_batch,
+ )
+ # Default HF loss handling (label smoothing) if no custom loss function
+ elif labels is not None:
unwrapped_model = self.accelerator.unwrap_model(model)
- if _is_peft_model(unwrapped_model):
- model_name = unwrapped_model.base_model.model._get_name()
- else:
- model_name = unwrapped_model._get_name()
- # User-defined compute_loss function
- if self.compute_loss_func is not None:
- loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
- elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
+ model_name = (
+ unwrapped_model.base_model.model._get_name()
+ if _is_peft_model(unwrapped_model)
+ else unwrapped_model._get_name()
+ )
+ if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
loss = self.label_smoother(outputs, labels, shift_labels=True)
else:
loss = self.label_smoother(outputs, labels)
@@ -4156,7 +4151,7 @@ def compute_loss(
and (self.model_accepts_loss_kwargs or self.compute_loss_func)
and num_items_in_batch is not None
):
- loss *= self.accelerator.num_processes
+ loss *= self.accelerator.num_processes if self.args.n_gpu <= 1 else self.args.n_gpu
return (loss, outputs) if return_outputs else loss
@@ -4208,9 +4203,7 @@ def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = Fa
elif (tp_size := getattr(self.model, "_tp_size", 0)) is not None and tp_size > 1:
self._save(output_dir)
elif self.is_fsdp_enabled:
- if ("FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)) and (
- version.parse(accelerate_version) > version.parse("0.24.1")
- ):
+ if "FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type):
state_dict = self.accelerator.get_state_dict(self.model)
if self.args.should_save:
self._save(output_dir, state_dict=state_dict)
@@ -4905,7 +4898,10 @@ def prediction_step(
else:
if has_labels or loss_without_labels:
with self.compute_loss_context_manager():
- loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+ num_items_in_batch = self._get_num_items_in_batch([inputs], self.args.device)
+ loss, outputs = self.compute_loss(
+ model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
+ )
loss = loss.detach().mean()
if isinstance(outputs, dict):
@@ -5594,21 +5590,16 @@ def _fsdp_qlora_plugin_updates(self):
self.model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True
)
- def get_batch_samples(
- self, epoch_iterator: Iterator, num_batches: int, device: torch.device
- ) -> tuple[list, Optional[Union[torch.Tensor, int]]]:
+ def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> int | None:
"""
- Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.
+ Counts the number of items in the batches to properly scale the loss.
+ Args:
+ batch_samples (`list`): List of batches
+ device (`torch.device`): The device on which the number of items in the batch should be.
+ Returns:
+ None if the number of items in the batch doesn't need to be computed else the number of items in the batch
"""
- batch_samples = []
num_items_in_batch = None
-
- for _ in range(num_batches):
- try:
- batch_samples.append(next(epoch_iterator))
- except StopIteration:
- break
-
count_num_items_in_batch = (
len(batch_samples) > 0
and "labels" in batch_samples[0]
@@ -5623,28 +5614,48 @@ def get_batch_samples(
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790
)
)
-
if count_num_items_in_batch:
# For now we don't support object detection
try:
- num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
+ num_items_in_batch = sum((batch["labels"].ne(-100)).sum() for batch in batch_samples)
except (TypeError, AttributeError):
pass
if num_items_in_batch is not None:
- if self.args.average_tokens_across_devices:
+ if self.args.average_tokens_across_devices and self.args.world_size >= 1:
num_items_in_batch = self.accelerator.gather(num_items_in_batch.to(device)).sum()
+ elif self.args.n_gpu >= 1:
+ # In DP case, if we don't average, we need to divide by the number of gpu. This is the simplest approximation.
+ # Otherwise, we would have to scatter labels and calculate num_items_in_batch for each gpu.
+ num_items_in_batch = num_items_in_batch // self.args.n_gpu
if torch.is_tensor(num_items_in_batch):
num_items_in_batch = num_items_in_batch.to(device)
if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0:
- # In the DataParallel case, convert the scalar tensor into a 1-dim tensor
- num_items_in_batch = num_items_in_batch.unsqueeze(0)
+ # In the DataParallel case, convert the scalar tensor into a 2-dim tensor with the same value repeated
+ num_items_in_batch = num_items_in_batch.unsqueeze(0).expand(self.args.n_gpu, -1)
# Divide by number of devices with the same batch
if pc := getattr(self.accelerator, "parallelism_config", None):
num_items_in_batch = num_items_in_batch // pc.non_data_parallel_size
+ return num_items_in_batch
+
+ def get_batch_samples(
+ self, epoch_iterator: Iterator, num_batches: int, device: torch.device
+ ) -> tuple[list, Optional[Union[torch.Tensor, int]]]:
+ """
+ Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.
+ """
+ batch_samples = []
+
+ for _ in range(num_batches):
+ try:
+ batch_samples.append(next(epoch_iterator))
+ except StopIteration:
+ break
+
+ num_items_in_batch = self._get_num_items_in_batch(batch_samples, device)
return batch_samples, num_items_in_batch
def set_initial_training_values(
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
index c32516b167fe..068ff81fd3cd 100644
--- a/src/transformers/trainer_pt_utils.py
+++ b/src/transformers/trainer_pt_utils.py
@@ -929,7 +929,7 @@ def _secs2timedelta(secs):
return f"{datetime.timedelta(seconds=int(secs))}.{msec:02d}"
-def metrics_format(self, metrics: dict[str, float]) -> dict[str, float]:
+def metrics_format(metrics: dict[str, float]) -> dict[str, float]:
"""
Reformat Trainer metrics values to a human-readable format.
@@ -1038,7 +1038,7 @@ def log_metrics(self, split, metrics):
return
print(f"***** {split} metrics *****")
- metrics_formatted = self.metrics_format(metrics)
+ metrics_formatted = metrics_format(metrics)
k_width = max(len(str(x)) for x in metrics_formatted)
v_width = max(len(str(x)) for x in metrics_formatted.values())
for key in sorted(metrics_formatted.keys()):
@@ -1285,7 +1285,7 @@ class AcceleratorConfig:
},
)
- non_blocking: Optional[bool] = field(
+ non_blocking: bool = field(
default=False,
metadata={
"help": "Whether to use non-blocking CUDA calls to help minimize synchronization during "
@@ -1349,7 +1349,7 @@ class LayerWiseDummyOptimizer(torch.optim.Optimizer):
https://github.com/hiyouga/LLaMA-Factory/commit/8664262cde3919e10eaecbd66e8c5d356856362e#diff-ebe08ab14496dfb9e06075f0fdd36799ef6d1535cc4dd4715b74c4e3e06fe3ba
"""
- def __init__(self, optimizer_dict=None, *args, **kwargs):
+ def __init__(self, optimizer_dict=None, **kwargs):
dummy_tensor = torch.randn(1, 1)
self.optimizer_dict = optimizer_dict
super().__init__([dummy_tensor], {"lr": kwargs.get("lr", 1e-03)})
diff --git a/src/transformers/trainer_utils.py b/src/transformers/trainer_utils.py
index e2a382db6c91..2e71367c70c7 100644
--- a/src/transformers/trainer_utils.py
+++ b/src/transformers/trainer_utils.py
@@ -24,7 +24,7 @@
import re
import threading
import time
-from typing import Any, NamedTuple, Optional, Union
+from typing import Any, Callable, NamedTuple, Optional, Union
import numpy as np
@@ -307,7 +307,7 @@ def default_hp_space_optuna(trial) -> dict[str, float]:
}
-def default_hp_space_ray(trial) -> dict[str, float]:
+def default_hp_space_ray(trial) -> dict[str, Any]:
from .integrations import is_ray_tune_available
assert is_ray_tune_available(), "This function needs ray installed: `pip install ray[tune]`"
@@ -334,7 +334,7 @@ def default_hp_space_sigopt(trial):
]
-def default_hp_space_wandb(trial) -> dict[str, float]:
+def default_hp_space_wandb(trial) -> dict[str, Any]:
from .integrations import is_wandb_available
if not is_wandb_available():
@@ -489,7 +489,7 @@ def __init__(self, skip_memory_metrics=False):
if self.skip_memory_metrics:
return
- import psutil # noqa
+ import psutil
if is_torch_cuda_available() or is_torch_mlu_available() or is_torch_musa_available():
import torch
@@ -793,14 +793,14 @@ def number_of_arguments(func):
def find_executable_batch_size(
- function: Optional[callable] = None, starting_batch_size: int = 128, auto_find_batch_size: bool = False
+ function: Optional[Callable] = None, starting_batch_size: int = 128, auto_find_batch_size: bool = False
):
"""
Args:
A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
CUDNN, the batch size is multiplied by 0.9 and passed to `function`. `function` must take in a `batch_size` parameter as
its first argument.
- function (`callable`, *optional*)
+ function (`Callable`, *optional*)
A function to wrap
starting_batch_size (`int`, *optional*)
The batch size to try and fit into memory
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
index 5219feb22023..91ab6f8ca4da 100644
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -41,7 +41,6 @@
is_accelerate_available,
is_apex_available,
is_ipex_available,
- is_safetensors_available,
is_sagemaker_dp_enabled,
is_sagemaker_mp_enabled,
is_torch_available,
@@ -395,7 +394,7 @@ class TrainingArguments:
Whether or not to use PyTorch jit trace for inference.
bf16 (`bool`, *optional*, defaults to `False`):
Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training. Requires Ampere or higher
- NVIDIA architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change.
+ NVIDIA architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU.
fp16 (`bool`, *optional*, defaults to `False`):
Whether to use fp16 16-bit (mixed) precision training instead of 32-bit training.
fp16_opt_level (`str`, *optional*, defaults to 'O1'):
@@ -409,7 +408,7 @@ class TrainingArguments:
requested backend.
bf16_full_eval (`bool`, *optional*, defaults to `False`):
Whether to use full bfloat16 evaluation instead of 32-bit. This will be faster and save memory but can harm
- metric values. This is an experimental API and it may change.
+ metric values.
fp16_full_eval (`bool`, *optional*, defaults to `False`):
Whether to use full float16 evaluation instead of 32-bit. This will be faster and save memory but can harm
metric values.
@@ -489,7 +488,7 @@ class TrainingArguments:
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.
- fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `''`):
+ fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `[]`):
Use PyTorch Distributed Parallel Training (in distributed training only).
A list of options along the following:
@@ -519,11 +518,9 @@ class TrainingArguments:
A list of options along the following:
- `"backward_pre"` : Prefetches the next set of parameters before the current set of parameter's
- gradient
- computation.
+ gradient computation.
- `"backward_post"` : This prefetches the next set of parameters after the current set of
- parameter’s
- gradient computation.
+ 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
@@ -633,6 +630,14 @@ class TrainingArguments:
`"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"neptune"`,
`"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations
installed, `"none"` for no integrations.
+ 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."
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.
@@ -697,7 +702,9 @@ class TrainingArguments:
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.
+ 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_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*):
@@ -760,11 +767,10 @@ class TrainingArguments:
Refer to the PyTorch doc for possible values and note that they may change across PyTorch versions.
This flag is experimental and subject to change in future releases.
- include_tokens_per_second (`bool`, *optional*):
+ include_tokens_per_second (`bool`, *optional*, defaults to `False`):
Whether or not to compute the number of tokens per second per device for training speed metrics.
This will iterate over the entire training dataloader once beforehand,
-
and will slow down the entire process.
include_num_input_tokens_seen (`bool`, *optional*):
@@ -783,7 +789,7 @@ class TrainingArguments:
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.
- batch_eval_metrics (`Optional[bool]`, defaults to `False`):
+ 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
@@ -887,7 +893,7 @@ class TrainingArguments:
metadata={"help": "Number of predictions steps to accumulate before moving the tensors to the CPU."},
)
- eval_delay: Optional[float] = field(
+ eval_delay: float = field(
default=0,
metadata={
"help": (
@@ -922,7 +928,7 @@ class TrainingArguments:
default="linear",
metadata={"help": "The scheduler type to use."},
)
- lr_scheduler_kwargs: Optional[Union[dict[str, Any], str]] = field(
+ lr_scheduler_kwargs: Union[dict[str, Any], str] = field(
default_factory=dict,
metadata={
"help": (
@@ -1005,7 +1011,7 @@ class TrainingArguments:
)
},
)
- save_safetensors: Optional[bool] = field(
+ save_safetensors: bool = field(
default=True,
metadata={
"help": "Use safetensors saving and loading for state dicts instead of default torch.load and torch.save."
@@ -1188,13 +1194,13 @@ class TrainingArguments:
default=None, metadata={"help": "Whether or not to disable the tqdm progress bars."}
)
- remove_unused_columns: Optional[bool] = field(
+ remove_unused_columns: bool = field(
default=True, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."}
)
label_names: Optional[list[str]] = field(
default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."}
)
- load_best_model_at_end: Optional[bool] = field(
+ load_best_model_at_end: bool = field(
default=False,
metadata={
"help": (
@@ -1218,8 +1224,8 @@ class TrainingArguments:
)
},
)
- fsdp: Optional[Union[list[FSDPOption], str]] = field(
- default="",
+ fsdp: Union[list[FSDPOption], str, bool] = field(
+ default_factory=list,
metadata={
"help": (
"Whether or not to use PyTorch Fully Sharded Data Parallel (FSDP) training (in distributed training"
@@ -1299,13 +1305,27 @@ class TrainingArguments:
default=False,
metadata={"help": "Whether or not to group samples of roughly the same length together when batching."},
)
- length_column_name: Optional[str] = field(
+ length_column_name: str = field(
default="length",
metadata={"help": "Column name with precomputed lengths to use when grouping by length."},
)
report_to: Union[None, str, list[str]] = field(
default=None, metadata={"help": "The list of integrations to report the results and logs to."}
)
+ project: str = field(
+ default="huggingface",
+ metadata={"help": "The name of the project to use for logging. Currenly, only used by Trackio."},
+ )
+ trackio_space_id: Optional[str] = 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."
+ },
+ )
ddp_find_unused_parameters: Optional[bool] = field(
default=None,
metadata={
@@ -1366,7 +1386,10 @@ class TrainingArguments:
hub_private_repo: Optional[bool] = 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."
+ "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_always_push: bool = field(
@@ -1492,12 +1515,12 @@ class TrainingArguments:
},
)
- include_tokens_per_second: Optional[bool] = field(
+ include_tokens_per_second: bool = field(
default=False,
metadata={"help": "If set to `True`, the speed metrics will include `tgs` (tokens per second per device)."},
)
- include_num_input_tokens_seen: Optional[Union[str, bool]] = field(
+ include_num_input_tokens_seen: Union[str, bool] = field(
default=False,
metadata={
"help": (
@@ -1534,7 +1557,7 @@ class TrainingArguments:
},
)
- use_liger_kernel: Optional[bool] = field(
+ use_liger_kernel: bool = field(
default=False,
metadata={"help": "Whether or not to enable the Liger Kernel for model training."},
)
@@ -1552,14 +1575,14 @@ class TrainingArguments:
},
)
- eval_use_gather_object: Optional[bool] = field(
+ 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."
},
)
- average_tokens_across_devices: Optional[bool] = field(
+ 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 "
@@ -1688,10 +1711,7 @@ def __post_init__(self):
f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
)
- safetensors_available = is_safetensors_available()
- if self.save_safetensors and not safetensors_available:
- raise ValueError(f"--save_safetensors={self.save_safetensors} requires safetensors to be installed!")
- if not self.save_safetensors and safetensors_available:
+ if not self.save_safetensors:
logger.info(
f"Found safetensors installation, but --save_safetensors={self.save_safetensors}. "
f"Safetensors should be a preferred weights saving format due to security and performance reasons. "
@@ -1790,18 +1810,6 @@ def __post_init__(self):
if self.framework == "pt" and is_torch_available():
self.device
- # Disable average tokens when using single device
- if self.average_tokens_across_devices:
- try:
- if self.world_size == 1:
- logger.info(
- "average_tokens_across_devices is True but world size is 1. Setting it to False automatically."
- )
- self.average_tokens_across_devices = False
- except ImportError as e:
- logger.warning(f"Can not specify world size due to {e}. Turn average_tokens_across_devices to False.")
- self.average_tokens_across_devices = False
-
if self.torchdynamo is not None:
warnings.warn(
"`torchdynamo` is deprecated and will be removed in version 5 of 🤗 Transformers. Use"
@@ -1861,8 +1869,13 @@ def __post_init__(self):
torch.backends.cudnn.allow_tf32 = False
# no need to assert on else
- # NOTE: Mixed precision environment variable setting moved to after DeepSpeed processing
- # to ensure DeepSpeed config can override TrainingArguments defaults
+ # if training args is specified, it will override the one specified in the accelerate config
+ mixed_precision_dtype = os.environ.get("ACCELERATE_MIXED_PRECISION", "no")
+ if self.fp16:
+ mixed_precision_dtype = "fp16"
+ elif self.bf16:
+ mixed_precision_dtype = "bf16"
+ os.environ["ACCELERATE_MIXED_PRECISION"] = mixed_precision_dtype
if self.report_to is None:
logger.info(
@@ -2591,7 +2604,7 @@ def to_json_string(self):
def to_sanitized_dict(self) -> dict[str, Any]:
"""
- Sanitized serialization to use with TensorBoard’s hparams
+ 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}}
diff --git a/src/transformers/utils/__init__.py b/src/transformers/utils/__init__.py
index 1b4671a55e8c..084ff016d283 100644
--- a/src/transformers/utils/__init__.py
+++ b/src/transformers/utils/__init__.py
@@ -107,7 +107,6 @@
is_offline_mode,
is_remote_url,
list_repo_templates,
- send_example_telemetry,
try_to_load_from_cache,
)
from .import_utils import (
diff --git a/src/transformers/utils/auto_docstring.py b/src/transformers/utils/auto_docstring.py
index 0847859450ea..be1e6b5e0fa7 100644
--- a/src/transformers/utils/auto_docstring.py
+++ b/src/transformers/utils/auto_docstring.py
@@ -1215,8 +1215,7 @@ def get_checkpoint_from_config_class(config_class):
# For example, `('google-bert/bert-base-uncased', 'https://huggingface.co/google-bert/bert-base-uncased')`
for ckpt_name, ckpt_link in checkpoints:
# allow the link to end with `/`
- if ckpt_link.endswith("/"):
- ckpt_link = ckpt_link[:-1]
+ ckpt_link = ckpt_link.removesuffix("/")
# verify the checkpoint name corresponds to the checkpoint link
ckpt_link_from_name = f"https://huggingface.co/{ckpt_name}"
@@ -1227,7 +1226,7 @@ def get_checkpoint_from_config_class(config_class):
return checkpoint
-def add_intro_docstring(func, class_name, parent_class=None, indent_level=0):
+def add_intro_docstring(func, class_name, indent_level=0):
intro_docstring = ""
if func.__name__ == "forward":
intro_docstring = rf"""The [`{class_name}`] forward method, overrides the `__call__` special method.
@@ -1469,9 +1468,7 @@ def find_sig_line(lines, line_end):
return sig_line_end
-def _process_kwargs_parameters(
- sig, func, parent_class, model_name_lowercase, documented_kwargs, indent_level, undocumented_parameters
-):
+def _process_kwargs_parameters(sig, func, parent_class, documented_kwargs, indent_level, undocumented_parameters):
"""
Process **kwargs parameters if needed.
@@ -1479,7 +1476,6 @@ def _process_kwargs_parameters(
sig (`inspect.Signature`): Function signature
func (`function`): Function the parameters belong to
parent_class (`class`): Parent class of the function
- model_name_lowercase (`str`): Lowercase model name
documented_kwargs (`dict`): Dictionary of kwargs that are already documented
indent_level (`int`): Indentation level
undocumented_parameters (`list`): List to append undocumented parameters to
@@ -1510,7 +1506,7 @@ def _process_kwargs_parameters(
# Extract documentation for kwargs
kwargs_documentation = kwarg_param.annotation.__args__[0].__doc__
if kwargs_documentation is not None:
- documented_kwargs, _ = parse_docstring(kwargs_documentation)
+ documented_kwargs = parse_docstring(kwargs_documentation)[0]
# Process each kwarg parameter
for param_name, param_type_annotation in kwarg_param.annotation.__args__[0].__annotations__.items():
@@ -1597,7 +1593,7 @@ def _process_parameters_section(
# Process **kwargs parameters if needed
kwargs_docstring = _process_kwargs_parameters(
- sig, func, parent_class, model_name_lowercase, documented_kwargs, indent_level, undocumented_parameters
+ sig, func, parent_class, documented_kwargs, indent_level, undocumented_parameters
)
docstring += kwargs_docstring
@@ -1757,9 +1753,7 @@ def auto_method_docstring(
if not docstring.strip().endswith("\n"):
docstring += "\n"
else:
- docstring = add_intro_docstring(
- func, class_name=class_name, parent_class=parent_class, indent_level=indent_level
- )
+ docstring = add_intro_docstring(func, class_name=class_name, indent_level=indent_level)
# Process Parameters section
docstring += _process_parameters_section(
diff --git a/src/transformers/utils/backbone_utils.py b/src/transformers/utils/backbone_utils.py
index 29b20a813ba6..d2f6277282d9 100644
--- a/src/transformers/utils/backbone_utils.py
+++ b/src/transformers/utils/backbone_utils.py
@@ -76,7 +76,7 @@ def verify_out_features_out_indices(
def _align_output_features_output_indices(
out_features: Optional[list[str]],
- out_indices: Optional[Union[list[int], tuple[int]]],
+ out_indices: Optional[Union[list[int], tuple[int, ...]]],
stage_names: list[str],
):
"""
@@ -284,7 +284,7 @@ def out_indices(self):
return self._out_indices
@out_indices.setter
- def out_indices(self, out_indices: Union[tuple[int], list[int]]):
+ def out_indices(self, out_indices: Union[tuple[int, ...], list[int]]):
"""
Set the out_indices attribute. This will also update the out_features attribute to match the new out_indices.
"""
diff --git a/src/transformers/utils/chat_template_utils.py b/src/transformers/utils/chat_template_utils.py
index 36018c19ccc6..69b3ec977241 100644
--- a/src/transformers/utils/chat_template_utils.py
+++ b/src/transformers/utils/chat_template_utils.py
@@ -468,9 +468,9 @@ def render_jinja_template(
tools: Optional[list[Union[dict, Callable]]] = None,
documents: Optional[list[dict[str, str]]] = None,
chat_template: Optional[str] = None,
- return_assistant_tokens_mask: Optional[bool] = False,
- continue_final_message: Optional[bool] = False,
- add_generation_prompt: Optional[bool] = False,
+ return_assistant_tokens_mask: bool = False,
+ continue_final_message: bool = False,
+ add_generation_prompt: bool = False,
**kwargs,
) -> str:
if return_assistant_tokens_mask and not re.search(r"\{\%-?\s*generation\s*-?\%\}", chat_template):
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
index 94d842eee826..59688632280c 100644
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -48,7 +48,7 @@
if is_torch_available():
# required for @can_return_tuple decorator to work with torchdynamo
- import torch # noqa: F401
+ import torch
from ..model_debugging_utils import model_addition_debugger_context
@@ -380,6 +380,8 @@ def __post_init__(self):
# if we provided an iterator as first field and the iterator is a (key, value) iterator
# set the associated fields
if first_field_iterator:
+ # reset first field to None
+ setattr(self, class_fields[0].name, None)
for idx, element in enumerate(iterator):
if not isinstance(element, (list, tuple)) or len(element) != 2 or not isinstance(element[0], str):
if idx == 0:
@@ -440,7 +442,7 @@ def __reduce__(self):
args = tuple(getattr(self, field.name) for field in fields(self))
return callable, args, *remaining
- def to_tuple(self) -> tuple[Any]:
+ def to_tuple(self) -> tuple:
"""
Convert self to a tuple containing all the attributes/keys that are not `None`.
"""
@@ -937,7 +939,7 @@ class OutputRecorder:
"""
target_class: "type[torch.nn.Module]"
- index: Optional[int] = 0
+ index: int = 0
layer_name: Optional[str] = None
class_name: Optional[str] = None
@@ -985,7 +987,7 @@ def wrapper(self, *args, **kwargs):
}
# We let cross attentions to be saved separately because some models add `cross-attn` layer
- # when certain condtions are met. Let's output cross attention if attentions are requested (for BC)
+ # when certain conditions are met. Let's output cross attention if attentions are requested (for BC)
if "output_attentions" in recordable_keys:
recordable_keys["output_cross_attentions"] = recordable_keys["output_attentions"]
diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
index 4beacbe25aeb..9eb68e0eee6b 100644
--- a/src/transformers/utils/hub.py
+++ b/src/transformers/utils/hub.py
@@ -56,7 +56,6 @@
build_hf_headers,
get_session,
hf_raise_for_status,
- send_telemetry,
)
from requests.exceptions import HTTPError
@@ -155,6 +154,7 @@ def list_repo_templates(
local_files_only: bool,
revision: Optional[str] = None,
cache_dir: Optional[str] = None,
+ token: Union[bool, str, None] = None,
) -> list[str]:
"""List template files from a repo.
@@ -171,6 +171,7 @@ def list_repo_templates(
revision=revision,
path_in_repo=CHAT_TEMPLATE_DIR,
recursive=False,
+ token=token,
)
if entry.path.endswith(".jinja")
]
@@ -993,41 +994,6 @@ def push_to_hub(
)
-def send_example_telemetry(example_name, *example_args, framework="pytorch"):
- """
- Sends telemetry that helps tracking the examples use.
-
- Args:
- example_name (`str`): The name of the example.
- *example_args (dataclasses or `argparse.ArgumentParser`): The arguments to the script. This function will only
- try to extract the model and dataset name from those. Nothing else is tracked.
- framework (`str`, *optional*, defaults to `"pytorch"`): The framework for the example.
- """
- if is_offline_mode():
- return
-
- data = {"example": example_name, "framework": framework}
- for args in example_args:
- args_as_dict = {k: v for k, v in args.__dict__.items() if not k.startswith("_") and v is not None}
- if "model_name_or_path" in args_as_dict:
- model_name = args_as_dict["model_name_or_path"]
- # Filter out local paths
- if not os.path.isdir(model_name):
- data["model_name"] = args_as_dict["model_name_or_path"]
- if "dataset_name" in args_as_dict:
- data["dataset_name"] = args_as_dict["dataset_name"]
- elif "task_name" in args_as_dict:
- # Extract script name from the example_name
- script_name = example_name.replace("tf_", "").replace("flax_", "").replace("run_", "")
- script_name = script_name.replace("_no_trainer", "")
- data["dataset_name"] = f"{script_name}-{args_as_dict['task_name']}"
-
- # Send telemetry in the background
- send_telemetry(
- topic="examples", library_name="transformers", library_version=__version__, user_agent=http_user_agent(data)
- )
-
-
def convert_file_size_to_int(size: Union[int, str]):
"""
Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
@@ -1086,7 +1052,6 @@ def get_checkpoint_shard_files(
For the description of each arg, see [`PreTrainedModel.from_pretrained`]. `index_filename` is the full path to the
index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub).
"""
- import json
use_auth_token = deprecated_kwargs.pop("use_auth_token", None)
if use_auth_token is not None:
diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py
index 2f6dc0b8e714..2bf0464dbf42 100644
--- a/src/transformers/utils/import_utils.py
+++ b/src/transformers/utils/import_utils.py
@@ -184,7 +184,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[
_auto_awq_available = importlib.util.find_spec("awq") is not None
_quark_available = _is_package_available("quark")
_fp_quant_available, _fp_quant_version = _is_package_available("fp_quant", return_version=True)
-_qutlass_available = _is_package_available("qutlass")
+_qutlass_available, _qutlass_version = _is_package_available("qutlass", return_version=True)
_is_optimum_quanto_available = False
try:
importlib.metadata.version("optimum_quanto")
@@ -470,11 +470,7 @@ def is_torchvision_available() -> bool:
def is_torchvision_v2_available() -> bool:
- if not is_torchvision_available():
- return False
-
- # NOTE: We require torchvision>=0.15 as v2 transforms are available from this version: https://pytorch.org/vision/stable/transforms.html#v1-or-v2-which-one-should-i-use
- return version.parse(_torchvision_version) >= version.parse("0.15")
+ return is_torchvision_available()
def is_galore_torch_available() -> Union[tuple[bool, str], bool]:
@@ -844,7 +840,7 @@ def is_torch_npu_available(check_device=False) -> bool:
@lru_cache
-def is_torch_mlu_available(check_device=False) -> bool:
+def is_torch_mlu_available() -> bool:
"""
Checks if `mlu` is available via an `cndev-based` check which won't trigger the drivers and leave mlu
uninitialized.
@@ -993,7 +989,7 @@ def is_habana_gaudi1() -> bool:
if not is_torch_hpu_available():
return False
- import habana_frameworks.torch.utils.experimental as htexp # noqa: F401
+ import habana_frameworks.torch.utils.experimental as htexp
# Check if the device is Gaudi1 (vs Gaudi2, Gaudi3)
return htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi
@@ -1019,7 +1015,7 @@ def is_torchdynamo_compiling() -> Union[tuple[bool, str], bool]:
return torch.compiler.is_compiling()
except Exception:
try:
- import torch._dynamo as dynamo # noqa: F401
+ import torch._dynamo as dynamo
return dynamo.is_compiling()
except Exception:
@@ -1036,7 +1032,7 @@ def is_torchdynamo_exporting() -> bool:
return torch.compiler.is_exporting()
except Exception:
try:
- import torch._dynamo as dynamo # noqa: F401
+ import torch._dynamo as dynamo
return dynamo.is_exporting()
except Exception:
@@ -1375,12 +1371,12 @@ def is_quark_available() -> Union[tuple[bool, str], bool]:
return _quark_available
-def is_fp_quant_available() -> bool:
- return _fp_quant_available and version.parse(_fp_quant_version) >= version.parse("0.1.6")
+def is_fp_quant_available():
+ return _fp_quant_available and version.parse(_fp_quant_version) >= version.parse("0.2.0")
-def is_qutlass_available() -> Union[tuple[bool, str], bool]:
- return _qutlass_available
+def is_qutlass_available():
+ return _qutlass_available and version.parse(_qutlass_version) >= version.parse("0.1.0")
def is_compressed_tensors_available() -> bool:
diff --git a/src/transformers/utils/logging.py b/src/transformers/utils/logging.py
index 88a6a9769f65..e383653871bf 100644
--- a/src/transformers/utils/logging.py
+++ b/src/transformers/utils/logging.py
@@ -20,13 +20,13 @@
import threading
from logging import (
CRITICAL, # NOQA
- DEBUG, # NOQA
- ERROR, # NOQA
+ DEBUG,
+ ERROR,
FATAL, # NOQA
- INFO, # NOQA
+ INFO,
NOTSET, # NOQA
WARN, # NOQA
- WARNING, # NOQA
+ WARNING,
)
from logging import captureWarnings as _captureWarnings
from typing import Optional
diff --git a/src/transformers/utils/metrics.py b/src/transformers/utils/metrics.py
index 62b41995a6d9..3703ddaca1fb 100644
--- a/src/transformers/utils/metrics.py
+++ b/src/transformers/utils/metrics.py
@@ -105,8 +105,6 @@ def decorator(func):
if not _has_opentelemetry:
return func
- import functools
-
@functools.wraps(func)
def wrapper(*args, **kwargs):
instance = args[0] if args and (hasattr(func, "__self__") and func.__self__ is not None) else None
@@ -339,7 +337,7 @@ def record_kv_cache_memory_metrics(self, cache) -> None:
page_size = cache.head_dim * cache.num_key_value_heads
page_mem_in_bytes = page_size * cache.dtype.itemsize
# When a block is allocated, it is for both K and V, so we multiply by 2
- # It's also allocated accross all cache tensors, so we multiply by the nb of tensors: len(cache.key_cache)
+ # It's also allocated across all cache tensors, so we multiply by the nb of tensors: len(cache.key_cache)
block_mem_in_bytes = 2 * len(cache.key_cache) * cache.block_size * page_mem_in_bytes
# Retrieve the number of used and free blocks
diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py
index 037bf3ed73d4..82357e6f0fe2 100644
--- a/src/transformers/utils/quantization_config.py
+++ b/src/transformers/utils/quantization_config.py
@@ -845,7 +845,7 @@ def post_init(self):
"You current version of `optimum` does not support `modules_in_block_to_quantize` quantization argument, please upgrade `optimum` package to a version superior than 1.15.0 ."
)
- def to_dict(self):
+ def to_dict(self) -> dict[str, Any]:
config_dict = super().to_dict()
config_dict.pop("disable_exllama", None)
return config_dict
@@ -1557,7 +1557,7 @@ class FPQuantConfig(QuantizationConfigMixin):
FPQuantConfig is a configuration class for quantization using the FPQuant method.
Args:
- forward_dtype (`str`, *optional*, defaults to `"mxfp4"`):
+ forward_dtype (`str`, *optional*, defaults to `"nvfp4"`):
The dtype to use for the forward pass.
forward_method (`str`, *optional*, defaults to `"abs_max"`):
The scaling to use for the forward pass. Can be `"abs_max"` or `"quest"`. `"abs_max"` is better for PTQ, `"quest"` is better for QAT.
@@ -1565,10 +1565,11 @@ class FPQuantConfig(QuantizationConfigMixin):
The dtype to use for the backward pass.
store_master_weights (`bool`, *optional*, defaults to `False`):
Whether to store the master weights. Needed for QAT over layer weights.
- hadamard_group_size (`int`, *optional*, defaults to 32):
- The group size for the hadamard transform before quantization for `"quest"` it matches the MXFP4 group size (32).
+ hadamard_group_size (`int`, *optional*):
+ The group size for the hadamard transform before quantization for `"quest"` it matches the MXFP4 group size (32). If `None`, it will be set to 16 for `"nvfp4"` and 32 for `"mxfp4"`.
pseudoquantization (`bool`, *optional*, defaults to `False`):
Whether to use Triton-based pseudo-quantization. Is mandatory for non-Blackwell GPUs. Doesn't provide any speedup. For debugging purposes.
+ transform_init (`str`, *optional*, defaults to `"hadamard"`): a method to initialize the pre-processing matrix with. Can be `"hadamard"`, `"identity"` or `"gsr"`.
modules_to_not_convert (`list`, *optional*):
The list of modules to not quantize, useful for quantizing models that explicitly require to have
some modules left in their original precision.
@@ -1576,12 +1577,13 @@ class FPQuantConfig(QuantizationConfigMixin):
def __init__(
self,
- forward_dtype: str = "mxfp4",
+ forward_dtype: str = "nvfp4",
forward_method: str = "abs_max",
backward_dtype: str = "bf16",
store_master_weights: bool = False,
- hadamard_group_size: int = 32,
+ hadamard_group_size: Optional[int] = None,
pseudoquantization: bool = False,
+ transform_init: str = "hadamard",
modules_to_not_convert: Optional[list[str]] = None,
**kwargs,
):
@@ -1591,6 +1593,7 @@ def __init__(
self.store_master_weights = store_master_weights
self.hadamard_group_size = hadamard_group_size
self.pseudoquantization = pseudoquantization
+ self.transform_init = transform_init
self.modules_to_not_convert = modules_to_not_convert
self.quant_method = QuantizationMethod.FPQUANT
@@ -1600,14 +1603,35 @@ def post_init(self):
r"""
Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
"""
- if self.forward_dtype not in ["mxfp4"]:
- raise ValueError("Only 'mxfp4' is supported for forward_dtype for now.")
- if self.forward_method not in ["abs_max", "quest"]:
- raise ValueError("Only 'abs_max' and 'quest' are supported for forward_method for now.")
- if self.backward_dtype not in ["bf16"]:
+
+ if self.hadamard_group_size is None:
+ if self.forward_dtype == "nvfp4":
+ self.hadamard_group_size = 16
+ else:
+ self.hadamard_group_size = 32
+
+ if self.forward_dtype == "mxfp4":
+ if self.forward_method not in ["abs_max", "quest"]:
+ raise ValueError("Only 'abs_max' and 'quest' are supported for forward_method for 'mxfp4'.")
+ if self.hadamard_group_size is None:
+ self.hadamard_group_size = 32
+ if self.hadamard_group_size not in [32, 64, 128]:
+ raise ValueError("Only a `hadamard_group_size` of [32, 64, 128] is supported for 'mxfp4'.")
+ elif self.forward_dtype == "nvfp4":
+ if self.forward_method != "abs_max":
+ raise ValueError("Only 'abs_max' is supported for forward_method for 'nvfp4'.")
+ if self.hadamard_group_size is None:
+ self.hadamard_group_size = 16
+ if self.hadamard_group_size not in [16, 32, 64, 128]:
+ raise ValueError("Only a `hadamard_group_size` of [16, 32, 64, 128] is supported for 'nvfp4'.")
+ else:
+ raise ValueError("Only 'mxfp4' and 'nvfp4' are supported for forward_dtype for now.")
+
+ if self.backward_dtype != "bf16":
raise ValueError("Only 'bf16' is supported for backward_dtype for now.")
- if self.hadamard_group_size not in [32]:
- raise ValueError("Only a hadamard_group_size of 32 is supported for now.")
+ if self.transform_init not in ["hadamard", "identity", "gsr"]:
+ raise ValueError("Only 'hadamard', 'identity' and 'gsr' are supported for transform_init.")
+
if self.modules_to_not_convert is None:
self.modules_to_not_convert = ["lm_head"]
@@ -1882,9 +1906,9 @@ class BitNetQuantConfig(QuantizationConfigMixin):
def __init__(
self,
modules_to_not_convert: Optional[list] = None,
- linear_class: Optional[str] = "bitlinear",
- quantization_mode: Optional[str] = "offline",
- use_rms_norm: Optional[bool] = False,
+ linear_class: str = "bitlinear",
+ quantization_mode: str = "offline",
+ use_rms_norm: bool = False,
rms_norm_eps: Optional[float] = 1e-6,
**kwargs,
):
@@ -2002,7 +2026,7 @@ def post_init(self):
Safety checker that arguments are correct
"""
self.activation_scheme = self.activation_scheme.lower()
- if self.activation_scheme not in ["dynamic"]:
+ if self.activation_scheme != "dynamic":
raise ValueError(f"Activation scheme {self.activation_scheme} not supported")
if len(self.weight_block_size) != 2:
raise ValueError("weight_block_size must be a tuple of two integers")
diff --git a/src/transformers/video_processing_utils.py b/src/transformers/video_processing_utils.py
index 4d0e9c58f314..0bc81bf8eb28 100644
--- a/src/transformers/video_processing_utils.py
+++ b/src/transformers/video_processing_utils.py
@@ -46,7 +46,6 @@
is_remote_url,
is_torch_available,
is_torchcodec_available,
- is_torchvision_available,
is_torchvision_v2_available,
logging,
)
@@ -70,8 +69,6 @@
if is_torchvision_v2_available():
from torchvision.transforms.v2 import functional as F
-elif is_torchvision_available():
- from torchvision.transforms import functional as F
logger = logging.get_logger(__name__)
diff --git a/src/transformers/video_utils.py b/src/transformers/video_utils.py
index 1749b0b3b1c5..2ed5720a8e41 100644
--- a/src/transformers/video_utils.py
+++ b/src/transformers/video_utils.py
@@ -74,7 +74,7 @@
Path,
list[Path],
list[list[Path]],
-] # noqa
+]
@dataclass
@@ -100,7 +100,7 @@ def __setitem__(self, key, value):
return setattr(self, key, value)
@property
- def timestamps(self) -> float:
+ 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.")
@@ -196,7 +196,9 @@ def make_batched_videos(videos) -> list[Union[np.ndarray, "torch.Tensor", "URL",
return convert_pil_frames_to_video([videos])
# only one frame passed, thus we unsqueeze time dim
elif is_valid_image(videos):
- return [np.array(videos)[None, ...]]
+ 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"
@@ -329,7 +331,7 @@ 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.
@@ -345,7 +347,7 @@ 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:
+ tuple[`np.ndarray`, `VideoMetadata`]: A tuple containing:
- Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
- `VideoMetadata` object.
"""
@@ -546,8 +548,8 @@ def sample_indices_fn(metadata, **kwargs):
metadata.update(
{
"frames_indices": indices,
- "height": video.shape[1],
- "width": video.shape[2],
+ "height": video.shape[2],
+ "width": video.shape[3],
}
)
return video, metadata
@@ -620,7 +622,7 @@ def load_video(
backend: str = "pyav",
sample_indices_fn: Optional[Callable] = None,
**kwargs,
-) -> np.array:
+) -> np.ndarray:
"""
Loads `video` to a numpy array.
@@ -646,7 +648,7 @@ def sample_indices_fn(metadata, **kwargs):
return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)
Returns:
- tuple[`np.array`, Dict]: A tuple containing:
+ tuple[`np.ndarray`, Dict]: A tuple containing:
- Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
- Metadata dictionary.
"""
@@ -692,7 +694,7 @@ def sample_indices_fn_func(metadata, **fn_kwargs):
# 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 in ["opencv"]:
+ 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 (
@@ -714,24 +716,21 @@ def sample_indices_fn_func(metadata, **fn_kwargs):
def convert_to_rgb(
video: np.ndarray,
- data_format: Optional[ChannelDimension] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = 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.array`):
+ video (`np.ndarray`):
The video to convert.
- data_format (`ChannelDimension`, *optional*):
- The channel dimension format of the output video. If unset, will use the inferred format from the input.
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 leet's convert it
+ # 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)
@@ -845,7 +844,7 @@ def _expand_for_data_format(values):
def group_videos_by_shape(
videos: list["torch.Tensor"],
-) -> tuple[dict[tuple[int, int], list["torch.Tensor"]], dict[int, tuple[tuple[int, int], int]]]:
+) -> 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,
@@ -867,7 +866,8 @@ def group_videos_by_shape(
def reorder_videos(
- processed_videos: dict[tuple[int, int], "torch.Tensor"], grouped_videos_index: dict[int, tuple[int, int]]
+ 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.
diff --git a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
index 3cd69eb95630..8ada67913b03 100755
--- a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
+++ b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
@@ -46,7 +46,6 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import send_example_telemetry
logger = logging.getLogger(__name__)
@@ -220,10 +219,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_{{cookiecutter.example_shortcut}}", model_args, data_args)
-
# Detecting last checkpoint.
last_checkpoint = None
if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
@@ -545,7 +540,6 @@ def _mp_fn(index):
get_scheduler,
set_seed,
)
-from transformers.utils import send_example_telemetry
logger = logging.getLogger(__name__)
@@ -698,10 +692,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_{{cookiecutter.example_shortcut}", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
accelerator = Accelerator()
# Make one log on every process with the configuration for debugging.
diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py
index 8600f1dc265e..dc57c708829c 100644
--- a/tests/causal_lm_tester.py
+++ b/tests/causal_lm_tester.py
@@ -19,7 +19,9 @@
from parameterized import parameterized
from transformers import AutoModelForCausalLM, PretrainedConfig, set_seed
+from transformers.models.auto.auto_factory import getattribute_from_module
from transformers.testing_utils import (
+ _COMMON_MODEL_NAMES_MAP,
is_flaky,
require_flash_attn,
require_torch_gpu,
@@ -43,31 +45,99 @@
class CausalLMModelTester:
- _required_attributes = ("base_model_class", "config_class", "causal_lm_class")
- forced_config_args = [
- "pad_token_id"
- ] # Arguments that should be passed to the config class even if not in its signature
- config_class = None
+ # If the model follows the standard naming conventions, only `base_model_class` needs to be set (the others are
+ # inferred from available public classes).
base_model_class = None
+ # ⚠️ Don't set these unless the model does NOT follow the standard naming conventions ⚠️
+ config_class = None
causal_lm_class = None
+ question_answering_class = None
sequence_classification_class = None
token_classification_class = None
- question_answering_class = None
+ # These attributes are required after the initialization phase of the tester.
+ _required_attributes = ("base_model_class", "config_class", "causal_lm_class")
+
+ # Arguments that should be passed to the config class even if not in its signature
+ forced_config_args = ["pad_token_id"]
+
+ @classmethod
+ def _verify_and_infer_model_attributes(cls):
+ """
+ Verifies that the required tester attributes are set correctly, and infers unset tester attributes.
+ Intentionally nitpicks the tester class attributes, to prevent human errors.
+ """
+ # `base_model_class` is mandatory, and it must be a valid model class.
+ base_model_class = getattr(cls, "base_model_class")
+ if base_model_class is None or "PreTrainedModel" not in str(base_model_class.__mro__):
+ raise ValueError(
+ f"You have inherited from `CausalLMModelTester` but did not set the `base_model_class` "
+ f"attribute to a valid model class. (It's set to `{base_model_class}`)"
+ )
+
+ # Infers other model classes from the base class name and available public classes, if the corresponding
+ # attributes are not set explicitly. If they are set, they must be set to a valid class (config or model).
+ model_name = base_model_class.__name__.replace("Model", "")
+ base_class_module = ".".join(base_model_class.__module__.split(".")[:-1])
+ for tester_attribute_name, model_class_termination in _COMMON_MODEL_NAMES_MAP.items():
+ if getattr(cls, tester_attribute_name) is None:
+ try:
+ model_class = getattribute_from_module(base_class_module, model_name + model_class_termination)
+ setattr(cls, tester_attribute_name, model_class)
+ except ValueError:
+ pass
+ else:
+ if tester_attribute_name == "config_class":
+ if "PretrainedConfig" not in str(getattr(cls, tester_attribute_name).__mro__):
+ raise ValueError(
+ f"You have inherited from `CausalLMModelTester` but did not set the "
+ f"`{tester_attribute_name}` attribute to a valid config class. (It's set to "
+ f"`{getattr(cls, tester_attribute_name)}`). If the config class follows a standard "
+ f"naming convention, you should unset `{tester_attribute_name}`."
+ )
+ else:
+ if "PreTrainedModel" not in str(getattr(cls, tester_attribute_name).__mro__):
+ raise ValueError(
+ f"You have inherited from `CausalLMModelTester` but did not set the "
+ f"`{tester_attribute_name}` attribute to a valid model class. (It's set to "
+ f"`{getattr(cls, tester_attribute_name)}`). If the model class follows a standard "
+ f"naming convention, you should unset `{tester_attribute_name}`."
+ )
+
+ # After inferring, if we don't have the basic classes set, we raise an error.
+ for required_attribute in cls._required_attributes:
+ if getattr(cls, required_attribute) is None:
+ raise ValueError(
+ f"You have inherited from `CausalLMModelTester` but did not set the `{required_attribute}` "
+ "attribute. It can't be automatically inferred either -- this means it is not following a "
+ "standard naming convention. If this is intentional, please set the attribute explicitly."
+ )
- def _verify_model_attributes(self):
- for required_attribute in self._required_attributes:
- if getattr(self, required_attribute) is None:
+ # To prevent issues with typos, no other attributes can be set to a model class
+ for instance_attribute_name, instance_attribute in cls.__dict__.items():
+ if (
+ (
+ instance_attribute_name not in _COMMON_MODEL_NAMES_MAP
+ and instance_attribute_name != "base_model_class"
+ )
+ and isinstance(instance_attribute, type)
+ and "PreTrainedModel" in str(instance_attribute.__mro__)
+ ):
raise ValueError(
- f"You have inherited from CausalLMModelTester but did not set the {required_attribute} attribute."
+ f"You have inherited from `CausalLMModelTester` but set an unexpected attribute to a model class "
+ f"(`{instance_attribute_name}` is set to `{instance_attribute}`). "
+ f"Only the following attributes can be set to model classes: {_COMMON_MODEL_NAMES_MAP.keys()}."
)
@property
def all_model_classes(self):
+ # Models that set `all_model_classes` in their `XXXModelTest` class must have a new class that doesn't fit
+ # any of the common classes.
return [
model_class
for model_class in (
self.base_model_class,
self.causal_lm_class,
+ self.question_answering_class,
self.sequence_classification_class,
self.token_classification_class,
)
@@ -118,7 +188,7 @@ def __init__(
mamba_expand=2,
mamba_chunk_size=16,
):
- self._verify_model_attributes()
+ self._verify_and_infer_model_attributes()
self.parent = parent
self.batch_size = batch_size
self.seq_length = seq_length
@@ -210,16 +280,7 @@ def create_and_check_model(
self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
def prepare_config_and_inputs_for_common(self):
- config_and_inputs = self.prepare_config_and_inputs()
- (
- config,
- input_ids,
- token_type_ids,
- input_mask,
- sequence_labels,
- token_labels,
- choice_labels,
- ) = config_and_inputs
+ config, input_ids, _, input_mask, _, _, _ = self.prepare_config_and_inputs()
inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
return config, inputs_dict
@@ -316,6 +377,27 @@ def test_token_classification_model(self):
(self.model_tester.batch_size, self.model_tester.seq_length, self.model_tester.num_labels),
)
+ def test_question_answering_model(self):
+ if self.model_tester.question_answering_class is None:
+ self.skipTest("Model does not support question answering")
+ config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
+ config.num_labels = 3
+
+ input_ids = input_dict["input_ids"]
+ attention_mask = input_ids.ne(1).to(torch_device)
+ model = self.model_tester.question_answering_class(config=config)
+ model.to(torch_device)
+ model.eval()
+ result = model(input_ids, attention_mask=attention_mask)
+ self.assertEqual(
+ result.start_logits.shape,
+ (self.model_tester.batch_size, self.model_tester.seq_length),
+ )
+ self.assertEqual(
+ result.end_logits.shape,
+ (self.model_tester.batch_size, self.model_tester.seq_length),
+ )
+
@parameterized.expand([("linear",), ("dynamic",), ("yarn",)])
def test_model_rope_scaling_from_config(self, scaling_type):
"""
@@ -497,7 +579,7 @@ def _config_supports_rope_scaling(config: PretrainedConfig) -> bool:
# Has rope_theta (and no rope_scaling) -> probably an older model, but should support rope scaling as well
main_config_has_rope = hasattr(config, "rope_scaling") or hasattr(config, "rope_theta")
sub_config_has_rope = any(
- hasattr(config[sub_config], "rope_scaling") or hasattr(config[sub_config], "rope_theta")
+ hasattr(getattr(config, sub_config), "rope_scaling") or hasattr(getattr(config, sub_config), "rope_theta")
for sub_config in config.sub_configs.keys()
)
return main_config_has_rope or sub_config_has_rope
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index e3dc9fc08c99..99b1450a0d59 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -1431,50 +1431,3 @@ def test_clm_from_config_zero3_fp16(self):
with CaptureStderr() as cs:
execute_subprocess_async(cmd, env=self.get_env())
self.assertIn("Detected DeepSpeed ZeRO-3", cs.err)
-
-
-@require_deepspeed
-class TestDeepSpeedMixedPrecisionPrecedence(TestCasePlus):
- """Test DeepSpeed mixed precision precedence over Accelerate defaults."""
-
- def setUp(self):
- super().setUp()
- unset_hf_deepspeed_config()
-
- def tearDown(self):
- super().tearDown()
- unset_hf_deepspeed_config()
-
- def test_deepspeed_fp16_overrides_defaults(self):
- """Test that DeepSpeed fp16 config overrides TrainingArguments defaults"""
- from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
-
- args = TrainingArguments(output_dir="./test_output", fp16=False, bf16=False)
- ds_config = {"fp16": {"enabled": True}, "bf16": {"enabled": False}, "zero_optimization": {"stage": 2}}
- hf_ds_config = HfTrainerDeepSpeedConfig(ds_config)
- hf_ds_config.trainer_config_process(args)
- self.assertTrue(args.fp16)
- self.assertFalse(args.bf16)
-
- def test_deepspeed_bf16_overrides_defaults(self):
- """Test that DeepSpeed bf16 config overrides TrainingArguments defaults"""
- from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
-
- args = TrainingArguments(output_dir="./test_output", fp16=False, bf16=False)
- ds_config = {"fp16": {"enabled": False}, "bf16": {"enabled": True}, "zero_optimization": {"stage": 2}}
- hf_ds_config = HfTrainerDeepSpeedConfig(ds_config)
- hf_ds_config.trainer_config_process(args)
- self.assertTrue(args.bf16)
- self.assertFalse(args.fp16)
-
- def test_user_explicit_settings_preserved(self):
- """Test that explicit user settings are preserved over DeepSpeed config"""
- from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
-
- args = TrainingArguments(output_dir="./test_output", fp16=True, bf16=False) # User explicit
- ds_config = {"fp16": {"enabled": False}, "bf16": {"enabled": True}, "zero_optimization": {"stage": 2}}
- hf_ds_config = HfTrainerDeepSpeedConfig(ds_config)
- hf_ds_config.trainer_config_process(args)
- # User's explicit choice should be preserved
- self.assertTrue(args.fp16)
- self.assertFalse(args.bf16)
diff --git a/tests/extended/test_trainer_ext.py b/tests/extended/test_trainer_ext.py
index 82d3cdb9c3ce..707d35a73697 100644
--- a/tests/extended/test_trainer_ext.py
+++ b/tests/extended/test_trainer_ext.py
@@ -259,8 +259,8 @@ def train_and_return_metrics(optim: str) -> tuple[int, float]:
f" gpu_total_mem_bnb={gpu_total_mem_bnb}MB",
)
- self.assertEqual(
- loss_orig, loss_bnb, f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"
+ self.assertAlmostEqual(
+ loss_orig, loss_bnb, 5, f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"
)
def run_trainer(
diff --git a/tests/fixtures/parakeet/expected_results_batch.json b/tests/fixtures/parakeet/expected_results_batch.json
new file mode 100644
index 000000000000..2ca30b96d85a
--- /dev/null
+++ b/tests/fixtures/parakeet/expected_results_batch.json
@@ -0,0 +1 @@
+{"transcriptions": ["mister quilter is the apostle of the middle classes and we are glad to welcome his gospel", "nor is mister quilter's manner less interesting than his matter", "he tells us that at this festive season of the year with christmas and roast beef looming before us similes drawn from eating and its results occur most readily to the mind", "he has grave doubts whether sir frederick leighton's work is really greek after all and can discover in it but little of rocky ithaca", "linnell's pictures are a sort of up guards and adam paintings and mason's exquisite idylls are as national as a jingo poem mr burket foster's landscapes smile at one much in the same way that mr carker used to flash his teeth and mr john collier gives his sitter a cheerful slap on the back before he says like a shampooer in a turkish bath next man"], "token_ids": [[1024, 1024, 1024, 1024, 1024, 1024, 19, 37, 132, 1024, 1024, 264, 128, 1024, 1024, 1024, 132, 1024, 58, 1024, 5, 645, 1024, 1000, 82, 52, 1024, 34, 1024, 5, 19, 68, 1007, 52, 1024, 235, 1024, 388, 1024, 27, 1024, 25, 1024, 56, 1024, 103, 1024, 1024, 727, 112, 1024, 22, 1024, 56, 1006, 1009, 405, 1024, 1024, 217, 1024, 1024, 95, 1003, 1024, 133, 1006, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024], [1024, 1024, 1024, 1024, 1024, 1024, 1024, 42, 28, 1024, 1024, 58, 1024, 19, 37, 1024, 132, 1024, 264, 128, 1024, 1024, 132, 1024, 1019, 1003, 1024, 284, 1024, 896, 1024, 32, 154, 1024, 715, 1024, 1024, 1024, 1024, 21, 1024, 322, 1024, 1024, 1024, 217, 1024, 1024, 1024, 1024, 19, 1024, 710, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024], [1024, 1024, 1024, 1024, 1024, 1024, 1024, 67, 1024, 634, 1024, 1024, 1003, 1024, 208, 1024, 1024, 39, 1024, 1024, 124, 1024, 1024, 77, 1024, 1024, 1024, 20, 156, 1024, 1024, 171, 1024, 1024, 101, 1024, 667, 1024, 1024, 34, 1024, 5, 1024, 696, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 93, 1024, 1024, 1024, 1024, 121, 1004, 172, 1024, 1010, 43, 1024, 25, 1024, 343, 250, 1024, 1024, 1024, 50, 1024, 846, 1024, 1024, 304, 44, 1024, 1024, 21, 1024, 1024, 497, 1024, 1024, 208, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 596, 1024, 1024, 1024, 128, 1024, 1024, 27, 1024, 26, 96, 447, 1024, 176, 1024, 48, 1024, 1024, 599, 1024, 25, 1024, 525, 1024, 1024, 338, 1024, 411, 1003, 1024, 1024, 9, 1009, 1024, 1024, 1009, 83, 1024, 1024, 463, 1024, 788, 1024, 1024, 522, 1024, 22, 1024, 5, 1024, 19, 191, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024], [1024, 1024, 1024, 1024, 1024, 1024, 67, 1024, 1024, 244, 1024, 1024, 657, 1024, 47, 1024, 1024, 26, 13, 1016, 998, 1003, 1024, 789, 1024, 1024, 8, 94, 1024, 20, 265, 1024, 12, 1024, 363, 184, 120, 1024, 1024, 1024, 18, 1024, 1019, 1003, 337, 1024, 1024, 58, 1024, 1024, 254, 1024, 1024, 1024, 1024, 1024, 41, 302, 1018, 1024, 1024, 451, 1024, 1024, 1024, 1024, 142, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 25, 1024, 1024, 117, 1024, 1024, 1024, 321, 1024, 394, 1024, 71, 1024, 35, 1024, 45, 1024, 106, 1024, 1024, 1024, 401, 1024, 1024, 1024, 34, 1024, 1024, 1024, 343, 1024, 137, 1024, 1024, 1011, 1024, 45, 1005, 1024, 765, 1024, 1024, 999, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024], [1024, 1024, 1024, 1024, 1024, 1024, 32, 1024, 10, 728, 728, 30, 1024, 1024, 1019, 1003, 1024, 24, 433, 1024, 799, 1024, 1024, 103, 1024, 1024, 3, 1024, 903, 1024, 1024, 34, 1024, 1024, 1024, 1024, 1024, 190, 1024, 1024, 1024, 415, 203, 1024, 1003, 1003, 25, 1024, 273, 1024, 1024, 104, 1024, 1024, 1024, 24, 164, 1024, 1024, 467, 1003, 1024, 1024, 1024, 1024, 1024, 25, 1024, 1024, 19, 1024, 1024, 1024, 667, 1024, 1019, 1003, 1024, 146, 1024, 162, 37, 1024, 320, 1024, 4, 1007, 1011, 1011, 30, 1024, 1003, 1024, 103, 1024, 1024, 88, 1024, 1024, 1024, 42, 1024, 1024, 1024, 895, 1024, 88, 1024, 1024, 3, 1024, 92, 1024, 21, 1024, 1024, 1000, 1024, 1024, 325, 1024, 1024, 215, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 747, 1024, 1024, 1024, 16, 83, 1024, 1018, 1024, 63, 1024, 453, 1024, 82, 1024, 12, 1024, 1019, 1003, 32, 187, 1003, 1024, 1009, 354, 27, 1024, 1024, 1024, 1024, 524, 1024, 429, 1024, 1024, 124, 1024, 1024, 165, 1024, 1024, 1024, 1024, 417, 1024, 1024, 35, 5, 1024, 545, 1024, 1024, 317, 1024, 1024, 39, 1024, 747, 1024, 1024, 1024, 1024, 15, 1024, 475, 1024, 1024, 1024, 12, 1024, 1024, 713, 1024, 1024, 1024, 22, 1024, 428, 1024, 958, 1024, 1024, 217, 1024, 1024, 261, 63, 1005, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 25, 1024, 1024, 747, 1024, 1024, 1024, 1024, 494, 1005, 1002, 1024, 737, 1024, 1024, 1001, 1024, 12, 1024, 1024, 1024, 41, 300, 1024, 27, 1024, 217, 1024, 882, 1024, 1024, 132, 1024, 1024, 3, 1024, 1024, 681, 12, 1024, 1024, 535, 1024, 1024, 635, 1024, 354, 1024, 1024, 1024, 62, 1024, 5, 1024, 344, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 497, 1024, 1024, 67, 1024, 1024, 858, 1024, 1024, 1024, 1024, 144, 1024, 3, 1024, 1024, 1024, 100, 104, 1024, 1015, 1024, 127, 1024, 12, 1024, 35, 1024, 3, 1, 83, 1018, 1024, 391, 1024, 1024, 16, 563, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 608, 1024, 1024, 1024, 1024, 284, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024]]}
\ No newline at end of file
diff --git a/tests/fixtures/parakeet/expected_results_single.json b/tests/fixtures/parakeet/expected_results_single.json
new file mode 100644
index 000000000000..b6b686fa4223
--- /dev/null
+++ b/tests/fixtures/parakeet/expected_results_single.json
@@ -0,0 +1 @@
+{"transcriptions": ["mister quilter is the apostle of the middle classes and we are glad to welcome his gospel"], "scores": [-0.08922013640403748], "token_ids": [[1024, 1024, 1024, 1024, 1024, 1024, 19, 37, 132, 1024, 1024, 264, 128, 1024, 1024, 1024, 132, 1024, 58, 1024, 5, 645, 1024, 1000, 82, 52, 1024, 34, 1024, 5, 19, 68, 1007, 52, 1024, 235, 1024, 388, 1024, 27, 1024, 25, 1024, 56, 1024, 103, 1024, 1024, 727, 112, 1024, 22, 1024, 56, 1006, 1009, 405, 1024, 1024, 217, 1024, 1024, 95, 1003, 1024, 133, 1006, 1024, 1024, 1024, 1024, 1024, 1024, 1024]]}
\ No newline at end of file
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index a932a1fbac67..6a4060b0a731 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -88,22 +88,11 @@ def get_master_port(real_launcher=False):
if is_torch_available():
- from tests.trainer.test_trainer import ( # noqa
- RegressionModelConfig,
- RegressionPreTrainedModel,
- )
-
# hack to restore original logging level pre #21700
get_regression_trainer = partial(tests.trainer.test_trainer.get_regression_trainer, log_level="info")
-require_fsdp_version = require_fsdp
if is_accelerate_available():
- from accelerate.utils.constants import (
- FSDP_PYTORCH_VERSION,
- FSDP_SHARDING_STRATEGY,
- )
-
- require_fsdp_version = partial(require_fsdp, min_version=FSDP_PYTORCH_VERSION)
+ from accelerate.utils.constants import FSDP_SHARDING_STRATEGY
FSDP2_ACCELERATE_VERSION = "1.6.0"
@@ -142,7 +131,6 @@ def _parameterized_custom_name_func(func, param_num, param):
@require_accelerate
@require_torch_accelerator
-@require_fsdp_version
class TrainerIntegrationFSDP(TestCasePlus, TrainerIntegrationCommon):
def setUp(self):
super().setUp()
diff --git a/tests/generation/test_continuous_batching.py b/tests/generation/test_continuous_batching.py
index 3179479bdb11..943320bfe00b 100644
--- a/tests/generation/test_continuous_batching.py
+++ b/tests/generation/test_continuous_batching.py
@@ -20,6 +20,7 @@
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers.generation.continuous_batching.cache import group_layers_by_attn_type
+from transformers.generation.continuous_batching.continuous_api import build_attention_mask
from transformers.testing_utils import Expectations, require_kernels, require_torch_gpu, slow
@@ -88,6 +89,48 @@ def test_group_layers(
f"Test failed for: {layer_types_str = }, {sliding_window = }, {group_types = }",
)
+ @parameterized.expand(
+ [
+ ([0, 4], [0, 4], 1, ["1000", "1100", "1110", "1111"]),
+ ([0, 4], [0, 4], 2, ["1000", "1100", "0110", "0011"]),
+ ([0, 3], [0, 5], 1, ["11100", "11110", "11111"]),
+ ([0, 3], [0, 5], 3, ["11100", "01110", "00111"]),
+ ([0, 3, 6], [0, 3, 6], 1, ["100000", "110000", "111000", "000100", "000110", "000111"]),
+ ([0, 3, 6], [0, 3, 6], 2, ["100000", "110000", "011000", "000100", "000110", "000011"]),
+ ]
+ )
+ def test_attention_mask(
+ self,
+ cumulative_seqlens_q: list[int],
+ cumulative_seqlens_k: list[int],
+ sliding_window: int, # the sliding window size, 1 means no sliding window
+ str_expected_mask: list[str], # the attention mask, broken down by line as a string of 0s and 1s
+ ) -> None:
+ # Build expected mask
+ minus_inf = torch.finfo(torch.float32).min
+ expected_mask = torch.empty((cumulative_seqlens_q[-1], cumulative_seqlens_k[-1]), dtype=torch.float32)
+ for i, line in enumerate(str_expected_mask):
+ expected_mask[i, :] = torch.tensor([minus_inf if c == "0" else 0 for c in line])
+ # Build actual mask
+ actual_mask = torch.full_like(expected_mask, minus_inf) # function modifies in place
+ build_attention_mask(
+ actual_mask, torch.tensor(cumulative_seqlens_q), torch.tensor(cumulative_seqlens_k), sliding_window
+ )
+ # Check that the actual mask matches the expected mask
+ matches = (expected_mask == actual_mask).all()
+ # If it doesn't match, print the masks in a readable form and fail the test
+ if not matches:
+ str_mask = [
+ "".join("1" if x == 0 else "0" for x in token_attn_vector) for token_attn_vector in actual_mask
+ ]
+ str_mask = "\n".join(str_mask)
+ str_expected_mask = "\n".join(str_expected_mask)
+ self.fail(
+ f"Test failed for: {cumulative_seqlens_q = }, {cumulative_seqlens_k = }, {sliding_window = }\n"
+ f"Expected mask:\n{str_expected_mask}\n"
+ f"Actual mask:\n{str_mask}"
+ )
+
def _continuous_batching_parity(
self, model_id: str, attn_implementation: str, expected_outputs: dict[str, str]
) -> None:
diff --git a/tests/generation/test_flash_attention_parity.py b/tests/generation/test_flash_attention_parity.py
index bcf11b4dc4fc..969cdddcd38d 100644
--- a/tests/generation/test_flash_attention_parity.py
+++ b/tests/generation/test_flash_attention_parity.py
@@ -81,39 +81,31 @@ def _benchmark_generation(self, model, inputs, n_warmup=3, n_runs=5):
@slow
def test_flash_attention_2_3_parity(self):
model_id = "meta-llama/Llama-3.2-1B-Instruct"
- prompt = "The ETH AI Center is"
+ prompt = ["The ETH AI Center is", "What is life?"]
- # 1. Load FA2 model and tokenizer
- model_2 = AutoModelForCausalLM.from_pretrained(
+ # 1. Load model and tokenizer
+ model = AutoModelForCausalLM.from_pretrained(
model_id,
dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_id)
+ tokenizer.pad_token_id = tokenizer.eos_token_id
- # 2. Load FA3 model
- try:
- model_3 = AutoModelForCausalLM.from_pretrained(
- model_id,
- dtype=torch.bfloat16,
- attn_implementation="flash_attention_3",
- ).to("cuda")
- except (ValueError, ImportError) as e:
- pytest.skip(f"Could not load Flash Attention 3 model, skipping test. Error: {e}")
-
- # 3. Generate with both models
- inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+ # 2. Generate with both models
+ inputs = tokenizer(prompt, padding=True, padding_side="left", return_tensors="pt").to("cuda")
with torch.no_grad():
- output_2 = model_2.generate(
+ output_2 = model.generate(
**inputs, max_new_tokens=20, do_sample=False, output_scores=True, return_dict_in_generate=True
)
- output_3 = model_3.generate(
+ model.set_attn_implementation("flash_attention_3")
+ output_3 = model.generate(
**inputs, max_new_tokens=20, do_sample=False, output_scores=True, return_dict_in_generate=True
)
- # 4. Correctness check
- # 4a. Logits
+ # 3. Correctness check
+ # 3a. Logits
logits_2 = torch.stack(output_2.scores)
logits_3 = torch.stack(output_3.scores)
torch.testing.assert_close(logits_2, logits_3, atol=1e-3, rtol=1e-3)
@@ -121,22 +113,27 @@ def test_flash_attention_2_3_parity(self):
logprobs_3 = torch.nn.functional.log_softmax(logits_3, dim=-1)
max_logprob_diff = torch.max(torch.abs(logprobs_2 - logprobs_3)).item()
- # 4b. Generated text
- text_2 = tokenizer.decode(output_2.sequences[0], skip_special_tokens=True)
- text_3 = tokenizer.decode(output_3.sequences[0], skip_special_tokens=True)
- rouge_score = self._calculate_rouge_l([text_2], [text_3])[0]
- assert rouge_score > 0.99, f"Generated texts do not match (ROUGE-L: {rouge_score})"
+ # 3b. Generated text
+ text_2s, text_3s = [], []
+ for i in range(len(prompt)):
+ text_2s.append(tokenizer.decode(output_2.sequences[i], skip_special_tokens=True))
+ text_3s.append(tokenizer.decode(output_3.sequences[i], skip_special_tokens=True))
+
+ rouge_scores = self._calculate_rouge_l(text_2s, text_3s)
+ for i in range(len(rouge_scores)):
+ assert rouge_scores[i] > 0.99, f"Generated texts at prompt {i} do not match (ROUGE-L: {rouge_scores[i]})"
- # 5. Performance check
+ # 4. Performance check
with torch.no_grad():
- time_2 = self._benchmark_generation(model_2, inputs)
- time_3 = self._benchmark_generation(model_3, inputs)
+ time_3 = self._benchmark_generation(model, inputs)
+ model.set_attn_implementation("flash_attention_2")
+ time_2 = self._benchmark_generation(model, inputs)
print(f"\n--- Flash Attention {2, 3} Parity Test on {model_id} ---")
print(f"Prompt: '{prompt}'")
- print(f"Generated text with Flash Attention 2: {text_2}")
- print(f"Generated text with Flash Attention 3: {text_3}")
- print(f"ROUGE-L: {rouge_score}")
+ print(f"Generated text with Flash Attention 2: {text_2s}")
+ print(f"Generated text with Flash Attention 3: {text_3s}")
+ print(f"ROUGE-L: {rouge_scores}")
print(f"Max absolute difference in logprobs: {max_logprob_diff:.5e}")
print(f"Flash Attention 2 latency: {time_2:.2f} ms")
print(f"Flash Attention 3 latency: {time_3:.2f} ms")
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
index 680002d4600b..9f6f7a01347b 100644
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -23,6 +23,7 @@
import unittest
import warnings
from pathlib import Path
+from typing import Optional
import numpy as np
import pytest
@@ -927,32 +928,44 @@ def test_prompt_lookup_decoding_stops_at_eos(self):
self.assertTrue(output_prompt_lookup.shape[-1] == 10)
@pytest.mark.generate
- def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
+ def test_left_padding_compatibility(
+ self, unpadded_custom_inputs: Optional[dict] = None, padded_custom_inputs: Optional[dict] = None
+ ):
+ """
+ Tests that adding left-padding yields the same logits as the original input. Exposes arguments for custom
+ inputs for overwrites, to prevent full rewrites of the test when all we need is model-specific input handling.
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
+ ! If you overwrite this test, make sure to document why you need to overwrite it !
+
+ NOTE: left-padding results in small numerical differences. This is expected.
+ See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
+
+ Args:
+ unpadded_custom_inputs (`dict`, *optional*):
+ Used in test overwrites. Custom inputs to add/overwrite over the default test inputs.
+ padded_custom_inputs (`dict`, *optional*):
+ Used in test overwrites. Custom inputs to add/overwrite over the padded test input handcrafted in this
+ test. Commonly used e.g. with multimodal cross attention masks.
+ """
- # - The model must support padding
+ # First, filter out models that don't support left padding
+ # 1. The model must support padding
if not self.has_attentions:
self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
+ # 2. [encoder-decoder] The model must be a decoder-only architecture. Encoder-based architectures can use
+ # right-padding in their (encoder) inputs. Encoder-decoder may use left-padding on their decoder inputs
+ # [TODO: lift this restriction? technically, we can test padding the decoder inputs.]
decoder_only_classes = []
for model_class in self.all_generative_model_classes:
config, _ = self.prepare_config_and_inputs_for_generate()
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
continue
else:
decoder_only_classes.append(model_class)
if len(decoder_only_classes) == 0:
self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
+ # 3. [old models] Decoder-only architectures derived from encoder-decoder models could support it in theory,
+ # but we haven't added support for it yet. We skip these models for now.
has_encoder_attributes = any(
attr_name
for attr_name in config.to_dict()
@@ -963,48 +976,73 @@ def test_left_padding_compatibility(self):
reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
)
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
+ # Now we can start testing
+ unpadded_custom_inputs = unpadded_custom_inputs or {}
+ padded_custom_inputs = padded_custom_inputs or {}
+
+ def _prepare_model_kwargs(model_inputs, signature):
+ model_kwargs = {"input_ids": model_inputs["input_ids"], "attention_mask": model_inputs["attention_mask"]}
if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = torch.cumsum(model_inputs["attention_mask"], dim=-1) - 1
+ position_ids.masked_fill_(model_inputs["attention_mask"] == 0, 1)
model_kwargs["position_ids"] = position_ids
if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[1], device=torch_device)
+ cache_position = torch.arange(model_inputs["input_ids"].shape[1], device=torch_device)
model_kwargs["cache_position"] = cache_position
+ # forward all other inputs, if they are in the signature
+ model_kwargs.update({k: v for k, v in model_inputs.items() if k not in model_kwargs and k in signature})
return model_kwargs
for model_class in decoder_only_classes:
config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
- attention_mask = inputs_dict.get("attention_mask")
- if attention_mask is None:
- attention_mask = torch.ones_like(input_ids)
-
model = model_class(config).to(torch_device).eval()
signature = inspect.signature(model.forward).parameters.keys()
- # no cache as some models require special cache classes to be init outside forward
+ # No cache to simplify the test (some models need careful init)
model.generation_config.use_cache = False
+ inputs_dict.update(unpadded_custom_inputs)
+ # special case: an inexistent `attention_mask` is a full mask
+ inputs_dict["attention_mask"] = inputs_dict.get("attention_mask", None)
+ if inputs_dict["attention_mask"] is None:
+ inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["input_ids"])
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(**model_kwargs).logits[:, -1, :]
+ # Get output logits from inputs without padding
+ model_kwargs_wo_padding = _prepare_model_kwargs(inputs_dict, signature)
+ next_logits_wo_padding = model(**model_kwargs_wo_padding).logits[:, -1, :]
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
+ # Prepare padding on common inputs (pad length 32)
+ input_ids = inputs_dict["input_ids"]
+ attention_mask = inputs_dict["attention_mask"]
+ token_type_ids = inputs_dict.get("token_type_ids", None)
+ pad_token_id = getattr(config.get_text_config(decoder=True), "pad_token_id", None) or 0
pad_size = (input_ids.shape[0], 32, *input_ids.shape[2:])
padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
padded_input_ids = torch.cat((padding, input_ids), dim=1)
padded_attention_mask = torch.cat(
(torch.zeros(pad_size[:2], dtype=input_ids.dtype, device=torch_device), attention_mask), dim=1
)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(**model_kwargs).logits[:, -1, :]
+ if token_type_ids is not None:
+ padded_token_type_ids = torch.cat(
+ (
+ # Assumption: `0` is a good default value for padding token type ids
+ torch.zeros(pad_size[:2], dtype=input_ids.dtype, device=torch_device),
+ token_type_ids,
+ ),
+ dim=1,
+ )
+ else:
+ padded_token_type_ids = None
+
+ # Get output logits from inputs with left-padding (pad length 32)
+ padded_inputs_dict = copy.deepcopy(inputs_dict)
+ padded_inputs_dict["input_ids"] = padded_input_ids
+ padded_inputs_dict["attention_mask"] = padded_attention_mask
+ if padded_token_type_ids is not None:
+ padded_inputs_dict["token_type_ids"] = padded_token_type_ids
+ padded_inputs_dict.update(padded_custom_inputs)
+
+ model_kwargs_with_padding = _prepare_model_kwargs(padded_inputs_dict, signature)
+ next_logits_with_padding = model(**model_kwargs_with_padding).logits[:, -1, :]
# They should result in very similar logits
torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
@@ -1192,7 +1230,7 @@ def test_generate_from_inputs_embeds(self, _, num_beams):
# This test is for decoder-only models (encoder-decoder models have native input embeddings support in the
# decoder)
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
continue
config.is_decoder = True
@@ -1271,7 +1309,7 @@ def test_generate_from_inputs_embeds_with_static_cache(self):
config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
self.skipTest(reason="This model is encoder-decoder and has Encoder-Decoder Cache")
model = model_class(config).to(torch_device).eval()
@@ -1422,7 +1460,7 @@ def test_generate_continue_from_inputs_embeds(self):
if "token_type_ids" in inputs_dict:
del inputs_dict["token_type_ids"]
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
self.skipTest(reason="This model is encoder-decoder")
# TODO (joao, raushan): the correct line below is `if not hasattr(config.get_text_config(), "use_cache")`,
# but it breaks a few models. Fix and then apply `has_similar_generate_outputs` pattern
@@ -1495,7 +1533,7 @@ def test_generate_with_static_cache(self):
set_config_for_less_flaky_test(config)
main_input = inputs_dict[model_class.main_input_name]
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
self.skipTest(reason="This model is encoder-decoder and has Encoder-Decoder Cache")
config.is_decoder = True
@@ -1550,10 +1588,7 @@ def test_generate_with_quant_cache(self):
for model_class in self.all_generative_model_classes:
config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- if (
- config.get_text_config(decoder=True).is_encoder_decoder
- or not model_class._supports_default_dynamic_cache()
- ):
+ if config.is_encoder_decoder or not model_class._supports_default_dynamic_cache():
self.skipTest(reason="This model does not support the quantized cache format")
config.is_decoder = True
@@ -1653,7 +1688,7 @@ def test_generate_compile_model_forward_fullgraph(self):
if not has_defined_cache_implementation:
decoder_cache = (
gen_out.past_key_values.self_attention_cache
- if config.get_text_config(decoder=True).is_encoder_decoder
+ if config.is_encoder_decoder
else gen_out.past_key_values
)
self.assertTrue(isinstance(decoder_cache, DynamicCache))
@@ -1679,7 +1714,7 @@ def test_generate_compile_model_forward_fullgraph(self):
# sanity checks
decoder_cache = (
gen_out.past_key_values.self_attention_cache
- if config.get_text_config(decoder=True).is_encoder_decoder
+ if config.is_encoder_decoder
else gen_out.past_key_values
)
self.assertFalse(isinstance(decoder_cache, DynamicCache))
@@ -2387,6 +2422,7 @@ def _check_generate_outputs(self, output, config, use_cache=False, num_return_se
"zamba",
"zamba2",
"lfm2",
+ "lfm2-vl",
)
has_standard_cache = not any(
model_name in config.__class__.__name__.lower() for model_name in models_without_standard_cache
diff --git a/tests/models/aimv2/test_modeling_aimv2.py b/tests/models/aimv2/test_modeling_aimv2.py
index 524cdc5e3016..6c9752f6f049 100644
--- a/tests/models/aimv2/test_modeling_aimv2.py
+++ b/tests/models/aimv2/test_modeling_aimv2.py
@@ -39,7 +39,6 @@
from ...test_modeling_common import (
TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION,
ModelTesterMixin,
- _config_zero_init,
_test_eager_matches_sdpa_inference,
floats_tensor,
ids_tensor,
@@ -430,30 +429,6 @@ def test_model_get_set_embeddings(self):
def test_multi_gpu_data_parallel_forward(self):
pass
- # Override as the `logit_scale` parameter initialization is different for Aimv2
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_load_vision_text_config(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
diff --git a/tests/models/align/test_modeling_align.py b/tests/models/align/test_modeling_align.py
index 167cb1ff7c2e..afda3156da7c 100644
--- a/tests/models/align/test_modeling_align.py
+++ b/tests/models/align/test_modeling_align.py
@@ -469,7 +469,7 @@ def test_batching_equivalence(self, atol=3e-4, rtol=3e-4):
@unittest.skip(reason="Start to fail after using torch `cu118`.")
def test_multi_gpu_data_parallel_forward(self):
- super().test_multi_gpu_data_parallel_forward()
+ pass
@unittest.skip(reason="Hidden_states is tested in individual model tests")
def test_hidden_states_output(self):
@@ -491,35 +491,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `temperature` parameter initialization is different for ALIGN
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `temperature` is initialized as per the original implementation
- if name == "temperature":
- self.assertAlmostEqual(
- param.data.item(),
- 1.0,
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif name == "text_projection.weight":
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/altclip/test_modeling_altclip.py b/tests/models/altclip/test_modeling_altclip.py
index 2a36470051f8..7795f6883bb1 100755
--- a/tests/models/altclip/test_modeling_altclip.py
+++ b/tests/models/altclip/test_modeling_altclip.py
@@ -467,29 +467,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for AltCLIP
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/apertus/test_modeling_apertus.py b/tests/models/apertus/test_modeling_apertus.py
index 77769c430e08..30e7fdbf21f5 100644
--- a/tests/models/apertus/test_modeling_apertus.py
+++ b/tests/models/apertus/test_modeling_apertus.py
@@ -33,7 +33,6 @@
if is_torch_available():
from transformers import (
- ApertusConfig,
ApertusForCausalLM,
ApertusForTokenClassification,
ApertusModel,
@@ -42,23 +41,11 @@
class ApertusModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = ApertusConfig
base_model_class = ApertusModel
- causal_lm_class = ApertusForCausalLM
- token_class = ApertusForTokenClassification
@require_torch
class ApertusModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- ApertusModel,
- ApertusForCausalLM,
- ApertusForTokenClassification,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": ApertusModel,
@@ -68,8 +55,6 @@ class ApertusModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
model_tester_class = ApertusModelTester
# Need to use `0.8` instead of `0.9` for `test_cpu_offload`
diff --git a/tests/models/arcee/test_modeling_arcee.py b/tests/models/arcee/test_modeling_arcee.py
index 7c6096081ecd..a8b485fd7eb2 100644
--- a/tests/models/arcee/test_modeling_arcee.py
+++ b/tests/models/arcee/test_modeling_arcee.py
@@ -43,26 +43,11 @@
class ArceeModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = ArceeConfig
base_model_class = ArceeModel
- causal_lm_class = ArceeForCausalLM
- sequence_class = ArceeForSequenceClassification
- token_class = ArceeForTokenClassification
@require_torch
class ArceeModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- ArceeModel,
- ArceeForCausalLM,
- ArceeForSequenceClassification,
- ArceeForQuestionAnswering,
- ArceeForTokenClassification,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": ArceeModel,
@@ -75,8 +60,6 @@ class ArceeModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False
model_tester_class = ArceeModelTester
diff --git a/tests/models/aria/test_modeling_aria.py b/tests/models/aria/test_modeling_aria.py
index 17259a5effa8..8a469e081fe5 100644
--- a/tests/models/aria/test_modeling_aria.py
+++ b/tests/models/aria/test_modeling_aria.py
@@ -199,10 +199,6 @@ def setUp(self):
self.model_tester = AriaVisionText2TextModelTester(self)
self.config_tester = ConfigTester(self, config_class=AriaConfig, has_text_modality=False)
- @unittest.skip(reason="Unstable test")
- def test_initialization(self):
- pass
-
SKIP = False
torch_accelerator_module = getattr(torch, torch_device)
diff --git a/tests/models/auto/test_modeling_auto.py b/tests/models/auto/test_modeling_auto.py
index 7af5315f844c..10de9f157086 100644
--- a/tests/models/auto/test_modeling_auto.py
+++ b/tests/models/auto/test_modeling_auto.py
@@ -26,7 +26,7 @@
from huggingface_hub import Repository
import transformers
-from transformers import BertConfig, GPT2Model, is_safetensors_available, is_torch_available
+from transformers import BertConfig, GPT2Model, is_torch_available
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
from transformers.testing_utils import (
DUMMY_UNKNOWN_IDENTIFIER,
@@ -112,7 +112,7 @@ def test_model_from_pretrained(self):
self.assertEqual(len(loading_info["missing_keys"]), 0)
# When using PyTorch checkpoint, the expected value is `8`. With `safetensors` checkpoint (if it is
# installed), the expected value becomes `7`.
- EXPECTED_NUM_OF_UNEXPECTED_KEYS = 7 if is_safetensors_available() else 8
+ EXPECTED_NUM_OF_UNEXPECTED_KEYS = 7
self.assertEqual(len(loading_info["unexpected_keys"]), EXPECTED_NUM_OF_UNEXPECTED_KEYS)
self.assertEqual(len(loading_info["mismatched_keys"]), 0)
self.assertEqual(len(loading_info["error_msgs"]), 0)
diff --git a/tests/models/autoformer/test_modeling_autoformer.py b/tests/models/autoformer/test_modeling_autoformer.py
index 954f9f16622b..414670b8d919 100644
--- a/tests/models/autoformer/test_modeling_autoformer.py
+++ b/tests/models/autoformer/test_modeling_autoformer.py
@@ -288,7 +288,7 @@ def test_forward_signature(self):
"future_time_features",
]
- if model.__class__.__name__ in ["AutoformerForPrediction"]:
+ if model.__class__.__name__ == "AutoformerForPrediction":
expected_arg_names.append("future_observed_mask")
expected_arg_names.extend(
diff --git a/tests/models/aya_vision/test_modeling_aya_vision.py b/tests/models/aya_vision/test_modeling_aya_vision.py
index 436cba19c290..27a8d5159492 100644
--- a/tests/models/aya_vision/test_modeling_aya_vision.py
+++ b/tests/models/aya_vision/test_modeling_aya_vision.py
@@ -71,7 +71,7 @@ def __init__(
"vocab_size": 99,
"hidden_size": 128,
"intermediate_size": 37,
- "num_hidden_layers": 4,
+ "num_hidden_layers": 2,
"num_attention_heads": 4,
"output_channels": 64,
"hidden_act": "silu",
@@ -198,10 +198,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- @unittest.skip(reason="Siglip uses the same initialization scheme as the Flax original implementation")
- def test_initialization(self):
- pass
-
@unittest.skip(reason="Compile not yet supported because in LLava models")
@pytest.mark.torch_compile_test
def test_sdpa_can_compile_dynamic(self):
diff --git a/tests/models/bamba/test_modeling_bamba.py b/tests/models/bamba/test_modeling_bamba.py
index c2e7c435dbfa..2c25efc24325 100644
--- a/tests/models/bamba/test_modeling_bamba.py
+++ b/tests/models/bamba/test_modeling_bamba.py
@@ -41,7 +41,7 @@
from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -73,7 +73,7 @@ def __init__(
use_labels=True,
vocab_size=99,
hidden_size=32,
- num_hidden_layers=4,
+ num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
intermediate_size=64,
@@ -336,37 +336,6 @@ def test_decoder_model_past_with_large_inputs(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)
- def test_initialization(self):
- r"""
- Overriding the test_initialization test as the A_log and D params of the Bamba mixer are initialized differently
- """
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if "A_log" in name:
- A = torch.arange(1, config.mamba_n_heads + 1, dtype=torch.float32)
- torch.testing.assert_close(param.data, torch.log(A), rtol=1e-5, atol=1e-5)
- elif "D" in name:
- D = torch.ones(config.mamba_n_heads, dtype=torch.float32)
- torch.testing.assert_close(param.data, D, rtol=1e-5, atol=1e-5)
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
- def test_mismatched_shapes_have_properly_initialized_weights(self):
- r"""
- Overriding the test_mismatched_shapes_have_properly_initialized_weights test because A_log and D params of the
- Bamba mixer are initialized differently and we tested that in test_initialization
- """
- self.skipTest(reason="Cumbersome and redundant for Bamba")
-
def test_attention_outputs(self):
r"""
Overriding the test_attention_outputs test as the Bamba model outputs attention only for its attention layers
@@ -438,88 +407,11 @@ def test_batching_equivalence(self):
super().test_batching_equivalence()
self.model_tester.use_input_mask = orig
- # essentially the same test in test_utils, just adjustment for rtol for this model
@pytest.mark.generate
def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
-
- # - The model must support padding
- if not self.has_attentions:
- self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
- decoder_only_classes = []
- for model_class in self.all_generative_model_classes:
- config, _ = self.prepare_config_and_inputs_for_generate()
- if config.is_encoder_decoder:
- continue
- else:
- decoder_only_classes.append(model_class)
- if len(decoder_only_classes) == 0:
- self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
- has_encoder_attributes = any(
- attr_name
- for attr_name in config.to_dict()
- if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
- )
- if has_encoder_attributes:
- self.skipTest(
- reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
- )
-
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in decoder_only_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
-
- # - for left padding we absolutely need to use an all ones
- # attention mask, so we do not use the one in inputs_dict
- attention_mask = torch.ones_like(input_ids)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(**model_kwargs).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(**model_kwargs).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
+ # TODO: document why a random attention mask causes this test to fail, but a full mask doesn't
+ unpadded_custom_inputs = {"attention_mask": None}
+ super().test_left_padding_compatibility(unpadded_custom_inputs=unpadded_custom_inputs)
@unittest.skip(
"Bamba requires additionally specifying position_ids, seq_idx, and FlashAttentionKwargs for padding-free training."
diff --git a/tests/models/beit/test_modeling_beit.py b/tests/models/beit/test_modeling_beit.py
index 48bfda0a4d85..fdcd90774bd8 100644
--- a/tests/models/beit/test_modeling_beit.py
+++ b/tests/models/beit/test_modeling_beit.py
@@ -34,7 +34,7 @@
from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -383,24 +383,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- # we skip lambda parameters as these require special initial values
- # determined by config.layer_scale_init_value
- if "lambda" in name:
- continue
- if param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@slow
def test_model_from_pretrained(self):
model_name = "microsoft/beit-base-patch16-224"
diff --git a/tests/models/bit/test_modeling_bit.py b/tests/models/bit/test_modeling_bit.py
index aa616c85aa6c..e3b12e46b19a 100644
--- a/tests/models/bit/test_modeling_bit.py
+++ b/tests/models/bit/test_modeling_bit.py
@@ -28,7 +28,6 @@
if is_torch_available():
import torch
- from torch import nn
from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
@@ -201,22 +200,6 @@ def test_backbone(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_backbone(*config_and_inputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- for model_class in self.all_model_classes:
- model = model_class(config=config)
- for name, module in model.named_modules():
- if isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
- self.assertTrue(
- torch.all(module.weight == 1),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- self.assertTrue(
- torch.all(module.bias == 0),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_hidden_states_output(self):
def check_hidden_states_output(inputs_dict, config, model_class):
model = model_class(config)
diff --git a/tests/models/bitnet/test_modeling_bitnet.py b/tests/models/bitnet/test_modeling_bitnet.py
index 75d885ba4d51..19bc0c45eb2e 100644
--- a/tests/models/bitnet/test_modeling_bitnet.py
+++ b/tests/models/bitnet/test_modeling_bitnet.py
@@ -49,7 +49,7 @@ def __init__(
use_input_mask=True,
vocab_size=99,
hidden_size=64,
- num_hidden_layers=5,
+ num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
intermediate_size=37,
diff --git a/tests/models/blip/test_modeling_blip.py b/tests/models/blip/test_modeling_blip.py
index 189773afd399..a061cbde7ccc 100644
--- a/tests/models/blip/test_modeling_blip.py
+++ b/tests/models/blip/test_modeling_blip.py
@@ -459,41 +459,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for Blip
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- # See PR #38607 (to avoid flakiness)
- data = torch.flatten(param.data)
- n_elements = torch.numel(data)
- # skip 2.5% of elements on each side to avoid issues caused by `nn.init.trunc_normal_` described in
- # https://github.com/huggingface/transformers/pull/27906#issuecomment-1846951332
- n_elements_to_skip_on_each_side = int(n_elements * 0.025)
- data_to_check = torch.sort(data).values
- if n_elements_to_skip_on_each_side > 0:
- data_to_check = data_to_check[
- n_elements_to_skip_on_each_side:-n_elements_to_skip_on_each_side
- ]
- self.assertIn(
- ((data_to_check.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
@@ -990,30 +955,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- # override as the `logit_scale` parameter initialization is different for Blip
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
@@ -1208,30 +1149,6 @@ def test_training_gradient_checkpointing(self):
loss = model(**inputs).loss
loss.backward()
- # override as the `logit_scale` parameter initialization is different for Blip
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/blip_2/test_modeling_blip_2.py b/tests/models/blip_2/test_modeling_blip_2.py
index 0b3ab74d519c..9c6cd4ce1f42 100644
--- a/tests/models/blip_2/test_modeling_blip_2.py
+++ b/tests/models/blip_2/test_modeling_blip_2.py
@@ -13,13 +13,11 @@
# limitations under the License.
"""Testing suite for the PyTorch BLIP-2 model."""
-import copy
import inspect
import tempfile
import unittest
import numpy as np
-import pytest
import requests
from parameterized import parameterized
@@ -43,7 +41,6 @@
from ...test_modeling_common import (
TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION,
ModelTesterMixin,
- _config_zero_init,
floats_tensor,
ids_tensor,
random_attention_mask,
@@ -597,89 +594,6 @@ def _check_generate_outputs(self, output, config, use_cache=False, num_return_se
output, config, use_cache=use_cache, num_return_sequences=num_return_sequences, num_beams=num_beams
)
- # overwrite because BLIP2 cannot generate only from input ids, and requires pixel values in all cases to be present
- @pytest.mark.generate
- def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
-
- # - The model must support padding
- if not self.has_attentions:
- self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
- decoder_only_classes = []
- for model_class in self.all_generative_model_classes:
- config, _ = self.prepare_config_and_inputs_for_generate()
- if config.is_encoder_decoder:
- continue
- else:
- decoder_only_classes.append(model_class)
- if len(decoder_only_classes) == 0:
- self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
- has_encoder_attributes = any(
- attr_name
- for attr_name in config.to_dict()
- if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
- )
- if has_encoder_attributes:
- self.skipTest(
- reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
- )
-
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in decoder_only_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
- attention_mask = inputs_dict.get("attention_mask")
- pixel_values = inputs_dict["pixel_values"]
- if attention_mask is None:
- attention_mask = torch.ones_like(input_ids)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(**model_kwargs, pixel_values=pixel_values).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(**model_kwargs, pixel_values=pixel_values).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
-
# this class is based on `T5ModelTester` found in tests/models/t5/test_modeling_t5.py
class Blip2TextModelTester:
@@ -1075,23 +989,6 @@ def test_get_qformer_features(self):
(self.model_tester.vision_model_tester.batch_size, 10, config.vision_config.hidden_size),
)
- # override from common to deal with nested configurations (`vision_config`, `text_config` and `qformer_config`)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for key in ["vision_config", "qformer_config", "text_config"]:
- setattr(configs_no_init, key, _config_zero_init(getattr(configs_no_init, key)))
- for model_class in self.all_model_classes:
- model = model_class(config=copy.deepcopy(configs_no_init))
- for name, param in model.named_parameters():
- if param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@unittest.skip("T5 backbone deepcopies the configs, and fixing it would be more involved")
def test_internal_model_config_and_subconfig_are_same(self):
pass
@@ -1601,36 +1498,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif name == "temp":
- self.assertAlmostEqual(
- param.data.item(),
- 0.07,
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# We will verify our results on an image of cute cats
def prepare_img():
diff --git a/tests/models/bloom/test_modeling_bloom.py b/tests/models/bloom/test_modeling_bloom.py
index 9480aac9ae94..e8cdd43ff5f7 100644
--- a/tests/models/bloom/test_modeling_bloom.py
+++ b/tests/models/bloom/test_modeling_bloom.py
@@ -446,7 +446,7 @@ def test_batch_generation(self):
@slow
@require_torch_accelerator
- def test_batch_generation_padd(self):
+ def test_batch_generation_padding(self):
path_560m = "bigscience/bloom-560m"
model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device)
model = model.eval()
diff --git a/tests/models/blt/__init__.py b/tests/models/blt/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/blt/test_modeling_blt.py b/tests/models/blt/test_modeling_blt.py
new file mode 100644
index 000000000000..34aab8f179c9
--- /dev/null
+++ b/tests/models/blt/test_modeling_blt.py
@@ -0,0 +1,555 @@
+# 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.
+"""Testing suite for the PyTorch Blt model."""
+
+import unittest
+
+import pytest
+from parameterized import parameterized
+
+from transformers import AutoTokenizer, is_torch_available, set_seed
+from transformers.testing_utils import (
+ cleanup,
+ require_read_token,
+ require_torch,
+ require_torch_accelerator,
+ require_torch_bf16,
+ slow,
+ torch_device,
+)
+
+from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
+from ...test_modeling_common import (
+ TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION,
+ _test_eager_matches_sdpa_inference,
+ ids_tensor,
+)
+
+
+if is_torch_available():
+ import torch
+
+ from transformers import BltConfig, BltForCausalLM, BltModel
+from transformers.models.blt.modeling_blt import BltRotaryEmbedding
+
+
+class BltModelTester(CausalLMModelTester):
+ if is_torch_available():
+ base_model_class = BltModel
+
+ def __init__(
+ self,
+ parent,
+ ignore_index=-100,
+ seq_length=7,
+ is_training=True,
+ ):
+ super().__init__(parent)
+ self.parent = parent
+ self.ignore_index = ignore_index
+ self.seq_length = seq_length
+ self.is_training = is_training
+ self.batch_size = 3
+
+ # Common parameters for all configs
+ self.hidden_size = 16
+ self.num_hidden_layers = 1
+ self.num_attention_heads = 2
+ self.num_key_value_heads = 2
+ self.intermediate_size = 32
+ self.hidden_act = "silu"
+ self.max_position_embeddings = 32
+ self.vocab_size = 32
+ self.rope_theta = 500000.0
+ self.rope_scaling = {"rope_type": "default"}
+ self.rms_norm_eps = 1e-5
+ self.dropout = 0.0
+ self.encoder_hash_byte_group_size = [2, 3]
+ self.encoder_hash_byte_group_vocab = 64
+ self.encoder_hash_byte_group_nb_functions = 1
+ # Common parameters for all configs
+ self.patcher_config = {
+ "hidden_size": self.hidden_size,
+ "num_hidden_layers": self.num_hidden_layers,
+ "num_attention_heads": self.num_attention_heads,
+ "num_key_value_heads": self.num_key_value_heads,
+ "intermediate_size": self.intermediate_size,
+ "max_position_embeddings": self.max_position_embeddings,
+ "rope_theta": self.rope_theta,
+ "rope_scaling": self.rope_scaling,
+ "hidden_act": self.hidden_act,
+ "rms_norm_eps": self.rms_norm_eps,
+ "dropout": self.dropout,
+ }
+
+ self.encoder_config = {
+ "hidden_size": self.hidden_size,
+ "num_hidden_layers": self.num_hidden_layers,
+ "num_attention_heads": self.num_attention_heads,
+ "num_key_value_heads": self.num_key_value_heads,
+ "intermediate_size": self.intermediate_size,
+ "max_position_embeddings": self.max_position_embeddings,
+ "rope_theta": self.rope_theta,
+ "rope_scaling": self.rope_scaling,
+ "hidden_act": self.hidden_act,
+ "rms_norm_eps": self.rms_norm_eps,
+ "dropout": self.dropout,
+ }
+
+ self.decoder_config = {
+ "vocab_size": self.vocab_size,
+ "hidden_size": self.hidden_size,
+ "hidden_size_global": self.hidden_size * 2, # Must match global transformer output size
+ "num_hidden_layers": self.num_hidden_layers,
+ "num_attention_heads": self.num_attention_heads,
+ "num_key_value_heads": self.num_key_value_heads,
+ "intermediate_size": self.intermediate_size,
+ "max_position_embeddings": self.max_position_embeddings,
+ "rope_theta": self.rope_theta,
+ "rope_scaling": self.rope_scaling,
+ "hidden_act": self.hidden_act,
+ "rms_norm_eps": self.rms_norm_eps,
+ "dropout": self.dropout,
+ }
+
+ self.global_config = {
+ "hidden_size": self.hidden_size * 2, # Double the hidden size for global transformer
+ "num_hidden_layers": self.num_hidden_layers,
+ "num_attention_heads": self.num_attention_heads,
+ "num_key_value_heads": self.num_key_value_heads,
+ "intermediate_size": self.intermediate_size,
+ "max_position_embeddings": self.max_position_embeddings,
+ "rope_theta": self.rope_theta,
+ "rope_scaling": self.rope_scaling,
+ "hidden_act": self.hidden_act,
+ "rms_norm_eps": self.rms_norm_eps,
+ "dropout": self.dropout,
+ }
+
+ self.num_hidden_layers = self.encoder_config["num_hidden_layers"]
+
+ def get_config(self):
+ config = BltConfig(
+ vocab_size=self.vocab_size,
+ max_position_embeddings=self.max_position_embeddings,
+ patch_in_forward=False, # Disable patching for tests
+ patch_size=4,
+ patching_mode="entropy",
+ patching_threshold=1.335442066192627,
+ patching_batch_size=1,
+ max_patch_length=None,
+ cross_attn_k=2,
+ encoder_hash_byte_group_size=self.encoder_hash_byte_group_size,
+ encoder_hash_byte_group_vocab=self.encoder_hash_byte_group_vocab,
+ encoder_hash_byte_group_nb_functions=self.encoder_hash_byte_group_nb_functions,
+ patcher_config=self.patcher_config,
+ encoder_config=self.encoder_config,
+ decoder_config=self.decoder_config,
+ global_config=self.global_config,
+ rope_scaling=self.rope_scaling,
+ tie_word_embeddings=False,
+ )
+
+ config.num_attention_heads = config.decoder_config.num_attention_heads
+ config.num_hidden_layers = config.encoder_config.num_hidden_layers
+ config.hidden_size = config.decoder_config.hidden_size
+
+ return config
+
+
+@require_torch
+class BltModelTest(CausalLMModelTest, unittest.TestCase):
+ all_model_classes = (
+ (
+ BltModel,
+ BltForCausalLM,
+ )
+ if is_torch_available()
+ else ()
+ )
+ pipeline_model_mapping = (
+ {
+ "feature-extraction": BltModel,
+ "text-generation": BltForCausalLM,
+ }
+ if is_torch_available()
+ else {}
+ )
+ test_headmasking = False
+ test_pruning = False
+ fx_compatible = False
+ model_tester_class = BltModelTester
+ rotary_embedding_layer = BltRotaryEmbedding # Enables RoPE tests if set
+
+ # Need to use `0.8` instead of `0.9` for `test_cpu_offload`
+ # This is because we are hitting edge cases with the causal_mask buffer
+ model_split_percents = [0.5, 0.7, 0.8]
+
+ # used in `test_torch_compile_for_training`
+ _torch_compile_train_cls = BltForCausalLM if is_torch_available() else None
+
+ @pytest.mark.generate
+ @parameterized.expand([("greedy", 1), ("beam search", 2)])
+ @unittest.skip(
+ "Blt requires real token IDs for its hash-based embedding computation, making inputs_embeds generation incompatible with identical outputs"
+ )
+ def test_generate_from_inputs_embeds(self, _, num_beams):
+ pass
+
+ @pytest.mark.generate
+ @unittest.skip(
+ "Blt requires real token IDs for its hash-based embedding computation, making inputs_embeds generation incompatible with identical outputs"
+ )
+ def test_inputs_embeds_matches_input_ids(self):
+ pass
+
+ @parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
+ def test_eager_matches_sdpa_inference(
+ self,
+ name,
+ torch_dtype,
+ padding_side,
+ use_attention_mask,
+ output_attentions,
+ enable_kernels,
+ ):
+ "We need to relax a bit the `atols` for fp32 here due to the altup projections"
+ atols = {
+ ("cpu", False, torch.float32): 2e-2, # this was relaxed
+ ("cpu", False, torch.float16): 5e-3,
+ ("cpu", False, torch.bfloat16): 1e-2,
+ ("cpu", True, torch.float32): 2e-2, # this was relaxed
+ ("cpu", True, torch.float16): 5e-3,
+ ("cpu", True, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float32): 2e-2, # this was relaxed
+ ("cuda", False, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float16): 5e-3,
+ ("cuda", True, torch.float32): 2e-2, # this was relaxed
+ ("cuda", True, torch.bfloat16): 1e-2,
+ ("cuda", True, torch.float16): 5e-3,
+ }
+ _test_eager_matches_sdpa_inference(
+ self, name, torch_dtype, padding_side, use_attention_mask, output_attentions, enable_kernels, atols=atols
+ )
+
+ @parameterized.expand([("linear",), ("dynamic",), ("yarn",)])
+ def test_model_rope_scaling_from_config(self, scaling_type):
+ """Override rope scaling from config test to handle Blt's sub-config structure."""
+ if self.rotary_embedding_layer is None:
+ self.skipTest("Rotary embedding layer not set")
+ config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+ short_input = ids_tensor([1, 10], config.vocab_size)
+ long_input = ids_tensor([1, int(config.max_position_embeddings * 1.5)], config.vocab_size)
+
+ set_seed(42) # Fixed seed at init time so the two models get the same random weights
+ original_model = self.model_tester_class.base_model_class(config)
+ original_model.to(torch_device)
+ original_model.eval()
+ original_short_output = original_model(short_input).last_hidden_state
+ original_long_output = original_model(long_input).last_hidden_state
+
+ set_seed(42) # Fixed seed at init time so the two models get the same random weights
+ config.rope_scaling = {"rope_type": scaling_type, "factor": 10.0}
+ # Propagate rope_scaling to sub-configs for Blt
+ config.encoder_config.rope_scaling = config.rope_scaling
+ config.decoder_config.rope_scaling = config.rope_scaling
+ config.global_config.rope_scaling = config.rope_scaling
+ config.patcher_config.rope_scaling = config.rope_scaling
+
+ scaled_model = self.model_tester_class.base_model_class(config)
+ scaled_model.to(torch_device)
+ scaled_model.eval()
+ scaled_short_output = scaled_model(short_input).last_hidden_state
+ scaled_long_output = scaled_model(long_input).last_hidden_state
+
+ # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
+ # maximum sequence length, so the outputs for the short input should match.
+ if scaling_type == "dynamic":
+ torch.testing.assert_close(original_short_output, scaled_short_output, rtol=1e-5, atol=1e-5)
+ else:
+ self.assertFalse(torch.allclose(original_short_output, scaled_short_output, atol=1e-5))
+
+ self.assertFalse(torch.allclose(original_long_output, scaled_long_output, atol=1e-5))
+
+
+@require_torch_accelerator
+class BltIntegrationTest(unittest.TestCase):
+ def tearDown(self):
+ # TODO (joao): automatic compilation, i.e. compilation when `cache_implementation="static"` is used, leaves
+ # some memory allocated in the cache, which means some object is not being released properly. This causes some
+ # unoptimal memory usage, e.g. after certain tests a 7B model in FP16 no longer fits in a 24GB GPU.
+ # Investigate the root cause.
+ cleanup(torch_device, gc_collect=False)
+
+ @slow
+ @require_read_token
+ def test_model(self):
+ NUM_TOKENS_TO_GENERATE = 200
+ EXPECTED_TEXT = "my name is alex and i am a student at the university of michigan. i am a senior majoring in computer science and minoring in mathematics. i am also a member of the michigan math club and the michigan computer s"
+
+ prompt = "my name is"
+
+ model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf", device_map="auto", attn_implementation="sdpa")
+
+ tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+ generated_ids = model.generate(
+ **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, use_cache=False
+ )
+
+ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+ self.assertEqual(output_text, EXPECTED_TEXT)
+
+ @slow
+ @require_read_token
+ def test_model_logits(self):
+ EXPECTED_OUTPUT = torch.tensor(
+ [
+ [
+ -10.4948,
+ -10.7065,
+ -6.1813,
+ -10.5545,
+ -10.3428,
+ -9.1493,
+ -8.4937,
+ -8.6382,
+ -9.2159,
+ -9.5907,
+ -9.3679,
+ -8.4184,
+ -9.0655,
+ -3.4436,
+ 2.9616,
+ -10.3157,
+ -6.3723,
+ -6.0133,
+ -9.7100,
+ -9.2128,
+ -8.8064,
+ -9.8179,
+ -9.7516,
+ -9.4681,
+ -9.7715,
+ -9.4897,
+ -9.0491,
+ -9.8098,
+ -9.4648,
+ -9.3294,
+ ],
+ [
+ -13.3010,
+ -13.1910,
+ -5.7230,
+ -13.2895,
+ -13.4864,
+ -8.7140,
+ -7.0275,
+ -7.0182,
+ -10.1362,
+ -10.3762,
+ -9.9086,
+ -7.8049,
+ -8.8660,
+ -5.2711,
+ -3.5778,
+ -12.5346,
+ -9.1609,
+ -6.7925,
+ -10.3717,
+ -9.2650,
+ -10.6393,
+ -11.4807,
+ -11.2128,
+ -10.9615,
+ -10.5806,
+ -10.8873,
+ -11.0651,
+ -11.3471,
+ -10.5437,
+ -9.9688,
+ ],
+ ]
+ ).to(torch_device)
+
+ input_ids = [1, 42, 21, 12, 43, 23, 1, 4]
+
+ model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf", attn_implementation="sdpa", device_map="auto")
+
+ with torch.no_grad():
+ output = model(torch.tensor([input_ids]).to(torch_device))[0]
+
+ torch.testing.assert_close(EXPECTED_OUTPUT, output[0, :2, :30], rtol=1e-4, atol=1e-4)
+
+ @slow
+ @require_read_token
+ @require_torch_bf16
+ def test_model_bf16(self):
+ """Test Blt model with bfloat16 precision."""
+ NUM_TOKENS_TO_GENERATE = 200
+ EXPECTED_TEXT = "my name is alex and i am a student at the university of michigan. i am a senior majoring in computer science and minoring in mathematics. i am also a member of the michigan math club and the michigan computer s"
+
+ prompt = "my name is"
+
+ model = BltForCausalLM.from_pretrained(
+ "itazap/blt-1b-hf", device_map="auto", attn_implementation="sdpa", torch_dtype=torch.bfloat16
+ )
+
+ tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+ generated_ids = model.generate(
+ **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, use_cache=False
+ )
+
+ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+ self.assertEqual(output_text, EXPECTED_TEXT)
+
+ @slow
+ @require_read_token
+ @require_torch_bf16
+ def test_model_logits_bf16(self):
+ """Test Blt model logits with bfloat16 precision."""
+
+ EXPECTED_OUTPUT = torch.tensor(
+ [
+ [
+ -10.5000,
+ -10.6875,
+ -6.1875,
+ -10.5625,
+ -10.3125,
+ -9.1875,
+ -8.5000,
+ -8.6875,
+ -9.1875,
+ -9.5625,
+ -9.3750,
+ -8.5000,
+ -9.0625,
+ -3.4219,
+ 2.9531,
+ -10.3125,
+ -6.4062,
+ -6.0000,
+ -9.6875,
+ -9.1875,
+ -8.8125,
+ -9.8125,
+ -9.7500,
+ -9.4375,
+ -9.8125,
+ -9.5000,
+ -9.0000,
+ -9.8125,
+ -9.4375,
+ -9.3125,
+ ],
+ [
+ -13.2500,
+ -13.1875,
+ -5.6875,
+ -13.3125,
+ -13.5000,
+ -8.7500,
+ -7.0625,
+ -7.0312,
+ -10.1250,
+ -10.3750,
+ -9.8750,
+ -7.8438,
+ -8.8750,
+ -5.2812,
+ -3.5625,
+ -12.5000,
+ -9.1875,
+ -6.8125,
+ -10.3750,
+ -9.3125,
+ -10.6250,
+ -11.5000,
+ -11.2500,
+ -11.0000,
+ -10.5625,
+ -10.8750,
+ -11.0625,
+ -11.3750,
+ -10.5625,
+ -10.0000,
+ ],
+ ]
+ ).to(torch_device)
+
+ input_ids = [1, 42, 21, 12, 43, 23, 1, 4]
+
+ model = BltForCausalLM.from_pretrained(
+ "itazap/blt-1b-hf", device_map="auto", attn_implementation="sdpa", torch_dtype=torch.bfloat16
+ )
+
+ with torch.no_grad():
+ output = model(torch.tensor([input_ids]).to(torch_device))[0]
+
+ torch.testing.assert_close(EXPECTED_OUTPUT, output[0, :2, :30], rtol=1e-3, atol=1e-3)
+
+ @slow
+ @require_read_token
+ def test_model_eager(self):
+ """Test Blt model with bfloat16 precision using eager attention implementation."""
+ NUM_TOKENS_TO_GENERATE = 200
+ EXPECTED_TEXT = "my name is alex and i am a student at the university of michigan. i am a senior majoring in computer science and minoring in mathematics. i am also a member of the michigan math club and the michigan computer s"
+
+ prompt = "my name is"
+
+ model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf", device_map="auto", attn_implementation="eager")
+
+ tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+ generated_ids = model.generate(
+ **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, use_cache=False
+ )
+
+ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+ self.assertEqual(output_text, EXPECTED_TEXT)
+
+ @slow
+ @require_read_token
+ @require_torch_bf16
+ def test_model_bf16_static_cache(self):
+ """Test Blt model with bfloat16 precision and static cache."""
+ NUM_TOKENS_TO_GENERATE = 200
+ EXPECTED_TEXT = "my name is alex and i am a student at the university of michigan. i am a senior majoring in computer science and minoring in mathematics. i am also a member of the michigan math club and the michigan computer s"
+
+ prompt = "my name is"
+
+ model = BltForCausalLM.from_pretrained(
+ "itazap/blt-1b-hf", device_map="auto", attn_implementation="sdpa", torch_dtype=torch.bfloat16
+ )
+
+ model.generation_config.cache_implementation = "static"
+
+ tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+ generated_ids = model.generate(
+ **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, use_cache=False
+ )
+
+ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+ self.assertEqual(output_text, EXPECTED_TEXT)
diff --git a/tests/models/bridgetower/test_modeling_bridgetower.py b/tests/models/bridgetower/test_modeling_bridgetower.py
index 59147a9d26a8..078e5d1384ab 100644
--- a/tests/models/bridgetower/test_modeling_bridgetower.py
+++ b/tests/models/bridgetower/test_modeling_bridgetower.py
@@ -28,7 +28,6 @@
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
ModelTesterMixin,
- _config_zero_init,
floats_tensor,
ids_tensor,
random_attention_mask,
@@ -438,29 +437,6 @@ def test_retain_grad_hidden_states_attentions(self):
if self.has_attentions:
self.assertIsNotNone(attentions.grad)
- # override as the `logit_scale` parameter initialization is different for BRIDGE TOWER
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- config.logit_scale_init_value,
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@unittest.skip(reason="""Bridge Tower does not have input/output embeddings. So this test is not applicable.""")
def test_model_get_set_embeddings(self):
pass
diff --git a/tests/models/bros/test_modeling_bros.py b/tests/models/bros/test_modeling_bros.py
index 3a80497cafc6..681c1e98bdd8 100644
--- a/tests/models/bros/test_modeling_bros.py
+++ b/tests/models/bros/test_modeling_bros.py
@@ -49,7 +49,7 @@ def __init__(
use_labels=True,
vocab_size=99,
hidden_size=64,
- num_hidden_layers=5,
+ num_hidden_layers=2,
num_attention_heads=4,
intermediate_size=37,
hidden_act="gelu",
@@ -323,7 +323,7 @@ def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
dtype=torch.bool,
device=torch_device,
)
- elif model_class.__name__ in ["BrosSpadeEEForTokenClassification"]:
+ elif model_class.__name__ == "BrosSpadeEEForTokenClassification":
inputs_dict["initial_token_labels"] = torch.zeros(
(self.model_tester.batch_size, self.model_tester.seq_length),
dtype=torch.long,
diff --git a/tests/models/chameleon/test_modeling_chameleon.py b/tests/models/chameleon/test_modeling_chameleon.py
index fa9e45506929..ecf873182234 100644
--- a/tests/models/chameleon/test_modeling_chameleon.py
+++ b/tests/models/chameleon/test_modeling_chameleon.py
@@ -76,7 +76,7 @@ def __init__(
pad_token_id=0,
vq_num_embeds=5,
vq_embed_dim=5,
- vq_channel_multiplier=[1, 4],
+ vq_channel_multiplier=[1, 2],
vq_img_token_start_id=10, # has to be less than vocab size when added with vq_num_embeds
scope=None,
):
@@ -255,10 +255,6 @@ def test_model_rope_scaling(self, scaling_type):
def test_batching_equivalence(self):
pass
- @unittest.skip("Chameleon VQ model cannot be squishes more due to hardcoded layer params in model code")
- def test_model_is_small(self):
- pass
-
class ChameleonVision2SeqModelTester(ChameleonModelTester):
def __init__(self, parent, image_size=10, **kwargs):
@@ -321,10 +317,6 @@ def test_disk_offload_bin(self):
def test_disk_offload_safetensors(self):
pass
- @unittest.skip("Chameleon VQ model cannot be squishes more due to hardcoded layer params in model code")
- def test_model_is_small(self):
- pass
-
@unittest.skip("Chameleon applies key/query norm which doesn't work with packing")
def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
pass
diff --git a/tests/models/chinese_clip/test_modeling_chinese_clip.py b/tests/models/chinese_clip/test_modeling_chinese_clip.py
index dc8e9a145b08..b2508fb3fbb5 100644
--- a/tests/models/chinese_clip/test_modeling_chinese_clip.py
+++ b/tests/models/chinese_clip/test_modeling_chinese_clip.py
@@ -18,7 +18,6 @@
import tempfile
import unittest
-import numpy as np
import requests
from transformers import ChineseCLIPConfig, ChineseCLIPTextConfig, ChineseCLIPVisionConfig
@@ -580,33 +579,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for CHINESE_CLIP
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for sub_config_key in ("vision_config", "text_config"):
- sub_config = getattr(configs_no_init, sub_config_key, {})
- setattr(configs_no_init, sub_config_key, _config_zero_init(sub_config))
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/clap/test_modeling_clap.py b/tests/models/clap/test_modeling_clap.py
index 0dab34123de4..a83b363f4aa5 100644
--- a/tests/models/clap/test_modeling_clap.py
+++ b/tests/models/clap/test_modeling_clap.py
@@ -528,30 +528,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for CLAP
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if "logit_scale" in name:
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/clip/test_modeling_clip.py b/tests/models/clip/test_modeling_clip.py
index b352b8160468..8940eeb510b1 100644
--- a/tests/models/clip/test_modeling_clip.py
+++ b/tests/models/clip/test_modeling_clip.py
@@ -565,30 +565,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for CLIP
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
@@ -754,10 +730,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- @unittest.skip(reason="CLIP uses the same initialization scheme as the Flax original implementation")
- def test_initialization(self):
- pass
-
@parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
@slow
@is_flaky()
diff --git a/tests/models/clipseg/test_modeling_clipseg.py b/tests/models/clipseg/test_modeling_clipseg.py
index 08a21f9dcf3b..5db226b7276e 100644
--- a/tests/models/clipseg/test_modeling_clipseg.py
+++ b/tests/models/clipseg/test_modeling_clipseg.py
@@ -393,7 +393,7 @@ def create_and_check_model(self, config, input_ids, attention_mask, pixel_values
result.logits_per_text.shape, (self.text_model_tester.batch_size, self.vision_model_tester.batch_size)
)
- def create_and_check_model_for_image_segmentation(self, config, input_ids, attention_maks, pixel_values):
+ def create_and_check_model_for_image_segmentation(self, config, input_ids, attention_mask, pixel_values):
model = CLIPSegForImageSegmentation(config).to(torch_device).eval()
with torch.no_grad():
result = model(input_ids, pixel_values)
@@ -493,33 +493,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- # override as the some parameters require custom initialization
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if "logit_scale" in name:
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif "film" in name or "transposed_conv" in name or "reduce" in name:
- # those parameters use PyTorch' default nn.Linear initialization scheme
- pass
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/clvp/test_modeling_clvp.py b/tests/models/clvp/test_modeling_clvp.py
index a33d787dc7cc..2c9548616c5e 100644
--- a/tests/models/clvp/test_modeling_clvp.py
+++ b/tests/models/clvp/test_modeling_clvp.py
@@ -32,7 +32,6 @@
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
ModelTesterMixin,
- _config_zero_init,
ids_tensor,
random_attention_mask,
)
@@ -501,36 +500,6 @@ def test_inputs_embeds(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for Clvp
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- expected_value = np.log(1 / 0.07)
- returned_value = param.data.item()
-
- self.assertAlmostEqual(
- returned_value,
- expected_value,
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- expected_range = [0.0, 1.0]
- returned_range = ((param.data.mean() * 1e9).round() / 1e9).item()
-
- self.assertIn(
- returned_range,
- expected_range,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_load_speech_text_decoder_config(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
diff --git a/tests/models/codegen/test_modeling_codegen.py b/tests/models/codegen/test_modeling_codegen.py
index ee16a5347ad6..5f97cfad359d 100644
--- a/tests/models/codegen/test_modeling_codegen.py
+++ b/tests/models/codegen/test_modeling_codegen.py
@@ -379,7 +379,7 @@ def test_batch_generation(self):
model.config.pad_token_id = model.config.eos_token_id
# use different length sentences to test batching
- sentences = ["def hellow_world():", "def greet(name):"]
+ sentences = ["def hello_world():", "def greet(name):"]
inputs = tokenizer(sentences, return_tensors="pt", padding=True)
input_ids = inputs["input_ids"].to(torch_device)
@@ -415,7 +415,7 @@ def test_batch_generation(self):
padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
expected_output_sentence = [
- 'def hellow_world():\n print("Hello World")\n\nhellow_world()',
+ 'def hello_world():\n print("Hello World")\n\nhellow_world()',
'def greet(name):\n print(f"Hello {name}")\n\ng',
]
self.assertListEqual(expected_output_sentence, batch_out_sentence)
diff --git a/tests/models/cohere/test_modeling_cohere.py b/tests/models/cohere/test_modeling_cohere.py
index 427a7f447d74..436d1f9d4226 100644
--- a/tests/models/cohere/test_modeling_cohere.py
+++ b/tests/models/cohere/test_modeling_cohere.py
@@ -54,7 +54,7 @@ def __init__(
use_labels=True,
vocab_size=99,
hidden_size=32,
- num_hidden_layers=4,
+ num_hidden_layers=2,
num_attention_heads=4,
intermediate_size=37,
hidden_act="gelu",
diff --git a/tests/models/cohere2/test_modeling_cohere2.py b/tests/models/cohere2/test_modeling_cohere2.py
index cdb78895e866..4619c7a7f19d 100644
--- a/tests/models/cohere2/test_modeling_cohere2.py
+++ b/tests/models/cohere2/test_modeling_cohere2.py
@@ -241,7 +241,7 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):
self.skipTest("FlashAttention2 is required for this test.")
if torch_device == "xpu" and attn_implementation == "flash_attention_2":
- self.skipTest(reason="Intel XPU doesn't support falsh_attention_2 as of now.")
+ self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
model_id = "CohereForAI/c4ai-command-r7b-12-2024"
EXPECTED_COMPLETIONS = [
diff --git a/tests/models/cohere2_vision/test_modeling_cohere2_vision.py b/tests/models/cohere2_vision/test_modeling_cohere2_vision.py
index 96843faa95f7..50ce19fd98ba 100644
--- a/tests/models/cohere2_vision/test_modeling_cohere2_vision.py
+++ b/tests/models/cohere2_vision/test_modeling_cohere2_vision.py
@@ -65,7 +65,7 @@ def __init__(
"vocab_size": 99,
"hidden_size": 128,
"intermediate_size": 37,
- "num_hidden_layers": 4,
+ "num_hidden_layers": 2,
"num_attention_heads": 4,
"output_channels": 64,
"hidden_act": "silu",
@@ -170,10 +170,6 @@ def setUp(self):
def test_config(self):
self.config_tester.run_common_tests()
- @unittest.skip(reason="Siglip backbone uses the same initialization scheme as the Flax original implementation")
- def test_initialization(self):
- pass
-
@require_read_token
@require_torch
diff --git a/tests/models/colpali/test_modeling_colpali.py b/tests/models/colpali/test_modeling_colpali.py
index 7966e34ce323..602993b11f93 100644
--- a/tests/models/colpali/test_modeling_colpali.py
+++ b/tests/models/colpali/test_modeling_colpali.py
@@ -272,12 +272,6 @@ def test_training_gradient_checkpointing_use_reentrant_false(self):
def test_model_parallelism(self):
pass
- @unittest.skip(
- reason="PaliGemma's SigLip encoder uses the same initialization scheme as the Flax original implementation"
- )
- def test_initialization(self):
- pass
-
# TODO extend valid outputs to include this test @Molbap
@unittest.skip(reason="PaliGemma has currently one output format.")
def test_model_outputs_equivalence(self):
diff --git a/tests/models/conditional_detr/test_modeling_conditional_detr.py b/tests/models/conditional_detr/test_modeling_conditional_detr.py
index a2d962a85a0f..1715bde36bc4 100644
--- a/tests/models/conditional_detr/test_modeling_conditional_detr.py
+++ b/tests/models/conditional_detr/test_modeling_conditional_detr.py
@@ -22,7 +22,7 @@
from transformers.testing_utils import require_timm, require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -512,29 +512,6 @@ def test_hf_backbone(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- configs_no_init.init_xavier_std = 1e9
-
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if "bbox_attention" in name and "bias" not in name:
- self.assertLess(
- 100000,
- abs(param.data.max().item()),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
TOLERANCE = 1e-4
diff --git a/tests/models/csm/test_modeling_csm.py b/tests/models/csm/test_modeling_csm.py
index 19e0beb39cb9..d298d32493fa 100644
--- a/tests/models/csm/test_modeling_csm.py
+++ b/tests/models/csm/test_modeling_csm.py
@@ -39,7 +39,6 @@
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
ModelTesterMixin,
- _config_zero_init,
ids_tensor,
)
@@ -190,25 +189,6 @@ def _get_logits_processor_kwargs(self, do_sample=False, config=None):
return logits_processor_kwargs
- def test_initialization(self):
- """
- Overrides [ModelTesterMixin.test_initialization] because of specificities of Mimi codec model.
- See https://github.com/huggingface/transformers/blob/1077603410cd73ba71d64a522033574d66d64b55/tests/models/mimi/test_modeling_mimi.py#L384-L397
- """
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = ["conv", "input_proj", "output_proj"]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _check_similar_generate_outputs(self, output_1, output_2, atol=1e-5, rtol=1e-5):
"""
Overrides [GenerationTesterMixin._check_similar_generate_outputs] to handle third input_ids dimension.
@@ -362,7 +342,7 @@ def _load_conversation(self):
def test_1b_model_integration_generate(self):
"""
Tests the generated tokens match the ones from the original model implementation.
- Such tokens are to be retreived using https://gist.github.com/eustlb/d25577a357ddcf8f4a8cd0d00baca551, which is a script that infers the original model.
+ Such tokens are to be retrieved using https://gist.github.com/eustlb/d25577a357ddcf8f4a8cd0d00baca551, which is a script that infers the original model.
"""
processor = AutoProcessor.from_pretrained(self.model_checkpoint)
prompt = "<|begin_of_text|>[0]What are you working on?<|end_of_text|><|AUDIO|><|audio_eos|><|begin_of_text|>[1]I'm figuring out my budget.<|end_of_text|>"
@@ -406,7 +386,7 @@ def test_1b_model_integration_generate(self):
def test_1b_model_integration_generate_no_audio(self):
"""
Tests the generated tokens match the ones from the original model implementation.
- Such tokens are to be retreived using https://gist.github.com/eustlb/aed822f765e928b9612e01b0d8836d69, which is a script that infers the original model.
+ Such tokens are to be retrieved using https://gist.github.com/eustlb/aed822f765e928b9612e01b0d8836d69, which is a script that infers the original model.
"""
processor = AutoProcessor.from_pretrained(self.model_checkpoint)
@@ -467,7 +447,7 @@ def test_1b_model_integration_generate_no_audio(self):
def test_1b_model_integration_generate_multiple_audio(self):
"""
Test the generated tokens match the ones from the original model implementation.
- Such tokens are to be retreived using https://gist.github.com/eustlb/0c94de002e1325abb61d32217f74c0f8, which is a script that infers the original model.
+ Such tokens are to be retrieved using https://gist.github.com/eustlb/0c94de002e1325abb61d32217f74c0f8, which is a script that infers the original model.
"""
processor = AutoProcessor.from_pretrained(self.model_checkpoint)
@@ -526,7 +506,7 @@ def test_1b_model_integration_generate_multiple_audio(self):
def test_1b_model_integration_generate_batched(self):
"""
Test the generated tokens match the ones from the original model implementation.
- Such tokens are to be retreived using https://gist.github.com/eustlb/bcc532b53161bc31da3d66cb07ae193f, which is a script that infers the original model.
+ Such tokens are to be retrieved using https://gist.github.com/eustlb/bcc532b53161bc31da3d66cb07ae193f, which is a script that infers the original model.
"""
processor = AutoProcessor.from_pretrained(self.model_checkpoint)
diff --git a/tests/models/d_fine/test_modeling_d_fine.py b/tests/models/d_fine/test_modeling_d_fine.py
index 7c381b8f6ae4..f01415fa699a 100644
--- a/tests/models/d_fine/test_modeling_d_fine.py
+++ b/tests/models/d_fine/test_modeling_d_fine.py
@@ -48,7 +48,7 @@
from transformers import RTDetrImageProcessor
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -361,10 +361,6 @@ def test_model_common_attributes(self):
def test_resize_tokens_embeddings(self):
pass
- @unittest.skip(reason="Not relevant for the model")
- def test_can_init_all_missing_weights(self):
- pass
-
@unittest.skip(reason="Feed forward chunking is not implemented")
def test_feed_forward_chunking(self):
pass
@@ -633,58 +629,6 @@ def test_hf_backbone(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- configs_no_init.initializer_bias_prior_prob = 0.2
- bias_value = -1.3863 # log_e ((1 - 0.2) / 0.2)
-
- failed_cases = []
-
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- # Skip the check for the backbone
- for name, module in model.named_modules():
- if module.__class__.__name__ == "DFineConvEncoder":
- backbone_params = [f"{name}.{key}" for key in module.state_dict()]
- break
-
- for name, param in model.named_parameters():
- if param.requires_grad:
- if ("class_embed" in name and "bias" in name) or "enc_score_head.bias" in name:
- bias_tensor = torch.full_like(param.data, bias_value)
- try:
- torch.testing.assert_close(param.data, bias_tensor, atol=1e-4, rtol=1e-4)
- except AssertionError:
- failed_cases.append(
- f"Parameter {name} of model {model_class} seems not properly initialized. "
- f"Biases should be initialized to {bias_value}, got {param.data}"
- )
- elif (
- "level_embed" in name
- or "sampling_offsets.bias" in name
- or "value_proj" in name
- or "output_proj" in name
- or "reference_points" in name
- or "enc_score_head.weight" in name
- or ("class_embed" in name and "weight" in name)
- or name in backbone_params
- ):
- continue
- else:
- mean = param.data.mean()
- round_mean = (mean * 1e9).round() / 1e9
- round_mean = round_mean.item()
- if round_mean not in [0.0, 1.0]:
- failed_cases.append(
- f"Parameter {name} of model {model_class} seems not properly initialized. "
- f"Mean is {round_mean}, but should be in [0, 1]"
- )
-
- message = "\n" + "\n".join(failed_cases)
- self.assertTrue(not failed_cases, message)
-
@parameterized.expand(["float32", "float16", "bfloat16"])
@require_torch_accelerator
@slow
diff --git a/tests/models/dab_detr/test_modeling_dab_detr.py b/tests/models/dab_detr/test_modeling_dab_detr.py
index 6f437ce7692d..ae2bcffa03f1 100644
--- a/tests/models/dab_detr/test_modeling_dab_detr.py
+++ b/tests/models/dab_detr/test_modeling_dab_detr.py
@@ -22,7 +22,7 @@
from transformers.testing_utils import require_timm, require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -194,7 +194,7 @@ def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
if return_labels:
- if model_class.__name__ in ["DabDetrForObjectDetection"]:
+ if model_class.__name__ == "DabDetrForObjectDetection":
labels = []
for i in range(self.model_tester.batch_size):
target = {}
@@ -706,55 +706,6 @@ def test_different_timm_backbone(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, _ = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- configs_no_init.init_xavier_std = 1e9
- # Copied from RT-DETR
- configs_no_init.initializer_bias_prior_prob = 0.2
- bias_value = -1.3863 # log_e ((1 - 0.2) / 0.2)
-
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if "bbox_attention" in name and "bias" not in name:
- self.assertLess(
- 100000,
- abs(param.data.max().item()),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- # Modified from RT-DETR
- elif "class_embed" in name and "bias" in name:
- bias_tensor = torch.full_like(param.data, bias_value)
- torch.testing.assert_close(
- param.data,
- bias_tensor,
- atol=1e-4,
- rtol=1e-4,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif "activation_fn" in name and config.activation_function == "prelu":
- self.assertTrue(
- param.data.mean() == 0.25,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif "backbone.conv_encoder.model" in name:
- continue
- elif "self_attn.in_proj_weight" in name:
- self.assertIn(
- ((param.data.mean() * 1e2).round() / 1e2).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
TOLERANCE = 1e-4
CHECKPOINT = "IDEA-Research/dab-detr-resnet-50"
diff --git a/tests/models/dac/test_modeling_dac.py b/tests/models/dac/test_modeling_dac.py
index cb7d6b388c19..8d8a5a88cb96 100644
--- a/tests/models/dac/test_modeling_dac.py
+++ b/tests/models/dac/test_modeling_dac.py
@@ -354,22 +354,6 @@ def recursive_check(tuple_object, dict_object):
dict_inputs = self._prepare_for_class(inputs_dict, model_class)
check_equivalence(model, tuple_inputs, dict_inputs)
- # Ignore copy
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = ["conv", "in_proj", "out_proj", "codebook"]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_identity_shortcut(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs()
config.use_conv_shortcut = False
diff --git a/tests/models/data2vec/test_modeling_data2vec_audio.py b/tests/models/data2vec/test_modeling_data2vec_audio.py
index 630f6238e76e..3357d92ec8e0 100644
--- a/tests/models/data2vec/test_modeling_data2vec_audio.py
+++ b/tests/models/data2vec/test_modeling_data2vec_audio.py
@@ -24,7 +24,7 @@
from transformers.testing_utils import require_torch, require_torchcodec, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init
+from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -458,39 +458,6 @@ def test_retain_grad_hidden_states_attentions(self):
self.assertIsNotNone(hidden_states.grad)
self.assertIsNotNone(attentions.grad)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = [
- "conv.weight",
- "masked_spec_embed",
- "codevectors",
- "quantizer.weight_proj.weight",
- "project_hid.weight",
- "project_hid.bias",
- "project_q.weight",
- "project_q.bias",
- "feature_projection.projection.weight",
- "feature_projection.projection.bias",
- "objective.weight",
- ]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# overwrite from test_modeling_common
def _mock_init_weights(self, module):
if hasattr(module, "weight") and module.weight is not None:
diff --git a/tests/models/data2vec/test_modeling_data2vec_vision.py b/tests/models/data2vec/test_modeling_data2vec_vision.py
index aebbe183cacf..e8aa7fa28973 100644
--- a/tests/models/data2vec/test_modeling_data2vec_vision.py
+++ b/tests/models/data2vec/test_modeling_data2vec_vision.py
@@ -32,7 +32,7 @@
)
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -292,24 +292,6 @@ def test_training_gradient_checkpointing(self):
loss = model(**inputs).loss
loss.backward()
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- # we skip lambda parameters as these require special initial values
- # determined by config.layer_scale_init_value
- if "lambda" in name:
- continue
- if param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_for_image_classification(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_for_image_classification(*config_and_inputs)
diff --git a/tests/models/dbrx/test_modeling_dbrx.py b/tests/models/dbrx/test_modeling_dbrx.py
index d194c74e7b43..4dbe5d403dc5 100644
--- a/tests/models/dbrx/test_modeling_dbrx.py
+++ b/tests/models/dbrx/test_modeling_dbrx.py
@@ -17,7 +17,7 @@
from parameterized import parameterized
-from transformers import DbrxConfig, is_torch_available
+from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
@@ -30,10 +30,8 @@
class DbrxModelTester(CausalLMModelTester):
- config_class = DbrxConfig
if is_torch_available():
base_model_class = DbrxModel
- causal_lm_class = DbrxForCausalLM
def __init__(
self,
diff --git a/tests/models/deepseek_v2/test_modeling_deepseek_v2.py b/tests/models/deepseek_v2/test_modeling_deepseek_v2.py
index 930f2504dee8..dc9886c70a6e 100644
--- a/tests/models/deepseek_v2/test_modeling_deepseek_v2.py
+++ b/tests/models/deepseek_v2/test_modeling_deepseek_v2.py
@@ -18,7 +18,7 @@
import pytest
-from transformers import BitsAndBytesConfig, Cache, DeepseekV2Config, is_torch_available
+from transformers import BitsAndBytesConfig, Cache, is_torch_available
from transformers.testing_utils import require_read_token, require_torch, require_torch_accelerator, slow, torch_device
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
@@ -33,10 +33,7 @@
class DeepseekV2ModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = DeepseekV2Config
base_model_class = DeepseekV2Model
- causal_lm_class = DeepseekV2ForCausalLM
- sequence_class = DeepseekV2ForSequenceClassification
def __init__(
self,
@@ -57,15 +54,6 @@ def __init__(
@require_torch
class DeepseekV2ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- DeepseekV2ForCausalLM,
- DeepseekV2ForSequenceClassification,
- DeepseekV2Model,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": DeepseekV2Model,
@@ -76,8 +64,6 @@ class DeepseekV2ModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False
test_torchscript = False
test_all_params_have_gradient = False
diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py
index 9ed521509408..46a5cdd7bdd0 100644
--- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py
+++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py
@@ -25,7 +25,6 @@
require_read_token,
require_torch,
require_torch_accelerator,
- require_torch_gpu,
require_torch_large_accelerator,
slow,
torch_device,
@@ -65,7 +64,7 @@ def __init__(
hidden_size=32,
intermediate_size=37,
moe_intermediate_size=12,
- num_hidden_layers=5,
+ num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=4,
n_shared_experts=1,
@@ -326,7 +325,9 @@ def test_model_rope_scaling(self):
long_input_length = int(config.max_position_embeddings * 1.5)
# Inputs
- x = torch.randn(1, dtype=torch.float32, device=torch_device) # used exlusively to get the dtype and the device
+ x = torch.randn(
+ 1, dtype=torch.float32, device=torch_device
+ ) # used exclusively to get the dtype and the device
position_ids_short = torch.arange(short_input_length, dtype=torch.long, device=torch_device)
position_ids_short = position_ids_short.unsqueeze(0)
position_ids_long = torch.arange(long_input_length, dtype=torch.long, device=torch_device)
@@ -447,7 +448,7 @@ def test_eager_matches_sdpa_generate(self):
msg=f"\n{tokenizer.batch_decode(res_eager)} \nvs\n{tokenizer.batch_decode(res_sdpa)}",
)
- @require_torch_gpu
+ @require_torch_accelerator
def test_flex_attention_with_grads(self):
"""
Overwriting as the namings/functionality on the attention part are different; for now it's more of a unique model.
diff --git a/tests/models/deepseek_vl/test_modeling_deepseek_vl.py b/tests/models/deepseek_vl/test_modeling_deepseek_vl.py
index a2d1950dcdc4..8e7389fe4f68 100644
--- a/tests/models/deepseek_vl/test_modeling_deepseek_vl.py
+++ b/tests/models/deepseek_vl/test_modeling_deepseek_vl.py
@@ -187,11 +187,6 @@ def test_inputs_embeds_matches_input_ids(self):
out_embeds = model(inputs_embeds=inputs_embeds, **inputs)[0]
torch.testing.assert_close(out_embeds, out_ids)
- @unittest.skip(reason="Siglip uses the same initialization scheme as the Flax original implementation")
- # Copied from tests.models.siglip.test_modeling_siglip.SiglipVisionModelTest.test_initialization
- def test_initialization(self):
- pass
-
# Copied from tests.models.janus.test_modeling_janus.JanusVisionText2TextModelTest.test_sdpa_can_dispatch_composite_models
def test_sdpa_can_dispatch_composite_models(self):
for model_class in self.all_model_classes:
diff --git a/tests/models/deepseek_vl_hybrid/test_modeling_deepseek_vl_hybrid.py b/tests/models/deepseek_vl_hybrid/test_modeling_deepseek_vl_hybrid.py
index fbb904da735b..485c08cc6523 100644
--- a/tests/models/deepseek_vl_hybrid/test_modeling_deepseek_vl_hybrid.py
+++ b/tests/models/deepseek_vl_hybrid/test_modeling_deepseek_vl_hybrid.py
@@ -218,11 +218,6 @@ def test_inputs_embeds_matches_input_ids(self):
out_embeds = model(inputs_embeds=inputs_embeds, **inputs)[0]
torch.testing.assert_close(out_embeds, out_ids)
- @unittest.skip(reason="Siglip uses the same initialization scheme as the Flax original implementation")
- # Copied from tests.models.siglip.test_modeling_siglip.SiglipVisionModelTest.test_initialization
- def test_initialization(self):
- pass
-
def test_sdpa_can_dispatch_composite_models(self):
for model_class in self.all_model_classes:
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
diff --git a/tests/models/deformable_detr/test_modeling_deformable_detr.py b/tests/models/deformable_detr/test_modeling_deformable_detr.py
index 14fa0994ebee..72e5cb24f613 100644
--- a/tests/models/deformable_detr/test_modeling_deformable_detr.py
+++ b/tests/models/deformable_detr/test_modeling_deformable_detr.py
@@ -30,7 +30,7 @@
)
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -581,29 +581,6 @@ def test_hf_backbone(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- print("Model class:", model_class)
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if (
- "level_embed" in name
- or "sampling_offsets.bias" in name
- or "value_proj" in name
- or "output_proj" in name
- or "reference_points" in name
- ):
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_two_stage_training(self):
model_class = DeformableDetrForObjectDetection
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
diff --git a/tests/models/depth_pro/test_modeling_depth_pro.py b/tests/models/depth_pro/test_modeling_depth_pro.py
index 0e644c7c1892..ca5382619241 100644
--- a/tests/models/depth_pro/test_modeling_depth_pro.py
+++ b/tests/models/depth_pro/test_modeling_depth_pro.py
@@ -22,7 +22,7 @@
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -301,45 +301,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- non_uniform_init_parms = [
- # these encoders are vision transformers
- # any layer outside these encoders is either Conv2d or ConvTranspose2d
- # which use kaiming initialization
- "patch_encoder",
- "image_encoder",
- "fov_model.encoder",
- ]
- if param.requires_grad:
- if any(x in name for x in non_uniform_init_parms):
- # See PR #38607 (to avoid flakiness)
- data = torch.flatten(param.data)
- n_elements = torch.numel(data)
- # skip 2.5% of elements on each side to avoid issues caused by `nn.init.trunc_normal_` described in
- # https://github.com/huggingface/transformers/pull/27906#issuecomment-1846951332
- n_elements_to_skip_on_each_side = int(n_elements * 0.025)
- data_to_check = torch.sort(data).values
- if n_elements_to_skip_on_each_side > 0:
- data_to_check = data_to_check[
- n_elements_to_skip_on_each_side:-n_elements_to_skip_on_each_side
- ]
- self.assertIn(
- ((data_to_check.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# this started when switched from normal initialization to kaiming_normal initialization
# maybe because the magnitude of offset values from ViT-encoders increases when followed by many convolution layers
def test_batching_equivalence(self, atol=1e-4, rtol=1e-4):
diff --git a/tests/models/detr/test_modeling_detr.py b/tests/models/detr/test_modeling_detr.py
index dcad75307691..bd8d1cb694f3 100644
--- a/tests/models/detr/test_modeling_detr.py
+++ b/tests/models/detr/test_modeling_detr.py
@@ -22,7 +22,7 @@
from transformers.testing_utils import Expectations, require_timm, require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -532,29 +532,6 @@ def test_greyscale_images(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- configs_no_init.init_xavier_std = 1e9
-
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if "bbox_attention" in name and "bias" not in name:
- self.assertLess(
- 100000,
- abs(param.data.max().item()),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
TOLERANCE = 1e-4
diff --git a/tests/models/dinat/test_modeling_dinat.py b/tests/models/dinat/test_modeling_dinat.py
index 4ffe5f6cd692..7c8f2a0b0ae3 100644
--- a/tests/models/dinat/test_modeling_dinat.py
+++ b/tests/models/dinat/test_modeling_dinat.py
@@ -23,7 +23,7 @@
from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -327,20 +327,6 @@ def test_model_from_pretrained(self):
model = DinatModel.from_pretrained(model_name)
self.assertIsNotNone(model)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if "embeddings" not in name and param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@require_natten
@require_vision
diff --git a/tests/models/dinov2/test_modeling_dinov2.py b/tests/models/dinov2/test_modeling_dinov2.py
index 2377bc1d2ee2..f22e216bd3e7 100644
--- a/tests/models/dinov2/test_modeling_dinov2.py
+++ b/tests/models/dinov2/test_modeling_dinov2.py
@@ -18,7 +18,6 @@
from transformers import Dinov2Config
from transformers.testing_utils import (
- is_flaky,
require_torch,
require_vision,
slow,
@@ -238,10 +237,6 @@ def setUp(self):
self.model_tester = Dinov2ModelTester(self)
self.config_tester = ConfigTester(self, config_class=Dinov2Config, has_text_modality=False, hidden_size=37)
- @is_flaky(max_attempts=3, description="`torch.nn.init.trunc_normal_` is flaky.")
- def test_initialization(self):
- super().test_initialization()
-
def test_config(self):
self.config_tester.run_common_tests()
diff --git a/tests/models/dinov2_with_registers/test_modeling_dinov2_with_registers.py b/tests/models/dinov2_with_registers/test_modeling_dinov2_with_registers.py
index b9f0f5fecfe0..dece1475be50 100644
--- a/tests/models/dinov2_with_registers/test_modeling_dinov2_with_registers.py
+++ b/tests/models/dinov2_with_registers/test_modeling_dinov2_with_registers.py
@@ -27,7 +27,7 @@
from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -245,29 +245,6 @@ def setUp(self):
self, config_class=Dinov2WithRegistersConfig, has_text_modality=False, hidden_size=37
)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad and "register_tokens" not in name:
- # See PR #38607 (to avoid flakiness)
- data = torch.flatten(param.data)
- n_elements = torch.numel(data)
- # skip 2.5% of elements on each side to avoid issues caused by `nn.init.trunc_normal_` described in
- # https://github.com/huggingface/transformers/pull/27906#issuecomment-1846951332
- n_elements_to_skip_on_each_side = int(n_elements * 0.025)
- data_to_check = torch.sort(data).values
- if n_elements_to_skip_on_each_side > 0:
- data_to_check = data_to_check[n_elements_to_skip_on_each_side:-n_elements_to_skip_on_each_side]
- self.assertIn(
- ((data_to_check.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_config(self):
self.config_tester.run_common_tests()
diff --git a/tests/models/dinov3_vit/test_modeling_dinov3_vit.py b/tests/models/dinov3_vit/test_modeling_dinov3_vit.py
index f0b8c92d22a0..7263ecc709fc 100644
--- a/tests/models/dinov3_vit/test_modeling_dinov3_vit.py
+++ b/tests/models/dinov3_vit/test_modeling_dinov3_vit.py
@@ -21,7 +21,7 @@
from transformers.utils import is_torch_available, is_vision_available
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -161,29 +161,6 @@ def setUp(self):
self.model_tester = DINOv3ViTModelTester(self)
self.config_tester = ConfigTester(self, config_class=DINOv3ViTConfig, has_text_modality=False, hidden_size=37)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad and "register_tokens" not in name:
- # See PR #38607 (to avoid flakiness)
- data = torch.flatten(param.data)
- n_elements = torch.numel(data)
- # skip 2.5% of elements on each side to avoid issues caused by `nn.init.trunc_normal_` described in
- # https://github.com/huggingface/transformers/pull/27906#issuecomment-1846951332
- n_elements_to_skip_on_each_side = int(n_elements * 0.025)
- data_to_check = torch.sort(data).values
- if n_elements_to_skip_on_each_side > 0:
- data_to_check = data_to_check[n_elements_to_skip_on_each_side:-n_elements_to_skip_on_each_side]
- self.assertIn(
- ((data_to_check.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_config(self):
self.config_tester.run_common_tests()
diff --git a/tests/models/distilbert/test_modeling_distilbert.py b/tests/models/distilbert/test_modeling_distilbert.py
index db90233b438a..a22d229a0405 100644
--- a/tests/models/distilbert/test_modeling_distilbert.py
+++ b/tests/models/distilbert/test_modeling_distilbert.py
@@ -383,7 +383,7 @@ def test_flash_attn_2_inference_equivalence_right_padding(self):
@require_torch
-class DistilBertModelIntergrationTest(unittest.TestCase):
+class DistilBertModelIntegrationTest(unittest.TestCase):
@slow
def test_inference_no_head_absolute_embedding(self):
model = DistilBertModel.from_pretrained("distilbert-base-uncased")
diff --git a/tests/models/donut/test_modeling_donut_swin.py b/tests/models/donut/test_modeling_donut_swin.py
index 456da8500799..6c5f44ebe1f9 100644
--- a/tests/models/donut/test_modeling_donut_swin.py
+++ b/tests/models/donut/test_modeling_donut_swin.py
@@ -21,7 +21,7 @@
from transformers.utils import is_torch_available
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -359,17 +359,3 @@ def test_model_from_pretrained(self):
model_name = "naver-clova-ix/donut-base"
model = DonutSwinModel.from_pretrained(model_name)
self.assertIsNotNone(model)
-
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if "embeddings" not in name and param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
diff --git a/tests/models/dots1/test_modeling_dots1.py b/tests/models/dots1/test_modeling_dots1.py
index 78707e6518ff..65cb64ee24ff 100644
--- a/tests/models/dots1/test_modeling_dots1.py
+++ b/tests/models/dots1/test_modeling_dots1.py
@@ -16,7 +16,7 @@
import gc
import unittest
-from transformers import AutoTokenizer, Dots1Config, is_torch_available
+from transformers import AutoTokenizer, is_torch_available
from transformers.testing_utils import (
backend_empty_cache,
cleanup,
@@ -39,10 +39,8 @@
class Dots1ModelTester(CausalLMModelTester):
- config_class = Dots1Config
if is_torch_available():
base_model_class = Dots1Model
- causal_lm_class = Dots1ForCausalLM
def __init__(
self,
diff --git a/tests/models/dpt/test_modeling_dpt.py b/tests/models/dpt/test_modeling_dpt.py
index 1d693e7f408c..64bc817a20a5 100644
--- a/tests/models/dpt/test_modeling_dpt.py
+++ b/tests/models/dpt/test_modeling_dpt.py
@@ -23,7 +23,7 @@
from transformers.testing_utils import Expectations, require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -261,29 +261,6 @@ def test_training_gradient_checkpointing_use_reentrant_false(self):
def test_sdpa_can_compile_dynamic(self):
pass
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- # Skip the check for the backbone
- backbone_params = []
- for name, module in model.named_modules():
- if module.__class__.__name__ == "DPTViTHybridEmbeddings":
- backbone_params = [f"{name}.{key}" for key in module.state_dict()]
- break
-
- for name, param in model.named_parameters():
- if param.requires_grad:
- if name in backbone_params:
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_backbone_selection(self):
def _validate_backbone_init():
for model_class in self.all_model_classes:
diff --git a/tests/models/dpt/test_modeling_dpt_auto_backbone.py b/tests/models/dpt/test_modeling_dpt_auto_backbone.py
index 165da4be6be5..3371099d37a6 100644
--- a/tests/models/dpt/test_modeling_dpt_auto_backbone.py
+++ b/tests/models/dpt/test_modeling_dpt_auto_backbone.py
@@ -21,7 +21,7 @@
from transformers.utils.import_utils import get_torch_major_and_minor_version
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -195,29 +195,6 @@ def test_training_gradient_checkpointing(self):
loss = model(**inputs).loss
loss.backward()
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- # Skip the check for the backbone
- backbone_params = []
- for name, module in model.named_modules():
- if module.__class__.__name__ == "DPTViTHybridEmbeddings":
- backbone_params = [f"{name}.{key}" for key in module.state_dict()]
- break
-
- for name, param in model.named_parameters():
- if param.requires_grad:
- if name in backbone_params:
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@unittest.skip(reason="DPT with AutoBackbone does not have a base model and hence no input_embeddings")
def test_model_get_set_embeddings(self):
pass
diff --git a/tests/models/dpt/test_modeling_dpt_hybrid.py b/tests/models/dpt/test_modeling_dpt_hybrid.py
index e7a184c400a7..4de0b3139930 100644
--- a/tests/models/dpt/test_modeling_dpt_hybrid.py
+++ b/tests/models/dpt/test_modeling_dpt_hybrid.py
@@ -20,7 +20,7 @@
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -271,29 +271,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- # Skip the check for the backbone
- backbone_params = []
- for name, module in model.named_modules():
- if module.__class__.__name__ == "DPTViTHybridEmbeddings":
- backbone_params = [f"{name}.{key}" for key in module.state_dict()]
- break
-
- for name, param in model.named_parameters():
- if param.requires_grad:
- if name in backbone_params:
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@slow
def test_model_from_pretrained(self):
model_name = "Intel/dpt-hybrid-midas"
diff --git a/tests/models/edgetam/__init__.py b/tests/models/edgetam/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/edgetam/test_modeling_edgetam.py b/tests/models/edgetam/test_modeling_edgetam.py
new file mode 100644
index 000000000000..152f3132583b
--- /dev/null
+++ b/tests/models/edgetam/test_modeling_edgetam.py
@@ -0,0 +1,730 @@
+# coding=utf-8
+# 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.
+"""Testing suite for the PyTorch EDGETAM model."""
+
+import gc
+import tempfile
+import unittest
+
+import requests
+
+from transformers import (
+ EdgeTamConfig,
+ EdgeTamMaskDecoderConfig,
+ EdgeTamPromptEncoderConfig,
+ EdgeTamVisionConfig,
+ Sam2Processor,
+ pipeline,
+)
+from transformers.testing_utils import (
+ backend_empty_cache,
+ require_torch,
+ slow,
+ torch_device,
+)
+from transformers.utils import is_torch_available, is_vision_available
+from transformers.video_utils import load_video
+
+from ...test_configuration_common import ConfigTester
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
+from ...test_pipeline_mixin import PipelineTesterMixin
+
+
+if is_torch_available():
+ import torch
+
+ from transformers import AutoConfig, EdgeTamModel, Sam2Processor
+
+
+if is_vision_available():
+ from PIL import Image
+
+
+class EdgeTamPromptEncoderTester:
+ def __init__(
+ self,
+ hidden_size=32,
+ input_image_size=128,
+ patch_size=16,
+ mask_input_channels=8,
+ num_point_embeddings=4,
+ hidden_act="gelu",
+ ):
+ self.hidden_size = hidden_size
+ self.input_image_size = input_image_size
+ self.patch_size = patch_size
+ self.mask_input_channels = mask_input_channels
+ self.num_point_embeddings = num_point_embeddings
+ self.hidden_act = hidden_act
+
+ def get_config(self):
+ return EdgeTamPromptEncoderConfig(
+ image_size=self.input_image_size,
+ patch_size=self.patch_size,
+ mask_input_channels=self.mask_input_channels,
+ hidden_size=self.hidden_size,
+ num_point_embeddings=self.num_point_embeddings,
+ hidden_act=self.hidden_act,
+ )
+
+ def prepare_config_and_inputs(self):
+ dummy_points = floats_tensor([self.batch_size, 3, 2])
+ config = self.get_config()
+
+ return config, dummy_points
+
+
+class EdgeTamMaskDecoderTester:
+ def __init__(
+ self,
+ hidden_size=32,
+ hidden_act="relu",
+ mlp_dim=64,
+ num_hidden_layers=2,
+ num_attention_heads=4,
+ attention_downsample_rate=2,
+ num_multimask_outputs=3,
+ iou_head_depth=3,
+ iou_head_hidden_dim=32,
+ ):
+ self.hidden_size = hidden_size
+ self.hidden_act = hidden_act
+ self.mlp_dim = mlp_dim
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.attention_downsample_rate = attention_downsample_rate
+ self.num_multimask_outputs = num_multimask_outputs
+ self.iou_head_depth = iou_head_depth
+ self.iou_head_hidden_dim = iou_head_hidden_dim
+
+ def get_config(self):
+ return EdgeTamMaskDecoderConfig(
+ hidden_size=self.hidden_size,
+ hidden_act=self.hidden_act,
+ mlp_dim=self.mlp_dim,
+ num_hidden_layers=self.num_hidden_layers,
+ num_attention_heads=self.num_attention_heads,
+ attention_downsample_rate=self.attention_downsample_rate,
+ num_multimask_outputs=self.num_multimask_outputs,
+ iou_head_depth=self.iou_head_depth,
+ iou_head_hidden_dim=self.iou_head_hidden_dim,
+ )
+
+ def prepare_config_and_inputs(self):
+ config = self.get_config()
+
+ dummy_inputs = {
+ "image_embedding": floats_tensor([self.batch_size, self.hidden_size]),
+ }
+
+ return config, dummy_inputs
+
+
+class EdgeTamModelTester:
+ def __init__(
+ self,
+ parent,
+ num_channels=3,
+ image_size=128,
+ hidden_size=12,
+ patch_kernel_size=7,
+ patch_stride=4,
+ patch_padding=3,
+ dim_mul=2.0,
+ backbone_channel_list=[96, 48, 24, 12],
+ backbone_feature_sizes=[[32, 32], [16, 16], [8, 8]],
+ fpn_hidden_size=32,
+ memory_encoder_hidden_size=32,
+ batch_size=2,
+ is_training=False,
+ ):
+ self.parent = parent
+ self.image_size = image_size
+ self.hidden_size = hidden_size
+ self.patch_kernel_size = patch_kernel_size
+ self.patch_stride = patch_stride
+ self.patch_padding = patch_padding
+ self.dim_mul = dim_mul
+ self.backbone_channel_list = backbone_channel_list
+ self.backbone_feature_sizes = backbone_feature_sizes
+ self.fpn_hidden_size = fpn_hidden_size
+ self.batch_size = batch_size
+ self.num_channels = num_channels
+ self.is_training = is_training
+ self.memory_encoder_hidden_size = memory_encoder_hidden_size
+
+ self.prompt_encoder_tester = EdgeTamPromptEncoderTester()
+ self.mask_decoder_tester = EdgeTamMaskDecoderTester()
+
+ def prepare_config_and_inputs(self):
+ pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
+ config = self.get_config()
+
+ return config, pixel_values
+
+ def get_config(self):
+ vision_config = EdgeTamVisionConfig(
+ backbone_config=AutoConfig.from_pretrained(
+ "timm/repvit_m1.dist_in1k",
+ model_args={
+ "in_chans": 3,
+ "features_only": True,
+ "out_indices": (0, 1, 2, 3),
+ "embed_dim": self.backbone_channel_list[::-1],
+ },
+ ),
+ backbone_channel_list=self.backbone_channel_list,
+ backbone_feature_sizes=self.backbone_feature_sizes,
+ fpn_hidden_size=self.fpn_hidden_size,
+ )
+
+ prompt_encoder_config = self.prompt_encoder_tester.get_config()
+
+ mask_decoder_config = self.mask_decoder_tester.get_config()
+
+ return EdgeTamConfig(
+ vision_config=vision_config,
+ prompt_encoder_config=prompt_encoder_config,
+ mask_decoder_config=mask_decoder_config,
+ memory_attention_hidden_size=self.hidden_size,
+ memory_encoder_hidden_size=self.memory_encoder_hidden_size,
+ image_size=self.image_size,
+ mask_downsampler_embed_dim=32,
+ memory_fuser_embed_dim=32,
+ memory_attention_num_layers=1,
+ memory_attention_feed_forward_hidden_size=32,
+ )
+
+ def create_and_check_model(self, config, pixel_values):
+ model = EdgeTamModel(config=config)
+ model.to(torch_device)
+ model.eval()
+ with torch.no_grad():
+ result = model(pixel_values)
+ self.parent.assertEqual(result.iou_scores.shape, (self.batch_size, 1, 3))
+ self.parent.assertEqual(result.pred_masks.shape[:3], (self.batch_size, 1, 3))
+
+ def prepare_config_and_inputs_for_common(self):
+ config_and_inputs = self.prepare_config_and_inputs()
+ config, pixel_values = config_and_inputs
+ inputs_dict = {"pixel_values": pixel_values}
+ return config, inputs_dict
+
+
+@require_torch
+class EdgeTamModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
+ """
+ Here we also overwrite some of the tests of test_modeling_common.py, as SAM's vision encoder does not use input_ids, inputs_embeds,
+ attention_mask and seq_length.
+ """
+
+ all_model_classes = (EdgeTamModel,) if is_torch_available() else ()
+ pipeline_model_mapping = (
+ {"feature-extraction": EdgeTamModel, "mask-generation": EdgeTamModel} if is_torch_available() else {}
+ )
+ fx_compatible = False
+ test_pruning = False
+ test_resize_embeddings = False
+ test_head_masking = False
+ test_torchscript = False
+ _is_composite = True
+
+ def setUp(self):
+ self.model_tester = EdgeTamModelTester(self)
+ common_properties = ["initializer_range"]
+ self.config_tester = ConfigTester(
+ self, config_class=EdgeTamConfig, has_text_modality=False, common_properties=common_properties
+ )
+
+ def test_config(self):
+ self.config_tester.run_common_tests()
+
+ @unittest.skip(reason="Timm model does not use inputs_embeds")
+ def test_inputs_embeds(self):
+ pass
+
+ @unittest.skip(reason="Can't get or set embeddings for Timm model")
+ def test_model_get_set_embeddings(self):
+ pass
+
+ def test_model(self):
+ config_and_inputs = self.model_tester.prepare_config_and_inputs()
+ self.model_tester.create_and_check_model(*config_and_inputs)
+
+ # Override as EdgeTamModel doesn't have hidden states
+ def flash_attn_inference_equivalence(self, attn_implementation: str, padding_side: str):
+ r"""
+ Tests the equivalence between the eager and flash attention implementations.
+ This test is only for inference and runs with `torch_dtype=torch.bfloat16`.
+ """
+ if not self.has_attentions:
+ self.skipTest(reason="Model architecture does not support attentions")
+
+ for model_class in self.all_model_classes:
+ if (attn_implementation == "flash_attention_2" and not model_class._supports_flash_attn_2) or (
+ attn_implementation == "flash_attention_3" and not model_class._supports_flash_attn_3
+ ):
+ self.skipTest(f"{model_class.__name__} does not support {attn_implementation}")
+
+ config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+ model = model_class(config)
+
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ model.save_pretrained(tmpdirname)
+ model_fa = model_class.from_pretrained(
+ tmpdirname, torch_dtype=torch.bfloat16, attn_implementation=attn_implementation
+ )
+ model_fa.to(torch_device)
+
+ model = model_class.from_pretrained(tmpdirname, torch_dtype=torch.bfloat16)
+ model.to(torch_device)
+
+ dummy_input = inputs_dict[model.main_input_name][:1]
+ if dummy_input.dtype in [torch.float32, torch.float16]:
+ dummy_input = dummy_input.to(torch.bfloat16)
+
+ dummy_attention_mask = inputs_dict.get("attention_mask", None)
+
+ if dummy_attention_mask is not None:
+ dummy_attention_mask = dummy_attention_mask[:1]
+ if padding_side == "left":
+ dummy_attention_mask[:, 1:] = 1
+ dummy_attention_mask[:, :1] = 0
+ else:
+ dummy_attention_mask[:, :-1] = 1
+ dummy_attention_mask[:, -1:] = 0
+ if model.config.is_encoder_decoder:
+ decoder_input_ids = inputs_dict.get("decoder_input_ids", dummy_input)[:1]
+
+ outputs = model(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
+ outputs_fa = model_fa(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
+ else:
+ outputs = model(dummy_input, output_hidden_states=True)
+ outputs_fa = model_fa(dummy_input, output_hidden_states=True)
+
+ logits = outputs.vision_hidden_states[-1]
+ logits_fa = outputs_fa.vision_hidden_states[-1]
+
+ assert torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2)
+
+ if model.config.is_encoder_decoder:
+ other_inputs = {
+ "decoder_input_ids": decoder_input_ids,
+ "decoder_attention_mask": dummy_attention_mask,
+ "output_hidden_states": True,
+ }
+ if dummy_attention_mask is not None:
+ other_inputs["attention_mask"] = dummy_attention_mask
+
+ outputs = model(dummy_input, **other_inputs)
+ outputs_fa = model_fa(dummy_input, **other_inputs)
+ else:
+ other_inputs = {
+ "output_hidden_states": True,
+ }
+ if dummy_attention_mask is not None:
+ other_inputs["attention_mask"] = dummy_attention_mask
+
+ outputs = model(dummy_input, **other_inputs)
+ outputs_fa = model_fa(dummy_input, **other_inputs)
+
+ logits = outputs.vision_hidden_states[-1]
+ logits_fa = outputs_fa.vision_hidden_states[-1]
+
+ if padding_side == "left":
+ assert torch.allclose(logits_fa[1:], logits[1:], atol=4e-2, rtol=4e-2)
+
+ # check with inference + dropout
+ model.train()
+ _ = model_fa(dummy_input, **other_inputs)
+ else:
+ assert torch.allclose(logits_fa[:-1], logits[:-1], atol=4e-2, rtol=4e-2)
+
+ # Override as diffence slightly higher than the threshold
+ # def test_batching_equivalence(self, atol=5e-4, rtol=5e-4):
+ # super().test_batching_equivalence(atol=atol, rtol=rtol)
+
+ @unittest.skip(reason="TimmWrapperModel does not support an attention implementation")
+ def test_can_set_attention_dynamically_composite_model(self):
+ pass
+
+ @unittest.skip(reason="vision_hidden_states from TimmWrapperModel")
+ def test_hidden_states_output(self):
+ pass
+
+ @unittest.skip(reason="Timm weights cannot be fully constructed in _init_weights")
+ def test_can_init_all_missing_weights(self):
+ pass
+
+ @unittest.skip(
+ reason="TIMM's attention implementation is self configured and won't raise ValueError on global attention implementation."
+ )
+ def test_flash_attn_2_can_dispatch_composite_models(self):
+ pass
+
+ @unittest.skip("TimmWrapperModel cannot be tested with meta device")
+ def test_can_be_initialized_on_meta(self):
+ pass
+
+ @unittest.skip("TimmWrapperModel cannot be tested with meta device")
+ def test_can_load_with_meta_device_context_manager(self):
+ pass
+
+ ## Skip flash attention releated tests below
+ ## correct configuration:
+ ## from_pretrained(model_id, attn_implementation={"text_config": "flash_attention_2", "vision_config": "eager"}
+ @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_eager_matches_fa2_generate(self):
+ pass
+
+ @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_flash_attn_2_fp32_ln(self):
+ pass
+
+ @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_flash_attn_2_from_config(self):
+ pass
+
+ @unittest.skip("SDPA test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_eager_matches_sdpa_generate_with_dynamic_cache(self):
+ pass
+
+ @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_flash_attn_2_inference_equivalence_right_padding(self):
+ pass
+
+ @unittest.skip("SDPA test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_eager_matches_sdpa_generate(self):
+ pass
+
+ @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.")
+ def test_flash_attn_2_inference_equivalence(self):
+ pass
+
+ @unittest.skip("EdgeTAM does not have language_model, vision_tower, multi_modal_projector.")
+ def test_sdpa_can_dispatch_composite_models(self):
+ pass
+
+ @unittest.skip("Cannot set `output_attentions` for timm models.")
+ def test_attention_outputs(self):
+ pass
+
+ @unittest.skip("Cannot set `output_attentions` for timm models.")
+ def test_retain_grad_hidden_states_attentions(self):
+ pass
+
+ @unittest.skip("Cannot set `output_attentions` for timm models.")
+ def test_generate_compilation_all_outputs(self):
+ pass
+
+ @slow
+ def test_model_from_pretrained(self):
+ model_name = "yonigozlan/EdgeTAM-hf"
+ model = EdgeTamModel.from_pretrained(model_name)
+ self.assertIsNotNone(model)
+
+ def test_sdpa_can_compile_dynamic(self):
+ self.skipTest(reason="EDGETAM model can't be compiled dynamic yet")
+
+
+def prepare_image():
+ img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_groceries_image():
+ img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/groceries.jpg"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_dog_img():
+ img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/dog-sam.png"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_video():
+ video_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/bedroom.mp4"
+ raw_video, _ = load_video(video_url)
+ return raw_video
+
+
+@slow
+class EdgeTamModelIntegrationTest(unittest.TestCase):
+ def setUp(self):
+ super().setUp()
+ self.model = EdgeTamModel.from_pretrained("yonigozlan/EdgeTAM-hf").to(torch.float32)
+ self.processor = Sam2Processor.from_pretrained("yonigozlan/EdgeTAM-hf")
+ self.model.to(torch_device)
+ self.model.eval()
+
+ def tearDown(self):
+ super().tearDown()
+ # clean-up as much as possible GPU memory occupied by PyTorch
+ gc.collect()
+ backend_empty_cache(torch_device)
+
+ def test_inference_mask_generation_one_point_multimask(self):
+ raw_image = prepare_image()
+ input_points = [[[[500, 375]]]]
+ input_labels = [[[1]]]
+
+ inputs = self.processor(
+ images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(torch_device)
+
+ with torch.no_grad():
+ outputs = self.model(**inputs)
+ self.assertEqual(outputs.iou_scores.shape, (1, 1, 3))
+ self.assertEqual(outputs.pred_masks.shape, (1, 1, 3, 256, 256))
+ sorted_indices = torch.argsort(outputs.iou_scores.squeeze(), descending=True)
+ scores = outputs.iou_scores.squeeze()[sorted_indices]
+ masks_logits = outputs.pred_masks.squeeze()[sorted_indices][0, :3, :3]
+ torch.testing.assert_close(
+ scores, torch.tensor([0.7621, 0.4859, 0.0461]).to(torch_device), atol=1e-4, rtol=1e-4
+ )
+ torch.testing.assert_close(
+ masks_logits,
+ torch.tensor(
+ [[-19.5483, -22.3549, -26.0962], [-18.1821, -23.4761, -24.2262], [-20.3549, -24.5518, -22.7232]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_one_point_no_multimask(self):
+ raw_image = prepare_image()
+ input_points = [[[[500, 375]]]]
+ input_labels = [[[1]]]
+
+ inputs = self.processor(
+ images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(torch_device)
+
+ with torch.no_grad():
+ outputs = self.model(**inputs, multimask_output=False)
+ self.assertEqual(outputs.iou_scores.shape, (1, 1, 1))
+ self.assertEqual(outputs.pred_masks.shape, (1, 1, 1, 256, 256))
+ scores = outputs.iou_scores.squeeze((0, 1))
+ masks_logits = outputs.pred_masks.squeeze((0, 1))[0, :3, :3]
+ torch.testing.assert_close(scores, torch.tensor([0.7621]).to(torch_device), atol=1e-4, rtol=1e-4)
+ torch.testing.assert_close(
+ masks_logits,
+ torch.tensor(
+ [[-19.5483, -22.3549, -26.0962], [-18.1821, -23.4761, -24.2262], [-20.3549, -24.5518, -22.7232]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_batched_images_multi_points(self):
+ raw_image1 = prepare_image()
+ raw_image2 = prepare_dog_img()
+ input_points = [[[[500, 375]]], [[[770, 200], [730, 120]]]]
+ input_labels = [[[1]], [[1, 0]]]
+
+ inputs = self.processor(
+ images=[raw_image1, raw_image2], input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(torch_device)
+
+ with torch.no_grad():
+ outputs = self.model(**inputs)
+ self.assertEqual(outputs.iou_scores.shape, (2, 1, 3))
+ self.assertEqual(outputs.pred_masks.shape, (2, 1, 3, 256, 256))
+
+ sorted_indices = torch.argsort(outputs.iou_scores[0].squeeze(), descending=True)
+ scores1 = outputs.iou_scores[0].squeeze()[sorted_indices]
+ masks_logits1 = outputs.pred_masks[0].squeeze()[sorted_indices][0, :3, :3]
+ sorted_indices = torch.argsort(outputs.iou_scores[1].squeeze(), descending=True)
+ scores2 = outputs.iou_scores[1].squeeze()[sorted_indices]
+ masks_logits2 = outputs.pred_masks[1].squeeze()[sorted_indices][0, :3, :3]
+ torch.testing.assert_close(
+ scores1, torch.tensor([0.7490, 0.4685, 0.0463]).to(torch_device), atol=1e-4, rtol=1e-4
+ )
+ torch.testing.assert_close(
+ masks_logits1,
+ torch.tensor(
+ [[-19.1423, -21.6488, -25.6816], [-17.8018, -22.6512, -23.5699], [-19.9140, -23.6919, -22.3147]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ torch.testing.assert_close(
+ scores2, torch.tensor([0.7225, 0.6515, 0.6350]).to(torch_device), atol=1e-4, rtol=1e-4
+ )
+ torch.testing.assert_close(
+ masks_logits2,
+ torch.tensor([[-8.8259, -7.7961, -9.3665], [-8.2648, -8.7771, -9.1390], [-9.5951, -8.3995, -9.0599]]).to(
+ torch_device
+ ),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_batched_images_batched_points_multi_points(self):
+ raw_image1 = prepare_image()
+ raw_image2 = prepare_groceries_image()
+ input_points = [[[[500, 375]], [[650, 750]]], [[[400, 300]], [[630, 300], [550, 300]]]]
+ input_labels = [[[1], [1]], [[1], [1, 1]]]
+ inputs = self.processor(
+ images=[raw_image1, raw_image2], input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(torch_device)
+ with torch.no_grad():
+ outputs = self.model(**inputs, multimask_output=False)
+ self.assertEqual(outputs.iou_scores.shape, (2, 2, 1))
+ self.assertEqual(outputs.pred_masks.shape, (2, 2, 1, 256, 256))
+ torch.testing.assert_close(
+ outputs.iou_scores,
+ torch.tensor([[[0.7490], [0.9397]], [[0.7952], [0.8723]]]).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+ torch.testing.assert_close(
+ outputs.pred_masks[:, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-19.1423, -21.6488], [-17.8018, -22.6512]]], [[[-7.1591, -9.8201], [-7.4133, -9.2781]]]],
+ [[[[-16.7645, -15.2790], [-16.1805, -16.2937]]], [[[-8.5934, -8.4215], [-8.1873, -8.3722]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_batched_images_batched_boxes(self):
+ raw_image1 = prepare_image()
+ raw_image2 = prepare_groceries_image()
+ input_boxes = [
+ [[75, 275, 1725, 850], [425, 600, 700, 875], [1375, 550, 1650, 800], [1240, 675, 1400, 750]],
+ [[450, 170, 520, 350], [350, 190, 450, 350], [500, 170, 580, 350], [580, 170, 640, 350]],
+ ]
+ inputs = self.processor(images=[raw_image1, raw_image2], input_boxes=input_boxes, return_tensors="pt").to(
+ torch_device
+ )
+ with torch.no_grad():
+ outputs = self.model(**inputs, multimask_output=False)
+ self.assertEqual(outputs.iou_scores.shape, (2, 4, 1))
+ self.assertEqual(outputs.pred_masks.shape, (2, 4, 1, 256, 256))
+ torch.testing.assert_close(
+ outputs.iou_scores,
+ torch.tensor([[[0.9773], [0.9415], [0.9683], [0.8792]], [[0.9721], [0.9852], [0.9812], [0.9760]]]).to(
+ torch_device
+ ),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+ torch.testing.assert_close(
+ outputs.pred_masks[:, :, :, :2, :2],
+ torch.tensor(
+ [
+ [
+ [[[-12.6412, -12.0553], [-11.8415, -13.1696]]],
+ [[[-16.0378, -19.9641], [-15.4939, -19.0260]]],
+ [[[-18.8254, -23.6185], [-17.7889, -23.2116]]],
+ [[[-25.7024, -29.8722], [-22.9264, -30.0557]]],
+ ],
+ [
+ [[[-19.0264, -17.0396], [-16.9458, -16.3287]]],
+ [[[-20.9671, -19.2132], [-18.5827, -18.0511]]],
+ [[[-22.4642, -19.7389], [-19.4541, -19.4717]]],
+ [[[-21.9226, -18.6297], [-18.9272, -18.8151]]],
+ ],
+ ]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_from_existing_points_and_mask(self):
+ raw_image = prepare_image()
+ input_points = [[[[500, 375]]]]
+ input_labels = [[[1]]]
+ original_inputs = self.processor(
+ images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt"
+ ).to(torch_device)
+ with torch.no_grad():
+ outputs = self.model(**original_inputs)
+
+ # best mask to use as input for new points
+ mask_input = outputs.pred_masks[:, :, torch.argmax(outputs.iou_scores)]
+
+ new_input_points = [[[[500, 375], [1125, 625]]]]
+ new_input_labels = [[[1, 1]]]
+ inputs = self.processor(
+ input_points=new_input_points,
+ input_labels=new_input_labels,
+ original_sizes=original_inputs["original_sizes"],
+ return_tensors="pt",
+ ).to(torch_device)
+ with torch.no_grad():
+ outputs = self.model(
+ **inputs,
+ input_masks=mask_input,
+ image_embeddings=outputs.image_embeddings,
+ multimask_output=False,
+ )
+
+ self.assertEqual(outputs.iou_scores.shape, (1, 1, 1))
+ self.assertEqual(outputs.pred_masks.shape, (1, 1, 1, 256, 256))
+ scores = outputs.iou_scores.squeeze((0, 1))
+ masks_logits = outputs.pred_masks.squeeze((0, 1))[0, :3, :3]
+ torch.testing.assert_close(scores, torch.tensor([0.9431]).to(torch_device), atol=1e-4, rtol=1e-4)
+ torch.testing.assert_close(
+ masks_logits,
+ torch.tensor([[-4.1968, -4.9034, -6.0680], [-4.4053, -5.1200, -5.8580], [-4.3920, -5.5096, -5.8166]]).to(
+ torch_device
+ ),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # with negative point
+ new_input_points = [[[[500, 375], [1125, 625]]]]
+ new_input_labels = [[[1, 0]]]
+ inputs = self.processor(
+ input_points=new_input_points,
+ input_labels=new_input_labels,
+ original_sizes=original_inputs["original_sizes"],
+ return_tensors="pt",
+ ).to(torch_device)
+ with torch.no_grad():
+ outputs = self.model(
+ **inputs,
+ input_masks=mask_input,
+ image_embeddings=outputs.image_embeddings,
+ multimask_output=False,
+ )
+ self.assertEqual(outputs.iou_scores.shape, (1, 1, 1))
+ self.assertEqual(outputs.pred_masks.shape, (1, 1, 1, 256, 256))
+ scores = outputs.iou_scores.squeeze((0, 1))
+ masks_logits = outputs.pred_masks.squeeze((0, 1))[0, :3, :3]
+ torch.testing.assert_close(scores, torch.tensor([0.9695]).to(torch_device), atol=1e-4, rtol=1e-4)
+ torch.testing.assert_close(
+ masks_logits,
+ torch.tensor(
+ [[-14.3212, -15.4295, -17.4482], [-13.2246, -15.9468, -17.1341], [-15.1678, -16.4498, -14.7385]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_dummy_pipeline_generation(self):
+ generator = pipeline("mask-generation", model="yonigozlan/EdgeTAM-hf", device=torch_device)
+ raw_image = prepare_image()
+
+ _ = generator(raw_image, points_per_batch=64)
diff --git a/tests/models/edgetam_video/__init__.py b/tests/models/edgetam_video/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/edgetam_video/test_modeling_edgetam_video.py b/tests/models/edgetam_video/test_modeling_edgetam_video.py
new file mode 100644
index 000000000000..a2ad383351d2
--- /dev/null
+++ b/tests/models/edgetam_video/test_modeling_edgetam_video.py
@@ -0,0 +1,507 @@
+# coding=utf-8
+# 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.
+"""Testing suite for the PyTorch SAM2 model."""
+
+import gc
+import unittest
+
+import requests
+
+from transformers.testing_utils import (
+ backend_empty_cache,
+ slow,
+ torch_device,
+)
+from transformers.utils import is_torch_available, is_vision_available
+from transformers.video_utils import load_video
+
+
+if is_torch_available():
+ import torch
+
+ from transformers import EdgeTamVideoModel, Sam2VideoProcessor
+
+
+if is_vision_available():
+ from PIL import Image
+
+
+def prepare_image():
+ img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_groceries_image():
+ img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/groceries.jpg"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_dog_img():
+ img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/dog-sam.png"
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+ return raw_image
+
+
+def prepare_video():
+ video_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/bedroom.mp4"
+ raw_video, _ = load_video(video_url)
+ return raw_video
+
+
+@slow
+class EdgeTamVideoModelIntegrationTest(unittest.TestCase):
+ def setUp(self):
+ super().setUp()
+ self.video_model = EdgeTamVideoModel.from_pretrained("yonigozlan/EdgeTAM-hf").to(torch.float32)
+ self.processor = Sam2VideoProcessor.from_pretrained("yonigozlan/EdgeTAM-hf")
+ self.video_model.to(torch_device)
+ self.video_model.eval()
+
+ def tearDown(self):
+ super().tearDown()
+ # clean-up as much as possible GPU memory occupied by PyTorch
+ gc.collect()
+ backend_empty_cache(torch_device)
+
+ def test_inference_mask_generation_video_one_point(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_points=[[[[210, 350]]]],
+ input_labels=[[[1]]],
+ )
+ outputs = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = outputs.pred_masks
+ self.assertEqual(low_res_masks.shape, (1, 1, 256, 256))
+ video_res_masks = self.processor.post_process_masks([low_res_masks], [raw_video.shape[-3:-1]], binarize=False)[
+ 0
+ ]
+ self.assertEqual(video_res_masks.shape, (1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[0, 0, :3, :3],
+ torch.tensor(
+ [[-28.3880, -28.3880, -27.9277], [-27.5260, -27.5260, -27.2455], [-25.5902, -25.5902, -25.7136]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-28.3880, -28.3880], [-27.5260, -27.5260]]]],
+ [[[[-15.3350, -15.3350], [-15.0002, -15.0002]]]],
+ [[[[-14.8729, -14.8729], [-14.6724, -14.6724]]]],
+ ],
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_video_one_point_propagate_in_video_directly(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_points=[[[[210, 350]]]],
+ input_labels=[[[1]]],
+ )
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ print(f"VIDEO_TEST2 - ACTUAL frames[:3, :, :, :2, :2]: {frames[:3, :, :, :2, :2]}")
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-28.3880, -28.3880], [-27.5260, -27.5260]]]],
+ [[[[-15.3350, -15.3350], [-15.0002, -15.0002]]]],
+ [[[[-14.8729, -14.8729], [-14.6724, -14.6724]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_mask_generation_video_multi_points(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_points=[[[[210, 350], [250, 220]]]],
+ input_labels=[[[1, 1]]],
+ )
+ outputs = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = outputs.pred_masks
+ video_res_masks = self.processor.post_process_masks(
+ [outputs.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ self.assertEqual(low_res_masks.shape, (1, 1, 256, 256))
+ self.assertEqual(video_res_masks.shape, (1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[0, 0, :3, :3],
+ torch.tensor(
+ [[-17.3081, -17.3081, -16.9805], [-16.8430, -16.8430, -16.6766], [-15.7986, -15.7986, -15.9941]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ # higher tolerance due to errors propagating from frame to frame
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-17.3081, -17.3081], [-16.8430, -16.8430]]]],
+ [[[[-14.9302, -14.9302], [-14.8802, -14.8802]]]],
+ [[[[-14.4372, -14.4372], [-14.3697, -14.3697]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-2,
+ rtol=1e-2,
+ )
+
+ def test_inference_mask_generation_video_one_bb(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_boxes=[[[300, 0, 500, 400]]],
+ )
+ outputs = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = outputs.pred_masks
+ video_res_masks = self.processor.post_process_masks(
+ [outputs.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ self.assertEqual(low_res_masks.shape, (1, 1, 256, 256))
+ self.assertEqual(video_res_masks.shape, (1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[0, 0, :3, :3],
+ torch.tensor(
+ [[-17.3245, -17.3245, -16.9231], [-16.8773, -16.8773, -16.6082], [-15.8731, -15.8731, -15.9011]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ # higher tolerance due to errors propagating from frame to frame
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-17.3245, -17.3245], [-16.8773, -16.8773]]]],
+ [[[[-16.2826, -16.2826], [-15.9087, -15.9087]]]],
+ [[[[-15.8716, -15.8716], [-15.3992, -15.3992]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-2,
+ rtol=1e-2,
+ )
+
+ def test_inference_mask_generation_video_one_point_one_bb(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_boxes=[[[300, 0, 500, 400]]],
+ input_points=[[[[460, 60]]]],
+ input_labels=[[[1]]],
+ )
+ outputs = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = outputs.pred_masks
+ video_res_masks = self.processor.post_process_masks(
+ [outputs.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ self.assertEqual(low_res_masks.shape, (1, 1, 256, 256))
+ self.assertEqual(video_res_masks.shape, (1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[0, 0, :3, :3],
+ torch.tensor(
+ [[-13.9780, -13.9780, -13.7824], [-13.7642, -13.7642, -13.6000], [-13.2842, -13.2842, -13.1904]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ # higher tolerance due to errors propagating from frame to frame
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-13.9780, -13.9780], [-13.7642, -13.7642]]]],
+ [[[[-16.0142, -16.0142], [-15.5600, -15.5600]]]],
+ [[[[-16.7568, -16.7568], [-16.2460, -16.2460]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-2,
+ rtol=1e-2,
+ )
+
+ def test_inference_mask_generation_video_multi_objects_multi_points(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_ids = [2, 3] # give a unique id to each object we interact with (it can be any integers)
+
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_ids,
+ input_points=[[[[200, 300], [230, 250], [275, 175]], [[400, 150]]]],
+ input_labels=[[[1, 1, 0], [1]]],
+ )
+ outputs = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = outputs.pred_masks
+ video_res_masks = self.processor.post_process_masks(
+ [outputs.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ self.assertEqual(low_res_masks.shape, (2, 1, 256, 256))
+ self.assertEqual(video_res_masks.shape, (2, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[:, 0, :2, :2], # first object
+ torch.tensor(
+ [[[-12.6233, -12.6233], [-12.1809, -12.1809]], [[-13.4556, -13.4556], [-12.9549, -12.9549]]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 2, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-12.6233, -12.6233], [-12.1809, -12.1809]]], [[[-13.4556, -13.4556], [-12.9549, -12.9549]]]],
+ [[[[-12.5589, -12.5589], [-12.4450, -12.4450]]], [[[-12.2181, -12.2181], [-12.0188, -12.0188]]]],
+ [[[[-15.3170, -15.3170], [-15.0254, -15.0254]]], [[[-11.4912, -11.4912], [-11.3171, -11.3171]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_propagate_video_from_mask_input(self):
+ raw_video = prepare_video()
+ inference_session = self.processor.init_video_session(video=raw_video, inference_device=torch_device)
+ ann_frame_idx = 0 # the frame index we interact with
+ ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
+
+ # get input_mask
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_points=[[[[210, 350], [250, 220]]]],
+ input_labels=[[[1, 1]]],
+ )
+ sam2_video_output = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+
+ # set mask as input
+ self.processor.add_inputs_to_inference_session(
+ inference_session=inference_session,
+ frame_idx=ann_frame_idx,
+ obj_ids=ann_obj_id,
+ input_masks=self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0],
+ )
+ sam2_video_output = self.video_model(inference_session=inference_session, frame_idx=ann_frame_idx)
+ low_res_masks = sam2_video_output.pred_masks
+ self.assertEqual(low_res_masks.shape, (1, 1, 256, 256))
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ self.assertEqual(video_res_masks.shape, (1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ video_res_masks[0, 0, :3, :3],
+ torch.tensor(
+ [[-10.0000, -10.0000, -10.0000], [-10.0000, -10.0000, -10.0000], [-10.0000, -10.0000, -10.0000]]
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ # test propagate in video frames
+ frames = []
+ for sam2_video_output in self.video_model.propagate_in_video_iterator(
+ inference_session=inference_session,
+ start_frame_idx=ann_frame_idx,
+ max_frame_num_to_track=2,
+ ):
+ video_res_masks = self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], [raw_video.shape[-3:-1]], binarize=False
+ )[0]
+ frames.append(video_res_masks)
+ frames = torch.stack(frames, dim=0)
+ self.assertEqual(frames.shape, (3, 1, 1, raw_video.shape[-3], raw_video.shape[-2]))
+ torch.testing.assert_close(
+ frames[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-10.0000, -10.0000], [-10.0000, -10.0000]]]],
+ [[[[-17.4083, -17.4083], [-17.2256, -17.2256]]]],
+ [[[[-13.8533, -13.8533], [-13.7759, -13.7759]]]],
+ ],
+ ).to(torch_device),
+ atol=1e-4,
+ rtol=1e-4,
+ )
+
+ def test_inference_propagate_on_streamed_video(self):
+ raw_video = prepare_video()
+
+ inference_session = self.processor.init_video_session(inference_device=torch_device)
+ video_res_masks = []
+ max_frame_num_to_track = 3
+ for frame_idx, frame in enumerate(raw_video):
+ if frame_idx >= max_frame_num_to_track:
+ break
+ inputs = self.processor(images=frame, device=torch_device, return_tensors="pt")
+ if frame_idx == 0:
+ self.processor.add_inputs_to_inference_session(
+ inference_session,
+ frame_idx=0,
+ obj_ids=1,
+ input_points=[[[[210, 350], [250, 220]]]],
+ input_labels=[[[1, 1]]],
+ original_size=inputs.original_sizes[0],
+ )
+ sam2_video_output = self.video_model(inference_session=inference_session, frame=inputs.pixel_values[0])
+ video_res_masks.append(
+ self.processor.post_process_masks(
+ [sam2_video_output.pred_masks], inputs.original_sizes, binarize=False
+ )[0]
+ )
+
+ video_res_masks = torch.stack(video_res_masks, dim=0)
+ self.assertEqual(
+ video_res_masks.shape, (max_frame_num_to_track, 1, 1, raw_video.shape[-3], raw_video.shape[-2])
+ )
+ # higher tolerance due to errors propagating from frame to frame
+ print(f"VIDEO_TEST8 - ACTUAL video_res_masks[:3, :, :, :2, :2]: {video_res_masks[:3, :, :, :2, :2]}")
+ torch.testing.assert_close(
+ video_res_masks[:3, :, :, :2, :2],
+ torch.tensor(
+ [
+ [[[[-17.3081, -17.3081], [-16.8430, -16.8430]]]],
+ [[[[-14.9302, -14.9302], [-14.8802, -14.8802]]]],
+ [[[[-14.4372, -14.4372], [-14.3697, -14.3697]]]],
+ ]
+ ).to(torch_device),
+ atol=1e-2,
+ rtol=1e-2,
+ )
diff --git a/tests/models/efficientloftr/test_modeling_efficientloftr.py b/tests/models/efficientloftr/test_modeling_efficientloftr.py
index aef77ac85686..4ea8a4d823c5 100644
--- a/tests/models/efficientloftr/test_modeling_efficientloftr.py
+++ b/tests/models/efficientloftr/test_modeling_efficientloftr.py
@@ -23,7 +23,6 @@
require_vision,
set_config_for_less_flaky_test,
set_model_for_less_flaky_test,
- set_model_tester_for_less_flaky_test,
slow,
torch_device,
)
@@ -47,18 +46,18 @@ def __init__(
self,
parent,
batch_size=2,
- image_width=80,
- image_height=60,
- stage_num_blocks: list[int] = [1, 1, 1],
- out_features: list[int] = [32, 32, 128],
- stage_stride: list[int] = [2, 1, 2],
+ image_width=6, # need to be a multiple of `stage_stride[0] * stage_stride[1]`
+ image_height=4, # need to be a multiple of `stage_stride[0] * stage_stride[1]`
+ stage_num_blocks: list[int] = [1, 1],
+ out_features: list[int] = [16, 16], # need to be >= 2 to make `config.fine_fusion_dims > 0`
+ stage_stride: list[int] = [2, 1],
q_aggregation_kernel_size: int = 1,
kv_aggregation_kernel_size: int = 1,
q_aggregation_stride: int = 1,
kv_aggregation_stride: int = 1,
num_attention_layers: int = 2,
num_attention_heads: int = 8,
- hidden_size: int = 128,
+ hidden_size: int = 16,
coarse_matching_threshold: float = 0.0,
fine_kernel_size: int = 2,
coarse_matching_border_removal: int = 0,
@@ -360,8 +359,6 @@ def recursive_check(batched_object, single_row_object, model_name, key):
msg += str(e)
raise AssertionError(msg)
- set_model_tester_for_less_flaky_test(self)
-
config, batched_input = self.model_tester.prepare_config_and_inputs_for_common()
set_config_for_less_flaky_test(config)
diff --git a/tests/models/emu3/test_modeling_emu3.py b/tests/models/emu3/test_modeling_emu3.py
index 8975cfe4a0b4..b04d41242909 100644
--- a/tests/models/emu3/test_modeling_emu3.py
+++ b/tests/models/emu3/test_modeling_emu3.py
@@ -350,19 +350,11 @@ def test_disk_offload_bin(self):
def test_cpu_offload(self):
pass
- @unittest.skip("VQ-VAE module doesn't initialize weights properly")
- def test_initialization(self):
- pass
-
@pytest.mark.generate
@unittest.skip("Emu3 has dynamic control flow in vision backbone")
def test_generate_with_static_cache(self):
pass
- # @unittest.skip("Emu3 can't be smaller than currently if we want to downsample images")
- # def test_model_is_small(self):
- # pass
-
@require_torch
class Emu3IntegrationTest(unittest.TestCase):
diff --git a/tests/models/encodec/test_modeling_encodec.py b/tests/models/encodec/test_modeling_encodec.py
index 05e13f9482d9..407c19df8a9d 100644
--- a/tests/models/encodec/test_modeling_encodec.py
+++ b/tests/models/encodec/test_modeling_encodec.py
@@ -414,28 +414,6 @@ def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}):
dict_inputs = self._prepare_for_class(inputs_dict, model_class)
check_equivalence(model, tuple_inputs, dict_inputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = ["conv"]
- ignore_init = ["lstm"]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- elif not any(x in name for x in ignore_init):
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_identity_shortcut(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs()
config.use_conv_shortcut = False
diff --git a/tests/models/eomt/test_modeling_eomt.py b/tests/models/eomt/test_modeling_eomt.py
index 1c92692f2795..faf99cbe26d2 100644
--- a/tests/models/eomt/test_modeling_eomt.py
+++ b/tests/models/eomt/test_modeling_eomt.py
@@ -22,7 +22,7 @@
from transformers.utils import is_torch_available, is_vision_available
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -47,7 +47,7 @@ def __init__(
num_labels=4,
hidden_size=8,
num_attention_heads=2,
- num_hidden_layers=4,
+ num_hidden_layers=2,
):
self.parent = parent
self.batch_size = batch_size
@@ -233,40 +233,6 @@ def test_training(self):
loss = model(**inputs).loss
loss.backward()
- def test_initialization(self):
- # Apart from the below params, all other parameters are initialized using kaiming uniform.
- non_uniform_init_parms = [
- "layernorm.bias",
- "layernorm.weight",
- "norm1.bias",
- "norm1.weight",
- "norm2.bias",
- "norm2.weight",
- "layer_scale1.lambda1",
- "layer_scale2.lambda1",
- "register_tokens",
- "cls_token",
- ]
-
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if any(x in name for x in non_uniform_init_parms):
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@require_torch
class EomtForUniversalSegmentationIntegrationTest(unittest.TestCase):
diff --git a/tests/models/ernie4_5/test_modeling_ernie4_5.py b/tests/models/ernie4_5/test_modeling_ernie4_5.py
index 10c3287e5b80..1fb5969e900d 100644
--- a/tests/models/ernie4_5/test_modeling_ernie4_5.py
+++ b/tests/models/ernie4_5/test_modeling_ernie4_5.py
@@ -33,7 +33,6 @@
from transformers import (
AutoTokenizer,
- Ernie4_5Config,
Ernie4_5ForCausalLM,
Ernie4_5Model,
)
@@ -41,21 +40,11 @@
class Ernie4_5ModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = Ernie4_5Config
base_model_class = Ernie4_5Model
- causal_lm_class = Ernie4_5ForCausalLM
@require_torch
class Ernie4_5ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- Ernie4_5Model,
- Ernie4_5ForCausalLM,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Ernie4_5Model,
@@ -64,8 +53,6 @@ class Ernie4_5ModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False # Broken by attention refactor cc @Cyrilvallez
model_tester_class = Ernie4_5ModelTester
diff --git a/tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py b/tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
index 2e27bfc9332a..59839c0466c1 100644
--- a/tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
+++ b/tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
@@ -18,7 +18,7 @@
import pytest
-from transformers import Ernie4_5_MoeConfig, is_torch_available
+from transformers import is_torch_available
from transformers.testing_utils import (
cleanup,
is_flaky,
@@ -46,22 +46,12 @@
class Ernie4_5_MoeModelTester(CausalLMModelTester):
- config_class = Ernie4_5_MoeConfig
if is_torch_available():
base_model_class = Ernie4_5_MoeModel
- causal_lm_class = Ernie4_5_MoeForCausalLM
@require_torch
class Ernie4_5_MoeModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- Ernie4_5_MoeModel,
- Ernie4_5_MoeForCausalLM,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Ernie4_5_MoeModel,
@@ -71,8 +61,6 @@ class Ernie4_5_MoeModelTest(CausalLMModelTest, unittest.TestCase):
else {}
)
- test_headmasking = False
- test_pruning = False
test_all_params_have_gradient = False
model_tester_class = Ernie4_5_MoeModelTester
diff --git a/tests/models/esm/test_modeling_esmfold.py b/tests/models/esm/test_modeling_esmfold.py
index b13e7fe58b1d..84172447e24b 100644
--- a/tests/models/esm/test_modeling_esmfold.py
+++ b/tests/models/esm/test_modeling_esmfold.py
@@ -244,12 +244,6 @@ def test_model_outputs_equivalence(self):
def test_feed_forward_chunking(self):
pass
- @unittest.skip(
- reason="ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."
- )
- def test_initialization(self):
- pass
-
@unittest.skip(reason="ESMFold doesn't support torchscript compilation.")
def test_torchscript_output_attentions(self):
pass
diff --git a/tests/models/evolla/test_modeling_evolla.py b/tests/models/evolla/test_modeling_evolla.py
index 50574c7c5096..28370874fcaf 100644
--- a/tests/models/evolla/test_modeling_evolla.py
+++ b/tests/models/evolla/test_modeling_evolla.py
@@ -32,7 +32,6 @@
from ...test_modeling_common import (
TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION,
ModelTesterMixin,
- _config_zero_init,
ids_tensor,
random_attention_mask,
)
@@ -257,7 +256,7 @@ def test_generate_multiple_proteins(self):
def test_saprot_output(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
config.return_dict = True
- protein_informations = {
+ protein_information = {
"input_ids": inputs_dict["protein_input_ids"],
"attention_mask": inputs_dict["protein_attention_mask"],
}
@@ -267,13 +266,13 @@ def test_saprot_output(self):
model = model_class(config)
model.to(torch_device)
model.eval()
- protein_encoder_outputs = model.protein_encoder.model(**protein_informations, return_dict=True)
+ protein_encoder_outputs = model.protein_encoder.model(**protein_information, return_dict=True)
print(model_class, protein_encoder_outputs)
def test_protein_encoder_output(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
config.return_dict = True
- protein_informations = {
+ protein_information = {
"input_ids": inputs_dict["protein_input_ids"],
"attention_mask": inputs_dict["protein_attention_mask"],
}
@@ -283,7 +282,7 @@ def test_protein_encoder_output(self):
model = model_class(config)
model.to(torch_device)
model.eval()
- protein_encoder_outputs = model.protein_encoder(**protein_informations, return_dict=True)
+ protein_encoder_outputs = model.protein_encoder(**protein_information, return_dict=True)
print(model_class, protein_encoder_outputs)
def test_single_forward(self):
@@ -301,25 +300,6 @@ def test_single_forward(self):
outputs = model(**self._prepare_for_class(inputs_dict, model_class))
print(outputs)
- def test_initialization(self):
- # we skip the latents initialization test
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # skip latents
- if name.endswith("latents"):
- print(f"Skipping latents {name}")
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
@unittest.skip("Evolla requires both text and protein inputs which is currently not done in this test.")
def test_eager_matches_sdpa_inference(self):
diff --git a/tests/models/exaone4/test_modeling_exaone4.py b/tests/models/exaone4/test_modeling_exaone4.py
index 1045c025b159..c934821b4599 100644
--- a/tests/models/exaone4/test_modeling_exaone4.py
+++ b/tests/models/exaone4/test_modeling_exaone4.py
@@ -21,7 +21,6 @@
from transformers import (
AutoTokenizer,
- Exaone4Config,
GenerationConfig,
is_torch_available,
)
@@ -35,7 +34,6 @@
)
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
-from ...test_configuration_common import ConfigTester
if is_torch_available():
@@ -51,28 +49,12 @@
class Exaone4ModelTester(CausalLMModelTester):
- config_class = Exaone4Config
if is_torch_available():
base_model_class = Exaone4Model
- causal_lm_class = Exaone4ForCausalLM
- sequence_class = Exaone4ForSequenceClassification
- token_class = Exaone4ForTokenClassification
- question_answering_class = Exaone4ForQuestionAnswering
@require_torch
class Exaone4ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- Exaone4Model,
- Exaone4ForCausalLM,
- Exaone4ForSequenceClassification,
- Exaone4ForQuestionAnswering,
- Exaone4ForTokenClassification,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Exaone4Model,
@@ -85,20 +67,17 @@ class Exaone4ModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False # Broken by attention refactor cc @Cyrilvallez
model_tester_class = Exaone4ModelTester
model_split_percents = [0.5, 0.6]
- def setUp(self):
- self.model_tester = Exaone4ModelTester(self)
- self.config_tester = ConfigTester(self, config_class=Exaone4Config, hidden_size=37)
-
@require_torch
class Exaone4IntegrationTest(unittest.TestCase):
- TEST_MODEL_ID = "LGAI-EXAONE/EXAONE-4.0-Instruct" # dummy model id
+ TEST_MODEL_ID = "LGAI-EXAONE/EXAONE-4.0-32B"
+
+ def setUp(self):
+ cleanup(torch_device, gc_collect=True)
def tearDown(self):
# TODO (joao): automatic compilation, i.e. compilation when `cache_implementation="static"` is used, leaves
@@ -111,124 +90,40 @@ def tearDown(self):
def test_model_logits(self):
input_ids = [405, 7584, 79579, 76636, 2907, 94640, 373]
model = Exaone4ForCausalLM.from_pretrained(
- self.TEST_MODEL_ID, device_map="auto", dtype=torch.float16, attn_implementation="eager"
- )
- input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
- with torch.no_grad():
- out = model(input_ids).logits.float().cpu()
-
- EXPECTED_MEAN = torch.tensor([[13.9380, 12.9951, 12.9442, 10.6576, 11.0901, 12.1466, 9.2482]])
- EXPECTED_SLICE = torch.tensor(
- [
- 4.9180,
- 11.6406,
- 21.1250,
- 13.4062,
- 20.8438,
- 18.0625,
- 17.9688,
- 18.7812,
- 18.0156,
- 18.3594,
- 18.5000,
- 19.1719,
- 18.5156,
- 19.3438,
- 19.5000,
- 20.6406,
- 19.4844,
- 19.2812,
- 19.4688,
- 20.0156,
- 19.8438,
- 19.9531,
- 19.7188,
- 20.5938,
- 20.5312,
- 20.1250,
- 20.4062,
- 21.4062,
- 21.2344,
- 20.7656,
- ]
- )
-
- torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
- torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-4, rtol=1e-4)
- del model
- cleanup(torch_device, gc_collect=True)
-
- @slow
- def test_model_logits_bf16(self):
- input_ids = [405, 7584, 79579, 76636, 2907, 94640, 373]
- model = Exaone4ForCausalLM.from_pretrained(
- self.TEST_MODEL_ID, device_map="auto", dtype=torch.bfloat16, attn_implementation="eager"
+ self.TEST_MODEL_ID,
+ device_map="auto",
+ dtype=torch.bfloat16,
)
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
with torch.no_grad():
out = model(input_ids).logits.float().cpu()
- EXPECTED_MEAN = torch.tensor([[13.8797, 13.0799, 12.9665, 10.7712, 11.1006, 12.2406, 9.3248]])
+ EXPECTED_MEAN = torch.tensor([[22.1993, 8.5845, 10.0401, 12.4262, 9.3112, 29.7933, 8.2628]])
EXPECTED_SLICE = torch.tensor(
- [
- 4.8750,
- 11.6250,
- 21.0000,
- 13.3125,
- 20.8750,
- 18.0000,
- 18.0000,
- 18.7500,
- 18.0000,
- 18.3750,
- 18.5000,
- 19.1250,
- 18.5000,
- 19.3750,
- 19.5000,
- 20.6250,
- 19.5000,
- 19.2500,
- 19.5000,
- 20.0000,
- 19.8750,
- 19.8750,
- 19.7500,
- 20.6250,
- 20.5000,
- 20.1250,
- 20.3750,
- 21.3750,
- 21.2500,
- 20.7500,
- ]
+ [20.6250, 19.6250, 14.5000, 21.1250, 24.5000, 22.1250, 24.0000, 24.8750, 25.0000, 25.3750]
)
torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
- torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-4, rtol=1e-4)
- del model
- cleanup(torch_device, gc_collect=True)
+ torch.testing.assert_close(out[0, 0, :10], EXPECTED_SLICE, atol=1e-4, rtol=1e-4)
@slow
- def test_model_generation(self):
- EXPECTED_TEXT = "Tell me about the Miracle on the Han river.\n\nThe Miracle on the Han River is a story about the miracle of the Korean War Armistice. The story is told by a Korean soldier who is a witness to the armistice negotiations. He is reluctant to tell the story because he does not want to be a hypocrite, but he feels that everyone should know what really happened.\n\nThe Korean War began on June 25, 1950, when North Korean troops invaded South Korea. Soon the United Nations troops, primarily from South Korea, were in support of the United States. The war was still ongoing when North Korean troops stopped their advance"
+ def test_model_generation_eager(self):
+ EXPECTED_TEXT = "Tell me about the Miracle on the Han river.\n\nOkay, the Miracle on the Han River refers to the rapid industrialization and economic growth of South"
prompt = "Tell me about the Miracle on the Han river."
tokenizer = AutoTokenizer.from_pretrained(self.TEST_MODEL_ID)
model = Exaone4ForCausalLM.from_pretrained(
- self.TEST_MODEL_ID, device_map="auto", dtype=torch.float16, attn_implementation="eager"
+ self.TEST_MODEL_ID, device_map="auto", dtype=torch.bfloat16, attn_implementation="eager"
)
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device)
# greedy generation outputs
- generated_ids = model.generate(input_ids, max_new_tokens=128, temperature=0)
+ generated_ids = model.generate(input_ids, max_new_tokens=20, temperature=0)
text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
self.assertEqual(EXPECTED_TEXT, text)
- del model
- cleanup(torch_device, gc_collect=True)
@slow
- def test_model_generation_bf16_sdpa(self):
- EXPECTED_TEXT = "Tell me about the Miracle on the Han river.\n\nThe Miracle on the Han River is a story about the miracle of the Korean War Armistice.\n\nThe Korean War broke out in 35 years ago in 1950. The war was the result of the ideological conflict between the communist north and the capitalist south. The war was brought to a halt in 1953. There was to be peace talks but no peace treaty. As a result of the stalemate the Korean people have neither a peace treaty nor a reunification nor a democratization of Korea. The stalemate of 35 years has produced a people of 70 million"
+ def test_model_generation_sdpa(self):
+ EXPECTED_TEXT = "Tell me about the Miracle on the Han river.\n\nOkay, the Miracle on the Han River refers to the rapid industrialization and economic growth of South"
prompt = "Tell me about the Miracle on the Han river."
tokenizer = AutoTokenizer.from_pretrained(self.TEST_MODEL_ID)
model = Exaone4ForCausalLM.from_pretrained(
@@ -237,11 +132,9 @@ def test_model_generation_bf16_sdpa(self):
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device)
# greedy generation outputs
- generated_ids = model.generate(input_ids, max_new_tokens=128, temperature=0)
+ generated_ids = model.generate(input_ids, max_new_tokens=20, temperature=0)
text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
self.assertEqual(EXPECTED_TEXT, text)
- del model
- cleanup(torch_device, gc_collect=True)
@slow
@require_torch_accelerator
@@ -250,33 +143,27 @@ def test_model_generation_long_flash(self):
EXPECTED_OUTPUT_TOKEN_IDS = [433, 9055]
input_ids = [433, 9055] * 2048
model = Exaone4ForCausalLM.from_pretrained(
- self.TEST_MODEL_ID, device_map="auto", dtype=torch.float16, attn_implementation="flash_attention_2"
+ self.TEST_MODEL_ID, device_map="auto", dtype=torch.bfloat16, attn_implementation="flash_attention_2"
)
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
generated_ids = model.generate(input_ids, max_new_tokens=4, temperature=0)
self.assertEqual(EXPECTED_OUTPUT_TOKEN_IDS, generated_ids[0][-2:].tolist())
- del model
- cleanup(torch_device, gc_collect=True)
@slow
@require_torch_accelerator
def test_model_generation_beyond_sliding_window(self):
- EXPECTED_TEXT_COMPLETION = (
- " but I'm not sure if I'm going to be able to see it. I really enjoy the scenery, but I'm not sure if I"
- )
+ EXPECTED_TEXT_COMPLETION = " This is a nice place. I really enjoy the scenery, and the atmosphere is so relaxing. I'm grateful for the opportunity to experience this place. It"
tokenizer = AutoTokenizer.from_pretrained(self.TEST_MODEL_ID)
prompt = "This is a nice place. " * 700 + "I really enjoy the scenery,"
model = Exaone4ForCausalLM.from_pretrained(
- self.TEST_MODEL_ID, device_map="auto", dtype=torch.float16, attn_implementation="sdpa"
+ self.TEST_MODEL_ID, device_map="auto", dtype=torch.bfloat16, attn_implementation="sdpa"
)
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device)
- generated_ids = model.generate(input_ids, max_new_tokens=32, temperature=0)
+ generated_ids = model.generate(input_ids, max_new_tokens=20, temperature=0)
text = tokenizer.decode(generated_ids[0, -32:], skip_special_tokens=True)
self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
- del model
- cleanup(torch_device, gc_collect=True)
@pytest.mark.torch_export_test
@slow
@@ -290,9 +177,7 @@ def test_export_static_cache(self):
)
tokenizer = AutoTokenizer.from_pretrained(self.TEST_MODEL_ID, padding_side="right")
- EXPECTED_TEXT_COMPLETION = [
- "The Deep Learning is 100% free and easy to use.\n\n## How to use Deep Learning?\n\n"
- ]
+ EXPECTED_TEXT_COMPLETION = ["The Deep Learning is \n['Deep Learning',"]
max_generation_length = tokenizer(EXPECTED_TEXT_COMPLETION, return_tensors="pt", padding=True)[
"input_ids"
].shape[-1]
diff --git a/tests/models/falcon/test_modeling_falcon.py b/tests/models/falcon/test_modeling_falcon.py
index f15b86d425f1..cf025f463516 100644
--- a/tests/models/falcon/test_modeling_falcon.py
+++ b/tests/models/falcon/test_modeling_falcon.py
@@ -36,7 +36,6 @@
from transformers import (
FalconForCausalLM,
- FalconForQuestionAnswering,
FalconForSequenceClassification,
FalconForTokenClassification,
FalconModel,
@@ -45,11 +44,7 @@
class FalconModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = FalconConfig
base_model_class = FalconModel
- causal_lm_class = FalconForCausalLM
- sequence_class = FalconForSequenceClassification
- token_class = FalconForTokenClassification
def __init__(self, parent, new_decoder_architecture=True):
super().__init__(parent)
@@ -59,17 +54,6 @@ def __init__(self, parent, new_decoder_architecture=True):
@require_torch
class FalconModelTest(CausalLMModelTest, unittest.TestCase):
model_tester_class = FalconModelTester
- all_model_classes = (
- (
- FalconModel,
- FalconForCausalLM,
- FalconForSequenceClassification,
- FalconForTokenClassification,
- FalconForQuestionAnswering,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": FalconModel,
@@ -81,8 +65,6 @@ class FalconModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
# TODO (ydshieh): Check this. See https://app.circleci.com/pipelines/github/huggingface/transformers/79245/workflows/9490ef58-79c2-410d-8f51-e3495156cf9c/jobs/1012146
def is_pipeline_test_to_skip(
@@ -208,7 +190,7 @@ def test_falcon_alibi_sdpa_matches_eager(self):
config = FalconConfig(
vocab_size=1000,
hidden_size=64,
- num_hidden_layers=3,
+ num_hidden_layers=2,
num_attention_heads=4,
new_decoder_architecture=True,
alibi=True,
diff --git a/tests/models/falcon_h1/test_modeling_falcon_h1.py b/tests/models/falcon_h1/test_modeling_falcon_h1.py
index cc78f7bf7c1d..27eb8e32713b 100644
--- a/tests/models/falcon_h1/test_modeling_falcon_h1.py
+++ b/tests/models/falcon_h1/test_modeling_falcon_h1.py
@@ -14,7 +14,6 @@
# limitations under the License.
"""Testing suite for the PyTorch FalconH1 model."""
-import inspect
import unittest
import pytest
@@ -55,7 +54,7 @@ def __init__(
use_labels=True,
vocab_size=99,
hidden_size=32,
- num_hidden_layers=4,
+ num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
intermediate_size=64,
@@ -311,37 +310,6 @@ def test_decoder_model_past_with_large_inputs(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)
- # def test_initialization(self):
- # r"""
- # Overriding the test_initialization test as the A_log and D params of the FalconH1 mixer are initialized differently
- # """
- # config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- # configs_no_init = _config_zero_init(config)
- # for model_class in self.all_model_classes:
- # model = model_class(config=configs_no_init)
- # for name, param in model.named_parameters():
- # if param.requires_grad:
- # if "A_log" in name:
- # A = torch.arange(1, config.mamba_n_heads + 1, dtype=torch.float32)
- # torch.testing.assert_close(param.data, torch.log(A), rtol=1e-5, atol=1e-5)
- # elif "D" in name:
- # D = torch.ones(config.mamba_n_heads, dtype=torch.float32)
- # torch.testing.assert_close(param.data, D, rtol=1e-5, atol=1e-5)
- # else:
- # self.assertIn(
- # ((param.data.mean() * 1e9).round() / 1e9).item(),
- # [0.0, 1.0],
- # msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- # )
-
- def test_mismatched_shapes_have_properly_initialized_weights(self):
- r"""
- Overriding the test_mismatched_shapes_have_properly_initialized_weights test because A_log and D params of the
- FalconH1 mixer are initialized differently and we tested that in test_initialization
- """
- self.skipTest(reason="Cumbersome and redundant for FalconH1")
-
def test_attention_outputs(self):
r"""
Overriding the test_attention_outputs test as the FalconH1 model outputs attention only for its attention layers
@@ -413,88 +381,11 @@ def test_batching_equivalence(self):
super().test_batching_equivalence()
self.model_tester.use_input_mask = orig
- # essentially the same test in test_utils, just adjustment for rtol for this model
@pytest.mark.generate
def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
-
- # - The model must support padding
- if not self.has_attentions:
- self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
- decoder_only_classes = []
- for model_class in self.all_generative_model_classes:
- config, _ = self.prepare_config_and_inputs_for_generate()
- if config.is_encoder_decoder:
- continue
- else:
- decoder_only_classes.append(model_class)
- if len(decoder_only_classes) == 0:
- self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
- has_encoder_attributes = any(
- attr_name
- for attr_name in config.to_dict()
- if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
- )
- if has_encoder_attributes:
- self.skipTest(
- reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
- )
-
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in decoder_only_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
-
- # - for left padding we absolutely need to use an all ones
- # attention mask, so we do not use the one in inputs_dict
- attention_mask = torch.ones_like(input_ids)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(**model_kwargs).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(**model_kwargs).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
+ # TODO: document why a random attention mask causes this test to fail, but a full mask doesn't
+ unpadded_custom_inputs = {"attention_mask": None}
+ super().test_left_padding_compatibility(unpadded_custom_inputs=unpadded_custom_inputs)
@slow
diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
index a04660362813..09b263c8ffbd 100644
--- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
+++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
@@ -13,7 +13,6 @@
# limitations under the License.
-import math
import unittest
from unittest.util import safe_repr
@@ -34,7 +33,7 @@
from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -333,45 +332,6 @@ def test_falcon_mamba_lm_head_forward_and_backwards(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_falcon_mamba_lm_head_forward_and_backwards(*config_and_inputs)
- def test_initialization(self):
- config, _ = self.model_tester.prepare_config_and_inputs_for_common()
- config.rescale_prenorm_residual = True
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if "dt_proj.bias" in name:
- dt = torch.exp(
- torch.tensor([0, 1]) * (math.log(config.time_step_max) - math.log(config.time_step_min))
- + math.log(config.time_step_min)
- ).clamp(min=config.time_step_floor)
- inv_dt = dt + torch.log(-torch.expm1(-dt))
- if param.requires_grad:
- self.assertTrue(param.data.max().item() <= inv_dt[1])
- self.assertTrue(param.data.min().item() >= inv_dt[0])
- elif "A_log" in name:
- A = torch.arange(1, config.state_size + 1, dtype=torch.float32)[None, :]
- A = A.expand(config.intermediate_size, -1).contiguous()
- torch.testing.assert_close(param.data, torch.log(A), rtol=1e-5, atol=1e-5)
- elif "D" in name:
- if param.requires_grad:
- # check if it's a ones like
- torch.testing.assert_close(param.data, torch.ones_like(param.data), rtol=1e-5, atol=1e-5)
- else:
- if param.requires_grad:
- if (
- "mixer.conv1d.weight" in name
- or "mixer.dt_proj.weight" in name
- or "mixer.out_proj.weight" in name
- ):
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@slow
# Ignore copy
def test_model_from_pretrained(self):
diff --git a/tests/models/fastspeech2_conformer/test_modeling_fastspeech2_conformer.py b/tests/models/fastspeech2_conformer/test_modeling_fastspeech2_conformer.py
index 6ee0015b9a65..065a60cfb2b4 100644
--- a/tests/models/fastspeech2_conformer/test_modeling_fastspeech2_conformer.py
+++ b/tests/models/fastspeech2_conformer/test_modeling_fastspeech2_conformer.py
@@ -33,7 +33,7 @@
)
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, ids_tensor
if is_torch_available():
@@ -142,22 +142,6 @@ def test_model(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_model(*config_and_inputs)
- def test_initialization(self):
- config, _ = self.model_tester.prepare_config_and_inputs_for_common()
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- msg = f"Parameter {name} of model {model_class} seems not properly initialized"
- if "norm" in name:
- if "bias" in name:
- self.assertEqual(param.data.mean().item(), 0.0, msg=msg)
- if "weight" in name:
- self.assertEqual(param.data.mean().item(), 1.0, msg=msg)
- elif "conv" in name or "embed" in name:
- self.assertTrue(-1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0, msg=msg)
-
def test_duration_energy_pitch_output(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
config.return_dict = True
@@ -575,22 +559,6 @@ def test_model(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_model(*config_and_inputs)
- def test_initialization(self):
- config, _ = self.model_tester.prepare_config_and_inputs_for_common()
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- msg = f"Parameter {name} of model {model_class} seems not properly initialized"
- if "norm" in name:
- if "bias" in name:
- self.assertEqual(param.data.mean().item(), 0.0, msg=msg)
- if "weight" in name:
- self.assertEqual(param.data.mean().item(), 1.0, msg=msg)
- elif "conv" in name or "embed" in name:
- self.assertTrue(-1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0, msg=msg)
-
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
return inputs_dict
diff --git a/tests/models/flava/test_modeling_flava.py b/tests/models/flava/test_modeling_flava.py
index 896ce256955a..ac5e66349693 100644
--- a/tests/models/flava/test_modeling_flava.py
+++ b/tests/models/flava/test_modeling_flava.py
@@ -925,30 +925,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for FLAVA
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale" or name == "flava.logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/flex_olmo/test_modeling_flex_olmo.py b/tests/models/flex_olmo/test_modeling_flex_olmo.py
index b73807502873..15e4bb57c4af 100644
--- a/tests/models/flex_olmo/test_modeling_flex_olmo.py
+++ b/tests/models/flex_olmo/test_modeling_flex_olmo.py
@@ -18,7 +18,7 @@
import pytest
-from transformers import FlexOlmoConfig, is_torch_available
+from transformers import is_torch_available
from transformers.models.auto.tokenization_auto import AutoTokenizer
from transformers.testing_utils import (
Expectations,
@@ -43,9 +43,7 @@
class FlexOlmoModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = FlexOlmoConfig
base_model_class = FlexOlmoModel
- causal_lm_class = FlexOlmoForCausalLM
@require_torch
@@ -59,8 +57,6 @@ class FlexOlmoModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False
test_torchscript = False
test_all_params_have_gradient = False
diff --git a/tests/models/focalnet/test_modeling_focalnet.py b/tests/models/focalnet/test_modeling_focalnet.py
index f4dac79f9ca0..1198f6c52d57 100644
--- a/tests/models/focalnet/test_modeling_focalnet.py
+++ b/tests/models/focalnet/test_modeling_focalnet.py
@@ -23,7 +23,7 @@
from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -388,20 +388,6 @@ def test_model_from_pretrained(self):
model = FocalNetModel.from_pretrained(model_name)
self.assertIsNotNone(model)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if "embeddings" not in name and param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@require_vision
@require_torch
diff --git a/tests/models/gemma/test_modeling_gemma.py b/tests/models/gemma/test_modeling_gemma.py
index 06e4c0031f78..0557b89459b9 100644
--- a/tests/models/gemma/test_modeling_gemma.py
+++ b/tests/models/gemma/test_modeling_gemma.py
@@ -18,7 +18,7 @@
import pytest
from packaging import version
-from transformers import AutoModelForCausalLM, AutoTokenizer, GemmaConfig, is_torch_available
+from transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_available
from transformers.generation.configuration_utils import GenerationConfig
from transformers.testing_utils import (
DeviceProperties,
@@ -50,21 +50,12 @@
@require_torch
class GemmaModelTester(CausalLMModelTester):
- config_class = GemmaConfig
if is_torch_available():
base_model_class = GemmaModel
- causal_lm_class = GemmaForCausalLM
- sequence_classification_class = GemmaForSequenceClassification
- token_classification_class = GemmaForTokenClassification
@require_torch
class GemmaModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (GemmaModel, GemmaForCausalLM, GemmaForSequenceClassification, GemmaForTokenClassification)
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": GemmaModel,
diff --git a/tests/models/gemma2/test_modeling_gemma2.py b/tests/models/gemma2/test_modeling_gemma2.py
index 28ef2eeb8b57..c4479d900a89 100644
--- a/tests/models/gemma2/test_modeling_gemma2.py
+++ b/tests/models/gemma2/test_modeling_gemma2.py
@@ -20,7 +20,7 @@
from parameterized import parameterized
from pytest import mark
-from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, Gemma2Config, is_torch_available, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, is_torch_available, pipeline
from transformers.cache_utils import DynamicLayer, DynamicSlidingWindowLayer
from transformers.generation.configuration_utils import GenerationConfig
from transformers.testing_utils import (
@@ -39,7 +39,6 @@
)
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
-from ...test_configuration_common import ConfigTester
if is_torch_available():
@@ -55,31 +54,11 @@
class Gemma2ModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = Gemma2Config
base_model_class = Gemma2Model
- causal_lm_class = Gemma2ForCausalLM
- sequence_class = Gemma2ForSequenceClassification
- token_class = Gemma2ForTokenClassification
- pipeline_model_mapping = (
- {
- "feature-extraction": Gemma2Model,
- "text-classification": Gemma2ForSequenceClassification,
- "token-classification": Gemma2ForTokenClassification,
- "text-generation": Gemma2ForCausalLM,
- "zero-shot": Gemma2ForSequenceClassification,
- }
- if is_torch_available()
- else {}
- )
@require_torch
class Gemma2ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (Gemma2Model, Gemma2ForCausalLM, Gemma2ForSequenceClassification, Gemma2ForTokenClassification)
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Gemma2Model,
@@ -92,16 +71,10 @@ class Gemma2ModelTest(CausalLMModelTest, unittest.TestCase):
else {}
)
- test_headmasking = False
- test_pruning = False
_is_stateful = True
model_split_percents = [0.5, 0.6]
model_tester_class = Gemma2ModelTester
- def setUp(self):
- self.model_tester = Gemma2ModelTester(self)
- self.config_tester = ConfigTester(self, config_class=Gemma2Config, hidden_size=37)
-
@slow
@require_torch_accelerator
diff --git a/tests/models/gemma3/test_modeling_gemma3.py b/tests/models/gemma3/test_modeling_gemma3.py
index ddef6e0d6bc1..bef9cb870691 100644
--- a/tests/models/gemma3/test_modeling_gemma3.py
+++ b/tests/models/gemma3/test_modeling_gemma3.py
@@ -41,8 +41,8 @@
torch_device,
)
+from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
from ...generation.test_utils import GenerationTesterMixin
-from ...models.gemma.test_modeling_gemma import GemmaModelTester
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
@@ -62,28 +62,28 @@
from transformers.pytorch_utils import is_torch_greater_or_equal
-class Gemma3ModelTester(GemmaModelTester):
+class Gemma3TextModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = Gemma3TextConfig
- model_class = Gemma3TextModel
- for_causal_lm_class = Gemma3ForCausalLM
+ base_model_class = Gemma3TextModel
+ causal_lm_class = Gemma3ForCausalLM
+ sequence_classification_class = Gemma3TextForSequenceClassification
@require_torch
-class Gemma3ModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
- all_model_classes = (
- (Gemma3TextModel, Gemma3ForCausalLM, Gemma3TextForSequenceClassification) if is_torch_available() else ()
+class Gemma3TextModelTest(CausalLMModelTest, unittest.TestCase):
+ model_tester_class = Gemma3TextModelTester
+ pipeline_model_mapping = (
+ {
+ "feature-extraction": Gemma3TextModel,
+ "text-classification": Gemma3TextForSequenceClassification,
+ "text-generation": Gemma3ForCausalLM,
+ }
+ if is_torch_available()
+ else {}
)
- all_generative_model_classes = (Gemma3ForCausalLM,) if is_torch_available() else ()
- test_headmasking = False
- test_pruning = False
_is_stateful = True
model_split_percents = [0.5, 0.6]
- def setUp(self):
- self.model_tester = Gemma3ModelTester(self)
- self.config_tester = ConfigTester(self, config_class=Gemma3Config, hidden_size=37)
-
@unittest.skip("Gemma3 applies key/query norm which doesn't work with packing")
def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
pass
@@ -152,20 +152,10 @@ def test_generation_beyond_sliding_window_tiny_model(self):
EXPECTED_OUTPUT = torch.tensor([[90109, 90109, 90109, 83191, 83191], [246901, 69832, 69832, 69832, 62288]])
torch.testing.assert_close(generated_sequences, EXPECTED_OUTPUT)
- def test_gemma3_text_sequence_classification_model(self):
- """Test the text-only sequence classification model."""
- config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
- config.num_labels = 3
- input_ids = input_dict["input_ids"]
- attention_mask = input_ids.ne(1).to(torch_device)
- sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.num_labels)
-
- model = Gemma3TextForSequenceClassification(config)
- model.to(torch_device)
- model.eval()
-
- result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
- self.assertEqual(result.logits.shape, (self.model_tester.batch_size, config.num_labels))
+ @parameterized.expand([("linear",), ("dynamic",), ("yarn",)])
+ @unittest.skip("TODO (joao): check why this is failing")
+ def test_model_rope_scaling_from_config(self):
+ pass
class Gemma3Vision2TextModelTester:
@@ -201,7 +191,7 @@ def __init__(
self.image_token_index = image_token_index
self.boi_token_index = boi_token_index
self.eoi_token_index = eoi_token_index
- self.llm_tester = Gemma3ModelTester(self.parent)
+ self.llm_tester = Gemma3TextModelTester(self.parent)
self.text_config = self.llm_tester.get_config()
self.vision_config = vision_config
self.seq_length = seq_length
@@ -345,20 +335,10 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- @unittest.skip(
- reason="Siglip (vision backbone) uses the same initialization scheme as the Flax original implementation"
- )
- def test_initialization(self):
- pass
-
@unittest.skip("Loading nested configs with overwritten `kwargs` isn't supported yet, FIXME @raushan.")
def test_load_with_mismatched_shapes(self):
pass
- @unittest.skip("Loading nested configs with overwritten `kwargs` isn't supported yet, FIXME @raushan.")
- def test_mismatched_shapes_have_properly_initialized_weights(self):
- pass
-
def test_automodelforcausallm(self):
"""
Regression test for #36741/#36917 -- make sure `AutoModelForCausalLM` works with a Gemma3 config, i.e. that
@@ -814,8 +794,8 @@ def test_dynamic_sliding_window_is_default(self):
prompt = "What is the capital of France?"
model_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
- foward_outputs = model(**model_inputs)
- self.assertIn("DynamicSlidingWindowLayer", str(foward_outputs.past_key_values))
+ forward_outputs = model(**model_inputs)
+ self.assertIn("DynamicSlidingWindowLayer", str(forward_outputs.past_key_values))
generate_outputs = model.generate(
**model_inputs, max_new_tokens=2, do_sample=False, return_dict_in_generate=True
diff --git a/tests/models/gemma3n/test_modeling_gemma3n.py b/tests/models/gemma3n/test_modeling_gemma3n.py
index 5e4b774a8bd0..b70bc2669789 100644
--- a/tests/models/gemma3n/test_modeling_gemma3n.py
+++ b/tests/models/gemma3n/test_modeling_gemma3n.py
@@ -31,16 +31,17 @@
Gemma3nAudioConfig,
Gemma3nAudioFeatureExtractor,
Gemma3nConfig,
- Gemma3nTextConfig,
GenerationConfig,
StaticCache,
is_torch_available,
)
from transformers.testing_utils import (
+ Expectations,
cleanup,
+ require_deterministic_for_xpu,
require_read_token,
require_torch,
- require_torch_gpu,
+ require_torch_accelerator,
set_config_for_less_flaky_test,
set_model_for_less_flaky_test,
slow,
@@ -48,6 +49,7 @@
)
from transformers.utils import is_flash_attn_2_available
+from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
from ...generation.test_utils import GenerationTesterMixin, has_similar_generate_outputs
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
@@ -57,7 +59,6 @@
floats_tensor,
ids_tensor,
)
-from ..gemma.test_modeling_gemma import GemmaModelTester
if is_torch_available():
@@ -147,8 +148,6 @@ class Gemma3nAudioModelTest(ModelTesterMixin, unittest.TestCase):
is_generative = False
_is_stateful = True
main_input_name = "audio_mel"
- test_initialization = False
- test_can_init_all_missing_weights = False
def setUp(self):
self.model_tester = Gemma3nAudioModelTester(self)
@@ -218,8 +217,6 @@ def test_feature_extractor(self):
self.assertEqual(input_features.shape, self.expected_input_features_shape)
np.testing.assert_allclose(input_features[0, 0, :5], self.expected_input_features_slice, rtol=1e-5, atol=1e-5)
- print(input_features[0, 0, :5])
-
input_features_mask = audio_inputs["input_features_mask"]
self.assertEqual(input_features_mask.shape, self.expected_input_features_mask_shape)
# The second audio sample is shorter (22 frames vs 48), so its mask should become False at index 22
@@ -236,8 +233,6 @@ def test_audio_encoder(self):
with torch.no_grad():
encoder_output, encoder_mask = model(**inputs_dict)
- print(encoder_output[0, 0, :5])
-
# Check output encodings
self.assertEqual(encoder_output.shape, self.expected_encoder_output_shape)
torch.testing.assert_close(
@@ -251,9 +246,10 @@ def test_audio_encoder(self):
torch.testing.assert_close(encoder_mask[1, :], self.expected_encoder_mask_slice.to(torch_device))
-class Gemma3nTextModelTester(GemmaModelTester):
- activation_sparsity_pattern = None
- forced_config_args = ["activation_sparsity_pattern"]
+class Gemma3nTextModelTester(CausalLMModelTester):
+ if is_torch_available():
+ base_model_class = Gemma3nTextModel
+ causal_lm_class = Gemma3nForCausalLM
def __init__(
self,
@@ -292,7 +288,7 @@ def __init__(
eos_token_id=2,
is_decoder=False,
):
- self._verify_model_attributes()
+ self._verify_and_infer_model_attributes()
self.parent = parent
self.batch_size = batch_size
self.seq_length = seq_length
@@ -324,30 +320,21 @@ def __init__(
self.head_dim = self.hidden_size // self.num_attention_heads
self.is_decoder = is_decoder
- if is_torch_available():
- config_class = Gemma3nTextConfig
- model_class = Gemma3nTextModel
- for_causal_lm_class = Gemma3nForCausalLM
-
@require_torch
-class Gemma3nTextModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
- all_model_classes = (Gemma3nTextModel, Gemma3nForCausalLM) if is_torch_available() else ()
- all_generative_model_classes = (Gemma3nForCausalLM,) if is_torch_available() else ()
- test_headmasking = False
- test_pruning = False
+class Gemma3nTextModelTest(CausalLMModelTest, unittest.TestCase):
+ model_tester_class = Gemma3nTextModelTester
+ pipeline_model_mapping = (
+ {
+ "feature-extraction": Gemma3nTextModel,
+ "text-generation": Gemma3nForCausalLM,
+ }
+ if is_torch_available()
+ else {}
+ )
_is_stateful = True
model_split_percents = [0.5, 0.6]
- def setUp(self):
- self.model_tester = Gemma3nTextModelTester(self)
- self.config_tester = ConfigTester(
- self,
- config_class=Gemma3nConfig,
- hidden_size=37,
- text_config={"activation_sparsity_pattern": None},
- )
-
def _check_hidden_states_for_generate(
self, batch_size, hidden_states, prompt_length, output_length, config, use_cache=False
):
@@ -461,7 +448,7 @@ def test_generate_from_inputs_embeds_with_static_cache(self):
config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
self.skipTest(reason="This model is encoder-decoder and has Encoder-Decoder Cache")
model = model_class(config).to(torch_device).eval()
@@ -522,7 +509,7 @@ def test_generate_with_static_cache(self):
set_config_for_less_flaky_test(config)
main_input = inputs_dict[model_class.main_input_name]
- if config.get_text_config(decoder=True).is_encoder_decoder:
+ if config.is_encoder_decoder:
self.skipTest(reason="This model is encoder-decoder and has Encoder-Decoder Cache")
config.is_decoder = True
@@ -714,12 +701,6 @@ def test_training_gradient_checkpointing_use_reentrant(self):
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
- @unittest.skip(
- reason="Siglip (vision backbone) uses the same initialization scheme as the Flax original implementation"
- )
- def test_initialization(self):
- pass
-
@unittest.skip(
reason="Siglip has no FLEX attention, and we don't have a proper way to set/test attn in VLMs. TODO @raushan"
)
@@ -748,7 +729,7 @@ def test_automodelforcausallm(self):
@slow
-@require_torch_gpu
+@require_torch_accelerator
@require_read_token
class Gemma3nIntegrationTest(unittest.TestCase):
def setUp(self):
@@ -769,7 +750,7 @@ def setUp(self):
audio_ds = load_dataset(
"etechgrid/28.5k_wavfiles_dataset", "default", data_files="wav_dataset/103-1240-0000.wav"
)
- self.audio_file_path = audio_ds["train"][0]["audio"].metadata.path
+ self.audio_file_path = audio_ds["train"][0]["audio"]["path"]
cleanup(torch_device, gc_collect=True)
def tearDown(self):
@@ -790,7 +771,10 @@ def test_model_4b_bf16(self):
output = model.generate(**inputs, max_new_tokens=30, do_sample=False)
output_text = self.processor.batch_decode(output, skip_special_tokens=True)
- EXPECTED_TEXTS = ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly'] # fmt: skip
+ EXPECTED_TEXTS = Expectations({
+ ("cuda", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly'],
+ ("rocm", (9, 4)): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a turquoise ocean. The sky is blue with a few white clouds. The'],
+ }).get_expectation() # fmt: skip
self.assertEqual(output_text, EXPECTED_TEXTS)
def test_model_with_audio(self):
@@ -871,8 +855,11 @@ def test_model_4b_batch(self):
output = model.generate(**inputs, max_new_tokens=30, do_sample=False)
output_text = self.processor.batch_decode(output, skip_special_tokens=True)
-
- EXPECTED_TEXTS = ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly', "user\nYou are a helpful assistant.\n\n\n\n\n\n\n\n\n\nAre these images identical?\nmodel\nNo, the images are not identical. \n\nHere's a breakdown of the differences:\n\n* **Subject:** The first image features a cow"] # fmt: skip
+ EXPECTED_TEXTS = Expectations({
+ ("cuda", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly', "user\nYou are a helpful assistant.\n\n\n\n\n\n\n\n\n\nAre these images identical?\nmodel\nNo, the images are not identical. \n\nHere's a breakdown of the differences:\n\n* **Subject:** The first image features a cow"],
+ ("rocm", (9, 4)): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly', "user\nYou are a helpful assistant.\n\n\n\n\n\n\n\n\n\nAre these images identical?\nmodel\nNo, the images are not identical. \n\nHere's a breakdown of the differences:\n\n* **Subject Matter:** The first image shows a"],
+ ("xpu", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a turquoise ocean. The cow is facing the viewer with its head slightly turned', "user\nYou are a helpful assistant.\n\n\n\n\n\n\n\n\n\nAre these images identical?\nmodel\nNo, the images are not identical. \n\nHere's a breakdown of the differences:\n\n* **Subject:** The first image features a cow"],
+ }).get_expectation() # fmt: skip
self.assertEqual(output_text, EXPECTED_TEXTS)
def test_model_4b_image(self):
@@ -894,10 +881,15 @@ def test_model_4b_image(self):
output_text = self.processor.batch_decode(output, skip_special_tokens=True)
EXPECTED_NUM_IMAGES = 1 # Gemma3n does not support crops
- EXPECTED_TEXTS = ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly'] # fmt: skip
+ EXPECTED_TEXTS = Expectations({
+ ("cuda", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly'],
+ ("xpu", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a clear blue ocean. The cow is facing the viewer with its head slightly'],
+ ("rocm", (9, 4)): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat is shown in this image?\nmodel\nThe image shows a brown and white cow standing on a sandy beach next to a turquoise ocean. The sky is blue with a few white clouds. The'],
+ }).get_expectation() # fmt: skip
self.assertEqual(len(inputs["pixel_values"]), EXPECTED_NUM_IMAGES)
self.assertEqual(output_text, EXPECTED_TEXTS)
+ @require_deterministic_for_xpu
def test_model_4b_multiimage(self):
model_id = "Google/gemma-3n-E4B-it"
@@ -931,7 +923,11 @@ def test_model_4b_multiimage(self):
output = model.generate(**inputs, max_new_tokens=30, do_sample=False)
output_text = self.processor.batch_decode(output, skip_special_tokens=True)
- EXPECTED_TEXTS = ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat do you see here?\nmodel\nIn the image, I see a street scene in what appears to be a Chinatown district. Here are some key elements:\n\n* **A prominent red'] # fmt: skip
+ EXPECTED_TEXTS = Expectations({
+ ("cuda", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat do you see here?\nmodel\nIn the image, I see a street scene in what appears to be a Chinatown district. Here are some key elements:\n\n* **A prominent red'],
+ ("xpu", None): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat do you see here?\nmodel\nIn the image, I see a street scene in what appears to be a Chinatown district. Here are the key elements:\n\n* **A prominent red'],
+ ("rocm", (9, 4)): ['user\nYou are a helpful assistant.\n\n\n\n\n\nWhat do you see here?\nmodel\nIn the image, I see a street scene in what appears to be a Chinatown district. \n\nHere are some key elements:\n\n* **A'],
+ }).get_expectation() # fmt: skip
self.assertEqual(output_text, EXPECTED_TEXTS)
@unittest.skip("For now, using a gemma model with the 3n class is not supported")
@@ -981,6 +977,7 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):
EXPECTED_COMPLETIONS = [" and I think it's a nice place to visit. This is a nice place. This is", ", green, yellow, orange, purple, pink, brown, black, white.\n\nHere'"] # fmt: skip
self.assertEqual(output_text, EXPECTED_COMPLETIONS)
+ @require_deterministic_for_xpu
def test_generation_beyond_sliding_window_with_generation_config(self):
"""Same as `test_generation_beyond_sliding_window`, but passing a GenerationConfig. Regression test for #36684 --
ensures `cache_implementation='hybrid'` is correctly inherited from the base `model.generation_config`.
@@ -1006,5 +1003,10 @@ def test_generation_beyond_sliding_window_with_generation_config(self):
]
output_text = tokenizer.batch_decode(out)
- EXPECTED_COMPLETIONS = [" and I am glad to be here. This is a nice place. This is a nice place.", ", green, yellow, purple, orange, pink, brown, black, white.\n\nHere are"] # fmt: skip
+ EXPECTED_COMPLETIONS = Expectations({
+ # FIXME: This test is VERY flaky on ROCm
+ ("cuda", None): [" and I am glad to be here. This is a nice place. This is a nice place.", ", green, yellow, purple, orange, pink, brown, black, white.\n\nHere are"],
+ ("rocm", (9, 4)): [' and I think it makes this place special. This is a nice place. This is a nice place', ', green, yellow, purple, orange, pink, brown, black, white.\n\nHere are'],
+ ("xpu", None): [" and I think it is very nice. I think it is nice. This is a nice place.", ", green, yellow, purple, orange, pink, brown, black, white.\n\nHere are"],
+ }).get_expectation() # fmt: skip
self.assertEqual(output_text, EXPECTED_COMPLETIONS)
diff --git a/tests/models/glm/test_modeling_glm.py b/tests/models/glm/test_modeling_glm.py
index 96969438ee48..08609cd90ca6 100644
--- a/tests/models/glm/test_modeling_glm.py
+++ b/tests/models/glm/test_modeling_glm.py
@@ -17,7 +17,7 @@
import pytest
-from transformers import AutoModelForCausalLM, AutoTokenizer, GlmConfig, is_torch_available
+from transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_available
from transformers.testing_utils import (
Expectations,
require_flash_attn,
@@ -43,21 +43,12 @@
@require_torch
class GlmModelTester(CausalLMModelTester):
- config_class = GlmConfig
if is_torch_available():
base_model_class = GlmModel
- causal_lm_class = GlmForCausalLM
- sequence_class = GlmForSequenceClassification
- token_class = GlmForTokenClassification
@require_torch
class GlmModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (GlmModel, GlmForCausalLM, GlmForSequenceClassification, GlmForTokenClassification)
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": GlmModel,
@@ -69,8 +60,6 @@ class GlmModelTest(CausalLMModelTest, unittest.TestCase):
else {}
)
- test_headmasking = False
- test_pruning = False
model_tester_class = GlmModelTester
diff --git a/tests/models/glm4/test_modeling_glm4.py b/tests/models/glm4/test_modeling_glm4.py
index d04711c92242..b810bf6a6066 100644
--- a/tests/models/glm4/test_modeling_glm4.py
+++ b/tests/models/glm4/test_modeling_glm4.py
@@ -18,7 +18,7 @@
import pytest
-from transformers import AutoModelForCausalLM, AutoTokenizer, Glm4Config, is_torch_available
+from transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_available
from transformers.testing_utils import (
Expectations,
cleanup,
@@ -46,21 +46,12 @@
class Glm4ModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = Glm4Config
base_model_class = Glm4Model
- causal_lm_class = Glm4ForCausalLM
- sequence_classification_class = Glm4ForSequenceClassification
- token_classification_class = Glm4ForTokenClassification
@require_torch
class Glm4ModelTest(CausalLMModelTest, unittest.TestCase):
model_tester_class = Glm4ModelTester
- all_model_classes = (
- (Glm4Model, Glm4ForCausalLM, Glm4ForSequenceClassification, Glm4ForTokenClassification)
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Glm4Model,
@@ -72,8 +63,6 @@ class Glm4ModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
_is_stateful = True
model_split_percents = [0.5, 0.6]
diff --git a/tests/models/glm4_moe/test_modeling_glm4_moe.py b/tests/models/glm4_moe/test_modeling_glm4_moe.py
index fbf79524c618..5ddf7a90ed0a 100644
--- a/tests/models/glm4_moe/test_modeling_glm4_moe.py
+++ b/tests/models/glm4_moe/test_modeling_glm4_moe.py
@@ -33,14 +33,12 @@
if is_torch_available():
- from transformers import AutoTokenizer, Glm4MoeConfig, Glm4MoeForCausalLM, Glm4MoeModel
+ from transformers import AutoTokenizer, Glm4MoeForCausalLM, Glm4MoeModel
class Glm4MoeModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = Glm4MoeConfig
base_model_class = Glm4MoeModel
- causal_lm_class = Glm4MoeForCausalLM
def __init__(
self,
@@ -60,14 +58,6 @@ def __init__(
@require_torch
class Glm4MoeModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- Glm4MoeModel,
- Glm4MoeForCausalLM,
- )
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": Glm4MoeModel,
@@ -76,8 +66,6 @@ class Glm4MoeModelTest(CausalLMModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
fx_compatible = False
model_tester_class = Glm4MoeModelTester
# used in `test_torch_compile_for_training`. Skip as "Dynamic control flow in MoE"
diff --git a/tests/models/glm4v_moe/test_modeling_glm4v_moe.py b/tests/models/glm4v_moe/test_modeling_glm4v_moe.py
index 995b3c0723db..1881fffa9dd9 100644
--- a/tests/models/glm4v_moe/test_modeling_glm4v_moe.py
+++ b/tests/models/glm4v_moe/test_modeling_glm4v_moe.py
@@ -297,6 +297,7 @@ def test_inputs_embeds_matches_input_ids(self):
@require_torch
+@slow
class Glm4vMoeIntegrationTest(unittest.TestCase):
model = None
@@ -310,7 +311,8 @@ def get_model(cls):
@classmethod
def tearDownClass(cls):
- del cls.model
+ if hasattr(cls, "model"):
+ del cls.model
cleanup(torch_device, gc_collect=True)
def setUp(self):
@@ -364,7 +366,6 @@ def setUp(self):
def tearDown(self):
cleanup(torch_device, gc_collect=True)
- @slow
def test_small_model_integration_test(self):
inputs = self.processor.apply_chat_template(
self.message, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
@@ -386,7 +387,6 @@ def test_small_model_integration_test(self):
)
torch.testing.assert_close(expected_pixel_slice, inputs.pixel_values[:6, :3], atol=1e-4, rtol=1e-4)
- @slow
def test_small_model_integration_test_batch(self):
model = self.get_model()
batch_messages = [self.message, self.message2, self.message_wo_image]
@@ -414,7 +414,6 @@ def test_small_model_integration_test_batch(self):
EXPECTED_DECODED_TEXT,
)
- @slow
def test_small_model_integration_test_with_video(self):
processor = AutoProcessor.from_pretrained("zai-org/GLM-4.5V", max_image_size={"longest_edge": 50176})
model = self.get_model()
@@ -437,7 +436,6 @@ def test_small_model_integration_test_with_video(self):
)
@run_first
- @slow
@require_flash_attn
@require_torch_gpu
def test_small_model_integration_test_batch_flashatt2(self):
diff --git a/tests/models/got_ocr2/test_modeling_got_ocr2.py b/tests/models/got_ocr2/test_modeling_got_ocr2.py
index 59577106b069..be7f447c7918 100644
--- a/tests/models/got_ocr2/test_modeling_got_ocr2.py
+++ b/tests/models/got_ocr2/test_modeling_got_ocr2.py
@@ -25,7 +25,7 @@
from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -59,7 +59,7 @@ def __init__(
"vocab_size": 99,
"hidden_size": 128,
"intermediate_size": 37,
- "num_hidden_layers": 4,
+ "num_hidden_layers": 2,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"output_channels": 64,
@@ -163,20 +163,6 @@ def setUp(self):
def test_config(self):
self.config_tester.run_common_tests()
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@require_torch
class GotOcr2IntegrationTest(unittest.TestCase):
diff --git a/tests/models/gpt2/test_modeling_gpt2.py b/tests/models/gpt2/test_modeling_gpt2.py
index ae37e2432ddb..4065e7179f5b 100644
--- a/tests/models/gpt2/test_modeling_gpt2.py
+++ b/tests/models/gpt2/test_modeling_gpt2.py
@@ -12,12 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import math
import unittest
import pytest
-from transformers import DynamicCache, GPT2Config, is_torch_available
+from transformers import is_torch_available
from transformers.testing_utils import (
Expectations,
cleanup,
@@ -28,10 +27,8 @@
torch_device,
)
-from ...generation.test_utils import GenerationTesterMixin
-from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
-from ...test_pipeline_mixin import PipelineTesterMixin
+from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
+from ...test_modeling_common import floats_tensor, ids_tensor
if is_torch_available():
@@ -48,149 +45,88 @@
)
-class GPT2ModelTester:
+class GPT2ModelTester(CausalLMModelTester):
+ if is_torch_available():
+ base_model_class = GPT2Model
+ causal_lm_class = GPT2LMHeadModel
+
def __init__(
self,
parent,
- batch_size=14,
- seq_length=7,
- is_training=True,
use_token_type_ids=True,
- use_input_mask=True,
- use_labels=True,
- use_mc_token_ids=True,
- vocab_size=99,
- hidden_size=32,
- num_hidden_layers=2,
- num_attention_heads=4,
- intermediate_size=37,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=16,
- type_sequence_label_size=2,
- initializer_range=0.02,
- num_labels=3,
num_choices=4,
- scope=None,
+ **kwargs,
):
- self.parent = parent
- self.batch_size = batch_size
- self.seq_length = seq_length
- self.is_training = is_training
- self.use_token_type_ids = use_token_type_ids
- self.use_input_mask = use_input_mask
- self.use_labels = use_labels
- self.use_mc_token_ids = use_mc_token_ids
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.type_sequence_label_size = type_sequence_label_size
- self.initializer_range = initializer_range
- self.num_labels = num_labels
+ super().__init__(parent, use_token_type_ids=use_token_type_ids, **kwargs)
self.num_choices = num_choices
- self.scope = None
- self.bos_token_id = vocab_size - 1
- self.eos_token_id = vocab_size - 1
- self.pad_token_id = vocab_size - 1
-
- def get_large_model_config(self):
- return GPT2Config.from_pretrained("openai-community/gpt2")
def prepare_config_and_inputs(
- self, gradient_checkpointing=False, scale_attn_by_inverse_layer_idx=False, reorder_and_upcast_attn=False
+ self, extra_inputs=False, scale_attn_by_inverse_layer_idx=False, reorder_and_upcast_attn=False
):
- input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
-
- input_mask = None
- if self.use_input_mask:
- input_mask = random_attention_mask([self.batch_size, self.seq_length])
-
- token_type_ids = None
- if self.use_token_type_ids:
- token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
+ # Overwritten: `GPT2DoubleHeadsModel` uses extra inputs
+ (config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels) = (
+ super().prepare_config_and_inputs()
+ )
- mc_token_ids = None
- if self.use_mc_token_ids:
+ if extra_inputs:
mc_token_ids = ids_tensor([self.batch_size, self.num_choices], self.seq_length)
-
- sequence_labels = None
- token_labels = None
- choice_labels = None
- if self.use_labels:
- sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
- token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
- choice_labels = ids_tensor([self.batch_size], self.num_choices)
+ head_mask = ids_tensor([self.num_hidden_layers, self.num_attention_heads], 2)
+ config_and_inputs = (
+ config,
+ input_ids,
+ input_mask,
+ head_mask,
+ token_type_ids,
+ mc_token_ids,
+ sequence_labels,
+ token_labels,
+ choice_labels,
+ )
+ else:
+ config_and_inputs = (
+ config,
+ input_ids,
+ token_type_ids,
+ input_mask,
+ sequence_labels,
+ token_labels,
+ choice_labels,
+ )
config = self.get_config(
- gradient_checkpointing=gradient_checkpointing,
scale_attn_by_inverse_layer_idx=scale_attn_by_inverse_layer_idx,
reorder_and_upcast_attn=reorder_and_upcast_attn,
)
- head_mask = ids_tensor([self.num_hidden_layers, self.num_attention_heads], 2)
+ return config_and_inputs
- return (
- config,
- input_ids,
- input_mask,
- head_mask,
- token_type_ids,
- mc_token_ids,
- sequence_labels,
- token_labels,
- choice_labels,
- )
-
- def get_config(
- self, gradient_checkpointing=False, scale_attn_by_inverse_layer_idx=False, reorder_and_upcast_attn=False
- ):
- return GPT2Config(
- vocab_size=self.vocab_size,
- n_embd=self.hidden_size,
- n_layer=self.num_hidden_layers,
- n_head=self.num_attention_heads,
- n_inner=self.intermediate_size,
- activation_function=self.hidden_act,
- resid_pdrop=self.hidden_dropout_prob,
- attn_pdrop=self.attention_probs_dropout_prob,
- n_positions=self.max_position_embeddings,
- type_vocab_size=self.type_vocab_size,
- initializer_range=self.initializer_range,
- use_cache=True,
- bos_token_id=self.bos_token_id,
- eos_token_id=self.eos_token_id,
- pad_token_id=self.pad_token_id,
- gradient_checkpointing=gradient_checkpointing,
- scale_attn_by_inverse_layer_idx=scale_attn_by_inverse_layer_idx,
- reorder_and_upcast_attn=reorder_and_upcast_attn,
- )
-
- def get_pipeline_config(self):
- config = self.get_config()
- config.vocab_size = 300
+ def get_config(self, scale_attn_by_inverse_layer_idx=False, reorder_and_upcast_attn=False):
+ # Overwritten: `GPT2Config` has extra flags and we want to test them
+ config = super().get_config()
+ config.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
+ config.reorder_and_upcast_attn = reorder_and_upcast_attn
return config
+ def prepare_config_and_inputs_for_common(self):
+ # Overwritten: we want `token_type_ids` as part of the common inputs
+ config_and_inputs = self.prepare_config_and_inputs(extra_inputs=True)
+ config, input_ids, _, head_mask, token_type_ids, _, _, _, _ = config_and_inputs
+ inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "head_mask": head_mask}
+ return config, inputs_dict
+
def prepare_config_and_inputs_for_decoder(self):
+ # Extra function: used in `encoder_decoder` tests
(
config,
input_ids,
input_mask,
head_mask,
token_type_ids,
- mc_token_ids,
+ _,
sequence_labels,
token_labels,
choice_labels,
- ) = self.prepare_config_and_inputs()
+ ) = self.prepare_config_and_inputs(extra_inputs=True)
encoder_hidden_states = floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
encoder_attention_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
@@ -208,283 +144,10 @@ def prepare_config_and_inputs_for_decoder(self):
encoder_attention_mask,
)
- def create_and_check_gpt2_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
- model = GPT2Model(config=config)
- model.to(torch_device)
- model.eval()
-
- result = model(input_ids, token_type_ids=token_type_ids, head_mask=head_mask)
- result = model(input_ids, token_type_ids=token_type_ids)
- result = model(input_ids)
-
- self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
- self.parent.assertEqual(len(result.past_key_values), config.n_layer)
-
- def create_and_check_gpt2_model_past(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
- model = GPT2Model(config=config)
- model.to(torch_device)
- model.eval()
-
- # first forward pass
- outputs = model(input_ids, token_type_ids=token_type_ids, use_cache=True)
- outputs_use_cache_conf = model(input_ids, token_type_ids=token_type_ids)
- outputs_no_past = model(input_ids, token_type_ids=token_type_ids, use_cache=False)
-
- self.parent.assertTrue(len(outputs) == len(outputs_use_cache_conf))
- self.parent.assertTrue(len(outputs) == len(outputs_no_past) + 1)
-
- output, past = outputs.to_tuple()
-
- # create hypothetical next token and extent to next_input_ids
- next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)
- next_token_types = ids_tensor([self.batch_size, 1], self.type_vocab_size)
-
- # append to next input_ids and token_type_ids
- next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
- next_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)
-
- output_from_no_past = model(next_input_ids, token_type_ids=next_token_type_ids)["last_hidden_state"]
- output_from_past = model(next_tokens, token_type_ids=next_token_types, past_key_values=past)[
- "last_hidden_state"
- ]
-
- # select random slice
- random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
- output_from_no_past_slice = output_from_no_past[:, -1, random_slice_idx].detach()
- output_from_past_slice = output_from_past[:, 0, random_slice_idx].detach()
-
- # test that outputs are equal for slice
- self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
-
- def create_and_check_gpt2_model_attention_mask_past(
- self, config, input_ids, input_mask, head_mask, token_type_ids, *args
- ):
- model = GPT2Model(config=config)
- model.to(torch_device)
- model.eval()
-
- # create attention mask
- attn_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)
- half_seq_length = self.seq_length // 2
- attn_mask[:, half_seq_length:] = 0
-
- # first forward pass
- output, past = model(input_ids, attention_mask=attn_mask).to_tuple()
-
- # create hypothetical next token and extent to next_input_ids
- next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)
-
- # change a random masked slice from input_ids
- random_seq_idx_to_change = ids_tensor((1,), half_seq_length).item() + 1
- random_other_next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size).squeeze(-1)
- input_ids[:, -random_seq_idx_to_change] = random_other_next_tokens
-
- # append to next input_ids and attn_mask
- next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
- attn_mask = torch.cat(
- [attn_mask, torch.ones((attn_mask.shape[0], 1), dtype=torch.long, device=torch_device)],
- dim=1,
- )
-
- # get two different outputs
- output_from_no_past = model(next_input_ids, attention_mask=attn_mask)["last_hidden_state"]
- output_from_past = model(next_tokens, past_key_values=past, attention_mask=attn_mask)["last_hidden_state"]
-
- # select random slice
- random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
- output_from_no_past_slice = output_from_no_past[:, -1, random_slice_idx].detach()
- output_from_past_slice = output_from_past[:, 0, random_slice_idx].detach()
-
- # test that outputs are equal for slice
- self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
-
- def create_and_check_gpt2_model_past_large_inputs(
- self, config, input_ids, input_mask, head_mask, token_type_ids, *args
- ):
- model = GPT2Model(config=config)
- model.to(torch_device)
- model.eval()
-
- # first forward pass
- outputs = model(input_ids, token_type_ids=token_type_ids, attention_mask=input_mask, use_cache=True)
-
- output, past = outputs.to_tuple()
-
- # create hypothetical next token and extent to next_input_ids
- next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
- next_token_types = ids_tensor([self.batch_size, 3], self.type_vocab_size)
- next_mask = ids_tensor((self.batch_size, 3), vocab_size=2)
-
- # append to next input_ids and token_type_ids
- next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
- next_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)
- next_attention_mask = torch.cat([input_mask, next_mask], dim=-1)
-
- output_from_no_past = model(
- next_input_ids, token_type_ids=next_token_type_ids, attention_mask=next_attention_mask
- )["last_hidden_state"]
- output_from_past = model(
- next_tokens, token_type_ids=next_token_types, attention_mask=next_attention_mask, past_key_values=past
- )["last_hidden_state"]
- self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1])
-
- # select random slice
- random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
- output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
- output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()
-
- # test that outputs are equal for slice
- self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
-
- def create_and_check_lm_head_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
- model = GPT2LMHeadModel(config)
- model.to(torch_device)
- model.eval()
-
- result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
- self.parent.assertEqual(result.loss.shape, ())
- self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
-
- def create_and_check_forward_and_backwards(
- self, config, input_ids, input_mask, head_mask, token_type_ids, *args, gradient_checkpointing=False
- ):
- model = GPT2LMHeadModel(config)
- model.to(torch_device)
- if gradient_checkpointing:
- model.gradient_checkpointing_enable()
-
- result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
- self.parent.assertEqual(result.loss.shape, ())
- self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
- result.loss.backward()
-
- def create_and_check_double_lm_head_model(
- self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, *args
- ):
- model = GPT2DoubleHeadsModel(config)
- model.to(torch_device)
- model.eval()
-
- multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
- multiple_choice_input_mask = input_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
- multiple_choice_token_type_ids = token_type_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
-
- inputs = {
- "input_ids": multiple_choice_inputs_ids,
- "mc_token_ids": mc_token_ids,
- "attention_mask": multiple_choice_input_mask,
- "token_type_ids": multiple_choice_token_type_ids,
- "labels": multiple_choice_inputs_ids,
- }
-
- result = model(**inputs)
- self.parent.assertEqual(result.loss.shape, ())
- self.parent.assertEqual(
- result.logits.shape, (self.batch_size, self.num_choices, self.seq_length, self.vocab_size)
- )
- self.parent.assertEqual(result.mc_logits.shape, (self.batch_size, self.num_choices))
-
- def create_and_check_gpt2_for_question_answering(
- self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
- ):
- config.num_labels = self.num_labels
- model = GPT2ForQuestionAnswering(config)
- model.to(torch_device)
- model.eval()
- result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
- self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))
- self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))
-
- def create_and_check_gpt2_for_sequence_classification(
- self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
- ):
- config.num_labels = self.num_labels
- model = GPT2ForSequenceClassification(config)
- model.to(torch_device)
- model.eval()
- result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels)
- self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))
-
- def create_and_check_gpt2_for_token_classification(
- self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
- ):
- config.num_labels = self.num_labels
- model = GPT2ForTokenClassification(config)
- model.to(torch_device)
- model.eval()
- result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
- self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
-
- def create_and_check_gpt2_weight_initialization(self, config, *args):
- model = GPT2Model(config)
- model_std = model.config.initializer_range / math.sqrt(2 * model.config.n_layer)
- for key in model.state_dict():
- if "c_proj" in key and "weight" in key:
- self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]) - model_std), 0.001)
- self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]) - 0.0), 0.01)
-
- def create_and_check_cached_forward_with_and_without_attention_mask(self, config, input_ids, *args):
- # Relevant issue: https://github.com/huggingface/transformers/issues/31943
- model = GPT2Model(config)
- model.to(torch_device)
- model.eval()
-
- # We want this for SDPA, eager works with a `None` attention mask
- assert model.config._attn_implementation == "sdpa", (
- "This test assumes the model to have the SDPA implementation for its attention calculations."
- )
-
- # Prepare cache and non_cache input, needs a full attention mask
- cached_len = input_ids.shape[-1] // 2
- input_mask = torch.ones(size=input_ids.size()).to(torch_device)
- cache_inputs = {"input_ids": input_ids[:, :cached_len], "attention_mask": input_mask[:, :cached_len]}
- non_cache_inputs = {"input_ids": input_ids[:, cached_len:], "attention_mask": input_mask}
-
- # Cached forward once with the attention mask provided and the other time without it (which should assume full attention)
- cache_outputs = model(**cache_inputs)
- # Caches are mutable (unlike legacy tuples), so we need to copy them before using multiple times
- pkv_copy = DynamicCache(config=config)
- pkv_copy.update(
- cache_outputs.past_key_values.layers[0].keys, cache_outputs.past_key_values.layers[0].values, 0
- )
- pkv_copy.update(
- cache_outputs.past_key_values.layers[1].keys, cache_outputs.past_key_values.layers[1].values, 1
- )
- full_outputs_with_attention_mask = model(**non_cache_inputs, past_key_values=pkv_copy).last_hidden_state
- full_outputs_without_attention_mask = model(
- non_cache_inputs["input_ids"], past_key_values=cache_outputs.past_key_values
- ).last_hidden_state
-
- self.parent.assertTrue(
- torch.allclose(full_outputs_with_attention_mask, full_outputs_without_attention_mask, atol=1e-5)
- )
-
- def prepare_config_and_inputs_for_common(self):
- config_and_inputs = self.prepare_config_and_inputs()
-
- (
- config,
- input_ids,
- input_mask,
- head_mask,
- token_type_ids,
- mc_token_ids,
- sequence_labels,
- token_labels,
- choice_labels,
- ) = config_and_inputs
-
- inputs_dict = {
- "input_ids": input_ids,
- "token_type_ids": token_type_ids,
- "head_mask": head_mask,
- }
-
- return config, inputs_dict
-
@require_torch
-class GPT2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
+class GPT2ModelTest(CausalLMModelTest, unittest.TestCase):
+ # `all_model_classes` is overwritten because of `GPT2DoubleHeadsModel`
all_model_classes = (
(
GPT2Model,
@@ -513,9 +176,10 @@ class GPT2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
fx_compatible = False # Broken by attention refactor cc @Cyrilvallez
test_missing_keys = False
test_model_parallel = True
+ model_tester_class = GPT2ModelTester
- # special case for DoubleHeads model
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
+ # Overwritten: special case for DoubleHeads model
inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
if return_labels:
@@ -537,220 +201,91 @@ def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
)
return inputs_dict
- def setUp(self):
- self.model_tester = GPT2ModelTester(self)
- self.config_tester = ConfigTester(self, config_class=GPT2Config, n_embd=37)
-
- def tearDown(self):
- super().tearDown()
- # clean-up as much as possible GPU memory occupied by PyTorch
- cleanup(torch_device)
-
- def test_config(self):
- self.config_tester.run_common_tests()
-
- def test_gpt2_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_model(*config_and_inputs)
-
- def test_gpt2_model_past(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_model_past(*config_and_inputs)
-
- def test_gpt2_model_att_mask_past(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_model_attention_mask_past(*config_and_inputs)
-
- def test_gpt2_model_past_large_inputs(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_model_past_large_inputs(*config_and_inputs)
-
- def test_gpt2_lm_head_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_lm_head_model(*config_and_inputs)
-
def test_gpt2_double_lm_head_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_double_lm_head_model(*config_and_inputs)
-
- def test_gpt2_question_answering_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_for_question_answering(*config_and_inputs)
-
- def test_gpt2_sequence_classification_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_for_sequence_classification(*config_and_inputs)
-
- def test_gpt2_token_classification_model(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_for_token_classification(*config_and_inputs)
-
- def test_gpt2_gradient_checkpointing(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs, gradient_checkpointing=True)
-
- def test_gpt2_scale_attn_by_inverse_layer_idx(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs(scale_attn_by_inverse_layer_idx=True)
- self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs)
-
- def test_gpt2_reorder_and_upcast_attn(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs(reorder_and_upcast_attn=True)
- self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs)
-
- def test_gpt2_weight_initialization(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_gpt2_weight_initialization(*config_and_inputs)
-
- def test_cached_forward_with_and_without_attention_mask(self):
- config_and_inputs = self.model_tester.prepare_config_and_inputs()
- self.model_tester.create_and_check_cached_forward_with_and_without_attention_mask(*config_and_inputs)
-
- @unittest.skip(
- reason="This architecture seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
- )
- def test_training_gradient_checkpointing(self):
- pass
-
- @unittest.skip(
- reason="This architecture seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
- )
- def test_training_gradient_checkpointing_use_reentrant(self):
- pass
-
- @unittest.skip(
- reason="This architecture seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
- )
- def test_training_gradient_checkpointing_use_reentrant_false(self):
- pass
-
- @slow
- def test_batch_generation(self):
- model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
+ # extra test: model-specific class
+ config_and_inputs = self.model_tester.prepare_config_and_inputs(extra_inputs=True)
+ config, input_ids, input_mask, _, token_type_ids, mc_token_ids, _, _, _ = config_and_inputs
+ model = GPT2DoubleHeadsModel(config)
model.to(torch_device)
- tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
-
- tokenizer.padding_side = "left"
-
- # Define PAD Token = EOS Token = 50256
- tokenizer.pad_token = tokenizer.eos_token
- model.config.pad_token_id = model.config.eos_token_id
-
- # use different length sentences to test batching
- sentences = [
- "Hello, my dog is a little",
- "Today, I",
- ]
+ model.eval()
- inputs = tokenizer(sentences, return_tensors="pt", padding=True)
- input_ids = inputs["input_ids"].to(torch_device)
- token_type_ids = torch.cat(
- [
- input_ids.new_full((input_ids.shape[0], input_ids.shape[1] - 1), 0),
- input_ids.new_full((input_ids.shape[0], 1), 500),
- ],
- dim=-1,
+ multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1).contiguous()
+ multiple_choice_input_mask = input_mask.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1).contiguous()
+ multiple_choice_token_type_ids = (
+ token_type_ids.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1).contiguous()
)
- outputs = model.generate(
- input_ids=input_ids,
- attention_mask=inputs["attention_mask"].to(torch_device),
- max_length=20,
- )
+ inputs = {
+ "input_ids": multiple_choice_inputs_ids,
+ "mc_token_ids": mc_token_ids,
+ "attention_mask": multiple_choice_input_mask,
+ "token_type_ids": multiple_choice_token_type_ids,
+ "labels": multiple_choice_inputs_ids,
+ }
- outputs_tt = model.generate(
- input_ids=input_ids,
- attention_mask=inputs["attention_mask"].to(torch_device),
- token_type_ids=token_type_ids,
- max_length=20,
+ result = model(**inputs)
+ self.assertEqual(result.loss.shape, ())
+ self.assertEqual(
+ result.logits.shape,
+ (
+ self.model_tester.batch_size,
+ self.model_tester.num_choices,
+ self.model_tester.seq_length,
+ self.model_tester.vocab_size,
+ ),
)
+ self.assertEqual(result.mc_logits.shape, (self.model_tester.batch_size, self.model_tester.num_choices))
- inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
- output_non_padded = model.generate(input_ids=inputs_non_padded, max_length=20)
-
- num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().item()
- inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
- output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
-
- batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
- batch_out_sentence_tt = tokenizer.batch_decode(outputs_tt, skip_special_tokens=True)
- non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
- padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
-
- expected_output_sentence = [
- "Hello, my dog is a little bit of a mess. I'm not sure if he's going",
- "Today, I'm going to be doing a lot of research on this. I",
- ]
- self.assertListEqual(expected_output_sentence, batch_out_sentence)
- self.assertTrue(batch_out_sentence_tt != batch_out_sentence) # token_type_ids should change output
- self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])
+ def test_gpt2_scale_attn_by_inverse_layer_idx(self):
+ # extra test: model-specific flag
+ config_and_inputs = self.model_tester.prepare_config_and_inputs(scale_attn_by_inverse_layer_idx=True)
+ config, input_ids, token_type_ids, _, _, _, _ = config_and_inputs
- @slow
- def test_batch_generation_2heads(self):
- model = GPT2DoubleHeadsModel.from_pretrained("openai-community/gpt2")
+ model = GPT2LMHeadModel(config)
model.to(torch_device)
- tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
-
- tokenizer.padding_side = "left"
-
- # This tokenizer has no pad token, so we have to set it in some way
- # Define PAD Token = EOS Token = 50256
- tokenizer.pad_token = tokenizer.eos_token
- model.config.pad_token_id = model.config.eos_token_id
-
- # use different length sentences to test batching
- sentences = [
- "Hello, my dog is a little",
- "Today, I",
- ]
-
- inputs = tokenizer(sentences, return_tensors="pt", padding=True)
- input_ids = inputs["input_ids"].to(torch_device)
- token_type_ids = torch.cat(
- [
- input_ids.new_full((input_ids.shape[0], input_ids.shape[1] - 1), 0),
- input_ids.new_full((input_ids.shape[0], 1), 500),
- ],
- dim=-1,
+ result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
+ self.assertEqual(result.loss.shape, ())
+ self.assertEqual(
+ result.logits.shape,
+ (self.model_tester.batch_size, self.model_tester.seq_length, self.model_tester.vocab_size),
)
+ result.loss.backward()
- outputs = model.generate(
- input_ids=input_ids,
- attention_mask=inputs["attention_mask"].to(torch_device),
- max_length=20,
- )
+ def test_gpt2_reorder_and_upcast_attn(self):
+ # extra test: model-specific flag
+ config_and_inputs = self.model_tester.prepare_config_and_inputs(reorder_and_upcast_attn=True)
+ config, input_ids, token_type_ids, _, _, _, _ = config_and_inputs
- outputs_tt = model.generate(
- input_ids=input_ids,
- attention_mask=inputs["attention_mask"].to(torch_device),
- token_type_ids=token_type_ids,
- max_length=20,
+ model = GPT2LMHeadModel(config)
+ model.to(torch_device)
+ result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
+ self.assertEqual(result.loss.shape, ())
+ self.assertEqual(
+ result.logits.shape,
+ (self.model_tester.batch_size, self.model_tester.seq_length, self.model_tester.vocab_size),
)
+ result.loss.backward()
- inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
- output_non_padded = model.generate(input_ids=inputs_non_padded, max_length=20)
-
- num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().item()
- inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
- output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
-
- batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
- batch_out_sentence_tt = tokenizer.batch_decode(outputs_tt, skip_special_tokens=True)
- non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
- padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
+ def test_training_gradient_checkpointing(self):
+ # overwritten: GPT2DoubleHeadsModel fails this test, non-standard class
+ self.original_all_model_classes = self.all_model_classes
+ self.all_model_classes = (cls for cls in self.all_model_classes if cls.__name__ != "GPT2DoubleHeadsModel")
+ super().test_training_gradient_checkpointing()
+ self.all_model_classes = self.original_all_model_classes
- expected_output_sentence = [
- "Hello, my dog is a little bit of a mess. I'm not sure if he's going",
- "Today, I'm going to be doing a lot of research on this. I",
- ]
- self.assertListEqual(expected_output_sentence, batch_out_sentence)
- self.assertTrue(batch_out_sentence_tt != batch_out_sentence) # token_type_ids should change output
- self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])
+ def test_training_gradient_checkpointing_use_reentrant(self):
+ # overwritten: GPT2DoubleHeadsModel fails this test, non-standard class
+ self.original_all_model_classes = self.all_model_classes
+ self.all_model_classes = (cls for cls in self.all_model_classes if cls.__name__ != "GPT2DoubleHeadsModel")
+ super().test_training_gradient_checkpointing_use_reentrant()
+ self.all_model_classes = self.original_all_model_classes
- @slow
- def test_model_from_pretrained(self):
- model_name = "openai-community/gpt2"
- model = GPT2Model.from_pretrained(model_name)
- self.assertIsNotNone(model)
+ def test_training_gradient_checkpointing_use_reentrant_false(self):
+ # overwritten: GPT2DoubleHeadsModel fails this test, non-standard class
+ self.original_all_model_classes = self.all_model_classes
+ self.all_model_classes = (cls for cls in self.all_model_classes if cls.__name__ != "GPT2DoubleHeadsModel")
+ super().test_training_gradient_checkpointing_use_reentrant_false()
+ self.all_model_classes = self.original_all_model_classes
@require_torch
@@ -915,3 +450,126 @@ def test_flash_attn_2_generate_padding_left(self):
self.assertListEqual(output_native, output_fa_2)
self.assertListEqual(output_native, expected_output)
+
+ @slow
+ def test_batch_generation(self):
+ model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
+ model.to(torch_device)
+ tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
+
+ tokenizer.padding_side = "left"
+
+ # Define PAD Token = EOS Token = 50256
+ tokenizer.pad_token = tokenizer.eos_token
+ model.config.pad_token_id = model.config.eos_token_id
+
+ # use different length sentences to test batching
+ sentences = [
+ "Hello, my dog is a little",
+ "Today, I",
+ ]
+
+ inputs = tokenizer(sentences, return_tensors="pt", padding=True)
+ input_ids = inputs["input_ids"].to(torch_device)
+ token_type_ids = torch.cat(
+ [
+ input_ids.new_full((input_ids.shape[0], input_ids.shape[1] - 1), 0),
+ input_ids.new_full((input_ids.shape[0], 1), 500),
+ ],
+ dim=-1,
+ )
+
+ outputs = model.generate(
+ input_ids=input_ids,
+ attention_mask=inputs["attention_mask"].to(torch_device),
+ max_length=20,
+ )
+
+ outputs_tt = model.generate(
+ input_ids=input_ids,
+ attention_mask=inputs["attention_mask"].to(torch_device),
+ token_type_ids=token_type_ids,
+ max_length=20,
+ )
+
+ inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
+ output_non_padded = model.generate(input_ids=inputs_non_padded, max_length=20)
+
+ num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().item()
+ inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
+ output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
+
+ batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+ batch_out_sentence_tt = tokenizer.batch_decode(outputs_tt, skip_special_tokens=True)
+ non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
+ padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
+
+ expected_output_sentence = [
+ "Hello, my dog is a little bit of a mess. I'm not sure if he's going",
+ "Today, I'm going to be doing a lot of research on this. I",
+ ]
+ self.assertListEqual(expected_output_sentence, batch_out_sentence)
+ self.assertTrue(batch_out_sentence_tt != batch_out_sentence) # token_type_ids should change output
+ self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])
+
+ @slow
+ def test_batch_generation_2heads(self):
+ model = GPT2DoubleHeadsModel.from_pretrained("openai-community/gpt2")
+ model.to(torch_device)
+ tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
+
+ tokenizer.padding_side = "left"
+
+ # This tokenizer has no pad token, so we have to set it in some way
+ # Define PAD Token = EOS Token = 50256
+ tokenizer.pad_token = tokenizer.eos_token
+ model.config.pad_token_id = model.config.eos_token_id
+
+ # use different length sentences to test batching
+ sentences = [
+ "Hello, my dog is a little",
+ "Today, I",
+ ]
+
+ inputs = tokenizer(sentences, return_tensors="pt", padding=True)
+ input_ids = inputs["input_ids"].to(torch_device)
+ token_type_ids = torch.cat(
+ [
+ input_ids.new_full((input_ids.shape[0], input_ids.shape[1] - 1), 0),
+ input_ids.new_full((input_ids.shape[0], 1), 500),
+ ],
+ dim=-1,
+ )
+
+ outputs = model.generate(
+ input_ids=input_ids,
+ attention_mask=inputs["attention_mask"].to(torch_device),
+ max_length=20,
+ )
+
+ outputs_tt = model.generate(
+ input_ids=input_ids,
+ attention_mask=inputs["attention_mask"].to(torch_device),
+ token_type_ids=token_type_ids,
+ max_length=20,
+ )
+
+ inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
+ output_non_padded = model.generate(input_ids=inputs_non_padded, max_length=20)
+
+ num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().item()
+ inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
+ output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
+
+ batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+ batch_out_sentence_tt = tokenizer.batch_decode(outputs_tt, skip_special_tokens=True)
+ non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
+ padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
+
+ expected_output_sentence = [
+ "Hello, my dog is a little bit of a mess. I'm not sure if he's going",
+ "Today, I'm going to be doing a lot of research on this. I",
+ ]
+ self.assertListEqual(expected_output_sentence, batch_out_sentence)
+ self.assertTrue(batch_out_sentence_tt != batch_out_sentence) # token_type_ids should change output
+ self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])
diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py
index 35e8f707c4b8..ec5588b1b989 100644
--- a/tests/models/gpt_oss/test_modeling_gpt_oss.py
+++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py
@@ -27,7 +27,6 @@
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
- GptOssConfig,
is_torch_available,
)
from transformers.testing_utils import (
@@ -40,7 +39,6 @@
)
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
-from ...test_configuration_common import ConfigTester
if is_torch_available():
@@ -58,31 +56,11 @@
class GptOssModelTester(CausalLMModelTester):
if is_torch_available():
- config_class = GptOssConfig
base_model_class = GptOssModel
- causal_lm_class = GptOssForCausalLM
- sequence_class = GptOssForSequenceClassification
- token_class = GptOssForTokenClassification
-
- pipeline_model_mapping = (
- {
- "feature-extraction": GptOssModel,
- "text-classification": GptOssForSequenceClassification,
- "text-generation": GptOssForCausalLM,
- "token-classification": GptOssForTokenClassification,
- }
- if is_torch_available()
- else {}
- )
@require_torch
class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (GptOssModel, GptOssForCausalLM, GptOssForSequenceClassification, GptOssForTokenClassification)
- if is_torch_available()
- else ()
- )
pipeline_model_mapping = (
{
"feature-extraction": GptOssModel,
@@ -94,16 +72,10 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
else {}
)
- test_headmasking = False
- test_pruning = False
_is_stateful = True
model_split_percents = [0.5, 0.6]
model_tester_class = GptOssModelTester
- def setUp(self):
- self.model_tester = GptOssModelTester(self)
- self.config_tester = ConfigTester(self, config_class=GptOssConfig, hidden_size=37)
-
@unittest.skip("GptOss's forcefully disables sdpa due to Sink")
def test_sdpa_can_dispatch_non_composite_models(self):
pass
diff --git a/tests/models/granite_speech/test_modeling_granite_speech.py b/tests/models/granite_speech/test_modeling_granite_speech.py
index 1ea1f73b4344..4939aecf0d49 100644
--- a/tests/models/granite_speech/test_modeling_granite_speech.py
+++ b/tests/models/granite_speech/test_modeling_granite_speech.py
@@ -40,7 +40,6 @@
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
ModelTesterMixin,
- _config_zero_init,
floats_tensor,
ids_tensor,
)
@@ -127,7 +126,7 @@ def __init__(
self.audio_token_index = audio_token_index
self.tie_word_embeddings = tie_word_embeddings
self.initializer_range = initializer_range
- self.has_lora_adapater = has_lora_adapter
+ self.has_lora_adapter = has_lora_adapter
self.downsample_rate = downsample_rate
self.window_size = window_size
self.is_training = is_training
@@ -152,7 +151,7 @@ def get_config(self):
audio_token_index=self.audio_token_index,
tie_word_embeddings=self.tie_word_embeddings,
initializer_range=self.initializer_range,
- has_lora_adapter=self.has_lora_adapater,
+ has_lora_adapter=self.has_lora_adapter,
)
def prepare_config_and_inputs(self):
@@ -252,22 +251,6 @@ def test_inputs_embeds(self):
with torch.no_grad():
model(**inputs)
- def test_initialization(self):
- config, _ = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if name == "projector.query":
- continue
- elif param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_sdpa_can_dispatch_composite_models(self):
# overwrite because Granite Speech is audio+text model (not vision+text)
if not self.has_attentions:
diff --git a/tests/models/grounding_dino/test_modeling_grounding_dino.py b/tests/models/grounding_dino/test_modeling_grounding_dino.py
index 1821262b2dec..1be8a7735489 100644
--- a/tests/models/grounding_dino/test_modeling_grounding_dino.py
+++ b/tests/models/grounding_dino/test_modeling_grounding_dino.py
@@ -40,7 +40,7 @@
)
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -570,32 +570,6 @@ def test_hf_backbone(self):
self.assertTrue(outputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- if (
- "level_embed" in name
- or "sampling_offsets.bias" in name
- or "text_param" in name
- or "vision_param" in name
- or "value_proj" in name
- or "output_proj" in name
- or "reference_points" in name
- or "vision_proj" in name
- or "text_proj" in name
- ):
- continue
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# Copied from tests.models.deformable_detr.test_modeling_deformable_detr.DeformableDetrModelTest.test_two_stage_training with DeformableDetr->GroundingDino
def test_two_stage_training(self):
model_class = GroundingDinoForObjectDetection
diff --git a/tests/models/groupvit/test_modeling_groupvit.py b/tests/models/groupvit/test_modeling_groupvit.py
index a4d521ff2a7b..33252f5ef8ac 100644
--- a/tests/models/groupvit/test_modeling_groupvit.py
+++ b/tests/models/groupvit/test_modeling_groupvit.py
@@ -569,30 +569,6 @@ def test_retain_grad_hidden_states_attentions(self):
def test_model_get_set_embeddings(self):
pass
- # override as the `logit_scale` parameter initialization is different for GROUPVIT
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- # check if `logit_scale` is initialized as per the original implementation
- if name == "logit_scale":
- self.assertAlmostEqual(
- param.data.item(),
- np.log(1 / 0.07),
- delta=1e-3,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def _create_and_check_torchscript(self, config, inputs_dict):
if not self.test_torchscript:
self.skipTest(reason="test_torchscript is set to False")
diff --git a/tests/models/helium/test_modeling_helium.py b/tests/models/helium/test_modeling_helium.py
index 61639ac48918..47bdb6bfc948 100644
--- a/tests/models/helium/test_modeling_helium.py
+++ b/tests/models/helium/test_modeling_helium.py
@@ -15,7 +15,7 @@
import unittest
-from transformers import AutoModelForCausalLM, AutoTokenizer, HeliumConfig, is_torch_available
+from transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_available
from transformers.testing_utils import (
Expectations,
require_read_token,
@@ -24,8 +24,7 @@
torch_device,
)
-from ...test_configuration_common import ConfigTester
-from ..gemma.test_modeling_gemma import GemmaModelTest, GemmaModelTester
+from ..gemma.test_modeling_gemma import GemmaModelTester
if is_torch_available():
@@ -41,20 +40,11 @@
class HeliumModelTester(GemmaModelTester):
if is_torch_available():
- config_class = HeliumConfig
- model_class = HeliumModel
- for_causal_lm_class = HeliumForCausalLM
- for_sequence_class = HeliumForSequenceClassification
- for_token_class = HeliumForTokenClassification
+ base_model_class = HeliumModel
@require_torch
-class HeliumModelTest(GemmaModelTest, unittest.TestCase):
- all_model_classes = (
- (HeliumModel, HeliumForCausalLM, HeliumForSequenceClassification, HeliumForTokenClassification)
- if is_torch_available()
- else ()
- )
+class HeliumModelTest(CausalLMModelTest, unittest.TestCase):
pipeline_model_mapping = (
{
"feature-extraction": HeliumModel,
@@ -66,14 +56,10 @@ class HeliumModelTest(GemmaModelTest, unittest.TestCase):
if is_torch_available()
else {}
)
- test_headmasking = False
- test_pruning = False
_is_stateful = True
model_split_percents = [0.5, 0.6]
- def setUp(self):
- self.model_tester = HeliumModelTester(self)
- self.config_tester = ConfigTester(self, config_class=HeliumConfig, hidden_size=37)
+ model_tester_class = HeliumModelTester
@slow
diff --git a/tests/models/hgnet_v2/test_modeling_hgnet_v2.py b/tests/models/hgnet_v2/test_modeling_hgnet_v2.py
index 2dad713308b4..7186a7045786 100644
--- a/tests/models/hgnet_v2/test_modeling_hgnet_v2.py
+++ b/tests/models/hgnet_v2/test_modeling_hgnet_v2.py
@@ -16,7 +16,6 @@
import unittest
import torch
-from torch import nn
from transformers import HGNetV2Config
from transformers.testing_utils import require_torch, torch_device
@@ -189,10 +188,6 @@ class HGNetV2ForImageClassificationTest(ModelTesterMixin, PipelineTesterMixin, u
def setUp(self):
self.model_tester = HGNetV2ModelTester(self)
- @unittest.skip(reason="Does not work on the tiny model.")
- def test_model_parallelism(self):
- super().test_model_parallelism()
-
@unittest.skip(reason="HGNetV2 does not output attentions")
def test_attention_outputs(self):
pass
@@ -209,34 +204,10 @@ def test_inputs_embeds(self):
def test_model_common_attributes(self):
pass
- @unittest.skip(reason="HGNetV2 does not have a model")
- def test_model(self):
- pass
-
- @unittest.skip(reason="Not relevant for the model")
- def test_can_init_all_missing_weights(self):
- pass
-
def test_backbone(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_backbone(*config_and_inputs)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- for model_class in self.all_model_classes:
- model = model_class(config=config)
- for name, module in model.named_modules():
- if isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
- self.assertTrue(
- torch.all(module.weight == 1),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- self.assertTrue(
- torch.all(module.bias == 0),
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
def test_hidden_states_output(self):
def check_hidden_states_output(inputs_dict, config, model_class):
model = model_class(config)
diff --git a/tests/models/hubert/test_modeling_hubert.py b/tests/models/hubert/test_modeling_hubert.py
index feec7a1de48d..4b44f9015350 100644
--- a/tests/models/hubert/test_modeling_hubert.py
+++ b/tests/models/hubert/test_modeling_hubert.py
@@ -408,32 +408,6 @@ def test_retain_grad_hidden_states_attentions(self):
self.assertIsNotNone(hidden_states.grad)
self.assertIsNotNone(attentions.grad)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = [
- "conv.weight",
- "conv.parametrizations.weight",
- "masked_spec_embed",
- "quantizer.weight_proj.weight",
- ]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# Hubert cannot be TorchScripted because of torch.nn.utils.weight_norm
def _create_and_check_torch_fx_tracing(self, config, inputs_dict, output_loss=False):
# TODO: fix it
@@ -673,32 +647,6 @@ def test_retain_grad_hidden_states_attentions(self):
self.assertIsNotNone(hidden_states.grad)
self.assertIsNotNone(attentions.grad)
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- uniform_init_parms = [
- "conv.weight",
- "conv.parametrizations.weight",
- "masked_spec_embed",
- "quantizer.weight_proj.weight",
- ]
- if param.requires_grad:
- if any(x in name for x in uniform_init_parms):
- self.assertTrue(
- -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
- else:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
# overwrite from test_modeling_common
def _mock_init_weights(self, module):
if hasattr(module, "weight") and module.weight is not None:
diff --git a/tests/models/hunyuan_v1_dense/test_modeling_hunyuan_v1_dense.py b/tests/models/hunyuan_v1_dense/test_modeling_hunyuan_v1_dense.py
index d01a1022f342..edcf9cd21088 100644
--- a/tests/models/hunyuan_v1_dense/test_modeling_hunyuan_v1_dense.py
+++ b/tests/models/hunyuan_v1_dense/test_modeling_hunyuan_v1_dense.py
@@ -17,7 +17,7 @@
from parameterized import parameterized
-from transformers import HunYuanDenseV1Config, is_torch_available
+from transformers import is_torch_available
from transformers.testing_utils import (
cleanup,
require_torch,
@@ -36,26 +36,12 @@
class HunYuanDenseV1ModelTester(CausalLMModelTester):
- config_class = HunYuanDenseV1Config
if is_torch_available():
base_model_class = HunYuanDenseV1Model
- causal_lm_class = HunYuanDenseV1ForCausalLM
- sequence_class = HunYuanDenseV1ForSequenceClassification
@require_torch
class HunYuanDenseV1ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- HunYuanDenseV1Model,
- HunYuanDenseV1ForCausalLM,
- HunYuanDenseV1ForSequenceClassification,
- )
- if is_torch_available()
- else ()
- )
- test_headmasking = False
- test_pruning = False
model_tester_class = HunYuanDenseV1ModelTester
pipeline_model_mapping = (
{
diff --git a/tests/models/hunyuan_v1_moe/test_modeling_hunyuan_v1_moe.py b/tests/models/hunyuan_v1_moe/test_modeling_hunyuan_v1_moe.py
index 3738ebee75d1..b835f0677cfe 100644
--- a/tests/models/hunyuan_v1_moe/test_modeling_hunyuan_v1_moe.py
+++ b/tests/models/hunyuan_v1_moe/test_modeling_hunyuan_v1_moe.py
@@ -18,7 +18,7 @@
import pytest
from parameterized import parameterized
-from transformers import HunYuanMoEV1Config, is_torch_available
+from transformers import is_torch_available
from transformers.testing_utils import (
cleanup,
require_torch,
@@ -40,26 +40,12 @@
class HunYuanMoEV1ModelTester(CausalLMModelTester):
- config_class = HunYuanMoEV1Config
if is_torch_available():
base_model_class = HunYuanMoEV1Model
- causal_lm_class = HunYuanMoEV1ForCausalLM
- sequence_class = HunYuanMoEV1ForSequenceClassification
@require_torch
class HunYuanMoEV1ModelTest(CausalLMModelTest, unittest.TestCase):
- all_model_classes = (
- (
- HunYuanMoEV1Model,
- HunYuanMoEV1ForCausalLM,
- HunYuanMoEV1ForSequenceClassification,
- )
- if is_torch_available()
- else ()
- )
- test_headmasking = False
- test_pruning = False
test_all_params_have_gradient = False
model_tester_class = HunYuanMoEV1ModelTester
pipeline_model_mapping = (
diff --git a/tests/models/idefics/test_modeling_idefics.py b/tests/models/idefics/test_modeling_idefics.py
index 2cf220fd6dfd..a517d69e18a6 100644
--- a/tests/models/idefics/test_modeling_idefics.py
+++ b/tests/models/idefics/test_modeling_idefics.py
@@ -13,7 +13,6 @@
# limitations under the License.
"""Testing suite for the PyTorch Idefics model."""
-import inspect
import unittest
from functools import cached_property
@@ -67,7 +66,7 @@ def __init__(
use_labels=True,
vocab_size=99,
hidden_size=32,
- num_hidden_layers=5,
+ num_hidden_layers=2,
num_attention_heads=4,
intermediate_size=37,
hidden_act="gelu",
@@ -85,7 +84,7 @@ def __init__(
vision_patch_size=2,
vision_image_size=30,
vision_num_attention_heads=4,
- vision_num_hidden_layers=5,
+ vision_num_hidden_layers=2,
vision_intermediate_size=37,
perceiver_qk_layer_norms_perceiver=False,
perceiver_resampler_depth=2,
@@ -327,7 +326,6 @@ class IdeficsModelTest(ModelTesterMixin, PipelineTesterMixin, GenerationTesterMi
test_pruning = False
test_headmasking = False
test_torchscript = False
- has_attentions = False # only supports SDOA and thus no attention probs returned
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
@@ -594,6 +592,33 @@ def test_generate_from_random_inputs_embeds(
):
pass
+ @pytest.mark.generate
+ def test_left_padding_compatibility(self):
+ # Overwrite -- Idefics needs to prepare `image_attention_mask`, and it must be padded accordingly
+ _, inputs_dict = self.prepare_config_and_inputs_for_generate()
+ input_ids = inputs_dict["input_ids"]
+ image_attention_mask = inputs_dict["image_attention_mask"]
+
+ pad_size_img = (input_ids.shape[0], 32, image_attention_mask.shape[-1])
+ extra_img_mask = torch.zeros(pad_size_img, dtype=image_attention_mask.dtype, device=torch_device)
+ padded_image_attention_mask = torch.cat([extra_img_mask, image_attention_mask], dim=1)
+
+ # `image_attention_mask` is randomly generated in `prepare_config_and_inputs_for_generate`, and it must match
+ # its padded version for the test to be valid -- we need to pass both
+ unpadded_custom_inputs = {"image_attention_mask": image_attention_mask}
+ padded_custom_inputs = {"image_attention_mask": padded_image_attention_mask}
+ super().test_left_padding_compatibility(
+ unpadded_custom_inputs=unpadded_custom_inputs, padded_custom_inputs=padded_custom_inputs
+ )
+
+ @unittest.skip(reason="Idefics can't do text-only inference (test filters non-text inputs)")
+ def test_eager_padding_matches_padding_free_with_position_ids(self):
+ pass
+
+ @unittest.skip(reason="Idefics can't do text-only inference (test filters non-text inputs)")
+ def test_sdpa_padding_matches_padding_free_with_position_ids(self):
+ pass
+
@require_torch
class IdeficsForVisionText2TextTest(IdeficsModelTest, GenerationTesterMixin, unittest.TestCase):
@@ -613,66 +638,6 @@ def test_eager_matches_sdpa_inference(
):
pass
- @pytest.mark.generate
- def test_left_padding_compatibility(self):
- """Overwrite because IDEFICS needs image attention mask to be also padded"""
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- def _prepare_model_kwargs(input_ids, attention_mask, image_attention_mask, signature):
- model_kwargs = {
- "input_ids": input_ids,
- "attention_mask": attention_mask,
- "image_attention_mask": image_attention_mask,
- }
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in self.all_generative_model_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict.pop("input_ids")
- attention_mask = inputs_dict.pop("attention_mask")
- if attention_mask is None:
- attention_mask = torch.ones_like(input_ids)
- image_attention_mask = inputs_dict.pop("image_attention_mask", None)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, image_attention_mask, signature)
- next_logits_wo_padding = model(**model_kwargs, **inputs_dict).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
-
- pad_size_img = (input_ids.shape[0], 32, image_attention_mask.shape[-1])
- extra_img_mask = torch.zeros(pad_size_img, dtype=image_attention_mask.dtype, device=torch_device)
- padded_image_attention_mask = torch.cat([extra_img_mask, image_attention_mask], dim=1)
- model_kwargs = _prepare_model_kwargs(
- padded_input_ids, padded_attention_mask, padded_image_attention_mask, signature
- )
- next_logits_with_padding = model(**model_kwargs, **inputs_dict).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
-
@pytest.mark.generate
def test_generate_continue_from_past_key_values(self):
"""Overwrite because IDEFICS needs image attention mask to be also processed"""
diff --git a/tests/models/idefics2/test_modeling_idefics2.py b/tests/models/idefics2/test_modeling_idefics2.py
index a500d8bf4946..6603f3604e0b 100644
--- a/tests/models/idefics2/test_modeling_idefics2.py
+++ b/tests/models/idefics2/test_modeling_idefics2.py
@@ -86,7 +86,7 @@ def __init__(
"vocab_size": 100,
"hidden_size": 64,
"intermediate_size": 56,
- "num_hidden_layers": 3,
+ "num_hidden_layers": 2,
"num_attention_heads": 2,
"num_key_value_heads": 2,
"hidden_act": "silu",
diff --git a/tests/models/idefics3/test_modeling_idefics3.py b/tests/models/idefics3/test_modeling_idefics3.py
index b4434f34b81c..fe05eda8c0fb 100644
--- a/tests/models/idefics3/test_modeling_idefics3.py
+++ b/tests/models/idefics3/test_modeling_idefics3.py
@@ -74,7 +74,7 @@ def __init__(
"vocab_size": 100,
"hidden_size": 64,
"intermediate_size": 56,
- "num_hidden_layers": 3,
+ "num_hidden_layers": 2,
"num_attention_heads": 2,
"num_key_value_heads": 2,
"hidden_act": "silu",
diff --git a/tests/models/imagegpt/test_modeling_imagegpt.py b/tests/models/imagegpt/test_modeling_imagegpt.py
index 1c10ed0797db..9a43671ad975 100644
--- a/tests/models/imagegpt/test_modeling_imagegpt.py
+++ b/tests/models/imagegpt/test_modeling_imagegpt.py
@@ -316,10 +316,6 @@ def test_forward_signature(self):
expected_arg_names = ["input_ids"]
self.assertListEqual(arg_names[:1], expected_arg_names)
- @unittest.skip(reason="The model doesn't support left padding") # and it's not used enough to be worth fixing :)
- def test_left_padding_compatibility(self):
- pass
-
@unittest.skip(reason="Model inputs don't fit test pattern") # and it's not used enough to be worth fixing :)
def test_past_key_values_format(self):
pass
diff --git a/tests/models/instructblip/test_modeling_instructblip.py b/tests/models/instructblip/test_modeling_instructblip.py
index 3ce58e4cb24a..17a54da482a2 100644
--- a/tests/models/instructblip/test_modeling_instructblip.py
+++ b/tests/models/instructblip/test_modeling_instructblip.py
@@ -18,7 +18,6 @@
import unittest
import numpy as np
-import pytest
import requests
from transformers import (
@@ -566,94 +565,6 @@ def _check_generate_outputs(self, output, config, use_cache=False, num_return_se
output, config, use_cache=use_cache, num_return_sequences=num_return_sequences, num_beams=num_beams
)
- # overwrite because InstructBLIP cannot generate only from input ids, and requires `pixel` values and `qformer_input_ids` in all cases to be present
- @pytest.mark.generate
- def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
-
- # - The model must support padding
- if not self.has_attentions:
- self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
- decoder_only_classes = []
- for model_class in self.all_generative_model_classes:
- config, _ = self.prepare_config_and_inputs_for_generate()
- if config.is_encoder_decoder:
- continue
- else:
- decoder_only_classes.append(model_class)
- if len(decoder_only_classes) == 0:
- self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
- has_encoder_attributes = any(
- attr_name
- for attr_name in config.to_dict()
- if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
- )
- if has_encoder_attributes:
- self.skipTest(
- reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
- )
-
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in decoder_only_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
- attention_mask = inputs_dict.get("attention_mask")
- pixel_values = inputs_dict["pixel_values"]
- qformer_input_ids = inputs_dict["qformer_input_ids"]
- if attention_mask is None:
- attention_mask = torch.ones_like(input_ids)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(
- **model_kwargs, pixel_values=pixel_values, qformer_input_ids=qformer_input_ids
- ).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(
- **model_kwargs, pixel_values=pixel_values, qformer_input_ids=qformer_input_ids
- ).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
-
def test_sdpa_can_dispatch_composite_models(self):
"""
Tests if composite models dispatch correctly on SDPA/eager when requested so when loading the model.
diff --git a/tests/models/instructblipvideo/test_modeling_instructblipvideo.py b/tests/models/instructblipvideo/test_modeling_instructblipvideo.py
index a91d31082da9..d6336c8c6840 100644
--- a/tests/models/instructblipvideo/test_modeling_instructblipvideo.py
+++ b/tests/models/instructblipvideo/test_modeling_instructblipvideo.py
@@ -18,7 +18,6 @@
import unittest
import numpy as np
-import pytest
from huggingface_hub import hf_hub_download
from transformers import (
@@ -578,94 +577,6 @@ def _check_generate_outputs(self, output, config, use_cache=False, num_return_se
output, config, use_cache=use_cache, num_return_sequences=num_return_sequences, num_beams=num_beams
)
- # overwrite because InstructBLIPVideo cannot generate only from input ids, and requires `pixel` values and `qformer_input_ids` in all cases to be present
- @pytest.mark.generate
- def test_left_padding_compatibility(self):
- # NOTE: left-padding results in small numerical differences. This is expected.
- # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
-
- # First, filter out models that don't support left padding
- # - The model must have generative capabilities
- if len(self.all_generative_model_classes) == 0:
- self.skipTest(reason="No generative architecture available for this model.")
-
- # - The model must support padding
- if not self.has_attentions:
- self.skipTest(reason="This model doesn't support padding.")
-
- # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
- decoder_only_classes = []
- for model_class in self.all_generative_model_classes:
- config, _ = self.prepare_config_and_inputs_for_generate()
- if config.is_encoder_decoder:
- continue
- else:
- decoder_only_classes.append(model_class)
- if len(decoder_only_classes) == 0:
- self.skipTest(reason="No decoder-only architecture available for this model.")
-
- # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
- # added support for it yet. We skip these models for now.
- has_encoder_attributes = any(
- attr_name
- for attr_name in config.to_dict()
- if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
- )
- if has_encoder_attributes:
- self.skipTest(
- reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
- )
-
- # Then, test left-padding
- def _prepare_model_kwargs(input_ids, attention_mask, signature):
- model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
- if "position_ids" in signature:
- position_ids = torch.cumsum(attention_mask, dim=-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- model_kwargs["position_ids"] = position_ids
- if "cache_position" in signature:
- cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
- model_kwargs["cache_position"] = cache_position
- return model_kwargs
-
- for model_class in decoder_only_classes:
- config, inputs_dict = self.prepare_config_and_inputs_for_generate()
- input_ids = inputs_dict["input_ids"]
- attention_mask = inputs_dict.get("attention_mask")
- pixel_values = inputs_dict["pixel_values"]
- qformer_input_ids = inputs_dict["qformer_input_ids"]
- if attention_mask is None:
- attention_mask = torch.ones_like(input_ids)
-
- model = model_class(config).to(torch_device).eval()
- signature = inspect.signature(model.forward).parameters.keys()
-
- # no cache as some models require special cache classes to be init outside forward
- model.generation_config.use_cache = False
-
- # Without padding
- model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
- next_logits_wo_padding = model(
- **model_kwargs, pixel_values=pixel_values, qformer_input_ids=qformer_input_ids
- ).logits[:, -1, :]
-
- # With left-padding (length 32)
- # can hardcode pad_token to be 0 as we'll do attn masking anyway
- pad_token_id = (
- config.get_text_config().pad_token_id if config.get_text_config().pad_token_id is not None else 0
- )
- pad_size = (input_ids.shape[0], 32)
- padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * pad_token_id
- padded_input_ids = torch.cat((padding, input_ids), dim=1)
- padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
- model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
- next_logits_with_padding = model(
- **model_kwargs, pixel_values=pixel_values, qformer_input_ids=qformer_input_ids
- ).logits[:, -1, :]
-
- # They should result in very similar logits
- torch.testing.assert_close(next_logits_wo_padding, next_logits_with_padding, rtol=1e-5, atol=1e-5)
-
def test_sdpa_can_dispatch_composite_models(self):
"""
Tests if composite models dispatch correctly on SDPA/eager when requested so when loading the model.
diff --git a/tests/models/internvl/test_modeling_internvl.py b/tests/models/internvl/test_modeling_internvl.py
index 297dc6cffe85..17fcf9a4b338 100644
--- a/tests/models/internvl/test_modeling_internvl.py
+++ b/tests/models/internvl/test_modeling_internvl.py
@@ -41,7 +41,7 @@
from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -74,7 +74,7 @@ def __init__(
"vocab_size": 99,
"hidden_size": 128,
"intermediate_size": 37,
- "num_hidden_layers": 4,
+ "num_hidden_layers": 2,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"output_channels": 64,
@@ -208,20 +208,6 @@ def test_flex_attention_with_grads(self):
def test_config(self):
self.config_tester.run_common_tests()
- def test_initialization(self):
- config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
- configs_no_init = _config_zero_init(config)
- for model_class in self.all_model_classes:
- model = model_class(config=configs_no_init)
- for name, param in model.named_parameters():
- if param.requires_grad:
- self.assertIn(
- ((param.data.mean() * 1e9).round() / 1e9).item(),
- [0.0, 1.0],
- msg=f"Parameter {name} of model {model_class} seems not properly initialized",
- )
-
@unittest.skip(reason="Compile not yet supported because in LLava models")
@pytest.mark.torch_compile_test
def test_sdpa_can_compile_dynamic(self):
diff --git a/tests/models/internvl/test_processing_internvl.py b/tests/models/internvl/test_processing_internvl.py
index 76e91a50d3ed..bbb4df973da6 100644
--- a/tests/models/internvl/test_processing_internvl.py
+++ b/tests/models/internvl/test_processing_internvl.py
@@ -17,6 +17,8 @@
import tempfile
import unittest
+from parameterized import parameterized
+
from transformers import AutoProcessor, AutoTokenizer, InternVLProcessor
from transformers.testing_utils import require_av, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available
@@ -219,7 +221,7 @@ def test_apply_chat_template_video_frame_sampling(self):
{
"type": "video",
"url": url_to_local_path(
- "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/Big_Buck_Bunny_720_10s_10MB.mp4"
+ "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/tiny_video.mp4"
),
},
{"type": "text", "text": "What is shown in this video?"},
@@ -251,7 +253,7 @@ def test_apply_chat_template_video_frame_sampling(self):
return_tensors="pt",
)
self.assertTrue(self.videos_input_name in out_dict_with_video)
- self.assertEqual(len(out_dict_with_video[self.videos_input_name]), 300)
+ self.assertEqual(len(out_dict_with_video[self.videos_input_name]), 11)
# Load video as a list of frames (i.e. images). NOTE: each frame should have same size
# because we assume they come from one video
@@ -345,13 +347,14 @@ def _test_apply_chat_template(
for idx, url in enumerate(input_data[:batch_size]):
batch_messages[idx][0]["content"] = [batch_messages[idx][0]["content"][0], {"type": modality, "url": url}]
+ num_frames = 2 # by default no more than 2 frames, otherwise too slow
out_dict = processor.apply_chat_template(
batch_messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
- num_frames=2, # by default no more than 2 frames, otherwise too slow
+ num_frames=num_frames,
)
self.assertTrue(self.videos_input_name in out_dict)
self.assertEqual(len(out_dict["input_ids"]), batch_size)
@@ -359,11 +362,15 @@ def _test_apply_chat_template(
# InternVL internally collects frames from all the videos in a batch and flattens the batch dimension (B T C H W) -> (B*T C H W) then patches and removes the frames
# hence output length does not equal batch size
- # removed hardcoded video length check video_len = 2 if batch_size == 1 else 3
- # from experiment video_len looks like batch_size + 1
- # TODO: update expected video_len calculation based on the internal processing logic of InternVLProcessor
- output_len = batch_size + 1 if modality == "video" else batch_size
- self.assertEqual(len(out_dict[self.videos_input_name]), output_len)
+ num_pixel_planes = 0 # i.e. images + video frames
+ for message_thread in batch_messages:
+ for message in message_thread:
+ for content in message.get("content", []):
+ if (content_type := content.get("type")) == "image":
+ num_pixel_planes += 1
+ elif content_type == "video":
+ num_pixel_planes += num_frames
+ self.assertEqual(len(out_dict[self.videos_input_name]), num_pixel_planes)
for k in out_dict:
self.assertIsInstance(out_dict[k], torch.Tensor)
@@ -377,3 +384,25 @@ def _test_apply_chat_template(
continue_prompt = processor.apply_chat_template(batch_messages, continue_final_message=True, tokenize=False)
for prompt in continue_prompt:
self.assertTrue(prompt.endswith("It is the sound of")) # no `eos` token at the end
+
+ @parameterized.expand([(1,), (2,)])
+ @require_torch
+ def test_frames_binding(self, batch_size: int):
+ texts = [
+ "
@@ -32,13 +30,11 @@ You can find all the original TrOCR checkpoints under the [Microsoft](https://hu
alt="drawing" width="600"/>
TrOCR architecture. Taken from the original paper.
-
> [!TIP]
> This model was contributed by [nielsr](https://huggingface.co/nielsr).
>
> Click on the TrOCR models in the right sidebar for more examples of how to apply TrOCR to different image and text tasks.
-
The example below demonstrates how to perform optical character recognition (OCR) with the [`AutoModel`] class.
@@ -67,6 +66,7 @@ reconstruction = outputs.logits
## Notes
+
- ViTMAE is typically used in two stages. Self-supervised pretraining with [`ViTMAEForPreTraining`], and then discarding the decoder and fine-tuning the encoder. After fine-tuning, the weights can be plugged into a model like [`ViTForImageClassification`].
- Use [`ViTImageProcessor`] for input preparation.
diff --git a/docs/source/en/model_doc/vit_msn.md b/docs/source/en/model_doc/vit_msn.md
index 5b727f34256c..d7a8172a18f3 100644
--- a/docs/source/en/model_doc/vit_msn.md
+++ b/docs/source/en/model_doc/vit_msn.md
@@ -40,11 +40,11 @@ while producing representations of a high semantic level that perform competitiv
on ImageNet-1K, with only 5,000 annotated images, our base MSN model achieves 72.4% top-1 accuracy,
and with 1% of ImageNet-1K labels, we achieve 75.7% top-1 accuracy, setting a new state-of-the-art for self-supervised learning on this benchmark.*
-
+
MSN architecture. Taken from the original paper.
-This model was contributed by [sayakpaul](https://huggingface.co/sayakpaul). The original code can be found [here](https://github.com/facebookresearch/msn).
+This model was contributed by [sayakpaul](https://huggingface.co/sayakpaul). The original code can be found [here](https://github.com/facebookresearch/msn).
## Usage tips
@@ -58,16 +58,16 @@ labels when fine-tuned.
### Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import ViTMSNForImageClassification
model = ViTMSNForImageClassification.from_pretrained("facebook/vit-msn-base", attn_implementation="sdpa", dtype=torch.float16)
...
diff --git a/docs/source/en/model_doc/vitdet.md b/docs/source/en/model_doc/vitdet.md
index 539ae5e376c8..a1250f1bb909 100644
--- a/docs/source/en/model_doc/vitdet.md
+++ b/docs/source/en/model_doc/vitdet.md
@@ -40,4 +40,4 @@ Tips:
## VitDetModel
[[autodoc]] VitDetModel
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/vitmatte.md b/docs/source/en/model_doc/vitmatte.md
index 519a2dd74d66..0584df8e67a5 100644
--- a/docs/source/en/model_doc/vitmatte.md
+++ b/docs/source/en/model_doc/vitmatte.md
@@ -62,4 +62,4 @@ The model expects both the image and trimap (concatenated) as input. Use [`ViTMa
## VitMatteForImageMatting
[[autodoc]] VitMatteForImageMatting
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/vits.md b/docs/source/en/model_doc/vits.md
index 2c1777b77f18..96dc93892470 100644
--- a/docs/source/en/model_doc/vits.md
+++ b/docs/source/en/model_doc/vits.md
@@ -149,11 +149,10 @@ Audio(waveform, rate=model.config.sampling_rate)
## VitsTokenizer
[[autodoc]] VitsTokenizer
-- __call__
-- save_vocabulary
+ - __call__
+ - save_vocabulary
## VitsModel
[[autodoc]] VitsModel
-- forward
-
+ - forward
diff --git a/docs/source/en/model_doc/vivit.md b/docs/source/en/model_doc/vivit.md
index 041f80f61ae6..fc127fa6f595 100644
--- a/docs/source/en/model_doc/vivit.md
+++ b/docs/source/en/model_doc/vivit.md
@@ -32,16 +32,16 @@ This model was contributed by [jegormeister](https://huggingface.co/jegormeister
### Using Scaled Dot Product Attention (SDPA)
-PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
-encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
-[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
+PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function
+encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the
+[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention)
page for more information.
-SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
+SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set
`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
-```
+```py
from transformers import VivitModel
model = VivitModel.from_pretrained("google/vivit-b-16x2-kinetics400", attn_implementation="sdpa", dtype=torch.float16)
...
@@ -56,8 +56,6 @@ On a local benchmark (A100-40GB, PyTorch 2.3.0, OS Ubuntu 22.04) with `float32`
|---------------------:|-------------:|----------:|--------------:|----------------------:|---------------------:|-----------------:|
| 100 | 1 | True | 7.122 | 2575.28 | 5932.54 | 130.364 |
-
-
### Inference
| num_batches | batch_size | is cuda | is half | Speedup (%) | Mem eager (MB) | Mem BT (MB) | Mem saved (%) |
|---------------|--------------|-----------|-----------|---------------|------------------|---------------|-----------------|
@@ -65,7 +63,6 @@ On a local benchmark (A100-40GB, PyTorch 2.3.0, OS Ubuntu 22.04) with `float32`
| 20 | 2 | True | False | 17.146 | 1234.75 | 447.175 | 176.122 |
| 20 | 4 | True | False | 18.093 | 2275.82 | 709.864 | 220.6 |
| 20 | 8 | True | False | 19.284 | 4358.19 | 1233.24 | 253.393 |
-
## VivitConfig
diff --git a/docs/source/en/model_doc/vjepa2.md b/docs/source/en/model_doc/vjepa2.md
index 93960f051893..049c7ff98f21 100644
--- a/docs/source/en/model_doc/vjepa2.md
+++ b/docs/source/en/model_doc/vjepa2.md
@@ -15,7 +15,6 @@ rendered properly in your Markdown viewer.
-->
*This model was released on 2025-06-11 and added to Hugging Face Transformers on 2025-06-11.*
-
@@ -34,7 +33,6 @@ rendered properly in your Markdown viewer.
You can find all original V-JEPA2 checkpoints under the [V-JEPA 2](https://huggingface.co/collections/facebook/v-jepa-2-6841bad8413014e185b497a6) collection.
-
This model was contributed by [koustuvs](https://huggingface.co/koustuvs), [yonigozlan](https://huggingface.co/yonigozlan) and [qubvel](https://huggingface.co/qubvel-hf). The original code can be found [here](https://github.com/facebookresearch/vjepa2).
## Usage example
diff --git a/docs/source/en/model_doc/voxtral.md b/docs/source/en/model_doc/voxtral.md
index 71f0661c8276..3dd2fc9e0d31 100644
--- a/docs/source/en/model_doc/voxtral.md
+++ b/docs/source/en/model_doc/voxtral.md
@@ -22,6 +22,7 @@ Voxtral is an upgrade of [Ministral 3B and Mistral Small 3B](https://mistral.ai/
You can read more in Mistral's [realease blog post](https://mistral.ai/news/voxtral).
The model is available in two checkpoints:
+
- 3B: [mistralai/Voxtral-Mini-3B-2507](https://huggingface.co/mistralai/Voxtral-Mini-3B-2507)
- 24B: [mistralai/Voxtral-Small-24B-2507](https://huggingface.co/mistralai/Voxtral-Small-24B-2507)
@@ -43,6 +44,7 @@ Voxtral builds on Ministral-3B by adding audio processing capabilities:
The model supports audio-text instructions, including multi-turn and multi-audio interactions, all processed in batches.
➡️ audio + text instruction
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -78,7 +80,8 @@ print(decoded_outputs[0])
print("=" * 80)
```
-➡️ multi-audio + text instruction
+➡️ multi-audio + text instruction
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -119,6 +122,7 @@ print("=" * 80)
```
➡️ multi-turn:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -173,6 +177,7 @@ print("=" * 80)
```
➡️ text only:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -208,6 +213,7 @@ print("=" * 80)
```
➡️ audio only:
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device
@@ -243,6 +249,7 @@ print("=" * 80)
```
➡️ batched inference!
+
```python
import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor, infer_device()
diff --git a/docs/source/en/model_doc/wav2vec2-bert.md b/docs/source/en/model_doc/wav2vec2-bert.md
index 4edb67498aaa..23409b0898c3 100644
--- a/docs/source/en/model_doc/wav2vec2-bert.md
+++ b/docs/source/en/model_doc/wav2vec2-bert.md
@@ -31,7 +31,7 @@ The official results of the model can be found in Section 3.2.1 of the paper.
The abstract from the paper is the following:
-*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one’s voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
+*Recent advancements in automatic speech translation have dramatically expanded language coverage, improved multimodal capabilities, and enabled a wide range of tasks and functionalities. That said, large-scale automatic speech translation systems today lack key features that help machine-mediated communication feel seamless when compared to human-to-human dialogue. In this work, we introduce a family of models that enable end-to-end expressive and multilingual translations in a streaming fashion. First, we contribute an improved version of the massively multilingual and multimodal SeamlessM4T model—SeamlessM4T v2. This newer model, incorporating an updated UnitY2 framework, was trained on more low-resource language data. The expanded version of SeamlessAlign adds 114,800 hours of automatically aligned data for a total of 76 languages. SeamlessM4T v2 provides the foundation on which our two newest models, SeamlessExpressive and SeamlessStreaming, are initiated. SeamlessExpressive enables translation that preserves vocal styles and prosody. Compared to previous efforts in expressive speech research, our work addresses certain underexplored aspects of prosody, such as speech rate and pauses, while also preserving the style of one's voice. As for SeamlessStreaming, our model leverages the Efficient Monotonic Multihead Attention (EMMA) mechanism to generate low-latency target translations without waiting for complete source utterances. As the first of its kind, SeamlessStreaming enables simultaneous speech-to-speech/text translation for multiple source and target languages. To understand the performance of these models, we combined novel and modified versions of existing automatic metrics to evaluate prosody, latency, and robustness. For human evaluations, we adapted existing protocols tailored for measuring the most relevant attributes in the preservation of meaning, naturalness, and expressivity. To ensure that our models can be used safely and responsibly, we implemented the first known red-teaming effort for multimodal machine translation, a system for the detection and mitigation of added toxicity, a systematic evaluation of gender bias, and an inaudible localized watermarking mechanism designed to dampen the impact of deepfakes. Consequently, we bring major components from SeamlessExpressive and SeamlessStreaming together to form Seamless, the first publicly available system that unlocks expressive cross-lingual communication in real-time. In sum, Seamless gives us a pivotal look at the technical foundation needed to turn the Universal Speech Translator from a science fiction concept into a real-world technology. Finally, contributions in this work—including models, code, and a watermark detector—are publicly released and accessible at the link below.*
This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The original code can be found [here](https://github.com/facebookresearch/seamless_communication).
@@ -54,7 +54,6 @@ This model was contributed by [ylacombe](https://huggingface.co/ylacombe). The o
- [`Wav2Vec2BertForSequenceClassification`] can be used by adapting this [example script](https://github.com/huggingface/transformers/tree/main/examples/pytorch/audio-classification).
- See also: [Audio classification task guide](../tasks/audio_classification)
-
## Wav2Vec2BertConfig
[[autodoc]] Wav2Vec2BertConfig
diff --git a/docs/source/en/model_doc/wav2vec2-conformer.md b/docs/source/en/model_doc/wav2vec2-conformer.md
index e2a56b450df3..663b6163011b 100644
--- a/docs/source/en/model_doc/wav2vec2-conformer.md
+++ b/docs/source/en/model_doc/wav2vec2-conformer.md
@@ -38,7 +38,7 @@ Note: Meta (FAIR) released a new version of [Wav2Vec2-BERT 2.0](https://huggingf
- Wav2Vec2-Conformer follows the same architecture as Wav2Vec2, but replaces the *Attention*-block with a *Conformer*-block
as introduced in [Conformer: Convolution-augmented Transformer for Speech Recognition](https://huggingface.co/papers/2005.08100).
-- For the same number of layers, Wav2Vec2-Conformer requires more parameters than Wav2Vec2, but also yields
+- For the same number of layers, Wav2Vec2-Conformer requires more parameters than Wav2Vec2, but also yields
an improved word error rate.
- Wav2Vec2-Conformer uses the same tokenizer and feature extractor as Wav2Vec2.
- Wav2Vec2-Conformer can use either no relative position embeddings, Transformer-XL-like position embeddings, or
diff --git a/docs/source/en/model_doc/wav2vec2.md b/docs/source/en/model_doc/wav2vec2.md
index 6c4772f90bc8..1f5f4a905767 100644
--- a/docs/source/en/model_doc/wav2vec2.md
+++ b/docs/source/en/model_doc/wav2vec2.md
@@ -80,13 +80,10 @@ model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-large-960h-lv60-self",
Below is an expected speedup diagram comparing the pure inference time between the native implementation in transformers of the `facebook/wav2vec2-large-960h-lv60-self` model and the flash-attention-2 and sdpa (scale-dot-product-attention) versions. . We show the average speedup obtained on the `librispeech_asr` `clean` validation split:
-
diff --git a/docs/source/en/model_doc/xcodec.md b/docs/source/en/model_doc/xcodec.md
index c4a0b92a26f6..957a74093484 100644
--- a/docs/source/en/model_doc/xcodec.md
+++ b/docs/source/en/model_doc/xcodec.md
@@ -33,9 +33,10 @@ The X-Codec model is a neural audio codec that integrates semantic information f
The abstract of the paper states the following:
-*Recent advancements in audio generation have been significantly propelled by the capabilities of Large Language Models (LLMs). The existing research on audio LLM has primarily focused on enhancing the architecture and scale of audio language models, as well as leveraging larger datasets, and generally, acoustic codecs, such as EnCodec, are used for audio tokenization. However, these codecs were originally designed for audio compression, which may lead to suboptimal performance in the context of audio LLM. Our research aims to address the shortcomings of current audio LLM codecs, particularly their challenges in maintaining semantic integrity in generated audio. For instance, existing methods like VALL-E, which condition acoustic token generation on text transcriptions, often suffer from content inaccuracies and elevated word error rates (WER) due to semantic misinterpretations of acoustic tokens, resulting in word skipping and errors. To overcome these issues, we propose a straightforward yet effective approach called X-Codec. X-Codec incorporates semantic features from a pre-trained semantic encoder before the Residual Vector Quantization (RVQ) stage and introduces a semantic reconstruction loss after RVQ. By enhancing the semantic ability of the codec, X-Codec significantly reduces WER in speech synthesis tasks and extends these benefits to non-speech applications, including music and sound generation. Our experiments in text-to-speech, music continuation, and text-to-sound tasks demonstrate that integrating semantic information substantially improves the overall performance of language models in audio generation.*
+*Recent advancements in audio generation have been significantly propelled by the capabilities of Large Language Models (LLMs). The existing research on audio LLM has primarily focused on enhancing the architecture and scale of audio language models, as well as leveraging larger datasets, and generally, acoustic codecs, such as EnCodec, are used for audio tokenization. However, these codecs were originally designed for audio compression, which may lead to suboptimal performance in the context of audio LLM. Our research aims to address the shortcomings of current audio LLM codecs, particularly their challenges in maintaining semantic integrity in generated audio. For instance, existing methods like VALL-E, which condition acoustic token generation on text transcriptions, often suffer from content inaccuracies and elevated word error rates (WER) due to semantic misinterpretations of acoustic tokens, resulting in word skipping and errors. To overcome these issues, we propose a straightforward yet effective approach called X-Codec. X-Codec incorporates semantic features from a pre-trained semantic encoder before the Residual Vector Quantization (RVQ) stage and introduces a semantic reconstruction loss after RVQ. By enhancing the semantic ability of the codec, X-Codec significantly reduces WER in speech synthesis tasks and extends these benefits to non-speech applications, including music and sound generation. Our experiments in text-to-speech, music continuation, and text-to-sound tasks demonstrate that integrating semantic information substantially improves the overall performance of language models in audio generation.*
Model cards:
+
- [xcodec-hubert-librispeech](https://huggingface.co/hf-audio/xcodec-hubert-librispeech) (for speech)
- [xcodec-wavlm-mls](https://huggingface.co/hf-audio/xcodec-wavlm-mls) (for speech)
- [xcodec-wavlm-more-data](https://huggingface.co/hf-audio/xcodec-wavlm-more-data) (for speech)
@@ -46,12 +47,11 @@ This model was contributed by [Manal El Aidouni](https://huggingface.co/Manel).
Demos can be found on this [page](https://x-codec-audio.github.io/).
-
-## Usage example
+## Usage example
Here is a quick example of how to encode and decode an audio using this model:
-```python
+```python
from datasets import load_dataset, Audio
from transformers import XcodecModel, AutoFeatureExtractor
dummy_dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
@@ -75,6 +75,7 @@ audio_values = decoder_outputs.audio_values
audio_values = model(inputs["input_values"]).audio_values
```
+
To listen to the original and reconstructed audio, run the snippet below and then open the generated `original.wav` and `reconstruction.wav` files in your music player to compare.
```python
@@ -88,15 +89,13 @@ sf.write("original.wav", original, sampling_rate)
sf.write("reconstruction.wav", reconstruction.T, sampling_rate)
```
-
## XcodecConfig
[[autodoc]] XcodecConfig
-
## XcodecModel
[[autodoc]] XcodecModel
- decode
- encode
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/xglm.md b/docs/source/en/model_doc/xglm.md
index d82bba7d23f9..9372b52af1f7 100644
--- a/docs/source/en/model_doc/xglm.md
+++ b/docs/source/en/model_doc/xglm.md
@@ -44,7 +44,6 @@ showing in particular that it enables cross-lingual in-context learning on some
on surface form robustness and adaptation to tasks that do not have a natural cloze form. Finally, we evaluate our models
in social value tasks such as hate speech detection in five languages and find it has limitations similar to comparable sized GPT-3 models.*
-
This model was contributed by [Suraj](https://huggingface.co/valhalla). The original code can be found [here](https://github.com/pytorch/fairseq/tree/main/examples/xglm).
## Resources
diff --git a/docs/source/en/model_doc/xlm-prophetnet.md b/docs/source/en/model_doc/xlm-prophetnet.md
index 4dad4c0afa78..fbf47d8c422a 100644
--- a/docs/source/en/model_doc/xlm-prophetnet.md
+++ b/docs/source/en/model_doc/xlm-prophetnet.md
@@ -41,7 +41,6 @@ You can do so by running the following command: `pip install -U transformers==4.
**DISCLAIMER:** If you see something strange, file a [Github Issue](https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title) and assign
@patrickvonplaten
-
## Overview
The XLM-ProphetNet model was proposed in [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training,](https://huggingface.co/papers/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei
diff --git a/docs/source/en/model_doc/xlm-roberta-xl.md b/docs/source/en/model_doc/xlm-roberta-xl.md
index 988107fdacc6..97dc6f1a7445 100644
--- a/docs/source/en/model_doc/xlm-roberta-xl.md
+++ b/docs/source/en/model_doc/xlm-roberta-xl.md
@@ -77,6 +77,7 @@ predicted_token = tokenizer.decode(predicted_token_id)
print(f"The predicted token is: {predicted_token}")
```
+
YOLOS architecture. Taken from the original paper.
-
> [!TIP]
> This model wasa contributed by [nielsr](https://huggingface.co/nielsr).
> Click on the YOLOS models in the right sidebar for more examples of how to apply YOLOS to different object detection tasks.
@@ -98,8 +96,8 @@ for score, label, box in zip(filtered_scores, filtered_labels, pixel_boxes):
-
## Notes
+
- Use [`YolosImageProcessor`] for preparing images (and optional targets) for the model. Contrary to [DETR](./detr), YOLOS doesn't require a `pixel_mask`.
## Resources
diff --git a/docs/source/en/model_doc/yoso.md b/docs/source/en/model_doc/yoso.md
index f07e5aba0827..211b0dcf8091 100644
--- a/docs/source/en/model_doc/yoso.md
+++ b/docs/source/en/model_doc/yoso.md
@@ -26,20 +26,20 @@ rendered properly in your Markdown viewer.
The YOSO model was proposed in [You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling](https://huggingface.co/papers/2111.09714)
by Zhanpeng Zeng, Yunyang Xiong, Sathya N. Ravi, Shailesh Acharya, Glenn Fung, Vikas Singh. YOSO approximates standard softmax self-attention
via a Bernoulli sampling scheme based on Locality Sensitive Hashing (LSH). In principle, all the Bernoulli random variables can be sampled with
-a single hash.
+a single hash.
The abstract from the paper is the following:
-*Transformer-based models are widely used in natural language processing (NLP). Central to the transformer model is
-the self-attention mechanism, which captures the interactions of token pairs in the input sequences and depends quadratically
-on the sequence length. Training such models on longer sequences is expensive. In this paper, we show that a Bernoulli sampling
-attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity of such models to linear.
-We bypass the quadratic cost by considering self-attention as a sum of individual tokens associated with Bernoulli random
-variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant).
-This leads to an efficient sampling scheme to estimate self-attention which relies on specific modifications of
-LSH (to enable deployment on GPU architectures). We evaluate our algorithm on the GLUE benchmark with standard 512 sequence
-length where we see favorable performance relative to a standard pretrained Transformer. On the Long Range Arena (LRA) benchmark,
-for evaluating performance on long sequences, our method achieves results consistent with softmax self-attention but with sizable
+*Transformer-based models are widely used in natural language processing (NLP). Central to the transformer model is
+the self-attention mechanism, which captures the interactions of token pairs in the input sequences and depends quadratically
+on the sequence length. Training such models on longer sequences is expensive. In this paper, we show that a Bernoulli sampling
+attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity of such models to linear.
+We bypass the quadratic cost by considering self-attention as a sum of individual tokens associated with Bernoulli random
+variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant).
+This leads to an efficient sampling scheme to estimate self-attention which relies on specific modifications of
+LSH (to enable deployment on GPU architectures). We evaluate our algorithm on the GLUE benchmark with standard 512 sequence
+length where we see favorable performance relative to a standard pretrained Transformer. On the Long Range Arena (LRA) benchmark,
+for evaluating performance on long sequences, our method achieves results consistent with softmax self-attention but with sizable
speed-ups and memory savings and often outperforms other efficient self-attention methods. Our code is available at this https URL*
This model was contributed by [novice03](https://huggingface.co/novice03). The original code can be found [here](https://github.com/mlpen/YOSO).
@@ -50,12 +50,12 @@ This model was contributed by [novice03](https://huggingface.co/novice03). The o
in parallel on a GPU.
- The kernels provide a `fast_hash` function, which approximates the random projections of the queries and keys using the Fast Hadamard Transform. Using these
hash codes, the `lsh_cumulation` function approximates self-attention via LSH-based Bernoulli sampling.
-- To use the custom kernels, the user should set `config.use_expectation = False`. To ensure that the kernels are compiled successfully,
-the user must install the correct version of PyTorch and cudatoolkit. By default, `config.use_expectation = True`, which uses YOSO-E and
+- To use the custom kernels, the user should set `config.use_expectation = False`. To ensure that the kernels are compiled successfully,
+the user must install the correct version of PyTorch and cudatoolkit. By default, `config.use_expectation = True`, which uses YOSO-E and
does not require compiling CUDA kernels.
+alt="drawing" width="600"/>
YOSO Attention Algorithm. Taken from the original paper.
@@ -99,4 +99,4 @@ alt="drawing" width="600"/>
## YosoForQuestionAnswering
[[autodoc]] YosoForQuestionAnswering
- - forward
\ No newline at end of file
+ - forward
diff --git a/docs/source/en/model_doc/zamba.md b/docs/source/en/model_doc/zamba.md
index bb9740807703..847f0532e2a7 100644
--- a/docs/source/en/model_doc/zamba.md
+++ b/docs/source/en/model_doc/zamba.md
@@ -24,7 +24,6 @@ rendered properly in your Markdown viewer.
This model was contributed by [pglo](https://huggingface.co/pglo).
-
## Model details
Zamba-7B-v1 is a hybrid between state-space models (Specifically [Mamba](https://github.com/state-spaces/mamba)) and transformer, and was trained using next-token prediction. Zamba uses a shared transformer layer after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba-7B-v1 was pre-trained on 1T tokens of text and code data.
@@ -33,23 +32,24 @@ Zamba-7B-v1 is a hybrid between state-space models (Specifically [Mamba](https:/
## Quick start
-
### Presequities
Zamba requires you use `transformers` version 4.46.0 or higher:
+
```bash
pip install transformers>=4.45.0
```
In order to run optimized Mamba implementations, you first need to install `mamba-ssm` and `causal-conv1d`:
+
```bash
pip install mamba-ssm causal-conv1d>=1.2.0
```
+
You also have to have the model on a CUDA device.
You can run the model not using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify `use_mamba_kernels=False` when loading the model.
-
## Inference
```python
@@ -66,39 +66,33 @@ outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))
```
-
## Model card
The model cards can be found at:
-* [Zamba-7B](https://huggingface.co/Zyphra/Zamba-7B-v1)
+* [Zamba-7B](https://huggingface.co/Zyphra/Zamba-7B-v1)
## Issues
For issues with model output, or community discussion, please use the Hugging Face community [forum](https://huggingface.co/Zyphra/Zamba-7B-v1/discussions)
-
## License
The model weights are open-sourced via an Apache 2.0 license.
-
## ZambaConfig
[[autodoc]] ZambaConfig
-
## ZambaModel
[[autodoc]] ZambaModel
- forward
-
## ZambaForCausalLM
[[autodoc]] ZambaForCausalLM
- forward
-
## ZambaForSequenceClassification
[[autodoc]] transformers.ZambaForSequenceClassification
diff --git a/docs/source/en/model_doc/zamba2.md b/docs/source/en/model_doc/zamba2.md
index 1d911a59c277..c9d3d3d1de75 100644
--- a/docs/source/en/model_doc/zamba2.md
+++ b/docs/source/en/model_doc/zamba2.md
@@ -26,19 +26,18 @@ rendered properly in your Markdown viewer.
This model was contributed by [pglo](https://huggingface.co/pglo).
-
## Model details
-[Zamba2-1.2B](https://www.zyphra.com/post/zamba2-mini), [Zamba2-2.7B](https://www.zyphra.com/post/zamba2-small) and [Zamba2-7B](https://www.zyphra.com/post/zamba2-7b) are hybrid models combining state-space models (Specifically [Mamba](https://github.com/state-spaces/mamba)) and transformer, and were trained using next-token prediction. Zamba2 uses shared transformer layers after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba2-1.2B, Zamba2-2.7B and Zamba2-7B were pre-trained on 2T and 3T tokens, respectively.
+[Zamba2-1.2B](https://www.zyphra.com/post/zamba2-mini), [Zamba2-2.7B](https://www.zyphra.com/post/zamba2-small) and [Zamba2-7B](https://www.zyphra.com/post/zamba2-7b) are hybrid models combining state-space models (Specifically [Mamba2](https://github.com/state-spaces/mamba)) and transformer, and were trained using next-token prediction. Zamba2 uses shared transformer layers after every 6 mamba blocks. It uses the [Mistral v0.1 tokenizer](https://huggingface.co/mistralai/Mistral-7B-v0.1). We came to this architecture after a series of ablations at small scales. Zamba2-1.2B, Zamba2-2.7B and Zamba2-7B were pre-trained on 2T and 3T tokens, respectively.
@@ -97,6 +96,7 @@ Image.fromarray(depth.astype("uint8"))
## Notes
- In the [original implementation](https://github.com/isl-org/ZoeDepth/blob/edb6daf45458569e24f50250ef1ed08c015f17a7/zoedepth/models/depth_model.py#L131) ZoeDepth performs inference on both the original and flipped images and averages the results. The `post_process_depth_estimation` function handles this by passing the flipped outputs to the optional `outputs_flipped` argument as shown below.
+
```py
with torch.no_grad():
outputs = model(pixel_values)
@@ -107,8 +107,9 @@ Image.fromarray(depth.astype("uint8"))
outputs_flipped=outputs_flipped,
)
```
-
+
## Resources
+
- Refer to this [notebook](https://github.com/NielsRogge/Transformers-Tutorials/tree/master/ZoeDepth) for an inference example.
## ZoeDepthConfig
diff --git a/docs/source/en/model_memory_anatomy.md b/docs/source/en/model_memory_anatomy.md
index 7ef53f40566e..f0a215b05c1b 100644
--- a/docs/source/en/model_memory_anatomy.md
+++ b/docs/source/en/model_memory_anatomy.md
@@ -16,24 +16,23 @@ limitations under the License.
# Model training anatomy
-To understand performance optimization techniques that one can apply to improve efficiency of model training
-speed and memory utilization, it's helpful to get familiar with how GPU is utilized during training, and how compute
+To understand performance optimization techniques that one can apply to improve efficiency of model training
+speed and memory utilization, it's helpful to get familiar with how GPU is utilized during training, and how compute
intensity varies depending on an operation performed.
-Let's start by exploring a motivating example of GPU utilization and the training run of a model. For the demonstration,
-we'll need to install a few libraries:
+Let's start by exploring a motivating example of GPU utilization and the training run of a model. For the demonstration,
+we'll need to install a few libraries:
```bash
-pip install transformers datasets accelerate nvidia-ml-py3
+pip install transformers datasets accelerate nvidia-ml-py
```
-The `nvidia-ml-py3` library allows us to monitor the memory usage of the models from within Python. You might be familiar
+The `nvidia-ml-py` library allows us to monitor the memory usage of the models from within Python. You might be familiar
with the `nvidia-smi` command in the terminal - this library allows to access the same information in Python directly.
-Then, we create some dummy data: random token IDs between 100 and 30000 and binary labels for a classifier.
+Then, we create some dummy data: random token IDs between 100 and 30000 and binary labels for a classifier.
In total, we get 512 sequences each with length 512 and store them in a [`~datasets.Dataset`] with PyTorch format.
-
```py
>>> import numpy as np
>>> from datasets import Dataset
@@ -74,9 +73,9 @@ Let's verify that we start with a free GPU memory:
GPU memory occupied: 0 MB.
```
-That looks good: the GPU memory is not occupied as we would expect before we load any models. If that's not the case on
-your machine make sure to stop all processes that are using GPU memory. However, not all free GPU memory can be used by
-the user. When a model is loaded to the GPU the kernels are also loaded, which can take up 1-2GB of memory. To see how
+That looks good: the GPU memory is not occupied as we would expect before we load any models. If that's not the case on
+your machine make sure to stop all processes that are using GPU memory. However, not all free GPU memory can be used by
+the user. When a model is loaded to the GPU the kernels are also loaded, which can take up 1-2GB of memory. To see how
much it is we load a tiny tensor into the GPU which triggers the kernels to be loaded as well.
```py
@@ -92,10 +91,9 @@ We see that the kernels alone take up 1.3GB of GPU memory. Now let's see how muc
## Load Model
-First, we load the `google-bert/bert-large-uncased` model. We load the model weights directly to the GPU so that we can check
+First, we load the `google-bert/bert-large-uncased` model. We load the model weights directly to the GPU so that we can check
how much space just the weights use.
-
```py
>>> from transformers import AutoModelForSequenceClassification
@@ -105,17 +103,16 @@ how much space just the weights use.
GPU memory occupied: 2631 MB.
```
-We can see that the model weights alone take up 1.3 GB of GPU memory. The exact number depends on the specific
-GPU you are using. Note that on newer GPUs a model can sometimes take up more space since the weights are loaded in an
-optimized fashion that speeds up the usage of the model. Now we can also quickly check if we get the same result
+We can see that the model weights alone take up 1.3 GB of GPU memory. The exact number depends on the specific
+GPU you are using. Note that on newer GPUs a model can sometimes take up more space since the weights are loaded in an
+optimized fashion that speeds up the usage of the model. Now we can also quickly check if we get the same result
as with `nvidia-smi` CLI:
-
```bash
nvidia-smi
```
-```bash
+```text
Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03 Driver Version: 460.91.03 CUDA Version: 11.2 |
@@ -138,8 +135,8 @@ Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
```
-We get the same number as before and you can also see that we are using a V100 GPU with 16GB of memory. So now we can
-start training the model and see how the GPU memory consumption changes. First, we set up a few standard training
+We get the same number as before and you can also see that we are using a V100 GPU with 16GB of memory. So now we can
+start training the model and see how the GPU memory consumption changes. First, we set up a few standard training
arguments:
```py
@@ -154,7 +151,7 @@ default_args = {
@@ -158,7 +159,7 @@ another image to illustrate:
Photo by [Denys Nevozhai](https://unsplash.com/@dnevozhai).
-
+
Textual and image prompts can be passed to the model's processor as a single list to create appropriate inputs.
```py
@@ -178,12 +179,12 @@ This is an image of the Eiffel Tower in Paris, France.
## Few-shot prompting
-While IDEFICS demonstrates great zero-shot results, your task may require a certain format of the caption, or come with
+While IDEFICS demonstrates great zero-shot results, your task may require a certain format of the caption, or come with
other restrictions or requirements that increase task's complexity. Few-shot prompting can be used to enable in-context learning.
-By providing examples in the prompt, you can steer the model to generate results that mimic the format of given examples.
+By providing examples in the prompt, you can steer the model to generate results that mimic the format of given examples.
-Let's use the previous image of the Eiffel Tower as an example for the model and build a prompt that demonstrates to the model
-that in addition to learning what the object in an image is, we would also like to get some interesting information about it.
+Let's use the previous image of the Eiffel Tower as an example for the model and build a prompt that demonstrates to the model
+that in addition to learning what the object in an image is, we would also like to get some interesting information about it.
Then, let's see, if we can get the same response format for an image of the Statue of Liberty:
@@ -213,24 +214,24 @@ User: Describe this image.
Assistant: An image of the Statue of Liberty. Fun fact: the Statue of Liberty is 151 feet tall.
```
-Notice that just from a single example (i.e., 1-shot) the model has learned how to perform the task. For more complex tasks,
+Notice that just from a single example (i.e., 1-shot) the model has learned how to perform the task. For more complex tasks,
feel free to experiment with a larger number of examples (e.g., 3-shot, 5-shot, etc.).
## Visual question answering
-Visual Question Answering (VQA) is the task of answering open-ended questions based on an image. Similar to image
-captioning it can be used in accessibility applications, but also in education (reasoning about visual materials), customer
+Visual Question Answering (VQA) is the task of answering open-ended questions based on an image. Similar to image
+captioning it can be used in accessibility applications, but also in education (reasoning about visual materials), customer
service (questions about products based on images), and image retrieval.
-Let's get a new image for this task:
+Let's get a new image for this task:
-Photo by [Jarritos Mexican Soda](https://unsplash.com/@jarritos).
+Photo by [Jarritos Mexican Soda](https://unsplash.com/@jarritos).
-You can steer the model from image captioning to visual question answering by prompting it with appropriate instructions:
+You can steer the model from image captioning to visual question answering by prompting it with appropriate instructions:
```py
>>> prompt = [
@@ -251,11 +252,11 @@ Instruction: Provide an answer to the question. Use the image to answer.
## Image classification
-IDEFICS is capable of classifying images into different categories without being explicitly trained on data containing
-labeled examples from those specific categories. Given a list of categories and using its image and text understanding
-capabilities, the model can infer which category the image likely belongs to.
+IDEFICS is capable of classifying images into different categories without being explicitly trained on data containing
+labeled examples from those specific categories. Given a list of categories and using its image and text understanding
+capabilities, the model can infer which category the image likely belongs to.
-Say, we have this image of a vegetable stand:
+Say, we have this image of a vegetable stand:
@@ -286,10 +287,10 @@ In the example above we instruct the model to classify the image into a single c
## Image-guided text generation
-For more creative applications, you can use image-guided text generation to generate text based on an image. This can be
-useful to create descriptions of products, ads, descriptions of a scene, etc.
+For more creative applications, you can use image-guided text generation to generate text based on an image. This can be
+useful to create descriptions of products, ads, descriptions of a scene, etc.
-Let's prompt IDEFICS to write a story based on a simple image of a red door:
+Let's prompt IDEFICS to write a story based on a simple image of a red door:
@@ -333,14 +334,14 @@ Looks like IDEFICS noticed the pumpkin on the doorstep and went with a spooky Ha
-For longer outputs like this, you will greatly benefit from tweaking the text generation strategy. This can help
-you significantly improve the quality of the generated output. Check out [Text generation strategies](../generation_strategies)
-to learn more.
+For longer outputs like this, you will greatly benefit from tweaking the text generation strategy. This can help
+you significantly improve the quality of the generated output. Check out [Text generation strategies](../generation_strategies)
+to learn more.
## Running inference in batch mode
-All of the earlier sections illustrated IDEFICS for a single example. In a very similar fashion, you can run inference
+All of the earlier sections illustrated IDEFICS for a single example. In a very similar fashion, you can run inference
for a batch of examples by passing a list of prompts:
```py
@@ -375,13 +376,13 @@ This is an image of a vegetable stand.
## IDEFICS instruct for conversational use
-For conversational use cases, you can find fine-tuned instructed versions of the model on the 🤗 Hub:
+For conversational use cases, you can find fine-tuned instructed versions of the model on the 🤗 Hub:
`HuggingFaceM4/idefics-80b-instruct` and `HuggingFaceM4/idefics-9b-instruct`.
-These checkpoints are the result of fine-tuning the respective base models on a mixture of supervised and instruction
+These checkpoints are the result of fine-tuning the respective base models on a mixture of supervised and instruction
fine-tuning datasets, which boosts the downstream performance while making the models more usable in conversational settings.
-The use and prompting for the conversational use is very similar to using the base models:
+The use and prompting for the conversational use is very similar to using the base models:
```py
>>> import torch
diff --git a/docs/source/en/tasks/image_captioning.md b/docs/source/en/tasks/image_captioning.md
index f9716f29a204..4b4b3ba5fa36 100644
--- a/docs/source/en/tasks/image_captioning.md
+++ b/docs/source/en/tasks/image_captioning.md
@@ -14,7 +14,6 @@ rendered properly in your Markdown viewer.
-->
-
# Image captioning
[[open-in-colab]]
@@ -26,7 +25,7 @@ helps to improve content accessibility for people by describing images to them.
This guide will show you how to:
* Fine-tune an image captioning model.
-* Use the fine-tuned model for inference.
+* Use the fine-tuned model for inference.
Before you begin, make sure you have all the necessary libraries installed:
@@ -37,7 +36,6 @@ pip install jiwer -q
We encourage you to log in to your Hugging Face account so you can upload and share your model with the community. When prompted, enter your token to log in:
-
```python
from huggingface_hub import notebook_login
@@ -47,8 +45,7 @@ notebook_login()
## Load the Pokémon BLIP captions dataset
Use the 🤗 Dataset library to load a dataset that consists of {image-caption} pairs. To create your own image captioning dataset
-in PyTorch, you can follow [this notebook](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/GIT/Fine_tune_GIT_on_an_image_captioning_dataset.ipynb).
-
+in PyTorch, you can follow [this notebook](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/GIT/Fine_tune_GIT_on_an_image_captioning_dataset.ipynb).
```python
from datasets import load_dataset
@@ -56,6 +53,7 @@ from datasets import load_dataset
ds = load_dataset("lambdalabs/pokemon-blip-captions")
ds
```
+
```bash
DatasetDict({
train: Dataset({
@@ -69,12 +67,11 @@ The dataset has two features, `image` and `text`.
-Many image captioning datasets contain multiple captions per image. In those cases, a common strategy is to randomly sample a caption amongst the available ones during training.
+Many image captioning datasets contain multiple captions per image. In those cases, a common strategy is to randomly sample a caption amongst the available ones during training.
-Split the dataset’s train split into a train and test set with the [`~datasets.Dataset.train_test_split`] method:
-
+Split the dataset's train split into a train and test set with the [`~datasets.Dataset.train_test_split`] method:
```python
ds = ds["train"].train_test_split(test_size=0.1)
@@ -82,8 +79,7 @@ train_ds = ds["train"]
test_ds = ds["test"]
```
-Let's visualize a couple of samples from the training set.
-
+Let's visualize a couple of samples from the training set.
```python
from textwrap import wrap
@@ -106,7 +102,7 @@ sample_images_to_visualize = [np.array(train_ds[i]["image"]) for i in range(5)]
sample_captions = [train_ds[i]["text"] for i in range(5)]
plot_images(sample_images_to_visualize, sample_captions)
```
-
+
@@ -115,7 +111,7 @@ plot_images(sample_images_to_visualize, sample_captions)
Since the dataset has two modalities (image and text), the pre-processing pipeline will preprocess images and the captions.
-To do so, load the processor class associated with the model you are about to fine-tune.
+To do so, load the processor class associated with the model you are about to fine-tune.
```python
from transformers import AutoProcessor
@@ -124,7 +120,7 @@ checkpoint = "microsoft/git-base"
processor = AutoProcessor.from_pretrained(checkpoint)
```
-The processor will internally pre-process the image (which includes resizing, and pixel scaling) and tokenize the caption.
+The processor will internally pre-process the image (which includes resizing, and pixel scaling) and tokenize the caption.
```python
def transforms(example_batch):
@@ -139,13 +135,12 @@ train_ds.set_transform(transforms)
test_ds.set_transform(transforms)
```
-With the dataset ready, you can now set up the model for fine-tuning.
+With the dataset ready, you can now set up the model for fine-tuning.
## Load a base model
Load the ["microsoft/git-base"](https://huggingface.co/microsoft/git-base) into a [`AutoModelForCausalLM`](https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForCausalLM) object.
-
```python
from transformers import AutoModelForCausalLM
@@ -154,10 +149,9 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint)
## Evaluate
-Image captioning models are typically evaluated with the [Rouge Score](https://huggingface.co/spaces/evaluate-metric/rouge) or [Word Error Rate](https://huggingface.co/spaces/evaluate-metric/wer). For this guide, you will use the Word Error Rate (WER).
-
-We use the 🤗 Evaluate library to do so. For potential limitations and other gotchas of the WER, refer to [this guide](https://huggingface.co/spaces/evaluate-metric/wer).
+Image captioning models are typically evaluated with the [Rouge Score](https://huggingface.co/spaces/evaluate-metric/rouge) or [Word Error Rate](https://huggingface.co/spaces/evaluate-metric/wer). For this guide, you will use the Word Error Rate (WER).
+We use the 🤗 Evaluate library to do so. For potential limitations and other gotchas of the WER, refer to [this guide](https://huggingface.co/spaces/evaluate-metric/wer).
```python
from evaluate import load
@@ -177,11 +171,10 @@ def compute_metrics(eval_pred):
## Train!
-Now, you are ready to start fine-tuning the model. You will use the 🤗 [`Trainer`] for this.
+Now, you are ready to start fine-tuning the model. You will use the 🤗 [`Trainer`] for this.
First, define the training arguments using [`TrainingArguments`].
-
```python
from transformers import TrainingArguments, Trainer
@@ -208,7 +201,7 @@ training_args = TrainingArguments(
)
```
-Then pass them along with the datasets and the model to 🤗 Trainer.
+Then pass them along with the datasets and the model to 🤗 Trainer.
```python
trainer = Trainer(
@@ -222,7 +215,7 @@ trainer = Trainer(
To start training, simply call [`~Trainer.train`] on the [`Trainer`] object.
-```python
+```python
trainer.train()
```
@@ -230,7 +223,6 @@ You should see the training loss drop smoothly as training progresses.
Once training is completed, share your model to the Hub with the [`~Trainer.push_to_hub`] method so everyone can use your model:
-
```python
trainer.push_to_hub()
```
@@ -239,7 +231,6 @@ trainer.push_to_hub()
Take a sample image from `test_ds` to test the model.
-
```python
from PIL import Image
import requests
@@ -252,7 +243,7 @@ image
-
+
Prepare image for the model.
```python
@@ -263,13 +254,14 @@ inputs = processor(images=image, return_tensors="pt").to(device)
pixel_values = inputs.pixel_values
```
-Call [`generate`] and decode the predictions.
+Call [`generate`] and decode the predictions.
```python
generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(generated_caption)
```
+
```bash
a drawing of a pink and blue pokemon
```
diff --git a/docs/source/en/tasks/image_classification.md b/docs/source/en/tasks/image_classification.md
index 39b013f129cc..4754a91bd482 100644
--- a/docs/source/en/tasks/image_classification.md
+++ b/docs/source/en/tasks/image_classification.md
@@ -175,7 +175,6 @@ Your `compute_metrics` function is ready to go now, and you'll return to it when
## Train
-
If you aren't familiar with finetuning a model with the [`Trainer`], take a look at the basic tutorial [here](../training#train-with-pytorch-trainer)!
@@ -238,7 +237,6 @@ Once training is completed, share your model to the Hub with the [`~transformers
>>> trainer.push_to_hub()
```
-
For a more in-depth example of how to finetune a model for image classification, take a look at the corresponding [PyTorch notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/image_classification.ipynb).
diff --git a/docs/source/en/tasks/image_feature_extraction.md b/docs/source/en/tasks/image_feature_extraction.md
index 455a2b425d41..e08ba89e4dd8 100644
--- a/docs/source/en/tasks/image_feature_extraction.md
+++ b/docs/source/en/tasks/image_feature_extraction.md
@@ -27,7 +27,7 @@ In this guide, you will:
## Image Similarity using `image-feature-extraction` Pipeline
-We have two images of cats sitting on top of fish nets, one of them is generated.
+We have two images of cats sitting on top of fish nets, one of them is generated.
```python
from PIL import Image
@@ -66,7 +66,7 @@ print(outputs)
# [[[-0.03909236937761307, 0.43381670117378235, -0.06913255900144577,
```
-To get the similarity score, we need to pass them to a similarity function.
+To get the similarity score, we need to pass them to a similarity function.
```python
from torch.nn.functional import cosine_similarity
@@ -131,4 +131,3 @@ print(similarity_score)
# tensor([0.6061], device='cuda:0', grad_fn=)
```
-
diff --git a/docs/source/en/tasks/image_text_to_text.md b/docs/source/en/tasks/image_text_to_text.md
index b34f4edf90f6..8820a534030c 100644
--- a/docs/source/en/tasks/image_text_to_text.md
+++ b/docs/source/en/tasks/image_text_to_text.md
@@ -23,6 +23,7 @@ Image-text-to-text models, also known as vision language models (VLMs), are lang
In this guide, we provide a brief overview of VLMs and show how to use them with Transformers for inference.
To begin with, there are multiple types of VLMs:
+
- base models used for fine-tuning
- chat fine-tuned models for conversation
- instruction fine-tuned models
@@ -63,7 +64,6 @@ The image inputs look like the following.
-
```python
from PIL import Image
import requests
@@ -76,7 +76,6 @@ images = [Image.open(requests.get(img_urls[0], stream=True).raw),
Below is an example of the chat template. We can feed conversation turns and the last message as an input by appending it at the end of the template.
-
```python
messages = [
{
@@ -207,7 +206,6 @@ We can use [text streaming](./generation_strategies#streaming) for a better gene
Assume we have an application that keeps chat history and takes in the new user input. We will preprocess the inputs as usual and initialize [`TextIteratorStreamer`] to handle the generation in a separate thread. This allows you to stream the generated text tokens in real-time. Any generation arguments can be passed to [`TextIteratorStreamer`].
-
```python
import time
from transformers import TextIteratorStreamer
diff --git a/docs/source/en/tasks/image_to_image.md b/docs/source/en/tasks/image_to_image.md
index da6a57ac9aa9..55380e9b0d1e 100644
--- a/docs/source/en/tasks/image_to_image.md
+++ b/docs/source/en/tasks/image_to_image.md
@@ -18,9 +18,10 @@ rendered properly in your Markdown viewer.
[[open-in-colab]]
-Image-to-Image task is the task where an application receives an image and outputs another image. This has various subtasks, including image enhancement (super resolution, low light enhancement, deraining and so on), image inpainting, and more.
+Image-to-Image task is the task where an application receives an image and outputs another image. This has various subtasks, including image enhancement (super resolution, low light enhancement, deraining and so on), image inpainting, and more.
This guide will show you how to:
+
- Use an image-to-image pipeline for super resolution task,
- Run image-to-image models for same task without a pipeline.
@@ -32,7 +33,7 @@ Let's begin by installing the necessary libraries.
pip install transformers
```
-We can now initialize the pipeline with a [Swin2SR model](https://huggingface.co/caidas/swin2SR-lightweight-x2-64). We can then infer with the pipeline by calling it with an image. As of now, only [Swin2SR models](https://huggingface.co/models?sort=trending&search=swin2sr) are supported in this pipeline.
+We can now initialize the pipeline with a [Swin2SR model](https://huggingface.co/caidas/swin2SR-lightweight-x2-64). We can then infer with the pipeline by calling it with an image. As of now, only [Swin2SR models](https://huggingface.co/models?sort=trending&search=swin2sr) are supported in this pipeline.
```python
from transformers import pipeline, infer_device
@@ -53,19 +54,22 @@ image = Image.open(requests.get(url, stream=True).raw)
print(image.size)
```
+
```bash
# (532, 432)
```
+
-We can now do inference with the pipeline. We will get an upscaled version of the cat image.
+We can now do inference with the pipeline. We will get an upscaled version of the cat image.
```python
upscaled = pipe(image)
print(upscaled.size)
```
+
```bash
# (1072, 880)
```
@@ -79,7 +83,7 @@ model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-lightweig
processor = Swin2SRImageProcessor("caidas/swin2SR-lightweight-x2-64")
```
-`pipeline` abstracts away the preprocessing and postprocessing steps that we have to do ourselves, so let's preprocess the image. We will pass the image to the processor and then move the pixel values to GPU.
+`pipeline` abstracts away the preprocessing and postprocessing steps that we have to do ourselves, so let's preprocess the image. We will pass the image to the processor and then move the pixel values to GPU.
```python
pixel_values = processor(image, return_tensors="pt").pixel_values
@@ -96,9 +100,10 @@ import torch
with torch.no_grad():
outputs = model(pixel_values)
```
-Output is an object of type `ImageSuperResolutionOutput` that looks like below 👇
-```
+Output is an object of type `ImageSuperResolutionOutput` that looks like below 👇
+
+```text
(loss=None, reconstruction=tensor([[[[0.8270, 0.8269, 0.8275, ..., 0.7463, 0.7446, 0.7453],
[0.8287, 0.8278, 0.8283, ..., 0.7451, 0.7448, 0.7457],
[0.8280, 0.8273, 0.8269, ..., 0.7447, 0.7446, 0.7452],
@@ -108,6 +113,7 @@ Output is an object of type `ImageSuperResolutionOutput` that looks like below
[0.5927, 0.5914, 0.5922, ..., 0.0664, 0.0694, 0.0718]]]],
device='cuda:0'), hidden_states=None, attentions=None)
```
+
We need to get the `reconstruction` and post-process it for visualization. Let's see how it looks like.
```python
@@ -128,6 +134,7 @@ output = np.moveaxis(output, source=0, destination=-1)
output = (output * 255.0).round().astype(np.uint8)
Image.fromarray(output)
```
+
diff --git a/docs/source/en/tasks/keypoint_detection.md b/docs/source/en/tasks/keypoint_detection.md
index 3a5871d01a2b..c850c67ae153 100644
--- a/docs/source/en/tasks/keypoint_detection.md
+++ b/docs/source/en/tasks/keypoint_detection.md
@@ -18,7 +18,7 @@ rendered properly in your Markdown viewer.
[[open-in-colab]]
-Keypoint detection identifies and locates specific points of interest within an image. These keypoints, also known as landmarks, represent meaningful features of objects, such as facial features or object parts. These models take an image input and return the following outputs:
+Keypoint detection identifies and locates specific points of interest within an image. These keypoints, also known as landmarks, represent meaningful features of objects, such as facial features or object parts. These models take an image input and return the following outputs:
- **Keypoints and Scores**: Points of interest and their confidence scores.
- **Descriptors**: A representation of the image region surrounding each keypoint, capturing its texture, gradient, orientation and other properties.
@@ -36,15 +36,14 @@ model = SuperPointForKeypointDetection.from_pretrained("magic-leap-community/sup
Let's test the model on the images below.
-SAM serves as a powerful foundation model for segmentation as it has large data coverage. It is trained on
-[SA-1B](https://ai.meta.com/datasets/segment-anything/), a dataset with 1 million images and 1.1 billion masks.
+SAM serves as a powerful foundation model for segmentation as it has large data coverage. It is trained on
+[SA-1B](https://ai.meta.com/datasets/segment-anything/), a dataset with 1 million images and 1.1 billion masks.
In this guide, you will learn how to:
+
- Infer in segment everything mode with batching,
- Infer in point prompting mode,
- Infer in box prompting mode.
@@ -114,7 +116,6 @@ Below is the original image in grayscale with colorful maps overlaid. Very impre
-
## Model Inference
### Point Prompting
@@ -132,7 +133,7 @@ processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
To do point prompting, pass the input point to the processor, then take the processor output
and pass it to the model for inference. To post-process the model output, pass the outputs and
-`original_sizes` and `reshaped_input_sizes` we take from the processor's initial output. We need to pass these
+`original_sizes` and `reshaped_input_sizes` we take from the processor's initial output. We need to pass these
since the processor resizes the image, and the output needs to be extrapolated.
```python
@@ -143,6 +144,7 @@ with torch.no_grad():
outputs = model(**inputs)
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
```
+
We can visualize the three masks in the `masks` output.
```python
@@ -177,10 +179,9 @@ plt.show()
### Box Prompting
You can also do box prompting in a similar fashion to point prompting. You can simply pass the input box in the format of a list
-`[x_min, y_min, x_max, y_max]` format along with the image to the `processor`. Take the processor output and directly pass it
+`[x_min, y_min, x_max, y_max]` format along with the image to the `processor`. Take the processor output and directly pass it
to the model, then post-process the output again.
-
```python
# bounding box around the bee
box = [2350, 1600, 2850, 2100]
@@ -219,7 +220,7 @@ plt.show()
-You can see the inference output below.
+You can see the inference output below.
```python
fig, ax = plt.subplots()
@@ -233,4 +234,3 @@ plt.show()
@@ -286,10 +287,10 @@ In the example above we instruct the model to classify the image into a single c
## Image-guided text generation
-For more creative applications, you can use image-guided text generation to generate text based on an image. This can be
-useful to create descriptions of products, ads, descriptions of a scene, etc.
+For more creative applications, you can use image-guided text generation to generate text based on an image. This can be
+useful to create descriptions of products, ads, descriptions of a scene, etc.
-Let's prompt IDEFICS to write a story based on a simple image of a red door:
+Let's prompt IDEFICS to write a story based on a simple image of a red door:
@@ -333,14 +334,14 @@ Looks like IDEFICS noticed the pumpkin on the doorstep and went with a spooky Ha
-
- 
-
```python
import torch
from PIL import Image
@@ -93,7 +92,7 @@ image_sizes = [(image.size[1], image.size[0]) for image in images]
outputs = processor.post_process_keypoint_detection(outputs, image_sizes)
```
-The outputs are now a list of dictionaries where each dictionary is a processed output of keypoints, scores and descriptors.
+The outputs are now a list of dictionaries where each dictionary is a processed output of keypoints, scores and descriptors.
```python
[{'keypoints': tensor([[ 226, 57],
@@ -144,11 +143,10 @@ for i in range(len(images)):
Below you can see the outputs.
-
- 
- 
-
diff --git a/docs/source/en/tasks/keypoint_matching.md b/docs/source/en/tasks/keypoint_matching.md
index f7065f315211..7183c308c27a 100644
--- a/docs/source/en/tasks/keypoint_matching.md
+++ b/docs/source/en/tasks/keypoint_matching.md
@@ -34,15 +34,15 @@ model = AutoModelForKeypointMatching.from_pretrained("zju-community/matchanythin
Load two images that have the same object of interest. The second photo is taken a second apart, it's colors are edited, and it is further cropped and rotated.
-
-
- 
-```python
+```python
from transformers.image_utils import load_image
image1 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg")
image2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee_edited.jpg")
@@ -69,7 +69,7 @@ print(outputs)
Here's the outputs.
-```
+```text
[{'keypoints0': tensor([[4514, 550],
[4813, 683],
[1972, 1547],
@@ -82,16 +82,16 @@ Here's the outputs.
[1521, 2560]], dtype=torch.int32),
'matching_scores': tensor([0.2189, 0.2073, 0.2414, ...
])}]
-```
+```
We have trimmed the output but there's 401 matches!
```python
len(outputs[0]["keypoints0"])
# 401
-```
+```
-We can visualize them using the processor's [`~EfficientLoFTRImageProcessor.visualize_keypoint_matching`] method.
+We can visualize them using the processor's [`~EfficientLoFTRImageProcessor.visualize_keypoint_matching`] method.
```python
plot_images = processor.visualize_keypoint_matching(images, outputs)
@@ -100,7 +100,7 @@ plot_images
-Optionally, you can use the [`Pipeline`] API and set the task to `keypoint-matching`.
+Optionally, you can use the [`Pipeline`] API and set the task to `keypoint-matching`.
```python
from transformers import pipeline
diff --git a/docs/source/en/tasks/knowledge_distillation_for_image_classification.md b/docs/source/en/tasks/knowledge_distillation_for_image_classification.md
index 7c4a684d3c05..d4b3dd8511df 100644
--- a/docs/source/en/tasks/knowledge_distillation_for_image_classification.md
+++ b/docs/source/en/tasks/knowledge_distillation_for_image_classification.md
@@ -52,7 +52,6 @@ processed_datasets = dataset.map(process, batched=True)
Essentially, we want the student model (a randomly initialized MobileNet) to mimic the teacher model (fine-tuned vision transformer). To achieve this, we first get the logits output from the teacher and the student. Then, we divide each of them by the parameter `temperature` which controls the importance of each soft target. A parameter called `lambda` weighs the importance of the distillation loss. In this example, we will use `temperature=5` and `lambda=0.5`. We will use the Kullback-Leibler Divergence loss to compute the divergence between the student and teacher. Given two data P and Q, KL Divergence explains how much extra information we need to represent P using Q. If two are identical, their KL divergence is zero, as there's no other information needed to explain P from Q. Thus, in the context of knowledge distillation, KL divergence is useful.
-
```python
from transformers import TrainingArguments, Trainer, infer_device
import torch
diff --git a/docs/source/en/tasks/mask_generation.md b/docs/source/en/tasks/mask_generation.md
index 5f66e68c2452..817cb9819e7d 100644
--- a/docs/source/en/tasks/mask_generation.md
+++ b/docs/source/en/tasks/mask_generation.md
@@ -16,24 +16,26 @@ rendered properly in your Markdown viewer.
# Mask Generation
-Mask generation is the task of generating semantically meaningful masks for an image.
-This task is very similar to [image segmentation](semantic_segmentation), but many differences exist. Image segmentation models are trained on labeled datasets and are limited to the classes they have seen during training; they return a set of masks and corresponding classes, given an image.
+Mask generation is the task of generating semantically meaningful masks for an image.
+This task is very similar to [image segmentation](semantic_segmentation), but many differences exist. Image segmentation models are trained on labeled datasets and are limited to the classes they have seen during training; they return a set of masks and corresponding classes, given an image.
-Mask generation models are trained on large amounts of data and operate in two modes.
-- Prompting mode: In this mode, the model takes in an image and a prompt, where a prompt can be a 2D point location (XY coordinates) in the image within an object or a bounding box surrounding an object. In prompting mode, the model only returns the mask over the object
-that the prompt is pointing out.
-- Segment Everything mode: In segment everything, given an image, the model generates every mask in the image. To do so, a grid of points is generated and overlaid on the image for inference.
+Mask generation models are trained on large amounts of data and operate in two modes.
-Mask generation task is supported by [Segment Anything Model (SAM)](model_doc/sam). It's a powerful model that consists of a Vision Transformer-based image encoder, a prompt encoder, and a two-way transformer mask decoder. Images and prompts are encoded, and the decoder takes these embeddings and generates valid masks.
+- Prompting mode: In this mode, the model takes in an image and a prompt, where a prompt can be a 2D point location (XY coordinates) in the image within an object or a bounding box surrounding an object. In prompting mode, the model only returns the mask over the object
+that the prompt is pointing out.
+- Segment Everything mode: In segment everything, given an image, the model generates every mask in the image. To do so, a grid of points is generated and overlaid on the image for inference.
+
+Mask generation task is supported by [Segment Anything Model (SAM)](model_doc/sam). It's a powerful model that consists of a Vision Transformer-based image encoder, a prompt encoder, and a two-way transformer mask decoder. Images and prompts are encoded, and the decoder takes these embeddings and generates valid masks.
-
In the original implementation ZoeDepth model performs inference on both the original and flipped images and averages out the results. The post_process_depth_estimation function can handle this for us by passing the flipped outputs to the optional outputs_flipped argument:
>>> with torch.no_grad():
+>>> with torch.no_grad():
... outputs = model(pixel_values)
... outputs_flipped = model(pixel_values=torch.flip(inputs.pixel_values, dims=[3]))
>>> post_processed_output = image_processor.post_process_depth_estimation(
diff --git a/docs/source/en/tasks/multiple_choice.md b/docs/source/en/tasks/multiple_choice.md
index 3f4c9d4637fb..d35f108ecce5 100644
--- a/docs/source/en/tasks/multiple_choice.md
+++ b/docs/source/en/tasks/multiple_choice.md
@@ -113,6 +113,7 @@ To apply the preprocessing function over the entire dataset, use 🤗 Datasets [
```
To create a batch of examples, it's more efficient to *dynamically pad* the sentences to the longest length in a batch during collation, instead of padding the whole dataset to the maximum length. [`DataCollatorForMultipleChoice`] flattens all the model inputs, applies padding, and then unflattens the results.
+
```py
>>> from transformers import DataCollatorForMultipleChoice
>>> collator = DataCollatorForMultipleChoice(tokenizer=tokenizer)
@@ -197,7 +198,6 @@ Once training is completed, share your model to the Hub with the [`~transformers
>>> trainer.push_to_hub()
```
-
For a more in-depth example of how to finetune a model for multiple choice, take a look at the corresponding
diff --git a/docs/source/en/tasks/object_detection.md b/docs/source/en/tasks/object_detection.md
index 394e77104b74..ef2a86190bbc 100644
--- a/docs/source/en/tasks/object_detection.md
+++ b/docs/source/en/tasks/object_detection.md
@@ -121,6 +121,7 @@ To get familiar with the data, explore what the examples look like.
```
The examples in the dataset have the following fields:
+
- `image_id`: the example image id
- `image`: a `PIL.Image.Image` object containing the image
- `width`: width of the image
@@ -171,11 +172,11 @@ To get an even better understanding of the data, visualize an example in the dat
>>> image
```
+
-
To visualize the bounding boxes with associated labels, you can get the labels from the dataset's metadata, specifically
the `category` field.
You'll also want to create dictionaries that map a label id to a label class (`id2label`) and the other way around (`label2id`).
@@ -216,6 +217,7 @@ Instantiate the image processor from the same checkpoint as the model you want t
```
Before passing the images to the `image_processor`, apply two preprocessing transformations to the dataset:
+
- Augmenting images
- Reformatting annotations to meet DETR expectations
@@ -505,6 +507,7 @@ The images in this dataset are still quite large, even after resizing. This mean
require at least one GPU.
Training involves the following steps:
+
1. Load the model with [`AutoModelForObjectDetection`] using the same checkpoint as in the preprocessing.
2. Define your training hyperparameters in [`TrainingArguments`].
3. Pass the training arguments to [`Trainer`] along with the model, dataset, image processor, and data collator.
@@ -527,9 +530,10 @@ and `id2label` maps that you created earlier from the dataset's metadata. Additi
In the [`TrainingArguments`] use `output_dir` to specify where to save your model, then configure hyperparameters as you see fit. For `num_train_epochs=30` training will take about 35 minutes in Google Colab T4 GPU, increase the number of epoch to get better results.
Important notes:
- - Do not remove unused columns because this will drop the image column. Without the image column, you
+
+- Do not remove unused columns because this will drop the image column. Without the image column, you
can't create `pixel_values`. For this reason, set `remove_unused_columns` to `False`.
- - Set `eval_do_concat_batches=False` to get proper evaluation results. Images have different number of target boxes, if batches are concatenated we will not be able to determine which boxes belongs to particular image.
+- Set `eval_do_concat_batches=False` to get proper evaluation results. Images have different number of target boxes, if batches are concatenated we will not be able to determine which boxes belongs to particular image.
If you wish to share your model by pushing to the Hub, set `push_to_hub` to `True` (you must be signed in to Hugging
Face to upload your model).
@@ -576,6 +580,7 @@ Finally, bring everything together, and call [`~transformers.Trainer.train`]:
>>> trainer.train()
```
+
@@ -1487,6 +1492,7 @@ Now that you have finetuned a model, evaluated it, and uploaded it to the Huggin
```
Load model and image processor from the Hugging Face Hub (skip to use already trained in this session):
+
```py
>>> from transformers import infer_device
diff --git a/docs/source/en/tasks/prompting.md b/docs/source/en/tasks/prompting.md
index eb8e61d67aaf..2678792c5f3d 100644
--- a/docs/source/en/tasks/prompting.md
+++ b/docs/source/en/tasks/prompting.md
@@ -80,7 +80,7 @@ This section covers a few prompting techniques.
### Few-shot prompting
-Few-shot prompting improves accuracy and performance by including specific examples of what a model should generate given an input. The explicit examples give the model a better understanding of the task and the output format you’re looking for. Try experimenting with different numbers of examples (2, 4, 8, etc.) to see how it affects performance. The example below provides the model with 1 example (1-shot) of the output format (a date in MM/DD/YYYY format) it should return.
+Few-shot prompting improves accuracy and performance by including specific examples of what a model should generate given an input. The explicit examples give the model a better understanding of the task and the output format you're looking for. Try experimenting with different numbers of examples (2, 4, 8, etc.) to see how it affects performance. The example below provides the model with 1 example (1-shot) of the output format (a date in MM/DD/YYYY format) it should return.
```python
from transformers import pipeline
@@ -127,7 +127,6 @@ for output in outputs:
print(f"Result: {output['generated_text']}")
```
-
While the basic few-shot prompting approach embedded examples within a single text string, the chat template format offers the following benefits.
- The model may have a potentially improved understanding because it can better recognize the pattern and the expected roles of user input and assistant output.
diff --git a/docs/source/en/tasks/semantic_segmentation.md b/docs/source/en/tasks/semantic_segmentation.md
index 5d3c8e70aa1f..de88a0af6866 100644
--- a/docs/source/en/tasks/semantic_segmentation.md
+++ b/docs/source/en/tasks/semantic_segmentation.md
@@ -23,6 +23,7 @@ rendered properly in your Markdown viewer.
Image segmentation models separate areas corresponding to different areas of interest in an image. These models work by assigning a label to each pixel. There are several types of segmentation: semantic segmentation, instance segmentation, and panoptic segmentation.
In this guide, we will:
+
1. [Take a look at different types of segmentation](#types-of-segmentation).
2. [Have an end-to-end fine-tuning example for semantic segmentation](#fine-tuning-a-model-for-segmentation).
@@ -69,6 +70,7 @@ results
```
The segmentation pipeline output includes a mask for every predicted class.
+
```bash
[{'score': None,
'label': 'road',
@@ -107,6 +109,7 @@ Taking a look at the mask for the car class, we can see every car is classified
```python
results[-1]["mask"]
```
+
@@ -135,11 +138,13 @@ As you can see below, there are multiple cars classified, and there's no classif
'label': 'person',
'mask': }]
```
+
Checking out one of the car masks below.
```python
results[2]["mask"]
```
+
@@ -151,6 +156,7 @@ panoptic_segmentation = pipeline("image-segmentation", "facebook/mask2former-swi
results = panoptic_segmentation(image)
results
```
+
As you can see below, we have more classes. We will later illustrate to see that every pixel is classified into one of the classes.
```bash
@@ -206,7 +212,6 @@ To see all architectures and checkpoints compatible with this task, we recommend
-
### Load SceneParse150 dataset
Start by loading a smaller subset of the SceneParse150 dataset from the 🤗 Datasets library. This'll give you a chance to experiment and make sure everything works before spending more time training on the full dataset.
@@ -473,7 +478,6 @@ Reload the dataset and load an image for inference.
-
We will now see how to infer without a pipeline. Process the image with an image processor and place the `pixel_values` on a GPU:
```py
@@ -503,7 +507,6 @@ Next, rescale the logits to the original image size:
>>> pred_seg = upsampled_logits.argmax(dim=1)[0]
```
-
To visualize the results, load the [dataset color palette](https://github.com/tensorflow/models/blob/3f1ca33afe3c1631b733ea7e40c294273b9e406d/research/deeplab/utils/get_dataset_colormap.py#L51) as `ade_palette()` that maps each class to their RGB values.
```py
diff --git a/docs/source/en/tasks/summarization.md b/docs/source/en/tasks/summarization.md
index c57097421fbc..b2f2beebc806 100644
--- a/docs/source/en/tasks/summarization.md
+++ b/docs/source/en/tasks/summarization.md
@@ -213,7 +213,6 @@ Once training is completed, share your model to the Hub with the [`~transformers
>>> trainer.push_to_hub()
```
-
For a more in-depth example of how to finetune a model for summarization, take a look at the corresponding
diff --git a/docs/source/en/tasks/token_classification.md b/docs/source/en/tasks/token_classification.md
index 49b0fcf216b8..5096298affd1 100644
--- a/docs/source/en/tasks/token_classification.md
+++ b/docs/source/en/tasks/token_classification.md
@@ -242,7 +242,6 @@ Before you start training your model, create a map of the expected ids to their
... }
```
-
If you aren't familiar with finetuning a model with the [`Trainer`], take a look at the basic tutorial [here](../training#train-with-pytorch-trainer)!
@@ -298,7 +297,6 @@ Once training is completed, share your model to the Hub with the [`~transformers
>>> trainer.push_to_hub()
```
-
For a more in-depth example of how to finetune a model for token classification, take a look at the corresponding
diff --git a/docs/source/en/tasks/video_classification.md b/docs/source/en/tasks/video_classification.md
index b387a8320dfc..bae638bd84ed 100644
--- a/docs/source/en/tasks/video_classification.md
+++ b/docs/source/en/tasks/video_classification.md
@@ -363,7 +363,6 @@ Leverage [`Trainer`](https://huggingface.co/docs/transformers/main_classes/train
Most of the training arguments are self-explanatory, but one that is quite important here is `remove_unused_columns=False`. This one will drop any features not used by the model's call function. By default it's `True` because usually it's ideal to drop unused feature columns, making it easier to unpack inputs into the model's call function. But, in this case, you need the unused features ('video' in particular) in order to create `pixel_values` (which is a mandatory key our model expects in its inputs).
-
```py
>>> from transformers import TrainingArguments, Trainer
@@ -477,7 +476,6 @@ The simplest way to try out your fine-tuned model for inference is to use it in
You can also manually replicate the results of the `pipeline` if you'd like.
-
```py
>>> def run_inference(model, video):
... # (num_frames, num_channels, height, width)
diff --git a/docs/source/en/tasks/video_text_to_text.md b/docs/source/en/tasks/video_text_to_text.md
index 0e0191af5884..58ca97e9a56c 100644
--- a/docs/source/en/tasks/video_text_to_text.md
+++ b/docs/source/en/tasks/video_text_to_text.md
@@ -18,13 +18,14 @@ rendered properly in your Markdown viewer.
[[open-in-colab]]
-Video-text-to-text models, also known as video language models or vision language models with video input, are language models that take a video input. These models can tackle various tasks, from video question answering to video captioning.
+Video-text-to-text models, also known as video language models or vision language models with video input, are language models that take a video input. These models can tackle various tasks, from video question answering to video captioning.
-These models have nearly the same architecture as [image-text-to-text](../image_text_to_text) models except for some changes to accept video data, since video data is essentially image frames with temporal dependencies. Some image-text-to-text models take in multiple images, but this alone is inadequate for a model to accept videos. Moreover, video-text-to-text models are often trained with all vision modalities. Each example might have videos, multiple videos, images and multiple images. Some of these models can also take interleaved inputs. For example, you can refer to a specific video inside a string of text by adding a video token in text like "What is happening in this video? `
@@ -158,6 +152,7 @@ class MyConfig(PretrainedConfig):
"norm": (["hidden_states"], ["hidden_states"]),
}
```
+
### Multimodal models
@@ -200,8 +195,8 @@ class MyMultimodalModelForConditionalGeneration(MyMultimodalPreTrainedModel, Gen
self.model = MyMultimodalModel(config)
self.lm_head = nn.Linear(hidden_dim, vocab_size)
```
-
+
2. A multimodal model config must be nested with the following fields.
* text_config: decoder language model config
@@ -210,7 +205,7 @@ class MyMultimodalModelForConditionalGeneration(MyMultimodalPreTrainedModel, Gen
3. A multimodal model's processing class must have the `self.image_token` and `self.image_token_ids` attributes. These are placeholder tokens used to indicate image positions in the input. The placeholder token is the same token used in the input prompt and to mask scatter image features.
- The processing class also needs ` self._get_num_multimodal_tokens` method to compute the number of placeholder tokens needed for multimodal inputs with given sizes and to return a [`MultiModalData`] object. The placeholder for row and column tokens don't count as image placeholders. Only the tokens that are actually replaced by image features are computed.
+ The processing class also needs `self._get_num_multimodal_tokens` method to compute the number of placeholder tokens needed for multimodal inputs with given sizes and to return a [`MultiModalData`] object. The placeholder for row and column tokens don't count as image placeholders. Only the tokens that are actually replaced by image features are computed.
Finally, when `return_mm_token_type_ids=True`, the class has to return `mm_token_type_ids` to indicate whether each position is a text token (`0`) or image placeholder token (`1`). Each image's token type IDs must be contiguous with no breaks between consecutive ones.
@@ -246,6 +241,7 @@ class MyMultimodalProcessor(ProcessorMixin):
vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
return MultiModalData(**vision_data)
```
+
## Resources
diff --git a/docs/source/en/troubleshooting.md b/docs/source/en/troubleshooting.md
index 7998881d3648..0cc5829d2e8d 100644
--- a/docs/source/en/troubleshooting.md
+++ b/docs/source/en/troubleshooting.md
@@ -34,12 +34,11 @@ Sometimes errors occur, but we are here to help! This guide covers some of the m
For more details about troubleshooting and getting help, take a look at [Chapter 8](https://huggingface.co/course/chapter8/1?fw=pt) of the Hugging Face course.
-
## Firewalled environments
Some GPU instances on cloud and intranet setups are firewalled to external connections, resulting in a connection error. When your script attempts to download model weights or datasets, the download will hang and then timeout with the following message:
-```
+```text
ValueError: Connection error, and we cannot find the requested files in the cached path.
Please try again or make sure your Internet connection is on.
```
@@ -50,7 +49,7 @@ In this case, you should try to run 🤗 Transformers on [offline mode](installa
Training large models with millions of parameters can be challenging without the appropriate hardware. A common error you may encounter when the GPU runs out of memory is:
-```
+```text
CUDA out of memory. Tried to allocate 256.00 MiB (GPU 0; 11.17 GiB total capacity; 9.70 GiB already allocated; 179.81 MiB free; 9.85 GiB reserved in total by PyTorch)
```
@@ -69,7 +68,7 @@ Refer to the Performance [guide](performance) for more details about memory-savi
Another common error you may encounter, especially if it is a newly released model, is `ImportError`:
-```
+```text
ImportError: cannot import name 'ImageGPTImageProcessor' from 'transformers' (unknown location)
```
@@ -83,7 +82,7 @@ pip install transformers --upgrade
Sometimes you may run into a generic CUDA error about an error in the device code.
-```
+```text
RuntimeError: CUDA error: device-side assert triggered
```
diff --git a/docs/source/en/video_processors.md b/docs/source/en/video_processors.md
index 4f44914c8cfc..2b26d9f9fc7f 100644
--- a/docs/source/en/video_processors.md
+++ b/docs/source/en/video_processors.md
@@ -14,17 +14,16 @@ rendered properly in your Markdown viewer.
-->
-
# Video Processor
-A **Video Processor** is a utility responsible for preparing input features for video models, as well as handling the post-processing of their outputs. It provides transformations such as resizing, normalization, and conversion into PyTorch.
+A **Video Processor** is a utility responsible for preparing input features for video models, as well as handling the post-processing of their outputs. It provides transformations such as resizing, normalization, and conversion into PyTorch.
The video processor extends the functionality of image processors by allowing the models to handle videos with a distinct set of arguments compared to images. It serves as the bridge between raw video data and the model, ensuring that input features are optimized for the VLM.
Use [`~BaseVideoProcessor.from_pretrained`] to load a video processors configuration (image size, whether to normalize and rescale, etc.) from a video model on the Hugging Face [Hub](https://hf.co) or local directory. The configuration for each pretrained model should be saved in a [video_preprocessor_config.json] file but older models might have the config saved in [preprocessor_config.json](https://huggingface.co/llava-hf/llava-onevision-qwen2-0.5b-ov-hf/blob/main/preprocessor_config.json) file. Note that the latter is less preferred and will be removed in the future.
+## Usage Example
-### Usage Example
Here's an example of how to load a video processor with [`llava-hf/llava-onevision-qwen2-0.5b-ov-hf`](https://huggingface.co/llava-hf/llava-onevision-qwen2-0.5b-ov-hf) model:
```python
diff --git a/docs/source/ko/_toctree.yml b/docs/source/ko/_toctree.yml
index df2d53c49a96..2412e497556f 100644
--- a/docs/source/ko/_toctree.yml
+++ b/docs/source/ko/_toctree.yml
@@ -607,6 +607,8 @@
title: LED
- local: in_translation
title: LFM2
+ - local: in_translation
+ title: LFM2-VL
- local: model_doc/llama
title: LLaMA
- local: model_doc/llama2
diff --git a/docs/source/zh/main_classes/deepspeed.md b/docs/source/zh/main_classes/deepspeed.md
index 7cdf3b62e427..a8863896235f 100644
--- a/docs/source/zh/main_classes/deepspeed.md
+++ b/docs/source/zh/main_classes/deepspeed.md
@@ -236,7 +236,7 @@ deepspeed --num_gpus=1 examples/pytorch/translation/run_translation.py \
}
```
-这会启用`optimizer offload `和一些其他重要功能。您可以尝试不同的buffer大小，有关详细信息，请参见下面的讨论。
+这会启用`optimizer offload`和一些其他重要功能。您可以尝试不同的buffer大小，有关详细信息，请参见下面的讨论。
关于这种启用类型的实际使用示例，请参阅 [此帖](https://github.com/huggingface/transformers/issues/8771#issuecomment-759176685)。
diff --git a/docs/source/zh/pipeline_tutorial.md b/docs/source/zh/pipeline_tutorial.md
index 92fbcbba31e4..7c497c6f1c65 100644
--- a/docs/source/zh/pipeline_tutorial.md
+++ b/docs/source/zh/pipeline_tutorial.md
@@ -306,5 +306,5 @@ pipe = pipeline(model="facebook/opt-1.3b", device_map="auto", model_kwargs={"loa
output = pipe("This is a cool example!", do_sample=True, top_p=0.95)
```
-请注意，您可以将`checkpoint `替换为任何支持大模型加载的Hugging Face模型，比如BLOOM！
+请注意，您可以将`checkpoint`替换为任何支持大模型加载的Hugging Face模型，比如BLOOM！
diff --git a/docs/source/zh/tasks/asr.md b/docs/source/zh/tasks/asr.md
index 3b66888bc107..228ba55c0d0e 100644
--- a/docs/source/zh/tasks/asr.md
+++ b/docs/source/zh/tasks/asr.md
@@ -83,7 +83,7 @@ DatasetDict({
})
```
-虽然数据集包含 `lang_id `和 `english_transcription` 等许多有用的信息，但在本指南中，
+虽然数据集包含 `lang_id` 和 `english_transcription` 等许多有用的信息，但在本指南中，
您将专注于 `audio` 和 `transcription`。使用 [`~datasets.Dataset.remove_columns`] 方法删除其他列：
```py
@@ -167,7 +167,7 @@ Wav2Vec2 分词器仅训练了大写字符，因此您需要确保文本与分
它还会动态地将您的文本和标签填充到其批次中最长元素的长度（而不是整个数据集），以使它们具有统一的长度。
虽然可以通过在 `tokenizer` 函数中设置 `padding=True` 来填充文本，但动态填充更有效。
-与其他数据整理器不同，这个特定的数据整理器需要对 `input_values` 和 `labels `应用不同的填充方法：
+与其他数据整理器不同，这个特定的数据整理器需要对 `input_values` 和 `labels` 应用不同的填充方法：
```py
>>> import torch
diff --git a/examples/legacy/pytorch-lightning/run_ner.py b/examples/legacy/pytorch-lightning/run_ner.py
index 144759d36aac..6cbb138f023f 100644
--- a/examples/legacy/pytorch-lightning/run_ner.py
+++ b/examples/legacy/pytorch-lightning/run_ner.py
@@ -72,12 +72,12 @@ def prepare_data(self):
self.labels,
args.max_seq_length,
self.tokenizer,
- cls_token_at_end=bool(self.config.model_type in ["xlnet"]),
+ cls_token_at_end=bool(self.config.model_type == "xlnet"),
cls_token=self.tokenizer.cls_token,
- cls_token_segment_id=2 if self.config.model_type in ["xlnet"] else 0,
+ cls_token_segment_id=2 if self.config.model_type == "xlnet" else 0,
sep_token=self.tokenizer.sep_token,
sep_token_extra=False,
- pad_on_left=bool(self.config.model_type in ["xlnet"]),
+ pad_on_left=bool(self.config.model_type == "xlnet"),
pad_token=self.tokenizer.pad_token_id,
pad_token_segment_id=self.tokenizer.pad_token_type_id,
pad_token_label_id=self.pad_token_label_id,
diff --git a/examples/legacy/run_chinese_ref.py b/examples/legacy/run_chinese_ref.py
index e63096d05244..7cb6caccefe1 100755
--- a/examples/legacy/run_chinese_ref.py
+++ b/examples/legacy/run_chinese_ref.py
@@ -55,7 +55,7 @@ def get_chinese_word(tokens: list[str]):
def add_sub_symbol(bert_tokens: list[str], chinese_word_set: set()):
if not chinese_word_set:
return bert_tokens
- max_word_len = max([len(w) for w in chinese_word_set])
+ max_word_len = max(len(w) for w in chinese_word_set)
bert_word = bert_tokens
start, end = 0, len(bert_word)
diff --git a/examples/legacy/token-classification/utils_ner.py b/examples/legacy/token-classification/utils_ner.py
index 0c1725b59b4e..833984bc0ec3 100644
--- a/examples/legacy/token-classification/utils_ner.py
+++ b/examples/legacy/token-classification/utils_ner.py
@@ -251,10 +251,10 @@ def __init__(
labels,
max_seq_length,
tokenizer,
- cls_token_at_end=bool(model_type in ["xlnet"]),
+ cls_token_at_end=bool(model_type == "xlnet"),
# xlnet has a cls token at the end
cls_token=tokenizer.cls_token,
- cls_token_segment_id=2 if model_type in ["xlnet"] else 0,
+ cls_token_segment_id=2 if model_type == "xlnet" else 0,
sep_token=tokenizer.sep_token,
sep_token_extra=False,
# roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805
diff --git a/examples/pytorch/audio-classification/run_audio_classification.py b/examples/pytorch/audio-classification/run_audio_classification.py
index ad82f4c401e8..bd190e801520 100644
--- a/examples/pytorch/audio-classification/run_audio_classification.py
+++ b/examples/pytorch/audio-classification/run_audio_classification.py
@@ -48,7 +48,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -218,10 +218,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_audio_classification", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/continuous_batching.py b/examples/pytorch/continuous_batching.py
index 2b0d506eb895..cf5379fc619c 100644
--- a/examples/pytorch/continuous_batching.py
+++ b/examples/pytorch/continuous_batching.py
@@ -40,7 +40,8 @@ def generate_simple(
attn_impl = {
"sdpa_paged": "sdpa",
"eager_paged": "eager",
- "flash_paged": "flash_attention_2",
+ "paged_attention": "eager", # TODO: this does not work on AMD docker
+ "flash_paged": "flash_attention_2", # TODO: this does not work on AMD docker
}[attn_impl]
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, attn_implementation=attn_impl)
diff --git a/examples/pytorch/contrastive-image-text/run_clip.py b/examples/pytorch/contrastive-image-text/run_clip.py
index 8b0b42252a2e..461062f6849b 100644
--- a/examples/pytorch/contrastive-image-text/run_clip.py
+++ b/examples/pytorch/contrastive-image-text/run_clip.py
@@ -56,7 +56,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -247,10 +247,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_clip", model_args, data_args)
-
# 2. Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/image-classification/run_image_classification.py b/examples/pytorch/image-classification/run_image_classification.py
index 9693d4b1c84a..8b498b545c45 100755
--- a/examples/pytorch/image-classification/run_image_classification.py
+++ b/examples/pytorch/image-classification/run_image_classification.py
@@ -59,7 +59,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -201,10 +201,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_image_classification", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/image-classification/run_image_classification_no_trainer.py b/examples/pytorch/image-classification/run_image_classification_no_trainer.py
index 656310424c17..cc8814305389 100644
--- a/examples/pytorch/image-classification/run_image_classification_no_trainer.py
+++ b/examples/pytorch/image-classification/run_image_classification_no_trainer.py
@@ -56,7 +56,7 @@
import transformers
from transformers import AutoConfig, AutoImageProcessor, AutoModelForImageClassification, SchedulerType, get_scheduler
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -234,10 +234,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_image_classification_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/image-pretraining/run_mae.py b/examples/pytorch/image-pretraining/run_mae.py
index d0ea39e780b5..2d92d8ab434d 100644
--- a/examples/pytorch/image-pretraining/run_mae.py
+++ b/examples/pytorch/image-pretraining/run_mae.py
@@ -42,7 +42,7 @@
ViTMAEForPreTraining,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -193,10 +193,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_mae", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/image-pretraining/run_mim.py b/examples/pytorch/image-pretraining/run_mim.py
index 746126596fbe..5a636bbad58b 100644
--- a/examples/pytorch/image-pretraining/run_mim.py
+++ b/examples/pytorch/image-pretraining/run_mim.py
@@ -45,7 +45,7 @@
TrainingArguments,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -257,10 +257,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_mim", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/image-pretraining/run_mim_no_trainer.py b/examples/pytorch/image-pretraining/run_mim_no_trainer.py
index 92c4d2242f76..1c5636088632 100644
--- a/examples/pytorch/image-pretraining/run_mim_no_trainer.py
+++ b/examples/pytorch/image-pretraining/run_mim_no_trainer.py
@@ -49,7 +49,7 @@
SchedulerType,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -384,10 +384,6 @@ def collate_fn(examples):
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_mim_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/instance-segmentation/run_instance_segmentation.py b/examples/pytorch/instance-segmentation/run_instance_segmentation.py
index 992d9854d078..ddfd05e0f661 100644
--- a/examples/pytorch/instance-segmentation/run_instance_segmentation.py
+++ b/examples/pytorch/instance-segmentation/run_instance_segmentation.py
@@ -50,7 +50,7 @@
from transformers.image_processing_utils import BatchFeature
from transformers.trainer import EvalPrediction
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -367,10 +367,6 @@ def main():
training_args.batch_eval_metrics = True
training_args.remove_unused_columns = False
- # # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_instance_segmentation", args)
-
# Setup logging and log on each process the small summary:
setup_logging(training_args)
logger.warning(
diff --git a/examples/pytorch/instance-segmentation/run_instance_segmentation_no_trainer.py b/examples/pytorch/instance-segmentation/run_instance_segmentation_no_trainer.py
index c538508b7b74..bff3abb32715 100644
--- a/examples/pytorch/instance-segmentation/run_instance_segmentation_no_trainer.py
+++ b/examples/pytorch/instance-segmentation/run_instance_segmentation_no_trainer.py
@@ -56,7 +56,7 @@
get_scheduler,
)
from transformers.image_processing_utils import BatchFeature
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -413,10 +413,6 @@ def handle_repository_creation(accelerator: Accelerator, args: argparse.Namespac
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_instance_segmentation_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/language-modeling/run_clm.py b/examples/pytorch/language-modeling/run_clm.py
index 69099bb79306..8c677b404630 100755
--- a/examples/pytorch/language-modeling/run_clm.py
+++ b/examples/pytorch/language-modeling/run_clm.py
@@ -64,7 +64,7 @@
)
from transformers.testing_utils import CaptureLogger
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -292,10 +292,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_clm", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/language-modeling/run_clm_no_trainer.py b/examples/pytorch/language-modeling/run_clm_no_trainer.py
index 874d95393f70..c750d9274a36 100755
--- a/examples/pytorch/language-modeling/run_clm_no_trainer.py
+++ b/examples/pytorch/language-modeling/run_clm_no_trainer.py
@@ -66,7 +66,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -268,10 +268,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_clm_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/language-modeling/run_fim.py b/examples/pytorch/language-modeling/run_fim.py
index 46b759e03002..134d741f6b6c 100644
--- a/examples/pytorch/language-modeling/run_fim.py
+++ b/examples/pytorch/language-modeling/run_fim.py
@@ -67,7 +67,7 @@
from transformers.integrations import is_deepspeed_zero3_enabled
from transformers.testing_utils import CaptureLogger
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -319,10 +319,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_fim", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/language-modeling/run_fim_no_trainer.py b/examples/pytorch/language-modeling/run_fim_no_trainer.py
index 67a94f1fae30..693f4d44b781 100644
--- a/examples/pytorch/language-modeling/run_fim_no_trainer.py
+++ b/examples/pytorch/language-modeling/run_fim_no_trainer.py
@@ -69,7 +69,7 @@
is_torch_xla_available,
)
from transformers.integrations import is_deepspeed_zero3_enabled
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -328,10 +328,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_fim_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/language-modeling/run_mlm.py b/examples/pytorch/language-modeling/run_mlm.py
index 5ba9262f451b..9c0bf50ede28 100755
--- a/examples/pytorch/language-modeling/run_mlm.py
+++ b/examples/pytorch/language-modeling/run_mlm.py
@@ -63,7 +63,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -264,10 +264,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_mlm", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/language-modeling/run_mlm_no_trainer.py b/examples/pytorch/language-modeling/run_mlm_no_trainer.py
index 501da0cff932..59ee11926c40 100755
--- a/examples/pytorch/language-modeling/run_mlm_no_trainer.py
+++ b/examples/pytorch/language-modeling/run_mlm_no_trainer.py
@@ -66,7 +66,7 @@
SchedulerType,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -275,10 +275,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_mlm_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/language-modeling/run_plm.py b/examples/pytorch/language-modeling/run_plm.py
index fd29c6a630d7..86bc31beedf8 100755
--- a/examples/pytorch/language-modeling/run_plm.py
+++ b/examples/pytorch/language-modeling/run_plm.py
@@ -56,7 +56,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -244,10 +244,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_plm", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/multiple-choice/run_swag.py b/examples/pytorch/multiple-choice/run_swag.py
index 585ac54febb2..a8679f2b739c 100755
--- a/examples/pytorch/multiple-choice/run_swag.py
+++ b/examples/pytorch/multiple-choice/run_swag.py
@@ -53,7 +53,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
@@ -188,10 +188,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_swag", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/multiple-choice/run_swag_no_trainer.py b/examples/pytorch/multiple-choice/run_swag_no_trainer.py
index 5d19486da0e1..c77a10e990d0 100755
--- a/examples/pytorch/multiple-choice/run_swag_no_trainer.py
+++ b/examples/pytorch/multiple-choice/run_swag_no_trainer.py
@@ -61,7 +61,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
@@ -238,10 +238,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_swag_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/object-detection/run_object_detection.py b/examples/pytorch/object-detection/run_object_detection.py
index f615488c7099..ee0bd66cae99 100644
--- a/examples/pytorch/object-detection/run_object_detection.py
+++ b/examples/pytorch/object-detection/run_object_detection.py
@@ -52,7 +52,7 @@
from transformers.image_transforms import center_to_corners_format
from transformers.trainer import EvalPrediction
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -349,10 +349,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_object_detection", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/object-detection/run_object_detection_no_trainer.py b/examples/pytorch/object-detection/run_object_detection_no_trainer.py
index f90bf1bbd3c0..543f3d108742 100644
--- a/examples/pytorch/object-detection/run_object_detection_no_trainer.py
+++ b/examples/pytorch/object-detection/run_object_detection_no_trainer.py
@@ -58,7 +58,7 @@
)
from transformers.image_processing_utils import BatchFeature
from transformers.image_transforms import center_to_corners_format
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -411,10 +411,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_object_detection_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/question-answering/run_qa.py b/examples/pytorch/question-answering/run_qa.py
index 5a639696f6cd..be93a526b803 100755
--- a/examples/pytorch/question-answering/run_qa.py
+++ b/examples/pytorch/question-answering/run_qa.py
@@ -44,7 +44,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -237,10 +237,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_qa", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/question-answering/run_qa_beam_search.py b/examples/pytorch/question-answering/run_qa_beam_search.py
index b778d9fc67ee..4bcf4f9af8c8 100755
--- a/examples/pytorch/question-answering/run_qa_beam_search.py
+++ b/examples/pytorch/question-answering/run_qa_beam_search.py
@@ -42,7 +42,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -235,10 +235,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_qa_beam_search", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py b/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py
index 9fd3ce223220..c95a5b46030c 100644
--- a/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py
+++ b/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py
@@ -49,7 +49,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -299,10 +299,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_qa_beam_search_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/question-answering/run_qa_no_trainer.py b/examples/pytorch/question-answering/run_qa_no_trainer.py
index dc1b9743e634..a8e3c72de862 100755
--- a/examples/pytorch/question-answering/run_qa_no_trainer.py
+++ b/examples/pytorch/question-answering/run_qa_no_trainer.py
@@ -51,7 +51,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -338,10 +338,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_qa_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
@@ -954,7 +950,7 @@ def create_and_fill_np_array(start_or_end_logits, dataset, max_len):
all_start_logits.append(accelerator.gather_for_metrics(start_logits).cpu().numpy())
all_end_logits.append(accelerator.gather_for_metrics(end_logits).cpu().numpy())
- max_len = max([x.shape[1] for x in all_start_logits]) # Get the max_length of the tensor
+ max_len = max(x.shape[1] for x in all_start_logits) # Get the max_length of the tensor
# concatenate the numpy array
start_logits_concat = create_and_fill_np_array(all_start_logits, eval_dataset, max_len)
@@ -993,7 +989,7 @@ def create_and_fill_np_array(start_or_end_logits, dataset, max_len):
all_start_logits.append(accelerator.gather_for_metrics(start_logits).cpu().numpy())
all_end_logits.append(accelerator.gather_for_metrics(end_logits).cpu().numpy())
- max_len = max([x.shape[1] for x in all_start_logits]) # Get the max_length of the tensor
+ max_len = max(x.shape[1] for x in all_start_logits) # Get the max_length of the tensor
# concatenate the numpy array
start_logits_concat = create_and_fill_np_array(all_start_logits, predict_dataset, max_len)
end_logits_concat = create_and_fill_np_array(all_end_logits, predict_dataset, max_len)
diff --git a/examples/pytorch/question-answering/run_seq2seq_qa.py b/examples/pytorch/question-answering/run_seq2seq_qa.py
index 408d4d23f59c..ac3c8ef4ec62 100644
--- a/examples/pytorch/question-answering/run_seq2seq_qa.py
+++ b/examples/pytorch/question-answering/run_seq2seq_qa.py
@@ -40,7 +40,7 @@
set_seed,
)
from transformers.trainer_utils import EvalLoopOutput, EvalPrediction, get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -282,10 +282,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_seq2seq_qa", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/semantic-segmentation/run_semantic_segmentation.py b/examples/pytorch/semantic-segmentation/run_semantic_segmentation.py
index ea678c094aef..cc45239f75c0 100644
--- a/examples/pytorch/semantic-segmentation/run_semantic_segmentation.py
+++ b/examples/pytorch/semantic-segmentation/run_semantic_segmentation.py
@@ -53,7 +53,7 @@
default_data_collator,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -197,10 +197,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_semantic_segmentation", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
index 97a3a249d484..66b0af0ef635 100644
--- a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
+++ b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
@@ -57,7 +57,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -253,10 +253,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_semantic_segmentation_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py b/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py
index f30fd1676a3a..1840d7e4ed7a 100755
--- a/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py
+++ b/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py
@@ -53,7 +53,6 @@
set_seed,
)
from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices, _sample_negative_indices
-from transformers.utils import send_example_telemetry
logger = get_logger(__name__)
@@ -410,10 +409,6 @@ def main():
# We now keep distinct sets of args, for a cleaner separation of concerns.
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_wav2vec2_pretraining_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
accelerator = Accelerator()
logger.info(accelerator.state, main_process_only=False)
diff --git a/examples/pytorch/speech-recognition/README.md b/examples/pytorch/speech-recognition/README.md
index 2889919655f4..41df41880b5a 100644
--- a/examples/pytorch/speech-recognition/README.md
+++ b/examples/pytorch/speech-recognition/README.md
@@ -66,7 +66,7 @@ The following command shows how to fine-tune [XLSR-Wav2Vec2](https://huggingface
```bash
python run_speech_recognition_ctc.py \
- --dataset_name="common_voice" \
+ --dataset_name="mozilla-foundation/common_voice_17_0" \
--model_name_or_path="facebook/wav2vec2-large-xlsr-53" \
--dataset_config_name="tr" \
--output_dir="./wav2vec2-common_voice-tr-demo" \
@@ -102,7 +102,7 @@ The following command shows how to fine-tune [XLSR-Wav2Vec2](https://huggingface
```bash
torchrun \
--nproc_per_node 8 run_speech_recognition_ctc.py \
- --dataset_name="common_voice" \
+ --dataset_name="mozilla-foundation/common_voice_17_0" \
--model_name_or_path="facebook/wav2vec2-large-xlsr-53" \
--dataset_config_name="tr" \
--output_dir="./wav2vec2-common_voice-tr-demo-dist" \
@@ -149,7 +149,7 @@ However, the `--shuffle_buffer_size` argument controls how many examples we can
```bash
**torchrun \
--nproc_per_node 4 run_speech_recognition_ctc_streaming.py \
- --dataset_name="common_voice" \
+ --dataset_name="mozilla-foundation/common_voice_17_0" \
--model_name_or_path="facebook/wav2vec2-xls-r-300m" \
--tokenizer_name_or_path="anton-l/wav2vec2-tokenizer-turkish" \
--dataset_config_name="tr" \
@@ -314,7 +314,7 @@ below 27%.
For an example run, you can have a look at [`patrickvonplaten/wav2vec2-common_voice-tr-mms-demo`](https://huggingface.co/patrickvonplaten/wav2vec2-common_voice-tr-mms-demo).
-If you'd like to train another adapter model with the same base model, you can simply re-use the same `--output_dir`,
+If you'd like to train another adapter model with the same base model, you can simply reuse the same `--output_dir`,
but make sure to pass the `--output_dir` folder also to `--tokenizer_name_or_path` so that the vocabulary is not
overwritten but **extended**. Assuming you would like to train adapter weights on Swedish in addition to Turkish and save
the adapter weights in the same model repo, you can run:
diff --git a/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py b/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py
index 4532bc511e9f..c756a6666187 100755
--- a/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py
+++ b/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py
@@ -56,14 +56,17 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.57.0.dev0")
-require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")
+require_version(
+ "datasets>=1.18.0",
+ "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt",
+)
logger = logging.getLogger(__name__)
@@ -91,13 +94,16 @@ class ModelArguments:
metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
)
freeze_feature_encoder: bool = field(
- default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+ default=True,
+ metadata={"help": "Whether to freeze the feature encoder layers of the model."},
)
attention_dropout: float = field(
- default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
+ default=0.0,
+ metadata={"help": "The dropout ratio for the attention probabilities."},
)
activation_dropout: float = field(
- default=0.0, metadata={"help": "The dropout ratio for activations inside the fully connected layer."}
+ default=0.0,
+ metadata={"help": "The dropout ratio for activations inside the fully connected layer."},
)
feat_proj_dropout: float = field(default=0.0, metadata={"help": "The dropout ratio for the projected features."})
hidden_dropout: float = field(
@@ -140,7 +146,8 @@ class ModelArguments:
)
layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
ctc_loss_reduction: Optional[str] = field(
- default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+ default="mean",
+ metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."},
)
ctc_zero_infinity: Optional[bool] = field(
default=False,
@@ -169,10 +176,13 @@ class DataTrainingArguments:
"""
dataset_name: str = field(
- metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ metadata={"help": "Path or name of the dataset (cf `load_dataset` method of the Datasets library)."}
)
dataset_config_name: str = field(
- default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ default=None,
+ metadata={
+ "help": "The configuration name of the dataset to use (cf `load_dataset` method of the Datasets library)."
+ },
)
train_split_name: str = field(
default="train+validation",
@@ -198,7 +208,8 @@ class DataTrainingArguments:
metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
)
overwrite_cache: bool = field(
- default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+ default=False,
+ metadata={"help": "Overwrite the cached preprocessed datasets or not."},
)
preprocessing_num_workers: Optional[int] = field(
default=None,
@@ -240,7 +251,8 @@ class DataTrainingArguments:
},
)
min_duration_in_seconds: float = field(
- default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+ default=0.0,
+ metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
)
preprocessing_only: bool = field(
default=False,
@@ -383,7 +395,8 @@ def extract_all_chars(batch):
# take union of all unique characters in each dataset
vocab_set = functools.reduce(
- lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]), vocabs.values()
+ lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]),
+ vocabs.values(),
)
vocab_dict = {v: k for k, v in enumerate(sorted(vocab_set))}
@@ -416,10 +429,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_speech_recognition_ctc", model_args, data_args)
-
# Detecting last checkpoint.
last_checkpoint = None
if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
@@ -575,7 +584,7 @@ def remove_special_characters(batch):
# it is defined by `tokenizer_class` if present in config else by `model_type`
tokenizer_kwargs = {
"config": config if config.tokenizer_class is not None else None,
- "tokenizer_type": config.model_type if config.tokenizer_class is None else None,
+ "tokenizer_type": (config.model_type if config.tokenizer_class is None else None),
"unk_token": unk_token,
"pad_token": pad_token,
"word_delimiter_token": word_delimiter_token,
@@ -643,7 +652,8 @@ def remove_special_characters(batch):
dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
if dataset_sampling_rate != feature_extractor.sampling_rate:
raw_datasets = raw_datasets.cast_column(
- data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
+ data_args.audio_column_name,
+ datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate),
)
# derive max & min input length for sample rate & max duration
diff --git a/examples/pytorch/speech-recognition/run_speech_recognition_ctc_adapter.py b/examples/pytorch/speech-recognition/run_speech_recognition_ctc_adapter.py
index 884201d9d993..aaebf59c8660 100755
--- a/examples/pytorch/speech-recognition/run_speech_recognition_ctc_adapter.py
+++ b/examples/pytorch/speech-recognition/run_speech_recognition_ctc_adapter.py
@@ -59,14 +59,17 @@
)
from transformers.models.wav2vec2.modeling_wav2vec2 import WAV2VEC2_ADAPTER_SAFE_FILE
from transformers.trainer_utils import get_last_checkpoint, is_main_process
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.57.0.dev0")
-require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")
+require_version(
+ "datasets>=1.18.0",
+ "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt",
+)
logger = logging.getLogger(__name__)
@@ -127,7 +130,8 @@ class ModelArguments:
)
layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
ctc_loss_reduction: Optional[str] = field(
- default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+ default="mean",
+ metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."},
)
adapter_attn_dim: int = field(
default=16,
@@ -148,9 +152,9 @@ class DataTrainingArguments:
"""
dataset_name: str = field(
- metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ metadata={"help": "Path or name of the dataset (cf `load_dataset` method of the Datasets library)."}
)
- target_language: Optional[str] = field(
+ target_language: str = field(
metadata={
"help": (
"The target language on which the adapter attention layers"
@@ -162,7 +166,10 @@ class DataTrainingArguments:
},
)
dataset_config_name: str = field(
- default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ default=None,
+ metadata={
+ "help": "The configuration name of the dataset to use (cf `load_dataset` method of the Datasets library)."
+ },
)
train_split_name: str = field(
default="train+validation",
@@ -188,7 +195,8 @@ class DataTrainingArguments:
metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
)
overwrite_cache: bool = field(
- default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+ default=False,
+ metadata={"help": "Overwrite the cached preprocessed datasets or not."},
)
preprocessing_num_workers: Optional[int] = field(
default=None,
@@ -230,7 +238,8 @@ class DataTrainingArguments:
},
)
min_duration_in_seconds: float = field(
- default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+ default=0.0,
+ metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
)
preprocessing_only: bool = field(
default=False,
@@ -363,7 +372,8 @@ def extract_all_chars(batch):
# take union of all unique characters in each dataset
vocab_set = functools.reduce(
- lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]), vocabs.values()
+ lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]),
+ vocabs.values(),
)
vocab_dict = {v: k for k, v in enumerate(sorted(vocab_set))}
@@ -396,10 +406,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_speech_recognition_ctc_adapter", model_args, data_args)
-
# Detecting last checkpoint.
last_checkpoint = None
if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
@@ -582,7 +588,7 @@ def remove_special_characters(batch):
# it is defined by `tokenizer_class` if present in config else by `model_type`
tokenizer_kwargs = {
"config": config if config.tokenizer_class is not None else None,
- "tokenizer_type": config.model_type if config.tokenizer_class is None else None,
+ "tokenizer_type": (config.model_type if config.tokenizer_class is None else None),
"unk_token": unk_token,
"pad_token": pad_token,
"word_delimiter_token": word_delimiter_token,
@@ -654,7 +660,8 @@ def remove_special_characters(batch):
dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
if dataset_sampling_rate != feature_extractor.sampling_rate:
raw_datasets = raw_datasets.cast_column(
- data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
+ data_args.audio_column_name,
+ datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate),
)
# derive max & min input length for sample rate & max duration
diff --git a/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py b/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py
index aee6ae3b8bae..4b6cda49925b 100755
--- a/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py
+++ b/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py
@@ -55,14 +55,17 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.57.0.dev0")
-require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")
+require_version(
+ "datasets>=1.18.0",
+ "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt",
+)
logger = logging.getLogger(__name__)
@@ -77,13 +80,16 @@ class ModelArguments:
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
)
config_name: Optional[str] = field(
- default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+ default=None,
+ metadata={"help": "Pretrained config name or path if not the same as model_name"},
)
tokenizer_name: Optional[str] = field(
- default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+ default=None,
+ metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"},
)
feature_extractor_name: Optional[str] = field(
- default=None, metadata={"help": "feature extractor name or path if not the same as model_name"}
+ default=None,
+ metadata={"help": "feature extractor name or path if not the same as model_name"},
)
cache_dir: Optional[str] = field(
default=None,
@@ -117,10 +123,12 @@ class ModelArguments:
},
)
freeze_feature_encoder: bool = field(
- default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+ default=True,
+ metadata={"help": "Whether to freeze the feature encoder layers of the model."},
)
freeze_encoder: bool = field(
- default=False, metadata={"help": "Whether to freeze the entire encoder of the seq2seq model."}
+ default=False,
+ metadata={"help": "Whether to freeze the entire encoder of the seq2seq model."},
)
forced_decoder_ids: list[list[int]] = field(
default=None,
@@ -150,13 +158,17 @@ class DataTrainingArguments:
"""
dataset_name: str = field(
- default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+ metadata={"help": "Path or name of the dataset (cf `load_dataset` method of the Datasets library)."}
)
- dataset_config_name: Optional[str] = field(
- default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ dataset_config_name: str = field(
+ default=None,
+ metadata={
+ "help": "The configuration name of the dataset to use (cf `load_dataset` method of the Datasets library)."
+ },
)
overwrite_cache: bool = field(
- default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+ default=False,
+ metadata={"help": "Overwrite the cached training and evaluation sets"},
)
preprocessing_num_workers: Optional[int] = field(
default=None,
@@ -198,7 +210,8 @@ class DataTrainingArguments:
},
)
min_duration_in_seconds: float = field(
- default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+ default=0.0,
+ metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
)
preprocessing_only: bool = field(
default=False,
@@ -300,10 +313,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_speech_recognition_seq2seq", model_args, data_args)
-
# 2. Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -391,7 +400,7 @@ def main():
# Distributed training:
# The .from_pretrained methods guarantee that only one local process can concurrently
config = AutoConfig.from_pretrained(
- model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+ (model_args.config_name if model_args.config_name else model_args.model_name_or_path),
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
@@ -403,14 +412,14 @@ def main():
config.update({"apply_spec_augment": model_args.apply_spec_augment})
feature_extractor = AutoFeatureExtractor.from_pretrained(
- model_args.feature_extractor_name if model_args.feature_extractor_name else model_args.model_name_or_path,
+ (model_args.feature_extractor_name if model_args.feature_extractor_name else model_args.model_name_or_path),
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
trust_remote_code=model_args.trust_remote_code,
)
tokenizer = AutoTokenizer.from_pretrained(
- model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+ (model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path),
cache_dir=model_args.cache_dir,
use_fast=model_args.use_fast_tokenizer,
revision=model_args.model_revision,
@@ -469,7 +478,8 @@ def main():
dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
if dataset_sampling_rate != feature_extractor.sampling_rate:
raw_datasets = raw_datasets.cast_column(
- data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
+ data_args.audio_column_name,
+ datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate),
)
# 7. Preprocessing the datasets.
@@ -498,7 +508,9 @@ def prepare_dataset(batch):
# process audio
sample = batch[audio_column_name]
inputs = feature_extractor(
- sample["array"], sampling_rate=sample["sampling_rate"], return_attention_mask=forward_attention_mask
+ sample["array"],
+ sampling_rate=sample["sampling_rate"],
+ return_attention_mask=forward_attention_mask,
)
# process audio length
batch[model_input_name] = inputs.get(model_input_name)[0]
@@ -583,7 +595,7 @@ def compute_metrics(pred):
eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,
processing_class=feature_extractor,
data_collator=data_collator,
- compute_metrics=compute_metrics if training_args.predict_with_generate else None,
+ compute_metrics=(compute_metrics if training_args.predict_with_generate else None),
)
# 12. Training
@@ -625,7 +637,10 @@ def compute_metrics(pred):
trainer.save_metrics("eval", metrics)
# 14. Write Training Stats
- kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "automatic-speech-recognition"}
+ kwargs = {
+ "finetuned_from": model_args.model_name_or_path,
+ "tasks": "automatic-speech-recognition",
+ }
if data_args.dataset_name is not None:
kwargs["dataset_tags"] = data_args.dataset_name
if data_args.dataset_config_name is not None:
diff --git a/examples/pytorch/summarization/run_summarization.py b/examples/pytorch/summarization/run_summarization.py
index e3554ec85829..dd7dd083b49a 100755
--- a/examples/pytorch/summarization/run_summarization.py
+++ b/examples/pytorch/summarization/run_summarization.py
@@ -62,7 +62,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, is_offline_mode, send_example_telemetry
+from transformers.utils import check_min_version, is_offline_mode
from transformers.utils.versions import require_version
@@ -337,10 +337,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_summarization", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/summarization/run_summarization_no_trainer.py b/examples/pytorch/summarization/run_summarization_no_trainer.py
index 19366f7b7248..b24bbc773bc0 100644
--- a/examples/pytorch/summarization/run_summarization_no_trainer.py
+++ b/examples/pytorch/summarization/run_summarization_no_trainer.py
@@ -66,7 +66,7 @@
SchedulerType,
get_scheduler,
)
-from transformers.utils import check_min_version, is_offline_mode, send_example_telemetry
+from transformers.utils import check_min_version, is_offline_mode
from transformers.utils.versions import require_version
@@ -338,9 +338,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_summarization_no_trainer", args)
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
diff --git a/examples/pytorch/text-classification/run_classification.py b/examples/pytorch/text-classification/run_classification.py
index 17eaccd96baf..35413cd7875b 100755
--- a/examples/pytorch/text-classification/run_classification.py
+++ b/examples/pytorch/text-classification/run_classification.py
@@ -56,7 +56,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -296,10 +296,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_classification", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/text-classification/run_glue.py b/examples/pytorch/text-classification/run_glue.py
index 1c8df2d54daf..afa09d746041 100755
--- a/examples/pytorch/text-classification/run_glue.py
+++ b/examples/pytorch/text-classification/run_glue.py
@@ -58,7 +58,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -241,10 +241,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_glue", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/text-classification/run_glue_no_trainer.py b/examples/pytorch/text-classification/run_glue_no_trainer.py
index a706e003f69e..05c51eb8ae3a 100644
--- a/examples/pytorch/text-classification/run_glue_no_trainer.py
+++ b/examples/pytorch/text-classification/run_glue_no_trainer.py
@@ -58,7 +58,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -234,9 +234,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_glue_no_trainer", args)
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
diff --git a/examples/pytorch/text-classification/run_xnli.py b/examples/pytorch/text-classification/run_xnli.py
index beb7bb778b1d..3027da5feae6 100755
--- a/examples/pytorch/text-classification/run_xnli.py
+++ b/examples/pytorch/text-classification/run_xnli.py
@@ -57,7 +57,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -199,10 +199,6 @@ def main():
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_xnli", model_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/token-classification/run_ner.py b/examples/pytorch/token-classification/run_ner.py
index d5bdb9ee3662..7620d697c126 100755
--- a/examples/pytorch/token-classification/run_ner.py
+++ b/examples/pytorch/token-classification/run_ner.py
@@ -55,7 +55,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -238,10 +238,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_ner", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/token-classification/run_ner_no_trainer.py b/examples/pytorch/token-classification/run_ner_no_trainer.py
index 7d5256f48e81..0d31cf46ab8a 100755
--- a/examples/pytorch/token-classification/run_ner_no_trainer.py
+++ b/examples/pytorch/token-classification/run_ner_no_trainer.py
@@ -62,7 +62,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -284,10 +284,6 @@ def parse_args():
def main():
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_ner_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/pytorch/translation/run_translation.py b/examples/pytorch/translation/run_translation.py
index dcfe9a6ffe94..8e005e0d7323 100755
--- a/examples/pytorch/translation/run_translation.py
+++ b/examples/pytorch/translation/run_translation.py
@@ -61,7 +61,7 @@
set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -285,10 +285,6 @@ def main():
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_translation", model_args, data_args)
-
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
diff --git a/examples/pytorch/translation/run_translation_no_trainer.py b/examples/pytorch/translation/run_translation_no_trainer.py
index 871504bb9877..9dae2ec653f0 100644
--- a/examples/pytorch/translation/run_translation_no_trainer.py
+++ b/examples/pytorch/translation/run_translation_no_trainer.py
@@ -66,7 +66,7 @@
default_data_collator,
get_scheduler,
)
-from transformers.utils import check_min_version, send_example_telemetry
+from transformers.utils import check_min_version
from transformers.utils.versions import require_version
@@ -330,10 +330,6 @@ def main():
# Parse the arguments
args = parse_args()
- # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
- # information sent is the one passed as arguments along with your Python/PyTorch versions.
- send_example_telemetry("run_translation_no_trainer", args)
-
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
# If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
# in the environment
diff --git a/examples/quantization/custom_quantization_int8_example.py b/examples/quantization/custom_quantization_int8_example.py
index 4bf907b77fe5..884b943f696b 100644
--- a/examples/quantization/custom_quantization_int8_example.py
+++ b/examples/quantization/custom_quantization_int8_example.py
@@ -159,24 +159,13 @@ def _process_model_before_weight_loading(self, model, **kwargs):
pre_quantized=self.pre_quantized,
)
- def check_quantized_param(
- self,
- model,
- param_value: "torch.Tensor",
- param_name: str,
- state_dict: dict[str, Any],
- **kwargs,
- ):
+ def param_needs_quantization(self, model, param_name: str, **kwargs) -> bool:
module, tensor_name = get_module_from_name(model, param_name)
if isinstance(module, Int8SymmetricLinear):
if self.pre_quantized or tensor_name == "bias":
- if tensor_name == "weight" and param_value.dtype != torch.int8:
- raise ValueError("Expect quantized weights but got an unquantized weight")
return False
else:
- if tensor_name == "weight_scale":
- raise ValueError("Expect unquantized weights but got a quantized weight_scale")
return True
return False
@@ -186,12 +175,18 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: dict[str, Any],
- unexpected_keys: Optional[list[str]] = None,
+ **kwargs,
):
- """
- Quantizes weights to INT8 symmetric format.
- """
+ # Sanity check
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module, Int8SymmetricLinear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.int8:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ else:
+ if tensor_name == "weight_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+
abs_max_per_row = torch.max(torch.abs(param_value), dim=1, keepdim=True)[0].clamp(min=1e-5)
weight_scale = abs_max_per_row / 127.0
diff --git a/i18n/README_ar.md b/i18n/README_ar.md
index cdf813445d6f..a0c86c770600 100644
--- a/i18n/README_ar.md
+++ b/i18n/README_ar.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_bn.md b/i18n/README_bn.md
new file mode 100644
index 000000000000..354521ee7ba3
--- /dev/null
+++ b/i18n/README_bn.md
@@ -0,0 +1,335 @@
+
+
+
+ 
+
+
+
+ 
+ 
+ 
+ 
+ 
+
+
+ English | + 简体中文 | + 繁體中文 | + 한국어 | + Español | + 日本語 | + हिन्दी | + Русский | + Português | + తెలుగు | + Français | + Deutsch | + Italiano | + Tiếng Việt | + العربية | + اردو | + বাংলা | +
+ + ++
ইনফারেন্স ও ট্রেনিংয়ের জন্য আধুনিকতম (State-of-the-art) প্রি-ট্রেইন্ড মডেলসমূহ
+ + +
+
+
+
+
+**Transformers** হলো একটা ফ্রেমওয়ার্ক যেটা দিয়ে টেক্সট, কম্পিউটার ভিশন, অডিও, ভিডিও আর মাল্টিমোডাল—সব ধরনের মডেল তৈরি আর চালানো যায়। এটা ট্রেইনিং আর ইনফারেন্স – দুই কাজেই ব্যবহার করা হয়।
+
+Transformers মডেলের ডেফিনিশন এক জায়গায় রাখে। এর মানে হলো, একবার কোনো মডেল `transformers`-এ সাপোর্ট পেলেই সেটা সহজে বিভিন্ন ট্রেইনিং ফ্রেমওয়ার্ক (Axolotl, Unsloth, DeepSpeed, FSDP, PyTorch-Lightning ইত্যাদি), ইনফারেন্স ইঞ্জিন (vLLM, SGLang, TGI ইত্যাদি) আর অন্যান্য লাইব্রেরি (llama.cpp, mlx ইত্যাদি)-তে ব্যবহার করা যায়।
+
+আমরা চাই নতুন আর আধুনিক মডেলগুলো সবাই ব্যবহার করতে পারে। তাই মডেলের ডেফিনিশন রাখা হয়েছে সহজ, কাস্টমাইজযোগ্য আর পারফরম্যান্স-ফ্রেন্ডলি।
+
+এখন পর্যন্ত [Hugging Face Hub](https://huggingface.com/models)-এ ১০ লাখেরও বেশি Transformers [মডেল চেকপয়েন্ট](https://huggingface.co/models?library=transformers&sort=trending) আছে, যেগুলো যেকোনো সময় ব্যবহার করা যায়।
+
+আজই [Hub](https://huggingface.com/) থেকে একটা মডেল বেছে নিন আর Transformers দিয়ে শুরু করুন।
+
+
+## ইনস্টলেশন
+
+Transformers Python 3.9+ সহ কাজ করে, এবং সমর্থিত ফ্রেমওয়ার্কগুলো হলো [PyTorch](https://pytorch.org/get-started/locally/) 2.1+, [TensorFlow](https://www.tensorflow.org/install/pip) 2.6+, এবং [Flax](https://flax.readthedocs.io/en/latest/) 0.4.1+।
+
+[venv](https://docs.python.org/3/library/venv.html) বা [uv](https://docs.astral.sh/uv/) ব্যবহার করে একটি ভার্চুয়াল এনভায়রনমেন্ট তৈরি এবং সক্রিয় করুন।
+
+```py
+# venv
+python -m venv .my-env
+source .my-env/bin/activate
+# uv
+uv venv .my-env
+source .my-env/bin/activate
+```
+আপনার ভার্চুয়াল পরিবেশে Transformers ইনস্টল করুন।
+
+```py
+# pip
+pip install "transformers[torch]"
+
+# uv
+uv pip install "transformers[torch]"
+```
+যদি আপনি লাইব্রেরির সর্বশেষ পরিবর্তনগুলি চান বা অবদান রাখতে আগ্রহী হন তবে উৎস থেকে Transformers ইনস্টল করুন। তবে, সর্বশেষ সংস্করণটি স্থিতিশীল নাও হতে পারে। যদি আপনি কোনো ত্রুটির সম্মুখীন হন তবে নির্দ্বিধায় একটি [issue](https://github.com/huggingface/transformers/issues) খুলুন।
+
+```Shell
+git clone [https://github.com/huggingface/transformers.git](https://github.com/huggingface/transformers.git)
+cd transformers
+
+# pip
+pip install .[torch]
+
+# uv
+uv pip install .[torch]
+```
+
+## কুইকস্টার্ট
+
+Transformers ব্যবহার শুরু করুন এখনই [Pipeline](https://huggingface.co/docs/transformers/pipeline_tutorial) API দিয়ে। `Pipeline` হলো একটি হাই-লেভেল ইনফারেন্স ক্লাস, যা টেক্সট, অডিও, ভিশন এবং মাল্টিমোডাল টাস্ক সাপোর্ট করে। এটি ইনপুট প্রিপ্রসেসিং করে এবং সঠিক আউটপুট রিটার্ন করে।
+
+একটি পাইপলাইন তৈরি করুন এবং টেক্সট জেনারেশনের জন্য কোন মডেল ব্যবহার করবেন তা নির্দিষ্ট করুন। মডেলটি ডাউনলোড হয়ে ক্যাশে রাখা হবে, ফলে পরে সহজেই আবার ব্যবহার করতে পারবেন। সবশেষে, মডেলকে প্রম্পট করার জন্য কিছু টেক্সট দিন।
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="text-generation", model="Qwen/Qwen2.5-1.5B")
+pipeline("the secret to baking a really good cake is ")
+[{'generated_text': 'the secret to baking a really good cake is 1) to use the right ingredients and 2) to follow the recipe exactly. the recipe for the cake is as follows: 1 cup of sugar, 1 cup of flour, 1 cup of milk, 1 cup of butter, 1 cup of eggs, 1 cup of chocolate chips. if you want to make 2 cakes, how much sugar do you need? To make 2 cakes, you will need 2 cups of sugar.'}]
+```
+
+মডেলের সাথে চ্যাট করতে হলেও ব্যবহার প্যাটার্ন একই। শুধু পার্থক্য হলো, আপনাকে একটি চ্যাট হিস্ট্রি তৈরি করতে হবে (যা `Pipeline`-এ ইনপুট হিসেবে যাবে) আপনার আর সিস্টেমের মধ্যে।
+
+> [!TIP]
+> আপনি সরাসরি কমান্ড লাইন থেকেও একটি মডেলের সাথে চ্যাট করতে পারেন।
+> ```Shell
+> transformers chat Qwen/Qwen2.5-0.5B-Instruct
+> ```
+
+```Python
+import torch
+from transformers import pipeline
+
+chat = [
+ {"role": "system", "content": "You are a sassy, wise-cracking robot as imagined by Hollywood circa 1986."},
+ {"role": "user", "content": "Hey, can you tell me any fun things to do in New York?"}
+]
+
+pipeline = pipeline(task="text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct", dtype=torch.bfloat16, device_map="auto")
+response = pipeline(chat, max_new_tokens=512)
+print(response[0]["generated_text"][-1]["content"])
+
+বিভিন্ন মোডালিটি এবং কাজের জন্য Pipeline কিভাবে কাজ করে তা দেখতে নিচের উদাহরণগুলো সম্প্রসারণ করুন।
+```
+
+
+
+
+অটোমেটিক স্পিচ রিকগনিশন (ASR)
+ +```Python +from transformers import pipeline + +pipeline = pipeline(task="automatic-speech-recognition", model="openai/whisper-large-v3") +pipeline("[https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac](https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac)") +{'text': ' I have a dream that one day this nation will rise up and live out the true meaning of its creed.'} +``` + +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="image-classification", model="facebook/dinov2-small-imagenet1k-1-layer")
+pipeline("[https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png](https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png)")
+[{'label': 'macaw', 'score': 0.997848391532898},
+ {'label': 'sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita',
+ 'score': 0.0016551691805943847},
+ {'label': 'lorikeet', 'score': 0.00018523589824326336},
+ {'label': 'African grey, African gray, Psittacus erithacus',
+ 'score': 7.85409429227002e-05},
+ {'label': 'quail', 'score': 5.502637941390276e-05}]
+ ```
+
+
+ইমেজ ক্লাসিফিকেশন
+ +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="image-classification", model="facebook/dinov2-small-imagenet1k-1-layer")
+pipeline("[https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png](https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png)")
+[{'label': 'macaw', 'score': 0.997848391532898},
+ {'label': 'sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita',
+ 'score': 0.0016551691805943847},
+ {'label': 'lorikeet', 'score': 0.00018523589824326336},
+ {'label': 'African grey, African gray, Psittacus erithacus',
+ 'score': 7.85409429227002e-05},
+ {'label': 'quail', 'score': 5.502637941390276e-05}]
+ ```
+
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="visual-question-answering", model="Salesforce/blip-vqa-base")
+pipeline(
+ image="[https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg)",
+ question="What is in the image?",
+)
+[{'answer': 'statue of liberty'}]
+```
+
+
+## কেন Transformers ব্যবহার করবেন?
+
+1. সহজে ব্যবহারযোগ্য সর্বাধুনিক মডেল:
+
+ * ন্যাচারাল ল্যাঙ্গুয়েজ আন্ডারস্ট্যান্ডিং ও জেনারেশন, কম্পিউটার ভিশন, অডিও, ভিডিও এবং মাল্টিমোডাল টাস্কে উচ্চ পারফরম্যান্স।
+ * গবেষক, ইঞ্জিনিয়ার এবং ডেভেলপারদের জন্য সহজে শুরু করার সুযোগ।
+ * মাত্র তিনটি ক্লাস শিখলেই ব্যবহার করা যায়।
+ * সব প্রি-ট্রেইন্ড মডেলের জন্য একটি একীভূত API।
+
+2. কম কম্পিউট খরচ, ছোট কার্বন ফুটপ্রিন্ট:
+
+ * শূন্য থেকে ট্রেইন না করে ট্রেইন্ড মডেল শেয়ার করুন।
+ * কম্পিউট টাইম ও প্রোডাকশন খরচ কমান।
+ * সব ধরনের মোডালিটির জন্য ১০ লক্ষ+ প্রি-ট্রেইন্ড চেকপয়েন্টসহ ডজনখানেক মডেল আর্কিটেকচার।
+
+3. মডেলের লাইফসাইকেলের প্রতিটি ধাপে সঠিক ফ্রেমওয়ার্ক বেছে নিন:
+
+ * মাত্র ৩ লাইনের কোডে সর্বাধুনিক মডেল ট্রেইন করুন।
+ * সহজে PyTorch / JAX / TF2.0 এর মধ্যে মডেল স্থানান্তর করুন।
+ * ট্রেইনিং, ইভ্যালুয়েশন ও প্রোডাকশনের জন্য আলাদা ফ্রেমওয়ার্ক ব্যবহার করুন।
+
+4. সহজেই মডেল বা উদাহরণ কাস্টমাইজ করুন:
+
+ * প্রতিটি আর্কিটেকচারের জন্য এমন উদাহরণ দেওয়া আছে যা মূল লেখকদের প্রকাশিত ফলাফল পুনরুত্পাদন করতে সক্ষম।
+ * মডেলের অভ্যন্তরীণ অংশগুলো যতটা সম্ভব একভাবে এক্সপোজ করা হয়েছে।
+ * দ্রুত এক্সপেরিমেন্টের জন্য লাইব্রেরি ছাড়াও মডেল ফাইল ব্যবহার করা যায়।
+
+
+
+ভিজুয়াল কোয়েশ্চন আনসারিং
+ +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="visual-question-answering", model="Salesforce/blip-vqa-base")
+pipeline(
+ image="[https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg)",
+ question="What is in the image?",
+)
+[{'answer': 'statue of liberty'}]
+```
++ +## কেন Transformers ব্যবহার করবেন না? + +* এই লাইব্রেরি নিউরাল নেটওয়ার্কের জন্য ব্লক-মডিউল টুলবক্স নয়। মডেল ফাইলের কোডে অতিরিক্ত অ্যাবস্ট্র্যাকশন intentionally করা হয়নি, যাতে গবেষকরা দ্রুত প্রতিটি মডেলের উপর কাজ করতে পারে কোনো অতিরিক্ত ফাইল বা স্তরে না গিয়ে। +* ট্রেইনিং API মূলত Transformers-এর PyTorch মডেলের সাথে কাজ করার জন্য অপটিমাইজ করা হয়েছে। সাধারণ মেশিন লার্নিং লুপের জন্য, [Accelerate](https://huggingface.co/docs/accelerate) এর মতো অন্য লাইব্রেরি ব্যবহার করা উচিত। +* [উদাহরণ স্ক্রিপ্টগুলো](https://github.com/huggingface/transformers/tree/main/examples) শুধু *উদাহরণ*। এগুলো সরাসরি আপনার ব্যবহারের ক্ষেত্রে কাজ নাও করতে পারে, তাই কোড সামঞ্জস্য করতে হতে পারে। + +## Transformers দিয়ে ১০০টি প্রজেক্ট + +Transformers শুধু প্রি-ট্রেইন্ড মডেল ব্যবহার করার টুলকিট নয়, এটি একটি কমিউনিটি, যা Hugging Face Hub-এর চারপাশে তৈরি। আমরা চাই যে ডেভেলপার, গবেষক, শিক্ষার্থী, অধ্যাপক, ইঞ্জিনিয়ার বা যে কেউ তাদের স্বপ্নের প্রজেক্ট তৈরি করতে পারে। + +Transformers 100,000 স্টার উদযাপন করতে আমরা কমিউনিটিকে তুলে ধরতে [awesome-transformers](./awesome-transformers.md) পেজ তৈরি করেছি, যেখানে Transformers দিয়ে তৈরি ১০০টি অসাধারণ প্রজেক্ট তালিকাভুক্ত আছে। + +আপনার কোনো প্রজেক্ট আছে যা তালিকায় থাকা উচিত মনে করেন? তাহলে PR খুলে যুক্ত করুন। + +## উদাহরণ মডেল + +আপনি আমাদের অধিকাংশ মডেল সরাসরি তাদের [Hub মডেল পেজ](https://huggingface.co/models) থেকে পরীক্ষা করতে পারেন। + +নিচের প্রতিটি মোডালিটি এক্সপ্যান্ড করে বিভিন্ন ব্যবহার কেসের জন্য কয়েকটি উদাহরণ মডেল দেখুন। + + +
+
+
+অডিও
+ +* [Whisper](https://huggingface.co/openai/whisper-large-v3-turbo) দিয়ে অডিও ক্লাসিফিকেশন +* [Moonshine](https://huggingface.co/UsefulSensors/moonshine) দিয়ে অটোমেটিক স্পিচ রিকগনিশন +* [Wav2Vec2](https://huggingface.co/superb/wav2vec2-base-superb-ks) দিয়ে কীওয়ার্ড স্পটিং +* [Moshi](https://huggingface.co/kyutai/moshiko-pytorch-bf16) দিয়ে স্পিচ-টু-স্পিচ জেনারেশন +* [MusicGen](https://huggingface.co/facebook/musicgen-large) দিয়ে টেক্সট-টু-অডিও +* [Bark](https://huggingface.co/suno/bark) দিয়ে টেক্সট-টু-স্পিচ + + +
+
+
+কম্পিউটার ভিশন
+ +* [SAM](https://huggingface.co/facebook/sam-vit-base) দিয়ে স্বয়ংক্রিয় মাস্ক জেনারেশন +* [DepthPro](https://huggingface.co/apple/DepthPro-hf) দিয়ে গভীরতা অনুমান +* [DINO v2](https://huggingface.co/facebook/dinov2-base) দিয়ে চিত্র শ্রেণীকরণ +* [SuperPoint](https://huggingface.co/magic-leap-community/superpoint) দিয়ে কীপয়েন্ট সনাক্তকরণ +* [SuperGlue](https://huggingface.co/magic-leap-community/superglue_outdoor) দিয়ে কীপয়েন্ট ম্যাচিং +* [RT-DETRv2](https://huggingface.co/PekingU/rtdetr_v2_r50vd) দিয়ে অবজেক্ট সনাক্তকরণ +* [VitPose](https://huggingface.co/usyd-community/vitpose-base-simple) দিয়ে পোস অনুমান +* [OneFormer](https://huggingface.co/shi-labs/oneformer_ade20k_swin_large) দিয়ে ইউনিভার্সাল সেগমেন্টেশন +* [VideoMAE](https://huggingface.co/MCG-NJU/videomae-large) দিয়ে ভিডিও শ্রেণীকরণ + + +
+
+
+মাল্টিমোডাল
+ +* [Qwen2-Audio](https://huggingface.co/Qwen/Qwen2-Audio-7B) দিয়ে অডিও বা টেক্সট থেকে টেক্সট জেনারেশন +* [LayoutLMv3](https://huggingface.co/microsoft/layoutlmv3-base) দিয়ে ডকুমেন্ট প্রশ্নোত্তর +* [Qwen-VL](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) দিয়ে ইমেজ বা টেক্সট থেকে টেক্সট জেনারেশন +* [BLIP-2](https://huggingface.co/Salesforce/blip2-opt-2.7b) দিয়ে ইমেজ ক্যাপশনিং +* [GOT-OCR2](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) দিয়ে OCR-ভিত্তিক ডকুমেন্ট আন্ডারস্ট্যান্ডিং +* [TAPAS](https://huggingface.co/google/tapas-base) দিয়ে টেবিল প্রশ্নোত্তর +* [Emu3](https://huggingface.co/BAAI/Emu3-Gen) দিয়ে ইউনিফাইড মাল্টিমোডাল আন্ডারস্ট্যান্ডিং এবং জেনারেশন +* [Llava-OneVision](https://huggingface.co/llava-hf/llava-onevision-qwen2-0.5b-ov-hf) দিয়ে ভিশন থেকে টেক্সট +* [Llava](https://huggingface.co/llava-hf/llava-1.5-7b-hf) দিয়ে ভিজুয়াল কোয়েশ্চন আনসারিং +* [Kosmos-2](https://huggingface.co/microsoft/kosmos-2-patch14-224) দিয়ে ভিজুয়াল রেফারিং এক্সপ্রেশন সেগমেন্টেশন + + +
+
+
+## সাইটেশন
+আমাদের [একটি পেপার](https://www.aclweb.org/anthology/2020.emnlp-demos.6/) আছে যা আপনি 🤗 Transformers লাইব্রেরির জন্য রেফারেন্স হিসেবে ব্যবহার করতে পারেন।
+
+```bibtex
+@inproceedings{wolf-etal-2020-transformers,
+ title = "Transformers: State-of-the-Art Natural Language Processing",
+ author = "Thomas Wolf and Lysandre Debut and Victor Sanh and Julien Chaumond and Clement Delangue and Anthony Moi and Pierric Cistac and Tim Rault and Rémi Louf and Morgan Funtowicz and Joe Davison and Sam Shleifer and Patrick von Platen and Clara Ma and Yacine Jernite and Julien Plu and Canwen Xu and Teven Le Scao and Sylvain Gugger and Mariama Drame and Quentin Lhoest and Alexander M. Rush",
+ booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations",
+ month = oct,
+ year = "2020",
+ address = "Online",
+ publisher = "Association for Computational Linguistics",
+ url = "https://www.aclweb.org/anthology/2020.emnlp-demos.6",
+ pages = "38--45"
+}
+```
\ No newline at end of file
diff --git a/i18n/README_de.md b/i18n/README_de.md
index b913df894dc1..2c54965371c1 100644
--- a/i18n/README_de.md
+++ b/i18n/README_de.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_es.md b/i18n/README_es.md
index d31b7f5f76c3..1a7a2256424a 100644
--- a/i18n/README_es.md
+++ b/i18n/README_es.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_fr.md b/i18n/README_fr.md
index 6512b4af0700..17e6c0424269 100644
--- a/i18n/README_fr.md
+++ b/i18n/README_fr.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_hd.md b/i18n/README_hd.md
index 1eb220efadc0..6c441088834c 100644
--- a/i18n/README_hd.md
+++ b/i18n/README_hd.md
@@ -72,9 +72,11 @@ checkpoint: जाँच बिंदु
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_it.md b/i18n/README_it.md
new file mode 100644
index 000000000000..3b8d71bdb721
--- /dev/null
+++ b/i18n/README_it.md
@@ -0,0 +1,337 @@
+
+
+NLP
+ +* [ModernBERT](https://huggingface.co/answerdotai/ModernBERT-base) দিয়ে মাস্কড ওয়ার্ড কমপ্লিশন +* [Gemma](https://huggingface.co/google/gemma-2-2b) দিয়ে নাম্বড এন্টিটি রিকগনিশন +* [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) দিয়ে প্রশ্নোত্তর +* [BART](https://huggingface.co/facebook/bart-large-cnn) দিয়ে সারসংক্ষেপ (Summarization) +* [T5](https://huggingface.co/google-t5/t5-base) দিয়ে অনুবাদ +* [Llama](https://huggingface.co/meta-llama/Llama-3.2-1B) দিয়ে টেক্সট জেনারেশন +* [Qwen](https://huggingface.co/Qwen/Qwen2.5-0.5B) দিয়ে টেক্সট ক্লাসিফিকেশন + +
+
+
+
+
+
+
+
+
+
+
+
+ English | + 简体中文 | + 繁體中文 | + 한국어 | + Español | + 日本語 | + हिन्दी | + Русский | + Português | + తెలుగు | + Français | + Deutsch | + Italiano | + Tiếng Việt | + العربية | + اردو | + বাংলা | +
+ + ++
Modelli preaddestrati all'avanguardia per l'inferenza e l'addestramento
+ + +
+
+
+
+Transformers funge da framework di definizione dei modelli per modelli di machine learning all'avanguardia nei
+modelli di testo, visione artificiale, audio, video e multimodali, sia per l'inferenza che per l'addestramento.
+
+Centralizza la definizione del modello in modo che tale definizione sia concordata all'interno dell'ecosistema.
+`transformers` è il perno tra i framework: se una definizione di modello è supportata, sarà compatibile con la
+maggior parte dei framework di addestramento (Axolotl, Unsloth, DeepSpeed, FSDP, PyTorch-Lightning, ...), motori
+di inferenza (vLLM, SGLang, TGI, ...) e librerie di modellazione adiacenti (llama.cpp, mlx, ...) che sfruttano
+la definizione del modello da `transformers`.
+
+Ci impegniamo a sostenere nuovi modelli all'avanguardia e a democratizzarne l'utilizzo rendendo la loro definizione
+semplice, personalizzabile ed efficiente.
+
+Ci sono oltre 1 milione di Transformers [model checkpoint](https://huggingface.co/models?library=transformers&sort=trending) su [Hugging Face Hub](https://huggingface.com/models) che puoi utilizzare.
+
+Esplora oggi stesso l'[Hub](https://huggingface.com/) per trovare un modello e utilizzare Transformers per aiutarti a iniziare subito.
+
+## Installazione
+
+Transformers funziona con Python 3.9+ e [PyTorch](https://pytorch.org/get-started/locally/) 2.1+.
+
+Crea e attiva un ambiente virtuale con [venv](https://docs.python.org/3/library/venv.html) o [uv](https://docs.astral.sh/uv/), un pacchetto Python veloce basato su Rust e un gestore di progetti.
+
+```py
+# venv
+python -m venv .my-env
+source .my-env/bin/activate
+# uv
+uv venv .my-env
+source .my-env/bin/activate
+```
+
+Installa Transformers nel tuo ambiente virtuale.
+
+```py
+# pip
+pip install "transformers[torch]"
+
+# uv
+uv pip install "transformers[torch]"
+```
+
+Installa Transformers dal sorgente se desideri le ultime modifiche nella libreria o sei interessato a contribuire. Tuttavia, la versione *più recente* potrebbe non essere stabile. Non esitare ad aprire una [issue](https://github.com/huggingface/transformers/issues) se riscontri un errore.
+
+```shell
+git clone https://github.com/huggingface/transformers.git
+cd transformers
+
+# pip
+pip install .[torch]
+
+# uv
+uv pip install .[torch]
+```
+
+## Quickstart
+
+Inizia subito a utilizzare Transformers con l'API [Pipeline](https://huggingface.co/docs/transformers/pipeline_tutorial). Pipeline è una classe di inferenza di alto livello che supporta attività di testo, audio, visione e multimodali. Gestisce la pre-elaborazione dell'input e restituisce l'output appropriato.
+
+Istanziare una pipeline e specificare il modello da utilizzare per la generazione di testo. Il modello viene scaricato e memorizzato nella cache in modo da poterlo riutilizzare facilmente. Infine, passare del testo per attivare il modello.
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="text-generation", model="Qwen/Qwen2.5-1.5B")
+pipeline("il segreto per preparare una torta davvero buona è ")
+[{'generated_text': 'il segreto per preparare una torta davvero buona è 1) usare gli ingredienti giusti e 2) seguire alla lettera la ricetta. la ricetta della torta è la seguente: 1 tazza di zucchero, 1 tazza di farina, 1 tazza di latte, 1 tazza di burro, 1 tazza di uova, 1 tazza di gocce di cioccolato. se vuoi preparare 2 torte, quanto zucchero ti serve? Per preparare 2 torte, avrete bisogno di 2 tazze di zucchero.'}]
+```
+
+Per chattare con un modello, lo schema di utilizzo è lo stesso. L'unica differenza è che è necessario creare una cronologia delle chat (l'input per `Pipeline`) tra l'utente e il sistema.
+
+> [!TIP]
+> È anche possibile chattare con un modello direttamente dalla riga di comando.
+> ```shell
+> transformers chat Qwen/Qwen2.5-0.5B-Instruct
+> ```
+
+```py
+import torch
+from transformers import pipeline
+
+chat = [
+ {"role": "system", "content": "Sei un robot sfacciato e spiritoso, proprio come lo immaginava Hollywood nel 1986."},
+ {"role": "user", "content": "Ehi, mi puoi suggerire qualcosa di divertente da fare a New York?"}
+]
+
+pipeline = pipeline(task="text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct", dtype=torch.bfloat16, device_map="auto")
+response = pipeline(chat, max_new_tokens=512)
+print(response[0]["generated_text"][-1]["content"])
+```
+
+Espandi gli esempi riportati di seguito per vedere come funziona `Pipeline` per diverse modalità e attività.
+
+
+
+
+Riconoscimento vocale automatico
+ +```py +from transformers import pipeline + +pipeline = pipeline(task="automatic-speech-recognition", model="openai/whisper-large-v3") +pipeline("https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac") +{'text': ' Ho un sogno: che un giorno questa nazione si solleverà e vivrà il vero significato del suo credo.'} +``` + +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="image-classification", model="facebook/dinov2-small-imagenet1k-1-layer")
+pipeline("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png")
+[{'label': 'macaw', 'score': 0.997848391532898},
+ {'label': 'cacatua dal ciuffo giallo, Kakatoe galerita, Cacatua galerita',
+ 'score': 0.0016551691805943847},
+ {'label': 'lorichetto', 'score': 0.00018523589824326336},
+ {'label': 'Pappagallo grigio africano, Psittacus erithacus',
+ 'score': 7.85409429227002e-05},
+ {'label': 'quaglia', 'score': 5.502637941390276e-05}]
+```
+
+
+
+Classificazione delle immagini
+ +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="image-classification", model="facebook/dinov2-small-imagenet1k-1-layer")
+pipeline("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png")
+[{'label': 'macaw', 'score': 0.997848391532898},
+ {'label': 'cacatua dal ciuffo giallo, Kakatoe galerita, Cacatua galerita',
+ 'score': 0.0016551691805943847},
+ {'label': 'lorichetto', 'score': 0.00018523589824326336},
+ {'label': 'Pappagallo grigio africano, Psittacus erithacus',
+ 'score': 7.85409429227002e-05},
+ {'label': 'quaglia', 'score': 5.502637941390276e-05}]
+```
+
+
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="visual-question-answering", model="Salesforce/blip-vqa-base")
+pipeline(
+ image="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg",
+ question="Cosa c'è nell'immagine?",
+)
+[{'answer': 'statua della libertà'}]
+```
+
+
+
+## Perché dovrei usare Transformers?
+
+1. Modelli all'avanguardia facili da usare:
+ - Prestazioni elevate nella comprensione e generazione del linguaggio naturale, nella visione artificiale, nell'audio, nel video e nelle attività multimodali.
+ - Bassa barriera di ingresso per ricercatori, ingegneri e sviluppatori.
+ - Poche astrazioni rivolte all'utente con solo tre classi da imparare.
+ - Un'API unificata per l'utilizzo di tutti i nostri modelli preaddestrati.
+
+1. Riduzione dei costi di calcolo e dell'impronta di carbonio:
+ - Condivisione dei modelli addestrati invece di addestrarli da zero.
+ - Riduzione dei tempi di calcolo e dei costi di produzione.
+ - Decine di architetture di modelli con oltre 1 milione di checkpoint preaddestrati in tutte le modalità.
+
+1. Scegli il framework giusto per ogni fase del ciclo di vita di un modello:
+ - Addestra modelli all'avanguardia con sole 3 righe di codice.
+ - Sposta un singolo modello tra i framework PyTorch/JAX/TF2.0 a tuo piacimento.
+ - Scegli il framework giusto per l'addestramento, la valutazione e la produzione.
+
+1. Personalizza facilmente un modello o un esempio in base alle tue esigenze:
+ - Forniamo esempi per ogni architettura per riprodurre i risultati pubblicati dagli autori originali.
+ - Gli interni del modello sono esposti nel modo più coerente possibile.
+ - I file del modello possono essere utilizzati indipendentemente dalla libreria per esperimenti rapidi.
+
+
+ Risposta a domande visive
+ +
+
+
+
+```py
+from transformers import pipeline
+
+pipeline = pipeline(task="visual-question-answering", model="Salesforce/blip-vqa-base")
+pipeline(
+ image="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/idefics-few-shot.jpg",
+ question="Cosa c'è nell'immagine?",
+)
+[{'answer': 'statua della libertà'}]
+```
+
++ +## Perché non dovrei usare Transformers? + +- Questa libreria non è un toolbox modulare di blocchi costitutivi per reti neurali. Il codice nei file dei modelli non è stato rifattorizzato con ulteriori astrazioni di proposito, in modo che i ricercatori possano iterare rapidamente su ciascuno dei modelli senza dover approfondire ulteriori astrazioni/file. +- L'API di addestramento è ottimizzata per funzionare con i modelli PyTorch forniti da Transformers. Per i loop generici di machine learning, è necessario utilizzare un'altra libreria come [Accelerate](https://huggingface.co/docs/accelerate). +- Gli [script di esempio](https://github.com/huggingface/transformers/tree/main/examples) sono solo *esempi*. Potrebbero non funzionare immediatamente nel vostro caso specifico e potrebbe essere necessario adattare il codice affinché funzioni. + +## 100 progetti che usano Transformers + +Transformers è più di un semplice toolkit per l'utilizzo di modelli preaddestrati, è una comunità di progetti costruita attorno ad esso e all' +Hugging Face Hub. Vogliamo che Transformers consenta a sviluppatori, ricercatori, studenti, professori, ingegneri e chiunque altro +di realizzare i propri progetti dei sogni. + +Per celebrare le 100.000 stelle di Transformers, abbiamo voluto puntare i riflettori sulla +comunità con la pagina [awesome-transformers](./awesome-transformers.md), che elenca 100 +incredibili progetti realizzati con Transformers. + +Se possiedi o utilizzi un progetto che ritieni debba essere inserito nell'elenco, apri una PR per aggiungerlo! + +## Modelli di esempio + +È possibile testare la maggior parte dei nostri modelli direttamente sulle loro [pagine dei modelli Hub](https://huggingface.co/models). + +Espandi ciascuna modalità qui sotto per vedere alcuni modelli di esempio per vari casi d'uso. + +
+
+
+Audio
+ +- Classificazione audio con [Whisper](https://huggingface.co/openai/whisper-large-v3-turbo) +- Riconoscimento vocale automatico con [Moonshine](https://huggingface.co/UsefulSensors/moonshine) +- Individuazione delle keyword con [Wav2Vec2](https://huggingface.co/superb/wav2vec2-base-superb-ks) +- Generazione da discorso a discorso con [Moshi](https://huggingface.co/kyutai/moshiko-pytorch-bf16) +- Testo in audio con [MusicGen](https://huggingface.co/facebook/musicgen-large) +- Sintesi vocale con [Bark](https://huggingface.co/suno/bark) + +
+
+
+Visione artificiale
+ +- Generazione automatica di maschere con [SAM](https://huggingface.co/facebook/sam-vit-base) +- Stima della profondità con [DepthPro](https://huggingface.co/apple/DepthPro-hf) +- Classificazione delle immagini con [DINO v2](https://huggingface.co/facebook/dinov2-base) +- Rilevamento dei punti chiave con [SuperPoint](https://huggingface.co/magic-leap-community/superpoint) +- Corrispondenza dei punti chiave con [SuperGlue](https://huggingface.co/magic-leap-community/superglue_outdoor) +- Rilevamento degli oggetti con [RT-DETRv2](https://huggingface.co/PekingU/rtdetr_v2_r50vd) +- Stima della posa con [VitPose](https://huggingface.co/usyd-community/vitpose-base-simple) +- Segmentazione universale con [OneFormer](https://huggingface.co/shi-labs/oneformer_ade20k_swin_large) +- Classificazione dei video con [VideoMAE](https://huggingface.co/MCG-NJU/videomae-large) + +
+
+
+Multimodale
+ +- Audio or text to text with [Qwen2-Audio](https://huggingface.co/Qwen/Qwen2-Audio-7B) +- Document question answering with [LayoutLMv3](https://huggingface.co/microsoft/layoutlmv3-base) +- Image or text to text with [Qwen-VL](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) +- Image captioning [BLIP-2](https://huggingface.co/Salesforce/blip2-opt-2.7b) +- OCR-based document understanding with [GOT-OCR2](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) +- Table question answering with [TAPAS](https://huggingface.co/google/tapas-base) +- Unified multimodal understanding and generation with [Emu3](https://huggingface.co/BAAI/Emu3-Gen) +- Vision to text with [Llava-OneVision](https://huggingface.co/llava-hf/llava-onevision-qwen2-0.5b-ov-hf) +- Visual question answering with [Llava](https://huggingface.co/llava-hf/llava-1.5-7b-hf) +- Visual referring expression segmentation with [Kosmos-2](https://huggingface.co/microsoft/kosmos-2-patch14-224) + +
+
+
+## Citazione
+
+Ora abbiamo un [paper](https://www.aclweb.org/anthology/2020.emnlp-demos.6/) che puoi citare per la libreria 🤗 Transformers:
+```bibtex
+@inproceedings{wolf-etal-2020-transformers,
+ title = "Transformers: State-of-the-Art Natural Language Processing",
+ author = "Thomas Wolf and Lysandre Debut and Victor Sanh and Julien Chaumond and Clement Delangue and Anthony Moi and Pierric Cistac and Tim Rault and Rémi Louf and Morgan Funtowicz and Joe Davison and Sam Shleifer and Patrick von Platen and Clara Ma and Yacine Jernite and Julien Plu and Canwen Xu and Teven Le Scao and Sylvain Gugger and Mariama Drame and Quentin Lhoest and Alexander M. Rush",
+ booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations",
+ month = oct,
+ year = "2020",
+ address = "Online",
+ publisher = "Association for Computational Linguistics",
+ url = "https://www.aclweb.org/anthology/2020.emnlp-demos.6",
+ pages = "38--45"
+}
+```
diff --git a/i18n/README_ja.md b/i18n/README_ja.md
index 5d5db4993239..98ad2643d23c 100644
--- a/i18n/README_ja.md
+++ b/i18n/README_ja.md
@@ -82,9 +82,11 @@ user: ユーザ
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_ko.md b/i18n/README_ko.md
index fded56a37c9b..a3e6b95cecb5 100644
--- a/i18n/README_ko.md
+++ b/i18n/README_ko.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_pt-br.md b/i18n/README_pt-br.md
index e3c71c6a3f35..bdd464ad0664 100644
--- a/i18n/README_pt-br.md
+++ b/i18n/README_pt-br.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
diff --git a/i18n/README_ru.md b/i18n/README_ru.md
index c30237fef885..3bcaab10f20b 100644
--- a/i18n/README_ru.md
+++ b/i18n/README_ru.md
@@ -47,9 +47,11 @@ limitations under the License.
తెలుగు |
Français |
Deutsch |
+ Italiano |
Tiếng Việt |
العربية |
اردو |
+ বাংলা |
NLP
+ +- Completamento parole mascherate con [ModernBERT](https://huggingface.co/answerdotai/ModernBERT-base) +- Riconoscimento delle entità denominate con [Gemma](https://huggingface.co/google/gemma-2-2b) +- Risposte alle domande con [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) +- Sintesi con [BART](https://huggingface.co/facebook/bart-large-cnn) +- Traduzione con [T5](https://huggingface.co/google-t5/t5-base) +- Generazione di testo con [Llama](https://huggingface.co/meta-llama/Llama-3.2-1B) +- Classificazione del testo con [Qwen](https://huggingface.co/Qwen/Qwen2.5-0.5B) + +diff --git a/i18n/README_te.md b/i18n/README_te.md index aee579b52abd..225bd74bb025 100644 --- a/i18n/README_te.md +++ b/i18n/README_te.md @@ -49,9 +49,11 @@ limitations under the License. తెలుగు | Français | Deutsch | + Italiano | Tiếng Việt | العربية | اردو | + বাংলা |
diff --git a/i18n/README_ur.md b/i18n/README_ur.md index bba5988e7717..215191e4cbb2 100644 --- a/i18n/README_ur.md +++ b/i18n/README_ur.md @@ -47,8 +47,10 @@ limitations under the License. తెలుగు | Français | Deutsch | + Italiano | Tiếng Việt | العربية | + বাংলা | اردو | diff --git a/i18n/README_vi.md b/i18n/README_vi.md index f78e3b6d4e9b..3e0146c1ddb0 100644 --- a/i18n/README_vi.md +++ b/i18n/README_vi.md @@ -47,9 +47,11 @@ limitations under the License. తెలుగు | Français | Deutsch | + Italiano | Tiếng việt | العربية | اردو | + বাংলা | diff --git a/i18n/README_zh-hans.md b/i18n/README_zh-hans.md index 8220e403b8b2..4c5859592c89 100644 --- a/i18n/README_zh-hans.md +++ b/i18n/README_zh-hans.md @@ -72,9 +72,11 @@ checkpoint: 检查点 తెలుగు | Français | Deutsch | + Italiano | Tiếng Việt | العربية | اردو | + বাংলা | diff --git a/i18n/README_zh-hant.md b/i18n/README_zh-hant.md index da6ed40910ea..5842e57255c3 100644 --- a/i18n/README_zh-hant.md +++ b/i18n/README_zh-hant.md @@ -84,9 +84,11 @@ user: 使用者 తెలుగు | Français | Deutsch | + Italiano | Tiếng Việt | العربية | اردو | + বাংলা | diff --git a/notebooks/README.md b/notebooks/README.md index 4d31797104f8..aed435878804 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -22,7 +22,6 @@ Also, we would like to list here interesting content created by the community. If you wrote some notebook(s) leveraging 🤗 Transformers and would like to be listed here, please open a Pull Request so it can be included under the Community notebooks. - ## Hugging Face's notebooks 🤗 ### Documentation notebooks @@ -38,7 +37,6 @@ You can open any page of the documentation as a notebook in Colab (there is a bu | [Summary of the tokenizers](https://github.com/huggingface/notebooks/blob/main/transformers_doc/en/tokenizer_summary.ipynb) | The differences between the tokenizers algorithm |[](https://colab.research.google.com/github/huggingface/notebooks/blob/main/transformers_doc/en/tokenizer_summary.ipynb)| [](https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/transformers_doc/en/tokenizer_summary.ipynb)| | [Multilingual models](https://github.com/huggingface/notebooks/blob/main/transformers_doc/en/multilingual.ipynb) | How to use the multilingual models of the library |[](https://colab.research.google.com/github/huggingface/notebooks/blob/main/transformers_doc/en/multilingual.ipynb)| [](https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/transformers_doc/en/multilingual.ipynb)| - ### PyTorch Examples #### Natural Language Processing[[pytorch-nlp]] @@ -88,7 +86,6 @@ You can open any page of the documentation as a notebook in Colab (there is a bu | [How to fine-tune a Nucleotide Transformer model](https://github.com/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling.ipynb) | See how to tokenize DNA and fine-tune a large pre-trained DNA "language" model | [](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling.ipynb) | [](https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling.ipynb) | | [Fine-tune a Nucleotide Transformer model with LoRA](https://github.com/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling_with_peft.ipynb) | Train even larger DNA models in a memory-efficient way | [](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling_with_peft.ipynb) | [](https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/examples/nucleotide_transformer_dna_sequence_modelling_with_peft.ipynb) | - #### Other modalities[[pytorch-other]] | Notebook | Description | | | @@ -101,7 +98,6 @@ You can open any page of the documentation as a notebook in Colab (there is a bu |:----------|:-------------|:-------------|------:| | [How to export model to ONNX](https://github.com/huggingface/notebooks/blob/main/examples/onnx-export.ipynb)| Highlight how to export and run inference workloads through ONNX | [](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/onnx-export.ipynb)| [](https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/examples/onnx-export.ipynb)| - ### Optimum notebooks 🤗 [Optimum](https://github.com/huggingface/optimum) is an extension of 🤗 Transformers, providing a set of performance optimization tools enabling maximum efficiency to train and run models on targeted hardwares. diff --git a/pyproject.toml b/pyproject.toml index 5d3a9436eb3f..80983fd49703 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ exclude_lines = [ ] [tool.ruff] -target-version = "py39" +target-version = "py310" line-length = 119 [tool.ruff.lint] @@ -27,7 +27,10 @@ line-length = 119 # UP031: Use format specifiers instead of percent format # UP004: Class `XXX` inherits from `object` # UP028: Checks for for loops that can be replaced with yield from expressions -ignore = ["C901", "E501", "E741", "F402", "F823", "SIM1", "SIM300", "SIM212", "SIM905", "UP009", "UP015", "UP031", "UP028", "UP004"] +# UP045: Use `X | None` for type annotations +# UP007: Use `X | Y` for type annotations +# UP035: temporarily disabled to minimize upgrade changes +ignore = ["C901", "E501", "E741", "F402", "F823", "SIM1", "SIM300", "SIM212", "SIM905", "UP009", "UP015", "UP031", "UP028", "UP004", "UP045", "UP007", "UP035"] # RUF013: Checks for the use of implicit Optional # in type annotations when the default parameter value is None. select = ["C", "E", "F", "I", "W", "RUF013", "PERF102", "PLC1802", "PLC0208", "SIM", "UP"] diff --git a/setup.py b/setup.py index 9f3bb1750597..dbeb2af336be 100644 --- a/setup.py +++ b/setup.py @@ -138,7 +138,7 @@ "opencv-python", "optimum-benchmark>=0.3.0", "optuna", - "optax>=0.0.8,<=0.1.4", + "optax>=0.08,<=0.1.4", "pandas<2.3.0", # `datasets` requires `pandas` while `pandas==2.3.0` has issues with CircleCI on 2025/06/05 "packaging>=20.0", "parameterized>=0.9", # older version of parameterized cause pytest collection to fail on .expand @@ -160,7 +160,7 @@ "rhoknp>=1.1.0,<1.3.1", "rjieba", "rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1", - "ruff==0.11.2", + "ruff==0.13.1", # `sacrebleu` not used in `transformers`. However, it is needed in several tests, when a test calls # `evaluate.load("sacrebleu")`. This metric is used in the examples that we use to test the `Trainer` with, in the # `Trainer` tests (see references to `run_translation.py`). diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 2cf1d5970b54..29c26b996f3c 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -928,7 +928,6 @@ 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 logging as logging # bitsandbytes config from .utils.quantization_config import AqlmConfig as AqlmConfig diff --git a/src/transformers/activations.py b/src/transformers/activations.py index 8bfd517add9f..7642e8aa238a 100644 --- a/src/transformers/activations.py +++ b/src/transformers/activations.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import math from collections import OrderedDict @@ -26,7 +27,8 @@ logger = logging.get_logger(__name__) -class PytorchGELUTanh(nn.Module): +@use_kernel_forward_from_hub("GeluTanh") +class GELUTanh(nn.Module): """ A fast C implementation of the tanh approximation of the GeLU activation function. See https://huggingface.co/papers/1606.08415. @@ -35,8 +37,18 @@ class PytorchGELUTanh(nn.Module): match due to rounding errors. """ + def __init__(self, use_gelu_tanh_python: bool = False): + super().__init__() + if use_gelu_tanh_python: + self.act = self._gelu_tanh_python + else: + self.act = functools.partial(nn.functional.gelu, approximate="tanh") + + def _gelu_tanh_python(self, input: Tensor) -> Tensor: + return input * 0.5 * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) + def forward(self, input: Tensor) -> Tensor: - return nn.functional.gelu(input, approximate="tanh") + return self.act(input) @use_kernel_forward_from_hub("NewGELU") @@ -50,6 +62,7 @@ def forward(self, input: Tensor) -> Tensor: return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) +@use_kernel_forward_from_hub("GeLU") class GELUActivation(nn.Module): """ Original Implementation of the GELU activation function in Google BERT repo when initially created. For @@ -72,6 +85,20 @@ def forward(self, input: Tensor) -> Tensor: return self.act(input) +@use_kernel_forward_from_hub("SiLU") +class SiLUActivation(nn.Module): + """ + See Gaussian Error Linear Units (Hendrycks et al., https://arxiv.org/abs/1606.08415) where the SiLU (Sigmoid Linear + Unit) was originally introduced and coined, and see Sigmoid-Weighted Linear Units for Neural Network Function + Approximation in Reinforcement Learning (Elfwing et al., https://arxiv.org/abs/1702.03118) and Swish: a Self-Gated + Activation Function (Ramachandran et al., https://arxiv.org/abs/1710.05941v1) where the SiLU was experimented with + later. + """ + + def forward(self, input: Tensor) -> Tensor: + return nn.functional.silu(input) + + @use_kernel_forward_from_hub("FastGELU") class FastGELUActivation(nn.Module): """ @@ -290,7 +317,8 @@ def forward(self, input: Tensor) -> Tensor: "gelu_fast": FastGELUActivation, "gelu_new": NewGELUActivation, "gelu_python": (GELUActivation, {"use_gelu_python": True}), - "gelu_pytorch_tanh": PytorchGELUTanh, + "gelu_pytorch_tanh": GELUTanh, + "gelu_python_tanh": (GELUTanh, {"use_gelu_tanh_python": True}), "gelu_accurate": AccurateGELUActivation, "laplace": LaplaceActivation, "leaky_relu": nn.LeakyReLU, @@ -301,7 +329,7 @@ def forward(self, input: Tensor) -> Tensor: "relu2": ReLUSquaredActivation, "relu6": nn.ReLU6, "sigmoid": nn.Sigmoid, - "silu": nn.SiLU, + "silu": SiLUActivation, "swish": nn.SiLU, "tanh": nn.Tanh, "prelu": nn.PReLU, diff --git a/src/transformers/audio_utils.py b/src/transformers/audio_utils.py index e848f558738c..5de56618014e 100644 --- a/src/transformers/audio_utils.py +++ b/src/transformers/audio_utils.py @@ -23,8 +23,11 @@ import warnings from collections.abc import Sequence from io import BytesIO -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union + +if TYPE_CHECKING: + import torch import numpy as np import requests from packaging import version @@ -51,7 +54,7 @@ if is_torchcodec_available(): TORCHCODEC_VERSION = version.parse(importlib.metadata.version("torchcodec")) -AudioInput = Union[np.ndarray, "torch.Tensor", Sequence[np.ndarray], Sequence["torch.Tensor"]] # noqa: F821 +AudioInput = Union[np.ndarray, "torch.Tensor", Sequence[np.ndarray], Sequence["torch.Tensor"]] def load_audio(audio: Union[str, np.ndarray], sampling_rate=16000, timeout=None) -> np.ndarray: @@ -78,9 +81,7 @@ def load_audio(audio: Union[str, np.ndarray], sampling_rate=16000, timeout=None) audio = load_audio_torchcodec(audio, sampling_rate=sampling_rate) else: audio = load_audio_librosa(audio, sampling_rate=sampling_rate, timeout=timeout) - elif isinstance(audio, np.ndarray): - audio = audio - else: + elif not isinstance(audio, np.ndarray): raise TypeError( "Incorrect format used for `audio`. Should be an url linking to an audio, a local path, or numpy array." ) @@ -318,9 +319,7 @@ def mel_to_hertz(mels: Union[float, np.ndarray], mel_scale: str = "htk") -> Unio return freq -def hertz_to_octave( - freq: Union[float, np.ndarray], tuning: Optional[float] = 0.0, bins_per_octave: Optional[int] = 12 -): +def hertz_to_octave(freq: Union[float, np.ndarray], tuning: float = 0.0, bins_per_octave: int = 12): """ Convert frequency from hertz to fractional octave numbers. Adapted from *librosa*. @@ -370,7 +369,7 @@ def chroma_filter_bank( tuning: float = 0.0, power: Optional[float] = 2.0, weighting_parameters: Optional[tuple[float, float]] = (5.0, 2.0), - start_at_c_chroma: Optional[bool] = True, + start_at_c_chroma: bool = True, ): """ Creates a chroma filter bank, i.e a linear transformation to project spectrogram bins onto chroma bins. @@ -391,7 +390,7 @@ def chroma_filter_bank( weighting_parameters (`tuple[float, float]`, *optional*, defaults to `(5., 2.)`): If specified, apply a Gaussian weighting parameterized by the first element of the tuple being the center and the second element being the Gaussian half-width. - start_at_c_chroma (`float`, *optional*, defaults to `True`): + start_at_c_chroma (`bool`, *optional*, defaults to `True`): If True, the filter bank will start at the 'C' pitch class. Otherwise, it will start at 'A'. Returns: `np.ndarray` of shape `(num_frequency_bins, num_chroma)` @@ -586,7 +585,7 @@ def window_function( window = np.hamming(length) elif name in ["hann", "hann_window"]: window = np.hanning(length) - elif name in ["povey"]: + elif name == "povey": window = np.power(np.hanning(length), 0.85) else: raise ValueError(f"Unknown window function '{name}'") @@ -627,7 +626,7 @@ def spectrogram( reference: float = 1.0, min_value: float = 1e-10, db_range: Optional[float] = None, - remove_dc_offset: Optional[bool] = None, + remove_dc_offset: bool = False, dtype: np.dtype = np.float32, ) -> np.ndarray: """ @@ -838,7 +837,7 @@ def spectrogram_batch( reference: float = 1.0, min_value: float = 1e-10, db_range: Optional[float] = None, - remove_dc_offset: Optional[bool] = None, + remove_dc_offset: bool = False, dtype: np.dtype = np.float32, ) -> list[np.ndarray]: """ diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index e6f2645a766e..99beb0b610a1 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -395,7 +395,12 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states) - cache_position = cache_kwargs.get("cache_position") + # Some old models give None for `cache_position` or even omit passing `cache_kwargs` when used as cross-attention, + # in which case we should copy the whole Layer (key_states.shape[-2] == self.max_cache_len) + cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None + cache_position = ( + cache_position if cache_position is not None else torch.arange(key_states.shape[-2], device=self.device) + ) cumulative_length = self.cumulative_length is_full = cumulative_length >= self.max_cache_len @@ -790,7 +795,7 @@ def early_initialization( for layer in self.layers: layer.lazy_initialization(fake_keys_tensor) - def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + def get_seq_length(self, layer_idx: int = 0) -> int: """Returns the sequence length of the cache for the given layer.""" if layer_idx >= len(self.layers): return 0 @@ -955,17 +960,19 @@ def __init__( layers = [] # If a config is passed, use it to infer the layer types and initialize accordingly if config is not None: - config = config.get_text_config(decoder=True) - sliding_window = getattr(config, "sliding_window", None) or getattr(config, "attention_chunk_size", None) - layer_types = getattr(config, "layer_types", None) + decoder_config = config.get_text_config(decoder=True) + sliding_window = getattr(decoder_config, "sliding_window", None) or getattr( + decoder_config, "attention_chunk_size", None + ) + layer_types = getattr(decoder_config, "layer_types", None) if layer_types is None: layer_types = [ "sliding_attention" if sliding_window is not None else "full_attention" - for _ in range(config.num_hidden_layers) + for _ in range(decoder_config.num_hidden_layers) ] # Some models have shared layers thus no cache is needed for them (e.g. Gemma3n) - if hasattr(config, "num_kv_shared_layers"): - layer_types = layer_types[: -config.num_kv_shared_layers] + if hasattr(decoder_config, "num_kv_shared_layers"): + layer_types = layer_types[: -decoder_config.num_kv_shared_layers] for layer_type in layer_types: # From a cache point of view, both sliding and chunked are the same in how they should behave and how many @@ -1286,7 +1293,7 @@ def from_legacy_cache( cache.is_updated[layer_idx] = True return cache - def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + def get_seq_length(self, layer_idx: int = 0) -> int: """Returns the sequence length of the cached states. A layer index can be optionally passed.""" return self.self_attention_cache.get_seq_length(layer_idx) diff --git a/src/transformers/commands/add_new_model_like.py b/src/transformers/commands/add_new_model_like.py index ffff54df93ba..fce524d4a6c0 100644 --- a/src/transformers/commands/add_new_model_like.py +++ b/src/transformers/commands/add_new_model_like.py @@ -755,7 +755,7 @@ def register_subcommand(parser: ArgumentParser): ) add_new_model_like_parser.set_defaults(func=add_new_model_like_command_factory) - def __init__(self, path_to_repo=None, *args): + def __init__(self, path_to_repo=None, **kwargs): ( self.old_model_infos, self.new_lowercase_name, diff --git a/src/transformers/commands/chat.py b/src/transformers/commands/chat.py index 70ee41c0c514..6ddf90164ba7 100644 --- a/src/transformers/commands/chat.py +++ b/src/transformers/commands/chat.py @@ -40,6 +40,12 @@ from transformers.utils import is_rich_available, is_torch_available +try: + import readline # noqa importing this enables GNU readline capabilities +except ImportError: + # some platforms may not support readline: https://docs.python.org/3/library/readline.html + pass + if platform.system() != "Windows": import pwd @@ -53,9 +59,7 @@ from transformers import ( AutoModelForCausalLM, - AutoTokenizer, BitsAndBytesConfig, - GenerationConfig, ) ALLOWED_KEY_CHARS = set(string.ascii_letters + string.whitespace) @@ -437,8 +441,7 @@ def parse_generate_flags(self, generate_flags: list[str]) -> dict: # 2. b. strings should be quoted def is_number(s: str) -> bool: # handle negative numbers - if s.startswith("-"): - s = s[1:] + s = s.removeprefix("-") return s.replace(".", "", 1).isdigit() generate_flags_as_dict = {k: f'"{v}"' if not is_number(v) else v for k, v in generate_flags_as_dict.items()} @@ -528,7 +531,7 @@ def parse_eos_tokens( # ----------------------------------------------------------------------------------------------------------------- # Model loading and performance automation methods @staticmethod - def get_quantization_config(model_args: ChatArguments) -> Optional["BitsAndBytesConfig"]: + def get_quantization_config(model_args: ChatArguments) -> Optional[BitsAndBytesConfig]: if model_args.load_in_4bit: quantization_config = BitsAndBytesConfig( load_in_4bit=True, @@ -684,7 +687,6 @@ async def _inner_run(self): model = self.args.model_name_or_path + "@" + self.args.model_revision host = "http://localhost" if self.args.host == "localhost" else self.args.host - client = AsyncInferenceClient(f"{host}:{self.args.port}") args = self.args if args.examples_path is None: @@ -707,48 +709,47 @@ async def _inner_run(self): # Starts the session with a minimal help message at the top, so that a user doesn't get stuck interface.print_help(minimal=True) - while True: - try: - user_input = interface.input() - - # User commands - if user_input.startswith("!"): - # `!exit` is special, it breaks the loop - if user_input == "!exit": - break - else: - chat, valid_command, generation_config, model_kwargs = self.handle_non_exit_user_commands( - user_input=user_input, - args=args, - interface=interface, - examples=examples, - generation_config=generation_config, - model_kwargs=model_kwargs, - chat=chat, - ) - # `!example` sends a user message to the model - if not valid_command or not user_input.startswith("!example"): - continue - else: - chat.append({"role": "user", "content": user_input}) - - stream = client.chat_completion( - chat, - stream=True, - extra_body={ - "generation_config": generation_config.to_json_string(), - "model": model, - }, - ) - model_output = await interface.stream_output(stream) + async with AsyncInferenceClient(f"{host}:{self.args.port}") as client: + while True: + try: + user_input = interface.input() + + # User commands + if user_input.startswith("!"): + # `!exit` is special, it breaks the loop + if user_input == "!exit": + break + else: + chat, valid_command, generation_config, model_kwargs = self.handle_non_exit_user_commands( + user_input=user_input, + args=args, + interface=interface, + examples=examples, + generation_config=generation_config, + model_kwargs=model_kwargs, + chat=chat, + ) + # `!example` sends a user message to the model + if not valid_command or not user_input.startswith("!example"): + continue + else: + chat.append({"role": "user", "content": user_input}) + + stream = client.chat_completion( + chat, + stream=True, + extra_body={ + "generation_config": generation_config.to_json_string(), + "model": model, + }, + ) - chat.append({"role": "assistant", "content": model_output}) + model_output = await interface.stream_output(stream) - except KeyboardInterrupt: - break - finally: - await client.close() + chat.append({"role": "assistant", "content": model_output}) + except KeyboardInterrupt: + break if __name__ == "__main__": diff --git a/src/transformers/commands/env.py b/src/transformers/commands/env.py index 983a858cd952..e15a699e80f6 100644 --- a/src/transformers/commands/env.py +++ b/src/transformers/commands/env.py @@ -14,7 +14,6 @@ import contextlib -import importlib.util import io import os import platform @@ -27,7 +26,6 @@ from ..utils import ( is_accelerate_available, is_flax_available, - is_safetensors_available, is_tf_available, is_torch_available, is_torch_hpu_available, @@ -61,18 +59,13 @@ def __init__(self, accelerate_config_file, *args) -> None: self._accelerate_config_file = accelerate_config_file def run(self): - safetensors_version = "not installed" - if is_safetensors_available(): - import safetensors + import safetensors - safetensors_version = safetensors.__version__ - elif importlib.util.find_spec("safetensors") is not None: - import safetensors - - safetensors_version = f"{safetensors.__version__} but is ignored because of PyTorch version too old." + safetensors_version = safetensors.__version__ accelerate_version = "not installed" accelerate_config = accelerate_config_str = "not found" + if is_accelerate_available(): import accelerate from accelerate.commands.config import default_config_file, load_config_from_file diff --git a/src/transformers/commands/serving.py b/src/transformers/commands/serving.py index 33a48aed7e64..970d59c96e74 100644 --- a/src/transformers/commands/serving.py +++ b/src/transformers/commands/serving.py @@ -31,7 +31,7 @@ from dataclasses import dataclass, field from io import BytesIO from threading import Thread -from typing import Optional, Union +from typing import Optional, TypedDict, Union from huggingface_hub import model_info from huggingface_hub.constants import HF_HUB_OFFLINE @@ -141,7 +141,7 @@ class TransformersTranscriptionCreateParams(TranscriptionCreateParamsBase, total file: bytes # Overwritten -- pydantic isn't happy with `typing.IO[bytes]`, present in the original type generation_config: str - stream: Optional[bool] = False + stream: bool = False # Contrarily to OpenAI's output types, input types are `TypedDict`, which don't have built-in validation. response_validator = TypeAdapter(TransformersResponseCreateParamsStreaming) @@ -528,7 +528,7 @@ def __init__(self, args: ServeArguments): def _validate_request( self, request: dict, - schema: "_TypedDictMeta", # noqa: F821 + schema: TypedDict, validator: "TypeAdapter", unused_fields: set, ): @@ -538,7 +538,7 @@ def _validate_request( Args: request (`dict`): The request to validate. - schema (`_TypedDictMeta`): + schema (`TypedDict`): The schema of the request to validate. It is a `TypedDict` definition. validator (`TypeAdapter`): The validator to use to validate the request. Built from `schema`. @@ -600,7 +600,7 @@ def validate_transcription_request(self, request: dict): def build_chat_completion_chunk( self, - request_id: Optional[str] = "", + request_id: str = "", content: Optional[int] = None, model: Optional[str] = None, role: Optional[str] = None, @@ -1026,7 +1026,9 @@ def generate_chat_completion(self, req: dict) -> Generator[str, None, None]: last_kv_cache = None if self.is_continuation(req) and not must_discard_cache: - last_kv_cache = self.last_kv_cache + seq_len = self.last_kv_cache.get_seq_length() + if inputs["input_ids"].shape[-1] > seq_len: + last_kv_cache = self.last_kv_cache generation_kwargs = { **inputs, @@ -1064,8 +1066,7 @@ def generate_with_cache(**kwargs): for result in streamer: # Temporary hack for GPTOS 3: don't emit the final "<|return|>" if "gptoss" in model.config.architectures[0].lower(): - if result.endswith("<|return|>"): - result = result[: -len("<|return|>")] + result = result.removesuffix("<|return|>") results += result # (related to temporary hack 2) @@ -1213,7 +1214,9 @@ def generate_response(self, req: dict) -> Generator[str, None, None]: last_kv_cache = None if self.is_continuation(req) and not must_discard_cache: - last_kv_cache = self.last_kv_cache + seq_len = self.last_kv_cache.get_seq_length() + if inputs["input_ids"].shape[-1] > seq_len: + last_kv_cache = self.last_kv_cache generation_kwargs = { "inputs": inputs, @@ -1321,8 +1324,7 @@ def generate_with_cache(**kwargs): for result in streamer: # Temporary hack for GPTOS 3: don't emit the final "<|return|>" if "gptoss" in model.config.architectures[0].lower(): - if result.endswith("<|return|>"): - result = result[: -len("<|return|>")] + result = result.removesuffix("<|return|>") results += result # (related to temporary hack 2) diff --git a/src/transformers/convert_slow_tokenizer.py b/src/transformers/convert_slow_tokenizer.py index a9e7c9bff5bc..aa32734ffb38 100644 --- a/src/transformers/convert_slow_tokenizer.py +++ b/src/transformers/convert_slow_tokenizer.py @@ -1454,7 +1454,7 @@ def pre_tokenizer(self, replacement, add_prefix_space): class HeliumConverter(SpmConverter): handle_byte_fallback = True - def __init__(self, vocab_file=None, *args): + def __init__(self, vocab_file=None, **kwargs): requires_backends(self, "protobuf") Converter.__init__(self, vocab_file) @@ -1540,6 +1540,54 @@ def post_processor(self): ) +class ParakeetConverter(SpmConverter): + handle_byte_fallback = True + + def __init__(self, vocab_file=None, *args): + self.vocab_file = vocab_file + + requires_backends(self, "protobuf") + + Converter.__init__(self, vocab_file) + + model_pb2 = import_protobuf() + m = model_pb2.ModelProto() + with open(vocab_file, "rb") as f: + m.ParseFromString(f.read()) + self.proto = m + + def tokenizer(self, proto): + vocab_scores = self.vocab(proto) + + _, merges = self.SpmExtractor(self.vocab_file).extract(vocab_scores) + bpe_vocab = {word: i for i, (word, score) in enumerate(vocab_scores)} + tokenizer = Tokenizer( + BPE( + bpe_vocab, + merges, + unk_token=proto.trainer_spec.unk_piece, + fuse_unk=True, + byte_fallback=self.handle_byte_fallback, + dropout=None, + ) + ) + + # Add user defined symbols and control tokens from sentencepiece model + spm_added_tokens = [ + (id, p.piece, p.type == 3 or p.piece in self.special_tokens) + for id, p in enumerate(proto.pieces) + if p.type in [3, 4] + ] + tokenizer.add_tokens( + [ + AddedToken(token, normalized=False, special=special) + for id, token, special in sorted(spm_added_tokens, key=lambda x: x[0]) + ] + ) + + return tokenizer + + # Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode def bytes_to_unicode(): """ @@ -1576,10 +1624,8 @@ def __init__( pattern=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""", add_prefix_space=False, additional_special_tokens=None, - *args, **kwargs, ): - super().__init__(*args) self.vocab_file = vocab_file self.pattern = pattern self.add_prefix_space = add_prefix_space diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py index 10ee10e01950..3fa9cb72de9f 100644 --- a/src/transformers/data/data_collator.py +++ b/src/transformers/data/data_collator.py @@ -18,26 +18,25 @@ from collections.abc import Mapping from dataclasses import dataclass from random import randint -from typing import Any, Callable, NewType, Optional, Union +from typing import Any, Callable, Optional, Union import numpy as np -from ..models.bert import BertTokenizer, BertTokenizerFast from ..tokenization_utils_base import PreTrainedTokenizerBase from ..utils import PaddingStrategy -InputDataClass = NewType("InputDataClass", Any) +InputDataClass = Any """ A DataCollator is a function that takes a list of samples from a Dataset and collate them into a batch, as a dictionary of PyTorch/TensorFlow tensors or NumPy arrays. """ -DataCollator = NewType("DataCollator", Callable[[list[InputDataClass]], dict[str, Any]]) +DataCollator = Callable[[list[InputDataClass]], dict[str, Any]] class DataCollatorMixin: - def __call__(self, features, return_tensors=None): + def __call__(self, features, return_tensors: Optional[str] = None): if return_tensors is None: return_tensors = self.return_tensors if return_tensors == "tf": @@ -773,6 +772,8 @@ class DataCollatorForLanguageModeling(DataCollatorMixin): Whether or not to use masked language modeling. If set to `False`, the labels are the same as the inputs with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for non-masked tokens and the value to predict for the masked token. + whole_word_mask (`bool`, *optional*, defaults to `False`): + Whether or not to mask whole words instead of individual tokens. mlm_probability (`float`, *optional*, defaults to 0.15): The probability with which to (randomly) mask tokens in the input, when `mlm` is set to `True`. mask_replace_prob (`float`, *optional*, defaults to 0.8): @@ -824,6 +825,7 @@ class DataCollatorForLanguageModeling(DataCollatorMixin): tokenizer: PreTrainedTokenizerBase mlm: bool = True + whole_word_mask: bool = False mlm_probability: Optional[float] = 0.15 mask_replace_prob: float = 0.8 random_replace_prob: float = 0.1 @@ -842,6 +844,11 @@ def __post_init__(self): if self.mlm_probability is None or self.mlm_probability < 0 or self.mlm_probability > 1: raise ValueError("mlm_probability should be between 0 and 1.") self.mlm_probability = float(self.mlm_probability) + elif self.whole_word_mask: + raise ValueError( + "Whole word masking can only be used with mlm=True." + "If you want to use whole word masking, please set mlm=True." + ) if self.mask_replace_prob + self.random_replace_prob > 1: raise ValueError("The sum of mask_replace_prob and random_replace_prob should not exceed 1") if self.mask_replace_prob < 0 or self.mask_replace_prob > 1: @@ -856,6 +863,20 @@ def __post_init__(self): import tensorflow as tf self.tf_mask_tokens = tf.function(self.tf_mask_tokens, jit_compile=True) + if self.whole_word_mask: + if not self.tokenizer.is_fast: + warnings.warn( + "Whole word masking depends on offset mapping which is only natively available with fast tokenizers.", + UserWarning, + ) + + if self.mask_replace_prob < 1: + warnings.warn( + "Random token replacement is not supported with whole word masking.", + "Setting mask_replace_prob to 1.", + ) + self.mask_replace_prob = 1 + self.random_replace_prob = 0 self.generator = None @@ -869,8 +890,6 @@ def get_generator(self, seed): return tf.random.Generator.from_seed(seed) else: - import numpy as np - return np.random.default_rng(seed) def create_rng(self): @@ -1021,9 +1040,10 @@ def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> d # If special token mask has been preprocessed, pop it from the dict. special_tokens_mask = batch.pop("special_tokens_mask", None) + offset_mapping = batch.pop("offset_mapping", None) if self.mlm: batch["input_ids"], batch["labels"] = self.torch_mask_tokens( - batch["input_ids"], special_tokens_mask=special_tokens_mask + batch["input_ids"], special_tokens_mask=special_tokens_mask, offset_mapping=offset_mapping ) else: labels = batch["input_ids"].clone() @@ -1032,9 +1052,11 @@ def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> d batch["labels"] = labels return batch - def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> tuple[Any, Any]: + def torch_mask_tokens( + self, inputs: Any, special_tokens_mask: Optional[Any] = None, offset_mapping: Optional[Any] = None + ) -> tuple[Any, Any]: """ - Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. + Prepare masked tokens inputs/labels for masked language modeling. """ import torch @@ -1045,12 +1067,24 @@ def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = No special_tokens_mask = [ self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist() ] - special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) + + if self.whole_word_mask: + word_ids, no_mask_mask = self._calc_word_ids_and_prob_mask( + to_numpy(offset_mapping), to_numpy(special_tokens_mask) + ) + no_mask_mask = torch.tensor(no_mask_mask, dtype=torch.bool) else: - special_tokens_mask = special_tokens_mask.bool() + no_mask_mask = ( + special_tokens_mask.bool() + if isinstance(special_tokens_mask, torch.Tensor) + else torch.tensor(special_tokens_mask, dtype=torch.bool) + ) - probability_matrix.masked_fill_(special_tokens_mask, value=0.0) + probability_matrix.masked_fill_(no_mask_mask, value=0.0) masked_indices = torch.bernoulli(probability_matrix, generator=self.generator).bool() + if self.whole_word_mask: + masked_indices = torch.BoolTensor(self._whole_word_mask(word_ids, masked_indices)) + labels[~masked_indices] = -100 # We only compute loss on masked tokens # mask_replace_prob% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) @@ -1100,9 +1134,10 @@ def numpy_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> d # If special token mask has been preprocessed, pop it from the dict. special_tokens_mask = batch.pop("special_tokens_mask", None) + offset_mapping = batch.pop("offset_mapping", None) if self.mlm: batch["input_ids"], batch["labels"] = self.numpy_mask_tokens( - batch["input_ids"], special_tokens_mask=special_tokens_mask + batch["input_ids"], special_tokens_mask=special_tokens_mask, offset_mapping=offset_mapping ) else: labels = np.copy(batch["input_ids"]) @@ -1111,9 +1146,14 @@ def numpy_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> d batch["labels"] = labels return batch - def numpy_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> tuple[Any, Any]: + def numpy_mask_tokens( + self, + inputs: Any, + special_tokens_mask: Optional[Any] = None, + offset_mapping: Optional[Any] = None, + ) -> tuple[Any, Any]: """ - Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. + Prepare masked tokens inputs/labels for masked language modeling. """ labels = np.copy(inputs) # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`) @@ -1122,16 +1162,28 @@ def numpy_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = No special_tokens_mask = [ self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist() ] - special_tokens_mask = np.array(special_tokens_mask, dtype=bool) + + if self.whole_word_mask: + word_ids, no_mask_mask = self._calc_word_ids_and_prob_mask( + to_numpy(offset_mapping), to_numpy(special_tokens_mask) + ) else: - special_tokens_mask = special_tokens_mask.astype(bool) + no_mask_mask = ( + special_tokens_mask.astype(bool) + if isinstance(special_tokens_mask, np.ndarray) + else np.array(special_tokens_mask, dtype=bool) + ) - probability_matrix[special_tokens_mask] = 0 + probability_matrix[no_mask_mask] = 0 # Numpy doesn't have bernoulli, so we use a binomial with 1 trial if self.generator: masked_indices = self.generator.binomial(1, probability_matrix, size=probability_matrix.shape).astype(bool) else: masked_indices = np.random.binomial(1, probability_matrix, size=probability_matrix.shape).astype(bool) + + if self.whole_word_mask: + masked_indices = self._whole_word_mask(word_ids, masked_indices) + labels[~masked_indices] = -100 # We only compute loss on masked tokens # mask_replace_prob% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) @@ -1176,6 +1228,51 @@ def numpy_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = No # The rest of the time (10% of the time) we keep the masked input tokens unchanged return inputs, labels + @staticmethod + def _calc_word_ids_and_prob_mask( + offsets: np.ndarray[np.ndarray[tuple[int, int]]], special_tokens_mask: np.ndarray[np.ndarray[int]] + ) -> tuple[np.ndarray[np.ndarray[int]], np.ndarray[np.ndarray[int]]]: + """ + Map tokens to word ids and create mask of tokens to not mask. + Tokens that are part of the same word will have the same word id and we will only + set a mask probability for the first token of each word. + """ + + token_starts = offsets[:, :, 0] + token_ends = offsets[:, :, 1] + + prev_token_ends = np.roll(token_ends, 1, axis=1) + prev_token_ends[:, 0] = -1 # First token has no previous token + + prev_token_special = np.roll(special_tokens_mask, 1, axis=1) + prev_token_special[:, 0] = 0 + + # Not special token AND (gap from previous or previous token was special) + special_tokens_mask = special_tokens_mask.astype(bool) + is_new_word = (~special_tokens_mask) & ((token_starts != prev_token_ends) | (prev_token_special == 1)) + + word_ids = np.cumsum(is_new_word, axis=1) + word_ids[special_tokens_mask] = -1 + + prob_mask = ~is_new_word + + return word_ids, prob_mask + + @staticmethod + def _whole_word_mask(word_ids: np.ndarray[np.ndarray[int]], mask: Any) -> Any: + """ + Mask whole words based on word ids and mask. + """ + mask = to_numpy(mask) + + valid_ids = word_ids != -1 + + # Create 3D mask where [batch, token_i, token_j] is True if token_i and token_j are the same word + same_word = (word_ids[:, :, None] == word_ids[:, None, :]) & valid_ids[:, :, None] & valid_ids[:, None, :] + + # For each token, set True if any token in the same word is masked + return np.any(same_word & mask[:, None, :], axis=2) + @dataclass class DataCollatorForWholeWordMask(DataCollatorForLanguageModeling): @@ -1322,6 +1419,8 @@ def _whole_word_mask(self, input_tokens: list[str], max_predictions=512): """ Get 0/1 labels for masked tokens with whole word mask proxy """ + from transformers import BertTokenizer, BertTokenizerFast + if not isinstance(self.tokenizer, (BertTokenizer, BertTokenizerFast)): warnings.warn( "DataCollatorForWholeWordMask is only suitable for BertTokenizer-like tokenizers. " @@ -1539,8 +1638,18 @@ def numpy_mask_tokens(self, inputs: Any, mask_labels: Any) -> tuple[Any, Any]: # The rest of the time ((1-mask_replace_prob-random_replace_prob)% of the time) we keep the masked input tokens unchanged return inputs, labels + def __init__(self, *args, **kwargs): + warnings.warn( + "DataCollatorForWholeWordMask is deprecated and will be removed in a future version, you can now use " + "DataCollatorForLanguageModeling with whole_word_mask=True instead.", + FutureWarning, + ) + super().__init__(*args, **kwargs) + self.mlm = True # Force masked language modeling + self.whole_word_mask = True # Force whole word masking + -def tolist(x): +def tolist(x) -> list[Any]: if isinstance(x, list): return x elif hasattr(x, "numpy"): # Checks for TF tensors without needing the import @@ -1548,6 +1657,15 @@ def tolist(x): return x.tolist() +def to_numpy(x) -> np.ndarray[Any]: + if isinstance(x, np.ndarray): + return x + elif hasattr(x, "detach"): + return x.detach().cpu().numpy() + else: + return np.array(x) + + @dataclass class DataCollatorForSOP(DataCollatorForLanguageModeling): """ diff --git a/src/transformers/data/datasets/squad.py b/src/transformers/data/datasets/squad.py index fdee571e249b..d4f76a51f422 100644 --- a/src/transformers/data/datasets/squad.py +++ b/src/transformers/data/datasets/squad.py @@ -122,9 +122,9 @@ def __init__( tokenizer: PreTrainedTokenizer, limit_length: Optional[int] = None, mode: Union[str, Split] = Split.train, - is_language_sensitive: Optional[bool] = False, + is_language_sensitive: bool = False, cache_dir: Optional[str] = None, - dataset_format: Optional[str] = "pt", + dataset_format: str = "pt", ): self.args = args self.is_language_sensitive = is_language_sensitive diff --git a/src/transformers/data/metrics/squad_metrics.py b/src/transformers/data/metrics/squad_metrics.py index f83c23bdeecf..0ffc025b65a0 100644 --- a/src/transformers/data/metrics/squad_metrics.py +++ b/src/transformers/data/metrics/squad_metrics.py @@ -148,7 +148,7 @@ def find_best_thresh_v2(preds, scores, na_probs, qid_to_has_ans): best_score = cur_score best_thresh = 0.0 qid_list = sorted(na_probs, key=lambda k: na_probs[k]) - for i, qid in enumerate(qid_list): + for qid in qid_list: if qid not in scores: continue if qid_to_has_ans[qid]: diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py index ab6e747d14db..e30fc7410d80 100644 --- a/src/transformers/dependency_versions_table.py +++ b/src/transformers/dependency_versions_table.py @@ -46,7 +46,6 @@ "opencv-python": "opencv-python", "optimum-benchmark": "optimum-benchmark>=0.3.0", "optuna": "optuna", - "optax": "optax>=0.0.8,<=0.1.4", "pandas": "pandas<2.3.0", "packaging": "packaging>=20.0", "parameterized": "parameterized>=0.9", @@ -68,7 +67,7 @@ "rhoknp": "rhoknp>=1.1.0,<1.3.1", "rjieba": "rjieba", "rouge-score": "rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1", - "ruff": "ruff==0.11.2", + "ruff": "ruff==0.13.1", "sacrebleu": "sacrebleu>=1.4.12,<2.0.0", "sacremoses": "sacremoses", "safetensors": "safetensors>=0.4.3", diff --git a/src/transformers/dynamic_module_utils.py b/src/transformers/dynamic_module_utils.py index 5b541c076f63..6d4e2bf48921 100644 --- a/src/transformers/dynamic_module_utils.py +++ b/src/transformers/dynamic_module_utils.py @@ -285,8 +285,7 @@ def get_class_in_module( `typing.Type`: The class looked for. """ name = os.path.normpath(module_path) - if name.endswith(".py"): - name = name[:-3] + name = name.removesuffix(".py") name = name.replace(os.path.sep, ".") module_file: Path = Path(HF_MODULES_CACHE) / module_path with _HF_REMOTE_CODE_LOCK: @@ -396,7 +395,7 @@ def get_cached_module_file( if is_local: submodule = _sanitize_module_name(os.path.basename(pretrained_model_name_or_path)) else: - submodule = _sanitize_module_name(pretrained_model_name_or_path.replace("/", os.path.sep)) + submodule = os.path.sep.join(map(_sanitize_module_name, pretrained_model_name_or_path.split("/"))) cached_module = try_to_load_from_cache( pretrained_model_name_or_path, module_file, cache_dir=cache_dir, revision=_commit_hash, repo_type=repo_type ) diff --git a/src/transformers/feature_extraction_utils.py b/src/transformers/feature_extraction_utils.py index a9ff39b0cc19..e007e72d4761 100644 --- a/src/transformers/feature_extraction_utils.py +++ b/src/transformers/feature_extraction_utils.py @@ -48,13 +48,12 @@ if TYPE_CHECKING: - if is_torch_available(): - import torch # noqa + from .feature_extraction_sequence_utils import SequenceFeatureExtractor logger = logging.get_logger(__name__) -PreTrainedFeatureExtractor = Union["SequenceFeatureExtractor"] # noqa: F821 +PreTrainedFeatureExtractor = Union["SequenceFeatureExtractor"] # type hinting: specifying the type of feature extractor class that inherits from FeatureExtractionMixin SpecificFeatureExtractorType = TypeVar("SpecificFeatureExtractorType", bound="FeatureExtractionMixin") @@ -127,7 +126,7 @@ def _get_is_as_tensor_fns(self, tensor_type: Optional[Union[str, TensorType]] = elif 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 # noqa + import torch def as_tensor(value): if isinstance(value, (list, tuple)) and len(value) > 0: @@ -216,7 +215,7 @@ def to(self, *args, **kwargs) -> "BatchFeature": [`BatchFeature`]: The same instance after modification. """ requires_backends(self, ["torch"]) - import torch # noqa + import torch device = kwargs.get("device") non_blocking = kwargs.get("non_blocking", False) @@ -563,7 +562,9 @@ def get_feature_extractor_dict( return feature_extractor_dict, kwargs @classmethod - def from_dict(cls, feature_extractor_dict: dict[str, Any], **kwargs) -> PreTrainedFeatureExtractor: + def from_dict( + cls, feature_extractor_dict: dict[str, Any], **kwargs + ) -> Union["FeatureExtractionMixin", tuple["FeatureExtractionMixin", dict[str, Any]]]: """ Instantiates a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a Python dictionary of parameters. @@ -613,7 +614,7 @@ def to_dict(self) -> dict[str, Any]: return output @classmethod - def from_json_file(cls, json_file: Union[str, os.PathLike]) -> PreTrainedFeatureExtractor: + def from_json_file(cls, json_file: Union[str, os.PathLike]) -> "FeatureExtractionMixin": """ Instantiates a feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] from the path to a JSON file of parameters. diff --git a/src/transformers/generation/beam_search.py b/src/transformers/generation/beam_search.py index ba2820cb437a..8510a02c803a 100644 --- a/src/transformers/generation/beam_search.py +++ b/src/transformers/generation/beam_search.py @@ -165,10 +165,10 @@ def __init__( batch_size: int, num_beams: int, device: torch.device, - length_penalty: Optional[float] = 1.0, - do_early_stopping: Optional[Union[bool, str]] = False, - num_beam_hyps_to_keep: Optional[int] = 1, - num_beam_groups: Optional[int] = 1, + length_penalty: float = 1.0, + do_early_stopping: Union[bool, str] = False, + num_beam_hyps_to_keep: int = 1, + num_beam_groups: int = 1, max_length: Optional[int] = None, ): logger.warning_once( @@ -214,7 +214,7 @@ def __init__( @property def is_done(self) -> bool: - return self._done.all() + return self._done.all().item() def process( self, @@ -225,8 +225,8 @@ def process( pad_token_id: Optional[Union[int, torch.Tensor]] = None, eos_token_id: Optional[Union[int, list[int], torch.Tensor]] = None, beam_indices: Optional[torch.LongTensor] = None, - group_index: Optional[int] = 0, - decoder_prompt_len: Optional[int] = 0, + group_index: int = 0, + decoder_prompt_len: int = 0, ) -> dict[str, torch.Tensor]: # add up to the length which the next_scores is calculated on (including decoder prompt) cur_len = input_ids.shape[-1] + 1 @@ -331,7 +331,7 @@ def finalize( pad_token_id: Optional[Union[int, torch.Tensor]] = None, eos_token_id: Optional[Union[int, list[int], torch.Tensor]] = None, beam_indices: Optional[torch.LongTensor] = None, - decoder_prompt_len: Optional[int] = 0, + decoder_prompt_len: int = 0, ) -> tuple[torch.LongTensor]: batch_size = len(self._beam_hyps) // self.num_beam_groups @@ -460,9 +460,9 @@ def __init__( num_beams: int, constraints: list[Constraint], device: torch.device, - length_penalty: Optional[float] = 1.0, - do_early_stopping: Optional[Union[bool, str]] = False, - num_beam_hyps_to_keep: Optional[int] = 1, + length_penalty: float = 1.0, + do_early_stopping: Union[bool, str] = False, + num_beam_hyps_to_keep: int = 1, max_length: Optional[int] = None, ): logger.warning_once( @@ -495,7 +495,7 @@ def __init__( @property def is_done(self) -> bool: - return self._done.all() + return self._done.all().item() def make_constraint_states(self, n): return [ConstraintListState([constraint.copy() for constraint in self.constraints]) for _ in range(n)] @@ -515,7 +515,7 @@ def process( pad_token_id: Optional[Union[int, torch.Tensor]] = None, eos_token_id: Optional[Union[int, list[int], torch.Tensor]] = None, beam_indices: Optional[torch.LongTensor] = None, - decoder_prompt_len: Optional[int] = 0, + decoder_prompt_len: int = 0, ) -> tuple[torch.Tensor]: r""" Args: @@ -804,7 +804,7 @@ def finalize( pad_token_id: Optional[Union[int, torch.Tensor]] = None, eos_token_id: Optional[Union[int, list[int], torch.Tensor]] = None, beam_indices: Optional[torch.LongTensor] = None, - decoder_prompt_len: Optional[int] = 0, + decoder_prompt_len: int = 0, ) -> tuple[torch.LongTensor]: batch_size = len(self._beam_hyps) @@ -912,7 +912,9 @@ def finalize( class BeamHypotheses: - def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool, max_length: Optional[int] = None): + def __init__( + self, num_beams: int, length_penalty: float, early_stopping: Union[bool, str], max_length: Optional[int] = None + ): """ Initialize n-best list of hypotheses. """ @@ -963,7 +965,7 @@ def add( else: self.worst_score = min(score, self.worst_score) - def is_done(self, best_sum_logprobs: float, cur_len: int, decoder_prompt_len: Optional[int] = 0) -> bool: + def is_done(self, best_sum_logprobs: float, cur_len: int, decoder_prompt_len: int = 0) -> bool: """ If there are enough hypotheses and that none of the hypotheses being generated can become better than the worst one in the heap, then we are done with this sentence. diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index a455e69d03ff..cd42288aebfa 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -524,7 +524,7 @@ def get_candidates(self, input_ids: torch.LongTensor) -> tuple[torch.LongTensor, self.assistant_kwargs.pop("attention_mask", None) assistant_output = self.assistant_model.generate(**generation_args, **self.assistant_kwargs) - new_target_ids = self._process_assistant_outputs(input_ids, assistant_output.sequences, assistant_input_ids) + new_target_ids = self._process_assistant_outputs(input_ids, assistant_output.sequences) # Update state self.prev_target_ids_len = input_ids.shape[1] @@ -583,7 +583,7 @@ def _prepare_assistant_input_ids(self, input_ids: torch.LongTensor) -> tuple[tor return assistant_input_ids, remove_from_pkv def _process_assistant_outputs( - self, input_ids: torch.LongTensor, assistant_sequences: torch.LongTensor, assistant_input_ids: torch.LongTensor + self, input_ids: torch.LongTensor, assistant_sequences: torch.LongTensor ) -> torch.LongTensor: """Processes assistant outputs to obtain target input IDs.""" num_prev_assistant = self.prev_assistant_ids.shape[1] diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index 05caed152c6e..98a0d14ade1a 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -1282,11 +1282,11 @@ class WatermarkingConfig(BaseWatermarkingConfig): def __init__( self, - greenlist_ratio: Optional[float] = 0.25, - bias: Optional[float] = 2.0, - hashing_key: Optional[int] = 15485863, - seeding_scheme: Optional[str] = "lefthash", - context_width: Optional[int] = 1, + greenlist_ratio: float = 0.25, + bias: float = 2.0, + hashing_key: int = 15485863, + seeding_scheme: str = "lefthash", + context_width: int = 1, ): self.greenlist_ratio = greenlist_ratio self.bias = bias diff --git a/src/transformers/generation/continuous_batching/cache.py b/src/transformers/generation/continuous_batching/cache.py index 05de093f661f..8d6e057be84a 100644 --- a/src/transformers/generation/continuous_batching/cache.py +++ b/src/transformers/generation/continuous_batching/cache.py @@ -79,7 +79,7 @@ class PagedAttentionCache: layer group, and the shape of the cache tensor is `[num_blocks * block_size, num_heads, head_size]`. Grouping layers into groups is useful because when we allocate one block to a group N, the block allocated is the - same for all layers in group N, equivalently it is allocated accross all cache tensors. This allows us to + same for all layers in group N, equivalently it is allocated across all cache tensors. This allows us to efficiently allocate and free blocks, and to efficiently read and write key and value states. For instance, imagine we have 8 blocks of cache and a model with two layer groups: a full-attention group with 3 @@ -349,7 +349,7 @@ class PagedAttentionMemoryHandler: The memory footprint consists of three main components: - Cache memory: the space needed to store the cache tensors: 2 * layer_group_size * [num_pages, page_size] * cache_dtype - - Activation memory: the space temporarly taken by the largest activation during the model forward pass: + - Activation memory: the space temporarily taken by the largest activation during the model forward pass: peak_activation_per_token * max_tokens_per_batch * activation_dtype_size - Static tensors: the space taken by the input/output buffers and metadata tensors for batch processing, sum of: - inputs_ids + outputs_ids + position_ids + logits_indices: 4 * max_tokens_per_batch * int32_size diff --git a/src/transformers/generation/continuous_batching/continuous_api.py b/src/transformers/generation/continuous_batching/continuous_api.py index b00c0a4825c3..0d1801fa163e 100644 --- a/src/transformers/generation/continuous_batching/continuous_api.py +++ b/src/transformers/generation/continuous_batching/continuous_api.py @@ -42,7 +42,56 @@ def build_attention_mask( ) -> None: """Builds an attention mask inplace using the cumulative seqlens of the query and key. If given a sliding window, it will also apply a sliding window mask on top. The attention mask is not boolean, it uses zeroes and -inf (or its - equivalent) so it's more of an attention score bias tensor.""" + equivalent) so it's more of an attention score bias tensor. + The attention mask is a block-diagonal matrix, with each block an attention mask for a single query-key pair. + Each of those block is built from a causal mask and, if there is a sliding window, a sliding window mask. + + An example is represented below, with seqlen_k = 8, seqlen_q = 4 and sliding_window = 6: + + CAUSAL MASK: + + █ █ █ █ █ ░ ░ ░ + █ █ █ █ █ █ ░ ░ + █ █ █ █ █ █ █ ░ + █ █ █ █ █ █ █ █ + + SLIDING WINDOW MASK: + ┌──────────────────────── seqlen_k - seqlen_q - sliding_window = 8 - 4 - 6 = -2 offset to the right + <─┴─> + ░ █ | █ █ █ █ █ █ █ █ + ░ ░ | █ █ █ █ █ █ █ █ + ░ ░ | ░ █ █ █ █ █ █ █ + ░ ░ | ░ ░ █ █ █ █ █ █ + + ATTENTION MASK (sum of causal and sliding window masks): + + █ █ █ █ █ ░ ░ ░ + █ █ █ █ █ █ ░ ░ + ░ █ █ █ █ █ █ ░ + ░ ░ █ █ █ █ █ █ + + Another example with seqlen_k = 5, seqlen_q = 3 and sliding_window = 2: + + CAUSAL MASK: + + █ █ █ ░ ░ + █ █ █ █ ░ + █ █ █ █ █ + + SLIDING WINDOW MASK: + ┌──────────────────────── seqlen_k - seqlen_q - sliding_window = 5 - 3 - 2 = 0 offset to the right + <┴> + | ░ █ █ █ █ + | ░ ░ █ █ █ + | ░ ░ ░ █ █ + + ATTENTION MASK (sum of causal and sliding window masks): + + ░ █ █ ░ ░ + ░ ░ █ █ ░ + ░ ░ ░ █ █ + + """ min_value = torch.finfo(attention_mask.dtype).min for i in range(len(cumulative_seqlens_q) - 1): seqlen_q = cumulative_seqlens_q[i + 1] - cumulative_seqlens_q[i] @@ -63,8 +112,8 @@ def build_attention_mask( masked = torch.triu(minus_inf, diagonal=causal_diagonal) # Apply sliding window mask if needed if sliding_window > 1: - sliding_diagonal = seqlen_k - seqlen_q + sliding_window - masked = torch.tril(masked, diagonal=sliding_diagonal) + sliding_diagonal = seqlen_k - seqlen_q - sliding_window + masked += torch.tril(minus_inf, diagonal=sliding_diagonal) # Replace in attention mask attention_mask[..., query_range, key_range] = masked diff --git a/src/transformers/generation/logits_process.py b/src/transformers/generation/logits_process.py index f63d2246c6a9..7d81501a783d 100644 --- a/src/transformers/generation/logits_process.py +++ b/src/transformers/generation/logits_process.py @@ -369,7 +369,6 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to if scores.dim() == 3: if self.logits_indices is not None and self.cu_seq_lens_q is not None: - batch_size, seq_len, vocab_size = scores.shape last_positions = self.logits_indices last_scores = scores[0, last_positions, :] @@ -2289,7 +2288,7 @@ def __init__( model, unconditional_ids: Optional[torch.LongTensor] = None, unconditional_attention_mask: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = True, + use_cache: bool = True, ): self.guidance_scale = guidance_scale self.model = model diff --git a/src/transformers/generation/stopping_criteria.py b/src/transformers/generation/stopping_criteria.py index 2b9e57aacd8d..5a013a49723d 100644 --- a/src/transformers/generation/stopping_criteria.py +++ b/src/transformers/generation/stopping_criteria.py @@ -76,9 +76,9 @@ def __init__(self, max_length: int, max_position_embeddings: Optional[int] = Non def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: cur_len = input_ids.shape[1] is_done = cur_len >= self.max_length - if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings: + if self.max_position_embeddings is not None and not is_done and cur_len > self.max_position_embeddings: logger.warning_once( - "This is a friendly reminder - the current text generation call will exceed the model's predefined " + "This is a friendly reminder - the current text generation call has exceeded the model's predefined " f"maximum length ({self.max_position_embeddings}). Depending on the model, you may observe " "exceptions, performance degradation, or nothing at all." ) @@ -249,7 +249,7 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_strings: Union[str, token_list, token_indices, tokenizer ) - self.maximum_token_len = max([len(stop_string) for stop_string in self.stop_strings]) + self.maximum_token_len = max(len(stop_string) for stop_string in self.stop_strings) self.num_stop_strings = len(self.stop_strings) self.target_lens = torch.tensor([len(stop_string) for stop_string in stop_strings], dtype=torch.int32) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 2e312bcb3c79..f9d58dfdf4f6 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -22,7 +22,6 @@ import torch import torch.distributed as dist -from huggingface_hub import file_exists from packaging import version from torch import nn @@ -414,23 +413,20 @@ def load_custom_generate( Returns: A callable that can be used to generate text. """ - # Does `pretrained_model_name_or_path` have a `custom_generate` subdirectory? If not -> OSError - is_local_code = os.path.exists(pretrained_model_name_or_path) - has_custom_generate_folder = True - if is_local_code: - if not os.path.exists(os.path.join(pretrained_model_name_or_path, "custom_generate/generate.py")): - has_custom_generate_folder = False - else: - if not file_exists(pretrained_model_name_or_path, "custom_generate/generate.py"): - has_custom_generate_folder = False - - if not has_custom_generate_folder: + # Fetches the generate.py file from the model repo. If it doesn't exist, a file in `.no_exist` cache directory + # is created (preventing future hub requests), and an OSError is raised. + try: + module = get_cached_module_file( + pretrained_model_name_or_path, module_file="custom_generate/generate.py", **kwargs + ) + except OSError: raise OSError( f"`{pretrained_model_name_or_path}` does not contain a `custom_generate` subdirectory with a " "`generate.py` file, can't load the custom generate function." ) # Handle opt-in `trust_remote_code` and related exceptions + is_local_code = os.path.exists(pretrained_model_name_or_path) error_message = ( f"The repository `{pretrained_model_name_or_path}` contains custom generation code that will override " "the default `generate` method." @@ -447,9 +443,6 @@ def load_custom_generate( check_python_requirements( pretrained_model_name_or_path, requirements_file="custom_generate/requirements.txt", **kwargs ) - module = get_cached_module_file( - pretrained_model_name_or_path, module_file="custom_generate/generate.py", **kwargs - ) custom_generate_function = get_class_in_module("generate", module) return custom_generate_function @@ -912,7 +905,7 @@ def _prepare_decoder_input_ids_for_generation( self.config.model_type == "vision-encoder-decoder" and "donut" in self.config.encoder.model_type.lower() ): pass - elif self.config.model_type in ["whisper"]: + elif self.config.model_type == "whisper": pass # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust # decoder_attention_mask if provided) @@ -1018,7 +1011,7 @@ def _get_candidate_generator( input_ids: torch.LongTensor, inputs_tensor: torch.Tensor, logits_processor: LogitsProcessorList, - model_kwargs: dict, + model_kwargs: dict[str, Any], assistant_model: Optional["PreTrainedModel"] = None, target_tokenizer: Optional["PreTrainedTokenizerBase"] = None, assistant_tokenizer: Optional["PreTrainedTokenizerBase"] = None, @@ -1709,7 +1702,10 @@ def _prepare_generated_length( return generation_config def _prepare_generation_config( - self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: dict + self, + generation_config: Optional[GenerationConfig], + use_model_defaults: Optional[bool] = None, + **kwargs: Any, ) -> tuple[GenerationConfig, dict]: """ Prepares the base generation config, then applies any generation configuration options from kwargs. This @@ -1903,6 +1899,7 @@ def _supports_default_dynamic_cache(cls) -> bool: "minimax", "xlnet", "lfm2", + "lfm2-vl", ] ) @@ -2136,7 +2133,7 @@ def _tensor_or_none(token, device=None): generation_config._pad_token_tensor = pad_token_tensor generation_config._decoder_start_token_tensor = decoder_start_token_tensor - def _valid_auto_compile_criteria(self, model_kwargs: dict, generation_config: GenerationConfig) -> bool: + def _valid_auto_compile_criteria(self, model_kwargs: dict[str, Any], generation_config: GenerationConfig) -> bool: """ Determines whether to trigger auto-compilation of the model's forward pass at generation time. """ @@ -3453,7 +3450,7 @@ def _assisted_decoding( generation_config: GenerationConfig, synced_gpus: bool = False, streamer: Optional["BaseStreamer"] = None, - inputs_tensor: torch.FloatTensor = None, + inputs_tensor: Optional[torch.FloatTensor] = None, assistant_model: Optional["PreTrainedModel"] = None, assistant_tokenizer: Optional["PreTrainedTokenizerBase"] = None, tokenizer: Optional["PreTrainedTokenizerBase"] = None, diff --git a/src/transformers/generation/watermarking.py b/src/transformers/generation/watermarking.py index e62742ef7514..df8a6ef7d483 100644 --- a/src/transformers/generation/watermarking.py +++ b/src/transformers/generation/watermarking.py @@ -24,14 +24,9 @@ from torch.nn import BCELoss from ..modeling_utils import PreTrainedModel -from ..utils import ModelOutput, is_torch_available, logging +from ..utils import ModelOutput, logging from .configuration_utils import PretrainedConfig, WatermarkingConfig - - -if is_torch_available(): - import torch - - from .logits_process import SynthIDTextWatermarkLogitsProcessor, WatermarkLogitsProcessor +from .logits_process import SynthIDTextWatermarkLogitsProcessor, WatermarkLogitsProcessor logger = logging.get_logger(__name__) @@ -43,31 +38,31 @@ class WatermarkDetectorOutput: Outputs of a watermark detector. Args: - num_tokens_scored (np.array of shape (batch_size)): + num_tokens_scored (np.ndarray of shape (batch_size)): Array containing the number of tokens scored for each element in the batch. - num_green_tokens (np.array of shape (batch_size)): + num_green_tokens (np.ndarray of shape (batch_size)): Array containing the number of green tokens for each element in the batch. - green_fraction (np.array of shape (batch_size)): + green_fraction (np.ndarray of shape (batch_size)): Array containing the fraction of green tokens for each element in the batch. - z_score (np.array of shape (batch_size)): + z_score (np.ndarray of shape (batch_size)): Array containing the z-score for each element in the batch. Z-score here shows how many standard deviations away is the green token count in the input text from the expected green token count for machine-generated text. - p_value (np.array of shape (batch_size)): + p_value (np.ndarray of shape (batch_size)): Array containing the p-value for each batch obtained from z-scores. - prediction (np.array of shape (batch_size)), *optional*: + prediction (np.ndarray of shape (batch_size)), *optional*: Array containing boolean predictions whether a text is machine-generated for each element in the batch. - confidence (np.array of shape (batch_size)), *optional*: + confidence (np.ndarray of shape (batch_size)), *optional*: Array containing confidence scores of a text being machine-generated for each element in the batch. """ - num_tokens_scored: Optional[np.array] = None - num_green_tokens: Optional[np.array] = None - green_fraction: Optional[np.array] = None - z_score: Optional[np.array] = None - p_value: Optional[np.array] = None - prediction: Optional[np.array] = None - confidence: Optional[np.array] = None + num_tokens_scored: Optional[np.ndarray] = None + num_green_tokens: Optional[np.ndarray] = None + green_fraction: Optional[np.ndarray] = None + z_score: Optional[np.ndarray] = None + p_value: Optional[np.ndarray] = None + prediction: Optional[np.ndarray] = None + confidence: Optional[np.ndarray] = None class WatermarkDetector: @@ -179,7 +174,7 @@ def _score_ngrams_in_passage(self, input_ids: torch.LongTensor): ) return num_tokens_scored_batch, green_token_count_batch - def _compute_z_score(self, green_token_count: np.ndarray, total_num_tokens: np.ndarray) -> np.array: + def _compute_z_score(self, green_token_count: np.ndarray, total_num_tokens: np.ndarray) -> np.ndarray: expected_count = self.greenlist_ratio numer = green_token_count - expected_count * total_num_tokens denom = np.sqrt(total_num_tokens * expected_count * (1 - expected_count)) @@ -195,7 +190,7 @@ def __call__( input_ids: torch.LongTensor, z_threshold: float = 3.0, return_dict: bool = False, - ) -> Union[WatermarkDetectorOutput, np.array]: + ) -> Union[WatermarkDetectorOutput, np.ndarray]: """ Args: input_ids (`torch.LongTensor`): @@ -207,8 +202,8 @@ def __call__( Whether to return `~generation.WatermarkDetectorOutput` or not. If not it will return boolean predictions, ma Return: - [`~generation.WatermarkDetectorOutput`] or `np.array`: A [`~generation.WatermarkDetectorOutput`] - if `return_dict=True` otherwise a `np.array`. + [`~generation.WatermarkDetectorOutput`] or `np.ndarray`: A [`~generation.WatermarkDetectorOutput`] + if `return_dict=True` otherwise a `np.ndarray`. """ diff --git a/src/transformers/hf_argparser.py b/src/transformers/hf_argparser.py index be7f05344faf..503130ea651a 100644 --- a/src/transformers/hf_argparser.py +++ b/src/transformers/hf_argparser.py @@ -262,19 +262,6 @@ def _add_dataclass_arguments(self, dtype: DataClassType): "removing line of `from __future__ import annotations` which opts in Postponed " "Evaluation of Annotations (PEP 563)" ) - except TypeError as ex: - # Remove this block when we drop Python 3.9 support - if sys.version_info[:2] < (3, 10) and "unsupported operand type(s) for |" in str(ex): - python_version = ".".join(map(str, sys.version_info[:3])) - raise RuntimeError( - f"Type resolution failed for {dtype} on Python {python_version}. Try removing " - "line of `from __future__ import annotations` which opts in union types as " - "`X | Y` (PEP 604) via Postponed Evaluation of Annotations (PEP 563). To " - "support Python versions that lower than 3.10, you need to use " - "`typing.Union[X, Y]` instead of `X | Y` and `typing.Optional[X]` instead of " - "`X | None`." - ) from ex - raise for field in dataclasses.fields(dtype): if not field.init: diff --git a/src/transformers/image_processing_utils_fast.py b/src/transformers/image_processing_utils_fast.py index 983fd4e16953..4dfa7f08b0db 100644 --- a/src/transformers/image_processing_utils_fast.py +++ b/src/transformers/image_processing_utils_fast.py @@ -46,7 +46,6 @@ auto_docstring, is_torch_available, is_torchvision_available, - is_torchvision_v2_available, is_vision_available, logging, ) @@ -60,14 +59,13 @@ import torch if is_torchvision_available(): + from torchvision.transforms.v2 import functional as F + from .image_utils import pil_torch_interpolation_mapping + else: pil_torch_interpolation_mapping = None -if is_torchvision_v2_available(): - from torchvision.transforms.v2 import functional as F -elif is_torchvision_available(): - from torchvision.transforms import functional as F logger = logging.get_logger(__name__) @@ -85,7 +83,7 @@ def validate_fast_preprocess_arguments( size: Optional[SizeDict] = None, interpolation: Optional["F.InterpolationMode"] = None, return_tensors: Optional[Union[str, TensorType]] = None, - data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, + data_format: ChannelDimension = ChannelDimension.FIRST, ): """ Checks validity of typically used arguments in an `ImageProcessorFast` `preprocess` method. @@ -131,7 +129,7 @@ def max_across_indices(values: Iterable[Any]) -> list[Any]: return [max(values_i) for values_i in zip(*values)] -def get_max_height_width(images: list["torch.Tensor"]) -> tuple[int]: +def get_max_height_width(images: list["torch.Tensor"]) -> tuple[int, ...]: """ Get the maximum height and width across all images in a batch. """ @@ -142,8 +140,8 @@ def get_max_height_width(images: list["torch.Tensor"]) -> tuple[int]: def divide_to_patches( - image: Union[np.array, "torch.Tensor"], patch_size: int -) -> list[Union[np.array, "torch.Tensor"]]: + image: Union[np.ndarray, "torch.Tensor"], patch_size: int +) -> list[Union[np.ndarray, "torch.Tensor"]]: """ Divides an image into patches of a specified size. @@ -248,7 +246,7 @@ def pad( pad_size: SizeDict = None, fill_value: Optional[int] = 0, padding_mode: Optional[str] = "constant", - return_mask: Optional[bool] = False, + return_mask: bool = False, disable_grouping: Optional[bool] = False, **kwargs, ) -> "torch.Tensor": @@ -375,9 +373,13 @@ def compile_friendly_resize( A wrapper around `F.resize` so that it is compatible with torch.compile when the image is a uint8 tensor. """ if image.dtype == torch.uint8: - image = image.float() / 255 + # 256 is used on purpose instead of 255 to avoid numerical differences + # see https://github.com/huggingface/transformers/pull/38540#discussion_r2127165652 + image = image.float() / 256 image = F.resize(image, new_size, interpolation=interpolation, antialias=antialias) - image = image * 255 + image = image * 256 + # torch.where is used on purpose instead of torch.clamp to avoid bug in torch.compile + # see https://github.com/huggingface/transformers/pull/38540#discussion_r2126888471 image = torch.where(image > 255, 255, image) image = torch.where(image < 0, 0, image) image = image.round().to(torch.uint8) diff --git a/src/transformers/image_transforms.py b/src/transformers/image_transforms.py index f0aeae8985b7..c0158b7111b7 100644 --- a/src/transformers/image_transforms.py +++ b/src/transformers/image_transforms.py @@ -255,7 +255,7 @@ def get_size_with_aspect_ratio(image_size, size, max_size=None) -> tuple[int, in # Logic adapted from torchvision resizing logic: https://github.com/pytorch/vision/blob/511924c1ced4ce0461197e5caa64ce5b9e558aab/torchvision/transforms/functional.py#L366 def get_resize_output_image_size( input_image: np.ndarray, - size: Union[int, tuple[int, int], list[int], tuple[int]], + size: Union[int, tuple[int, int], list[int], tuple[int, ...]], default_to_square: bool = True, max_size: Optional[int] = None, input_data_format: Optional[Union[str, ChannelDimension]] = None, @@ -323,7 +323,7 @@ def get_resize_output_image_size( def resize( image: np.ndarray, size: tuple[int, int], - resample: "PILImageResampling" = None, + resample: Optional["PILImageResampling"] = None, reducing_gap: Optional[int] = None, data_format: Optional[ChannelDimension] = None, return_numpy: bool = True, diff --git a/src/transformers/image_utils.py b/src/transformers/image_utils.py index 2079c21f3b0c..c5f4d4a3fa4c 100644 --- a/src/transformers/image_utils.py +++ b/src/transformers/image_utils.py @@ -30,7 +30,6 @@ is_torch_available, is_torch_tensor, is_torchvision_available, - is_torchvision_v2_available, is_vision_available, logging, requires_backends, @@ -56,9 +55,7 @@ from torchvision.transforms import InterpolationMode pil_torch_interpolation_mapping = { - PILImageResampling.NEAREST: InterpolationMode.NEAREST_EXACT - if is_torchvision_v2_available() - else InterpolationMode.NEAREST, + PILImageResampling.NEAREST: InterpolationMode.NEAREST_EXACT, PILImageResampling.BOX: InterpolationMode.BOX, PILImageResampling.BILINEAR: InterpolationMode.BILINEAR, PILImageResampling.HAMMING: InterpolationMode.HAMMING, @@ -78,7 +75,7 @@ ImageInput = Union[ "PIL.Image.Image", np.ndarray, "torch.Tensor", list["PIL.Image.Image"], list[np.ndarray], list["torch.Tensor"] -] # noqa +] class ChannelDimension(ExplicitEnum): @@ -486,9 +483,7 @@ def load_image(image: Union[str, "PIL.Image.Image"], timeout: Optional[float] = raise ValueError( f"Incorrect image source. Must be a valid URL starting with `http://` or `https://`, a valid path to an image file, or a base64 encoded string. Got {image}. Failed with {e}" ) - elif isinstance(image, PIL.Image.Image): - image = image - else: + elif not isinstance(image, PIL.Image.Image): raise TypeError( "Incorrect format used for image. Should be an url linking to an image, a base64 string, a local path, or a PIL image." ) @@ -579,7 +574,7 @@ class ImageFeatureExtractionMixin: def _ensure_format_supported(self, image): if not isinstance(image, (PIL.Image.Image, np.ndarray)) and not is_torch_tensor(image): raise ValueError( - f"Got type {type(image)} which is not supported, only `PIL.Image.Image`, `np.array` and " + f"Got type {type(image)} which is not supported, only `PIL.Image.Image`, `np.ndarray` and " "`torch.Tensor` are." ) diff --git a/src/transformers/integrations/deepspeed.py b/src/transformers/integrations/deepspeed.py index 47d7a7ffcb5f..c5f9ecc03b53 100644 --- a/src/transformers/integrations/deepspeed.py +++ b/src/transformers/integrations/deepspeed.py @@ -130,58 +130,11 @@ def fill_match(self, ds_key_long, hf_val, hf_key=None, must_match=True): fill_only = partialmethod(fill_match, must_match=False) - def override_training_args_from_deepspeed(self, args): - """ - Override TrainingArguments based on DeepSpeed config values to ensure compatibility. - - This method ensures that the DeepSpeed config takes precedence over TrainingArguments - defaults when there are conflicts, particularly for mixed precision settings. - - Args: - args: TrainingArguments object to potentially modify - """ - # Check precision settings in DeepSpeed config and override TrainingArguments accordingly - # Only override defaults, not explicit user settings - - # Check if user explicitly set precision options (we assume defaults are False) - user_set_fp16 = args.fp16 is True - user_set_bf16 = args.bf16 is True - - if self.is_true("fp16.enabled"): - # DeepSpeed config explicitly enables fp16 - if not user_set_fp16 and not user_set_bf16: - # User didn't explicitly set either, so apply DeepSpeed config - args.fp16 = True - args.bf16 = False - elif user_set_bf16 and not user_set_fp16: - # User explicitly chose bf16, but DeepSpeed config wants fp16 - # This is a potential conflict - let user choice win but log a warning - pass # Keep user's bf16=True, fp16=False - elif self.is_true("bf16.enabled"): - # DeepSpeed config explicitly enables bf16 - if not user_set_fp16 and not user_set_bf16: - # User didn't explicitly set either, so apply DeepSpeed config - args.bf16 = True - args.fp16 = False - elif user_set_fp16 and not user_set_bf16: - # User explicitly chose fp16, but DeepSpeed config wants bf16 - # This is a potential conflict - let user choice win but log a warning - pass # Keep user's fp16=True, bf16=False - elif self.is_false("fp16.enabled") and self.is_false("bf16.enabled"): - # Both are explicitly disabled in DeepSpeed config - if not user_set_fp16 and not user_set_bf16: - # User didn't explicitly set either, so apply DeepSpeed config (fp32) - args.fp16 = False - args.bf16 = False - def trainer_config_process(self, args, auto_find_batch_size=False): """ Adjust the config with `TrainingArguments` values. This stage is run during `TrainingArguments` object creation. """ - # First, override TrainingArguments based on DeepSpeed config to ensure compatibility - self.override_training_args_from_deepspeed(args) - # DeepSpeed does: # train_batch_size = world_size * train_micro_batch_size_per_gpu * gradient_accumulation_steps train_batch_size = args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps @@ -268,17 +221,20 @@ def trainer_config_finalize(self, args, model, num_training_steps): hidden_size_auto_keys = [x for x in hidden_size_based_keys if self.is_auto(x)] if len(hidden_size_auto_keys) > 0: - if hasattr(model.config, "hidden_size"): - hidden_size = model.config.hidden_size - elif hasattr(model.config, "hidden_sizes"): - # if there are many hidden sizes pick the largest one - hidden_size = max(model.config.hidden_sizes) - elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_size"): - hidden_size = model.config.text_config.hidden_size - elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_sizes"): - # if there are many hidden sizes pick the largest one - hidden_size = max(model.config.text_config.hidden_sizes) - else: + hidden_size = None + if hasattr(model, "config"): + if hasattr(model.config, "hidden_size"): + hidden_size = model.config.hidden_size + elif hasattr(model.config, "hidden_sizes"): + # if there are many hidden sizes pick the largest one + hidden_size = max(model.config.hidden_sizes) + elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_size"): + hidden_size = model.config.text_config.hidden_size + elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_sizes"): + # if there are many hidden sizes pick the largest one + hidden_size = max(model.config.text_config.hidden_sizes) + + if hidden_size is None: raise ValueError( "The model's config file has neither `hidden_size` nor `hidden_sizes` entry, " "therefore it's not possible to automatically fill out the following `auto` entries " @@ -416,7 +372,7 @@ def deepspeed_optim_sched(trainer, hf_deepspeed_config, args, num_training_steps optimizer = None if "optimizer" in config: - if args.adafactor: + if args.optim == "adafactor": raise ValueError( "--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. " "Only one optimizer can be configured." diff --git a/src/transformers/integrations/flash_paged.py b/src/transformers/integrations/flash_paged.py index 329fab4c9323..1d1db72a7605 100644 --- a/src/transformers/integrations/flash_paged.py +++ b/src/transformers/integrations/flash_paged.py @@ -6,11 +6,21 @@ from ..utils import is_flash_attn_2_available +# For some reason, if we dont assign the function to a variable here, it will be garbage collected try: if is_flash_attn_2_available(): from flash_attn import flash_attn_varlen_func # noqa: F401 -except Exception: - pass + + FLASH_ATTN_VARLEN_FUNC = flash_attn_varlen_func + else: + raise RuntimeError( + "Flash Attention 2 is not installed. Please refer to https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install it" + ) +except Exception as e: + msg = repr(e) + + def FLASH_ATTN_VARLEN_FUNC(*args, **kwargs): + raise Exception(f"flash_attn_varlen_func is not available: {msg}") def paged_attention_forward( @@ -58,14 +68,13 @@ def paged_attention_forward( # Retrieve the cumulative sequence lengths for the current layer if isinstance(cu_seq_lens_k, dict): - cu_seq_lens_k = cu_seq_lens_k[layer_type].clone() + cu_seq_lens_k = cu_seq_lens_k[layer_type] max_seqlen_k = max_seqlen_k[layer_type] - else: - cu_seq_lens_k = cu_seq_lens_k.clone() - max_seqlen_k = max_seqlen_k if implementation is not None and hasattr(implementation, "flash_attn_varlen_func"): flash_attn_varlen_func = implementation.flash_attn_varlen_func + else: + flash_attn_varlen_func = FLASH_ATTN_VARLEN_FUNC custom_kwargs = {"s_aux": kwargs.get("s_aux")} if "s_aux" in kwargs else {} diff --git a/src/transformers/integrations/flex_attention.py b/src/transformers/integrations/flex_attention.py index 85ddc433e67a..2701936685dd 100644 --- a/src/transformers/integrations/flex_attention.py +++ b/src/transformers/integrations/flex_attention.py @@ -36,7 +36,7 @@ if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size # noqa: N811 + 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, flex_attention @@ -272,12 +272,9 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): score = score + score_mask[batch_idx][0][q_idx][kv_idx] if head_mask is not None: score = score + head_mask[batch_idx][head_idx][0][0] - if s_aux is not None: - logits_max = torch.max(score, dim=-1, keepdim=True).values - sinks = torch.exp(s_aux - logits_max) - unnormalized_scores = torch.exp(score - logits_max) - normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + sinks - score = unnormalized_scores / normalizer + # Note: attention sinks cannot be correctly implemented in score_mod + # because it requires operating on the full attention matrix before softmax. + # ==> this is done after flex attention return score enable_gqa = True @@ -293,6 +290,11 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it return_lse = query.device.type != "cpu" + if not return_lse and s_aux is not None: + raise ValueError( + "Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA" + ) + flex_attention_output = compile_friendly_flex_attention( query, key, @@ -311,6 +313,21 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): if return_lse: attention_output, lse = flex_attention_output # type: ignore[misc] lse = lse.to(value.dtype) + + if s_aux is not None: + # Apply attention sinks by renormalizing using LSE + batch_size, num_heads, seq_len_q, _ = attention_output.shape # batch, num_heads, seq_len, head_dim + sinks = s_aux.view(1, -1, 1, 1).expand(batch_size, num_heads, seq_len_q, 1) + + # We need to compute the normalization that includes the sinks + # since log(sum(exp(scores))) = lse, exp(log(sum(exp(scores)))) = exp(lse) + # NB: log(sum(exp(scores)) + exp(sink)) = log(exp(lse) + exp(sink)) + lse_expanded = lse.unsqueeze(-1) # [batch, num_heads, seq_len, 1] + combined_lse = torch.logsumexp(torch.cat([lse_expanded, sinks], dim=-1), dim=-1, keepdim=True) + + # Use new_norm / old_norm = exp(combined_lse - lse) to compute renorm and apply + renorm_factor = torch.exp(lse_expanded - combined_lse) + attention_output = attention_output * renorm_factor else: attention_output = flex_attention_output # type: ignore[assignment] lse = None diff --git a/src/transformers/integrations/fp_quant.py b/src/transformers/integrations/fp_quant.py index 89ebac7004ee..0ac441e36f93 100644 --- a/src/transformers/integrations/fp_quant.py +++ b/src/transformers/integrations/fp_quant.py @@ -28,6 +28,8 @@ def adapt_fp_quant_config(config: FPQuantConfig): if config.forward_dtype == "mxfp4": forward_dtype = FPQuantDtype.MXFP4 + elif config.forward_dtype == "nvfp4": + forward_dtype = FPQuantDtype.NVFP4 else: raise ValueError(f"Unsupported forward dtype: {config.forward_dtype}") @@ -43,5 +45,6 @@ def adapt_fp_quant_config(config: FPQuantConfig): store_master_weights=config.store_master_weights, hadamard_group_size=config.hadamard_group_size, pseudoquantization=config.pseudoquantization, + transform_init=config.transform_init, modules_to_not_convert=config.modules_to_not_convert, ) diff --git a/src/transformers/integrations/ggml.py b/src/transformers/integrations/ggml.py index 703fd0156365..d5600050188f 100644 --- a/src/transformers/integrations/ggml.py +++ b/src/transformers/integrations/ggml.py @@ -90,6 +90,19 @@ "expert_count": "num_experts", "expert_used_count": "num_experts_per_tok", }, + "lfm2": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + "shortconv.l_cache": "conv_L_cache", + }, "qwen3": { "context_length": "max_position_embeddings", "block_count": "num_hidden_layers", @@ -316,11 +329,11 @@ def _gguf_parse_value(_value, data_type): _value = int(_value[0]) elif data_type in [6, 12]: _value = float(_value[0]) - elif data_type in [7]: + elif data_type == 7: _value = bool(_value[0]) - elif data_type in [8]: + elif data_type == 8: _value = array("B", list(_value)).tobytes().decode() - elif data_type in [9]: + elif data_type == 9: _value = _gguf_parse_value(_value, array_data_type) return _value diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 5be21e2f9a51..6bf8dbcc0219 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -111,6 +111,27 @@ ) } }, + "SiLU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="Silu", version=">=0.1.0" + ) + } + }, + "GeLU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="Gelu", version=">=0.1.0" + ) + } + }, + "GeluTanh": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="GeluTanh", version=">=0.1.0" + ) + } + }, } register_kernel_mapping(_KERNEL_MAPPING) @@ -152,7 +173,10 @@ def load_and_register_kernel(attn_implementation: str) -> None: if not is_kernel(attn_implementation): return if not _kernels_available: - raise ImportError("`kernels` is not installed. Please install it with `pip install kernels`.") + raise ImportError( + "`kernels` is either not installed or uses an incompatible version. " + "Please install the latest version with `pip install -U kernels`." + ) # Need to be imported here as otherwise we have a circular import in `modeling_utils` from ..masking_utils import ALL_MASK_ATTENTION_FUNCTIONS @@ -188,7 +212,7 @@ def load_and_register_kernel(attn_implementation: str) -> None: if attention_wrapper is None: attention_wrapper = flash_attention_forward kernel_function = partial(attention_wrapper, implementation=kernel) - lazy_import_flash_attention(kernel) + lazy_import_flash_attention(kernel, force_import=True) elif kernel_name is not None: kernel_function = getattr(kernel, kernel_name) # Register the kernel as a valid attention diff --git a/src/transformers/integrations/integration_utils.py b/src/transformers/integrations/integration_utils.py index 5ef1123b8fce..b81d47831b6b 100755 --- a/src/transformers/integrations/integration_utils.py +++ b/src/transformers/integrations/integration_utils.py @@ -547,8 +547,6 @@ def run_hp_search_sigopt(trainer, n_trials: int, direction: str, **kwargs) -> Be def run_hp_search_wandb(trainer, n_trials: int, direction: str, **kwargs) -> BestRun: - from ..integrations import is_wandb_available - if not is_wandb_available(): raise ImportError("This function needs wandb installed: `pip install wandb`") import wandb @@ -686,7 +684,7 @@ def __init__(self, tb_writer=None): ) if has_tensorboard: try: - from torch.utils.tensorboard import SummaryWriter # noqa: F401 + from torch.utils.tensorboard import SummaryWriter self._SummaryWriter = SummaryWriter except ImportError: @@ -1092,19 +1090,28 @@ def setup(self, args, state, model, **kwargs): """ Setup the optional Trackio integration. - To customize the setup you can also override the following environment variables: - - Environment: - - **TRACKIO_PROJECT** (`str`, *optional*, defaults to `"huggingface"`): - The name of the project (can be an existing project to continue tracking or a new project to start tracking - from scratch). - - **TRACKIO_SPACE_ID** (`str`, *optional*, defaults to `None`): - If set, the project will be logged to a Hugging Face Space instead of a local directory. 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 the Space does not - exist, it will be created. If the Space already exists, the project will be logged to it. + To customize the setup you can also set the arguments `project`, `trackio_space_id` and `hub_private_repo` in + [`TrainingArguments`]. Please refer to the docstring of for more details. """ if state.is_world_process_zero: + if os.getenv("TRACKIO_PROJECT"): + logger.warning( + "The `TRACKIO_PROJECT` environment variable is deprecated and will be removed in a future " + "version. Use TrainingArguments.project instead." + ) + project = os.getenv("TRACKIO_PROJECT") + else: + project = args.project + + if os.getenv("TRACKIO_SPACE_ID"): + logger.warning( + "The `TRACKIO_SPACE_ID` environment variable is deprecated and will be removed in a future " + "version. Use TrainingArguments.trackio_space_id instead." + ) + space_id = os.getenv("TRACKIO_SPACE_ID") + else: + space_id = args.trackio_space_id + combined_dict = {**args.to_dict()} if hasattr(model, "config") and model.config is not None: @@ -1115,10 +1122,11 @@ def setup(self, args, state, model, **kwargs): combined_dict = {**{"peft_config": peft_config}, **combined_dict} self._trackio.init( - project=os.getenv("TRACKIO_PROJECT", "huggingface"), + project=project, name=args.run_name, - space_id=os.getenv("TRACKIO_SPACE_ID", None), + space_id=space_id, resume="allow", + private=args.hub_private_repo, ) # Add config parameters (run may have been created manually) diff --git a/src/transformers/integrations/mistral.py b/src/transformers/integrations/mistral.py index 78172329277e..cdf237645fc1 100644 --- a/src/transformers/integrations/mistral.py +++ b/src/transformers/integrations/mistral.py @@ -16,10 +16,8 @@ def __init__( pattern=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""", add_prefix_space=False, additional_special_tokens=None, - *args, **kwargs, ): - super().__init__(*args) self.vocab = vocab self.pattern = pattern self.add_prefix_space = add_prefix_space diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index c40b202c54e8..6a6ce1db17e7 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -23,6 +23,7 @@ from accelerate import init_empty_weights import re +from contextlib import contextmanager logger = logging.get_logger(__name__) @@ -47,6 +48,28 @@ ] +@contextmanager +def on_device(dev): + if is_torch_available(): + import torch + + if isinstance(dev, torch.Tensor): + dev = dev.device + elif isinstance(dev, str): + dev = torch.device(dev) + dev_type = getattr(dev, "type", None) + if dev_type == "cuda": + with torch.cuda.device(dev): + yield + return + if dev_type == "xpu" and hasattr(torch, "xpu"): + with torch.xpu.device(dev): + yield + return + # other: CPU + yield + + # Copied from GPT_OSS repo and vllm def quantize_to_mxfp4(w, triton_kernels_hub): downcast_to_mxfp_torch = triton_kernels_hub.numerics_details.mxfp.downcast_to_mxfp_torch @@ -173,7 +196,7 @@ def forward(self, hidden_states: torch.Tensor, routing_data, gather_idx, scatter ) swiglu_fn = triton_kernels_hub.swiglu.swiglu_fn - with torch.cuda.device(hidden_states.device): + with on_device(hidden_states.device): act = FusedActivation(FnSpecs("swiglu", swiglu_fn, ("alpha", "limit")), (self.alpha, self.limit), 2) intermediate_cache1 = matmul_ogs( @@ -214,7 +237,7 @@ def routing_torch_dist( triton_kernels_hub.routing.compute_expt_data_torch, ) - with torch.cuda.device(logits.device): + with on_device(logits.device): world_size = torch.distributed.get_world_size() rank = int(os.environ.get("LOCAL_RANK", "0")) replace_value = -1 @@ -281,7 +304,7 @@ def mlp_forward(self, hidden_states): hidden_states = hidden_states.reshape(-1, self.router.hidden_dim) router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias) - with torch.cuda.device(router_logits.device): + with on_device(router_logits.device): routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k) routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx) @@ -320,7 +343,6 @@ def dequantize(module, param_name, param_value, target_device, dq_param_name, ** to_contiguous, rank, device_mesh, - set_param=False, ) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" @@ -376,7 +398,7 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito target_device = "cuda" blocks = blocks.to(target_device).contiguous() scales = scales.to(target_device).contiguous() - with torch.cuda.device(target_device): + with on_device(target_device): triton_weight_tensor, weight_scale = swizzle_mxfp4( blocks.transpose(-2, -1), scales.transpose(-2, -1), triton_kernels_hub ) diff --git a/src/transformers/integrations/peft.py b/src/transformers/integrations/peft.py index 87dd6cffc2fa..22261eecad0b 100644 --- a/src/transformers/integrations/peft.py +++ b/src/transformers/integrations/peft.py @@ -15,7 +15,6 @@ import importlib import inspect import re -import warnings from typing import Any, Optional, Union from packaging import version @@ -70,14 +69,9 @@ class PeftAdapterMixin: more details about adapters and injecting them on a transformer-based model, check out the documentation of PEFT library: https://huggingface.co/docs/peft/index - Currently supported PEFT methods are all non-prefix tuning methods. Below is the list of supported PEFT methods - that anyone can load, train and run with this mixin class: - - Low Rank Adapters (LoRA): https://huggingface.co/docs/peft/conceptual_guides/lora - - IA3: https://huggingface.co/docs/peft/conceptual_guides/ia3 - - AdaLora: https://huggingface.co/papers/2303.10512 - - Other PEFT models such as prompt tuning, prompt learning are out of scope as these adapters are not "injectable" - into a torch module. For using these methods, please refer to the usage guide of PEFT library. + Currently supported PEFT methods are all non-prompt learning methods (LoRA, IA³, etc.). Other PEFT models such as + prompt tuning, prompt learning are out of scope as these adapters are not "injectable" into a torch module. For + using these methods, please refer to the usage guide of PEFT library. With this mixin, if the correct PEFT version is installed, it is possible to: @@ -96,7 +90,7 @@ def load_adapter( adapter_name: Optional[str] = None, revision: Optional[str] = None, token: Optional[str] = None, - device_map: Optional[str] = "auto", + device_map: str = "auto", max_memory: Optional[str] = None, offload_folder: Optional[str] = None, offload_index: Optional[int] = None, @@ -110,24 +104,21 @@ def load_adapter( Load adapter weights from file or remote Hub folder. If you are not familiar with adapters and PEFT methods, we invite you to read more about them on PEFT official documentation: https://huggingface.co/docs/peft - Requires peft as a backend to load the adapter weights. + Requires PEFT to be installed as a backend to load the adapter weights. Args: peft_model_id (`str`, *optional*): The identifier of the model to look for on the Hub, or a local path to the saved adapter config file and adapter weights. adapter_name (`str`, *optional*): - The adapter name to use. If not set, will use the default adapter. + The adapter name to use. If not set, will use the name "default". 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. -