diff --git a/.github/workflows/push-important-models.yml b/.github/workflows/push-important-models.yml
index 1887af0f4c5bac..7294777655e183 100644
--- a/.github/workflows/push-important-models.yml
+++ b/.github/workflows/push-important-models.yml
@@ -134,10 +134,3 @@ jobs:
slackChannel: ${{ secrets.SLACK_CIFEEDBACK_CHANNEL }}
slackToken: ${{ secrets.SLACK_CIFEEDBACK_BOT_TOKEN }}
waitForSSH: true
-
- benchmark:
- name: Benchmark workflow
- needs: get_modified_models
- if: ${{ needs.get_modified_models.outputs.matrix != '[]' && needs.get_modified_models.outputs.matrix != '' && fromJson(needs.get_modified_models.outputs.matrix)[0] != null }}
- uses: ./.github/workflows/benchmark.yml
- secrets: inherit
diff --git a/docs/source/ar/quicktour.md b/docs/source/ar/quicktour.md
index 9a99c28287d622..1795c3a5d74fcc 100644
--- a/docs/source/ar/quicktour.md
+++ b/docs/source/ar/quicktour.md
@@ -347,8 +347,8 @@ tensor([[0.0021, 0.0018, 0.0115, 0.2121, 0.7725],
```py
>>> from transformers import AutoModel
->>> tokenizer = AutoTokenizer.from_pretrained(tf_save_directory)
->>> pt_model = AutoModelForSequenceClassification.from_pretrained(tf_save_directory, from_tf=True)
+>>> tokenizer = AutoTokenizer.from_pretrained(pt_save_directory)
+>>> pt_model = AutoModelForSequenceClassification.from_pretrained(pt_save_directory, from_pt=True)
```
@@ -356,8 +356,8 @@ tensor([[0.0021, 0.0018, 0.0115, 0.2121, 0.7725],
```py
>>> from transformers import TFAutoModel
->>> tokenizer = AutoTokenizer.from_pretrained(pt_save_directory)
->>> tf_model = TFAutoModelForSequenceClassification.from_pretrained(pt_save_directory, from_pt=True)
+>>> tokenizer = AutoTokenizer.from_pretrained(tf_save_directory)
+>>> tf_model = TFAutoModelForSequenceClassification.from_pretrained(tf_save_directory, from_tf=True)
```
diff --git a/docs/source/de/quicktour.md b/docs/source/de/quicktour.md
index 01cd7200750c4d..c01609207fec2a 100644
--- a/docs/source/de/quicktour.md
+++ b/docs/source/de/quicktour.md
@@ -109,7 +109,7 @@ label: NEGATIVE, with score: 0.5309
Die [`pipeline`] kann auch über einen ganzen Datensatz iterieren. Starten wir mit der Installation der [🤗 Datasets](https://huggingface.co/docs/datasets/) Bibliothek:
```bash
-pip install datasets
+pip install datasets
```
Erstellen wir eine [`pipeline`] mit der Aufgabe die wir lösen und dem Modell welches wir nutzen möchten.
@@ -191,7 +191,7 @@ Wenn Sie kein Modell für Ihren Anwendungsfall finden können, müssen Sie ein v
-Unter der Haube arbeiten die Klassen [`AutoModelForSequenceClassification`] und [`AutoTokenizer`] zusammen, um die [`pipeline`] zu betreiben. Eine [`AutoClass`](./model_doc/auto) ist eine Abkürzung, die automatisch die Architektur eines trainierten Modells aus dessen Namen oder Pfad abruft. Sie müssen nur die passende `AutoClass` für Ihre Aufgabe und den zugehörigen Tokenizer mit [`AutoTokenizer`] auswählen.
+Unter der Haube arbeiten die Klassen [`AutoModelForSequenceClassification`] und [`AutoTokenizer`] zusammen, um die [`pipeline`] zu betreiben. Eine [`AutoClass`](./model_doc/auto) ist eine Abkürzung, die automatisch die Architektur eines trainierten Modells aus dessen Namen oder Pfad abruft. Sie müssen nur die passende `AutoClass` für Ihre Aufgabe und den zugehörigen Tokenizer mit [`AutoTokenizer`] auswählen.
Kehren wir zu unserem Beispiel zurück und sehen wir uns an, wie Sie die `AutoClass` verwenden können, um die Ergebnisse der [`pipeline`] zu replizieren.
@@ -281,7 +281,7 @@ Jetzt können Sie Ihren vorverarbeiteten Stapel von Eingaben direkt an das Model
```
Das Modell gibt die endgültigen Aktivierungen in dem Attribut "logits" aus. Wenden Sie die Softmax-Funktion auf die "logits" an, um die Wahrscheinlichkeiten zu erhalten:
-
+
```py
>>> from torch import nn
@@ -308,7 +308,7 @@ In der [Aufgabenzusammenfassung](./task_summary) steht, welche [AutoModel]-Klass
Jetzt können Sie Ihren vorverarbeiteten Stapel von Eingaben direkt an das Modell übergeben, indem Sie die Wörterbuchschlüssel direkt an die Tensoren übergeben:
-
+
```py
>>> tf_outputs = tf_model(tf_batch)
```
@@ -383,8 +383,8 @@ Ein besonders cooles 🤗 Transformers-Feature ist die Möglichkeit, ein Modell
```py
>>> from transformers import AutoModel
->>> tokenizer = AutoTokenizer.from_pretrained(tf_save_directory)
->>> pt_model = AutoModelForSequenceClassification.from_pretrained(tf_save_directory, from_tf=True)
+>>> tokenizer = AutoTokenizer.from_pretrained(pt_save_directory)
+>>> pt_model = AutoModelForSequenceClassification.from_pretrained(pt_save_directory, from_pt=True)
```
@@ -392,8 +392,8 @@ Ein besonders cooles 🤗 Transformers-Feature ist die Möglichkeit, ein Modell
```py
>>> from transformers import TFAutoModel
->>> tokenizer = AutoTokenizer.from_pretrained(pt_save_directory)
->>> tf_model = TFAutoModelForSequenceClassification.from_pretrained(pt_save_directory, from_pt=True)
+>>> tokenizer = AutoTokenizer.from_pretrained(tf_save_directory)
+>>> tf_model = TFAutoModelForSequenceClassification.from_pretrained(tf_save_directory, from_tf=True)
```
diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml
index c4707d5f20a027..435b482df599cf 100644
--- a/docs/source/en/_toctree.yml
+++ b/docs/source/en/_toctree.yml
@@ -396,6 +396,8 @@
title: ESM
- local: model_doc/falcon
title: Falcon
+ - local: model_doc/falcon3
+ title: Falcon3
- local: model_doc/falcon_mamba
title: FalconMamba
- local: model_doc/fastspeech2_conformer
@@ -834,6 +836,8 @@
title: CLIPSeg
- local: model_doc/clvp
title: CLVP
+ - local: model_doc/colpali
+ title: ColPali
- local: model_doc/data2vec
title: Data2Vec
- local: model_doc/deplot
diff --git a/docs/source/en/add_new_pipeline.md b/docs/source/en/add_new_pipeline.md
index 1e5b95e9b48cfc..e646f832831504 100644
--- a/docs/source/en/add_new_pipeline.md
+++ b/docs/source/en/add_new_pipeline.md
@@ -184,7 +184,7 @@ class PairClassificationPipeline(Pipeline):
```
The implementation is framework agnostic, and will work for PyTorch and TensorFlow models. If we have saved this in
-a file named `pair_classification.py`, we can then import it and register it like this:
+a file named `pair_classification.py`, we can then import it and register it like this. The [register_pipeline](https://github.com/huggingface/transformers/blob/9feae5fb0164e89d4998e5776897c16f7330d3df/src/transformers/pipelines/base.py#L1387) function registers the pipeline details (task type, pipeline class, supported backends) to a models `config.json` file.
```py
from pair_classification import PairClassificationPipeline
diff --git a/docs/source/en/index.md b/docs/source/en/index.md
index 49c44874e320ef..3bd1c286d43240 100644
--- a/docs/source/en/index.md
+++ b/docs/source/en/index.md
@@ -100,6 +100,7 @@ Flax), PyTorch, and/or TensorFlow.
| [CodeLlama](model_doc/code_llama) | ✅ | ❌ | ✅ |
| [Cohere](model_doc/cohere) | ✅ | ❌ | ❌ |
| [Cohere2](model_doc/cohere2) | ✅ | ❌ | ❌ |
+| [ColPali](model_doc/colpali) | ✅ | ❌ | ❌ |
| [Conditional DETR](model_doc/conditional_detr) | ✅ | ❌ | ❌ |
| [ConvBERT](model_doc/convbert) | ✅ | ✅ | ❌ |
| [ConvNeXT](model_doc/convnext) | ✅ | ✅ | ❌ |
@@ -140,6 +141,7 @@ Flax), PyTorch, and/or TensorFlow.
| [ESM](model_doc/esm) | ✅ | ✅ | ❌ |
| [FairSeq Machine-Translation](model_doc/fsmt) | ✅ | ❌ | ❌ |
| [Falcon](model_doc/falcon) | ✅ | ❌ | ❌ |
+| [Falcon3](model_doc/falcon3) | ✅ | ❌ | ✅ |
| [FalconMamba](model_doc/falcon_mamba) | ✅ | ❌ | ❌ |
| [FastSpeech2Conformer](model_doc/fastspeech2_conformer) | ✅ | ❌ | ❌ |
| [FLAN-T5](model_doc/flan-t5) | ✅ | ✅ | ✅ |
diff --git a/docs/source/en/model_doc/beit.md b/docs/source/en/model_doc/beit.md
index f7605ebcdf90d4..25b0eafb26a039 100644
--- a/docs/source/en/model_doc/beit.md
+++ b/docs/source/en/model_doc/beit.md
@@ -71,6 +71,43 @@ alt="drawing" width="600"/>
BEiT pre-training. Taken from the original paper.
+### 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)
+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
+`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
+
+```
+from transformers import BeitForImageClassification
+model = BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224", attn_implementation="sdpa", torch_dtype=torch.float16)
+...
+```
+
+For the best speedups, we recommend loading the model in half-precision (e.g. `torch.float16` or `torch.bfloat16`).
+
+On a local benchmark (NVIDIA GeForce RTX 2060-8GB, PyTorch 2.5.1, OS Ubuntu 20.04) with `float16` and
+`microsoft/beit-base-patch16-224` model, we saw the following improvements during training and inference:
+
+#### Training
+
+| num_training_steps | batch_size | image_size | is_cuda | Time per batch (eager - s) | Time per batch (sdpa - s) | Speedup (%) | Eager peak mem (MB) | SDPA peak mem (MB) | Mem saving (%) |
+|--------------------|------------|--------------|---------|----------------------------|---------------------------|-------------|----------------------|--------------------|----------------|
+| 50 | 2 | (1048, 640) | True | 0.984 | 0.746 | 31.975 | 6738.915 | 4319.886 | 55.998 |
+
+#### Inference
+
+| Image batch size | Eager (s/iter) | Eager CI, % | Eager memory (MB) | SDPA (s/iter) | SDPA CI, % | SDPA memory (MB) | SDPA speedup | SDPA memory saved (%) |
+|-------------------:|-----------------:|:--------------|--------------------:|----------------:|:-------------|-------------------:|---------------:|----------------------:|
+| 1 | 0.012 | ±0.3% | 3.76657e+08 | 0.011 | ±0.5% | 3.75739e+08 | 1.05 | 0.244 |
+| 4 | 0.013 | ±0.1% | 4.03147e+08 | 0.011 | ±0.2% | 3.90554e+08 | 1.178 | 3.225 |
+| 16 | 0.045 | ±0.1% | 4.96697e+08 | 0.035 | ±0.1% | 4.51232e+08 | 1.304 | 10.076 |
+| 32 | 0.088 | ±0.1% | 6.24417e+08 | 0.066 | ±0.1% | 5.33488e+08 | 1.325 | 17.044 |
+
## Resources
A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with BEiT.
diff --git a/docs/source/en/model_doc/cohere2.md b/docs/source/en/model_doc/cohere2.md
index 4d3a1f0cb0929f..33e67d48fb0e8b 100644
--- a/docs/source/en/model_doc/cohere2.md
+++ b/docs/source/en/model_doc/cohere2.md
@@ -1,5 +1,12 @@
# Cohere
+## Overview
+[C4AI Command R7B](https://cohere.com/blog/command-r7b) is an open weights research release of a 7B billion parameter model developed by Cohere and Cohere For AI. It has advanced capabilities optimized for various use cases, including reasoning, summarization, question answering, and code. The model is trained to perform sophisticated tasks including Retrieval Augmented Generation (RAG) and tool use. The model also has powerful agentic capabilities that can use and combine multiple tools over multiple steps to accomplish more difficult tasks. It obtains top performance on enterprise-relevant code use cases. C4AI Command R7B is a multilingual model trained on 23 languages.
+
+The model features three layers with sliding window attention (window size 4096) 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.
+
+The model has been trained on 23 languages: English, French, Spanish, Italian, German, Portuguese, Japanese, Korean, Arabic, Chinese, Russian, Polish, Turkish, Vietnamese, Dutch, Czech, Indonesian, Ukrainian, Romanian, Greek, Hindi, Hebrew, and Persian.
+
## Usage tips
The model and tokenizer can be loaded via:
@@ -20,7 +27,7 @@ gen_tokens = model.generate(
max_new_tokens=100,
do_sample=True,
temperature=0.3,
- )
+)
gen_text = tokenizer.decode(gen_tokens[0])
print(gen_text)
diff --git a/docs/source/en/model_doc/colpali.md b/docs/source/en/model_doc/colpali.md
new file mode 100644
index 00000000000000..d47f0aa072262c
--- /dev/null
+++ b/docs/source/en/model_doc/colpali.md
@@ -0,0 +1,95 @@
+
+
+# ColPali
+
+## Overview
+
+The ColPali model was proposed in [ColPali: Efficient Document Retrieval with Vision Language Models](https://doi.org/10.48550/arXiv.2407.01449) by **Manuel Faysse***, **Hugues Sibille***, **Tony Wu***, Bilel Omrani, Gautier Viaud, Céline Hudelot, Pierre Colombo (* denotes equal contribution).
+
+With our new model *ColPali*, we propose to leverage VLMs to construct efficient multi-vector embeddings in the visual space for document retrieval. By feeding the ViT output patches from PaliGemma-3B to a linear projection, we create a multi-vector representation of documents. We train the model to maximize the similarity between these document embeddings and the query embeddings, following the ColBERT method.
+
+Using ColPali removes the need for potentially complex and brittle layout recognition and OCR pipelines with a single model that can take into account both the textual and visual content (layout, charts, ...) of a document. ColPali is also highly interpretable: similarity maps can be obtained between patches and query tokens. These maps highlight ColPali’s strong OCR capabilities and chart understanding.
+
+**Paper abstract:**
+
+> Documents are visually rich structures that convey information through text, but also figures, page layouts, tables, or even fonts. Since modern retrieval systems mainly rely on the textual information they extract from document pages to index documents -often through lengthy and brittle processes-, they struggle to exploit key visual cues efficiently. This limits their capabilities in many practical document retrieval applications such as Retrieval Augmented Generation (RAG). To benchmark current systems on visually rich document retrieval, we introduce the Visual Document Retrieval Benchmark *ViDoRe*, composed of various page-level retrieval tasks spanning multiple domains, languages, and practical settings. The inherent complexity and performance shortcomings of modern systems motivate a new concept; doing document retrieval by directly embedding the images of the document pages. We release *ColPali*, a Vision Language Model trained to produce high-quality multi-vector embeddings from images of document pages. Combined with a late interaction matching mechanism, *ColPali* largely outperforms modern document retrieval pipelines while being drastically simpler, faster and end-to-end trainable.
+>
+> We release models, data, code and benchmarks under open licenses at [https://huggingface.co/vidore](https://huggingface.co/vidore).
+
+## Resources
+
+- The official blog post detailing ColPali can be found [here](https://huggingface.co/blog/manu/colpali). 📝
+- The original model implementation code for the ColPali model and for the `colpali-engine` package can be found [here](https://github.com/illuin-tech/colpali). 🌎
+- Cookbooks for learning to use the transformers-native version of ColPali, fine-tuning, and similarity maps generation can be found [here](https://github.com/tonywu71/colpali-cookbooks). 📚
+
+This model was contributed by [@tonywu71](https://huggingface.co/tonywu71) and [@yonigozlan](https://huggingface.co/yonigozlan).
+
+## Usage
+
+This example demonstrates how to use ColPali to embed both queries and images, calculate their similarity scores, and identify the most relevant matches. For a specific query, you can retrieve the top-k most similar images by selecting the ones with the highest similarity scores.
+
+```python
+import torch
+from PIL import Image
+
+from transformers import ColPaliForRetrieval, ColPaliProcessor
+
+model_name = "vidore/colpali-v1.2-hf"
+
+model = ColPaliForRetrieval.from_pretrained(
+ model_name,
+ torch_dtype=torch.bfloat16,
+ device_map="cuda:0", # or "mps" if on Apple Silicon
+).eval()
+
+processor = ColPaliProcessor.from_pretrained(model_name)
+
+# Your inputs (replace dummy images with screenshots of your documents)
+images = [
+ Image.new("RGB", (32, 32), color="white"),
+ Image.new("RGB", (16, 16), color="black"),
+]
+queries = [
+ "What is the organizational structure for our R&D department?",
+ "Can you provide a breakdown of last year’s financial performance?",
+]
+
+# Process the inputs
+batch_images = processor(images=images).to(model.device)
+batch_queries = processor(text=queries).to(model.device)
+
+# Forward pass
+with torch.no_grad():
+ image_embeddings = model(**batch_images)
+ query_embeddings = model(**batch_queries)
+
+# Score the queries against the images
+scores = processor.score_retrieval(query_embeddings, image_embeddings)
+```
+
+## ColPaliConfig
+
+[[autodoc]] ColPaliConfig
+
+## ColPaliProcessor
+
+[[autodoc]] ColPaliProcessor
+
+## ColPaliForRetrieval
+
+[[autodoc]] ColPaliForRetrieval
+ - forward
diff --git a/docs/source/en/model_doc/data2vec.md b/docs/source/en/model_doc/data2vec.md
index 517a51ce46a3a4..cb1dc675caa55e 100644
--- a/docs/source/en/model_doc/data2vec.md
+++ b/docs/source/en/model_doc/data2vec.md
@@ -48,6 +48,46 @@ The original code for vision can be found [here](https://github.com/facebookrese
- For Data2VecText, preprocessing is identical to [`RobertaModel`], including tokenization.
- For Data2VecVision, preprocessing is identical to [`BeitModel`], including feature extraction.
+### 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)
+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
+`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used.
+
+The SDPA implementation is currently available for the Data2VecAudio and Data2VecVision models.
+
+```
+from transformers import Data2VecVisionForImageClassification
+model = Data2VecVisionForImageClassification.from_pretrained("facebook/data2vec-vision-base", attn_implementation="sdpa", torch_dtype=torch.float16)
+...
+```
+
+For the best speedups, we recommend loading the model in half-precision (e.g. `torch.float16` or `torch.bfloat16`).
+
+For the Data2VecVision model, on a local benchmark (NVIDIA GeForce RTX 2060-8GB, PyTorch 2.5.1, OS Ubuntu 20.04)
+with `float16` and `facebook/data2vec-vision-base` model, we saw the following improvements during training and
+inference:
+
+#### Training
+
+| num_training_steps | batch_size | image_size | is_cuda | Time per batch (eager - s) | Time per batch (sdpa - s) | Speedup (%) | Eager peak mem (MB) | SDPA peak mem (MB) | Mem saving (%) |
+|--------------------|------------|--------------|---------|----------------------------|---------------------------|-------------|----------------------|--------------------|----------------|
+| 50 | 2 | (1048, 640) | True | 0.996 | 0.754 | 32.147 | 6722.198 | 4264.653 | 57.626 |
+
+#### Inference
+
+| Image batch size | Eager (s/iter) | Eager CI, % | Eager memory (MB) | SDPA (s/iter) | SDPA CI, % | SDPA memory (MB) | SDPA speedup | SDPA memory saved |
+|-------------------:|-----------------:|:--------------|--------------------:|----------------:|:-------------|-------------------:|---------------:|--------------------:|
+| 1 | 0.011 | ±0.3% | 3.76143e+08 | 0.01 | ±0.3% | 3.74397e+08 | 1.101 | 0.466 |
+| 4 | 0.014 | ±0.1% | 4.02756e+08 | 0.012 | ±0.2% | 3.91373e+08 | 1.219 | 2.909 |
+| 16 | 0.046 | ±0.3% | 4.96482e+08 | 0.035 | ±0.2% | 4.51017e+08 | 1.314 | 10.081 |
+| 32 | 0.088 | ±0.1% | 6.23903e+08 | 0.067 | ±0.1% | 5.32974e+08 | 1.33 | 17.061 |
+
## Resources
A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with Data2Vec.
diff --git a/docs/source/en/model_doc/falcon3.md b/docs/source/en/model_doc/falcon3.md
new file mode 100644
index 00000000000000..813533dd7f4d0a
--- /dev/null
+++ b/docs/source/en/model_doc/falcon3.md
@@ -0,0 +1,29 @@
+
+
+# Falcon3
+
+## Overview
+
+Falcon3 represents a natural evolution from previous releases, emphasizing expanding the models' science, math, and code capabilities. This iteration includes five base models: Falcon3-1B-Base, Falcon3-3B-Base, Falcon3-Mamba-7B-Base, Falcon3-7B-Base, and Falcon3-10B-Base. In developing these models, we incorporated several key innovations aimed at improving the models' performances while reducing training costs:
+
+One pre-training: We conducted a single large-scale pretraining run on the 7B model, using 2048 H100 GPU chips, leveraging 14 trillion tokens featuring web, code, STEM, and curated high-quality and multilingual data.
+Depth up-scaling for improved reasoning: Building on recent studies on the effects of model depth, we upscaled the 7B model to a 10B parameters model by duplicating the redundant layers and continuing pre-training with 2TT of high-quality data. This yielded Falcon3-10B-Base which achieves state-of-the-art zero-shot and few-shot performance for models under 13B parameters.
+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/llava.md b/docs/source/en/model_doc/llava.md
index dec19ca5ef45db..e883572995e924 100644
--- a/docs/source/en/model_doc/llava.md
+++ b/docs/source/en/model_doc/llava.md
@@ -131,7 +131,7 @@ prompt_2 = processor.apply_chat_template(conversation_2, add_generation_prompt=T
prompts = [prompt_1, prompt_2]
# We can simply feed images in the order they have to be used in the text prompt
-inputs = processor(images=[image_stop, image_cats, image_snowman], text=prompts, padding=True, return_tensors="pt").to(model.device, torch.float16)
+inputs = processor(images=[image_stop, image_cats], text=prompts, padding=True, return_tensors="pt").to(model.device, torch.float16)
# Generate
generate_ids = model.generate(**inputs, max_new_tokens=30)
diff --git a/docs/source/en/perf_infer_gpu_multi.md b/docs/source/en/perf_infer_gpu_multi.md
index 9975094411527a..ea9421747c13df 100644
--- a/docs/source/en/perf_infer_gpu_multi.md
+++ b/docs/source/en/perf_infer_gpu_multi.md
@@ -64,5 +64,5 @@ You can benefit from considerable speedups for inference, especially for inputs
For a single forward pass on [Llama](https://huggingface.co/docs/transformers/model_doc/llama#transformers.LlamaModel) with a sequence length of 512 and various batch sizes, the expected speedup is as follows:
-
+
-
+
",
+ image_seq_len=3,
+ image_token="",
+ )
+ The output will be:
+ "Initial str"
+ Args:
+ prompt (`List[Union[str, ImageInput]]`): The input prompt.
+ bos_token (`str`): The beginning of sentence token.
+ image_seq_len (`int`): The length of the image sequence.
+ image_token (`str`): The image token.
+ num_images (`int`): Number of images in the prompt.
+ """
+ return f"{image_token * image_seq_len * num_images}{bos_token}{prompt}\n"
+
+
+def make_batched_images(images) -> List[List[ImageInput]]:
+ """
+ Accepts images in list or nested list format, and makes a list of images for preprocessing.
+
+ Args:
+ images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):
+ The input image.
+
+ Returns:
+ list: A list of images.
+ """
+ if isinstance(images, (list, tuple)) and isinstance(images[0], (list, tuple)) and is_valid_image(images[0][0]):
+ return [img for img_list in images for img in img_list]
+
+ elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
+ return images
+
+ elif is_valid_image(images):
+ return [images]
+
+ raise ValueError(f"Could not make batched video from {images}")
+
+
+class ColPaliProcessor(ProcessorMixin):
+ r"""
+ Constructs a ColPali processor which wraps a PaliGemmaProcessor and special methods to process images and queries, as
+ well as to compute the late-interaction retrieval score.
+
+ [`ColPaliProcessor`] offers all the functionalities of [`PaliGemmaProcessor`]. See the [`~PaliGemmaProcessor.__call__`]
+ for more information.
+
+ Args:
+ image_processor ([`SiglipImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`LlamaTokenizerFast`], *optional*):
+ 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.
+ """
+
+ attributes = ["image_processor", "tokenizer"]
+ valid_kwargs = ["chat_template"]
+ image_processor_class = "SiglipImageProcessor"
+ tokenizer_class = ("GemmaTokenizer", "GemmaTokenizerFast")
+
+ visual_prompt_prefix: ClassVar[str] = "Describe the image."
+ query_prefix: ClassVar[str] = "Question: "
+
+ def __init__(
+ self,
+ image_processor=None,
+ tokenizer=None,
+ chat_template=None,
+ **kwargs,
+ ):
+ if image_processor is None:
+ raise ValueError("You need to specify an `image_processor`.")
+ if tokenizer is None:
+ raise ValueError("You need to specify a `tokenizer`.")
+ if not hasattr(image_processor, "image_seq_length"):
+ raise ValueError("Image processor is missing an `image_seq_length` attribute.")
+
+ self.image_seq_length = image_processor.image_seq_length
+
+ if not hasattr(tokenizer, "image_token"):
+ image_token = AddedToken(IMAGE_TOKEN, normalized=False, special=True)
+ tokens_to_add = {"additional_special_tokens": [image_token]}
+ tokenizer.add_special_tokens(tokens_to_add)
+ self.image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+ else:
+ self.image_token_id = tokenizer.image_token_id
+
+ tokenizer.add_tokens(EXTRA_TOKENS)
+ tokenizer.add_bos_token = False
+ tokenizer.add_eos_token = False
+
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ def __call__(
+ self,
+ images: ImageInput = None,
+ text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+ audio=None,
+ videos=None,
+ **kwargs: Unpack[ColPaliProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model either (1) one or several texts, either (2) one or several image(s). This method is custom
+ wrapper around the PaliGemmaProcessor's [`~PaliGemmaProcessor.__call__`] method adapted for the ColPali model. It cannot process
+ both text and images at the same time.
+
+ When preparing the the text(s), this method forwards the `text` and `kwargs` arguments to LlamaTokenizerFast's
+ [`~LlamaTokenizerFast.__call__`].
+ When preparing the the image(s), this method forwards the `images` and `kwargs` arguments to SiglipImageProcessor's
+ [`~SiglipImageProcessor.__call__`].
+ Please refer to the doctsring of the above two methods for more information.
+
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
+ number of channels, H and W are image height and width.
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+
+ - `'tf'`: Return TensorFlow `tf.constant` objects.
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+ - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ ColPaliProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ suffix = output_kwargs["text_kwargs"].pop("suffix", None)
+
+ return_token_type_ids = True if suffix is not None else False
+
+ if text is None and images is None:
+ raise ValueError("Either text or images must be provided")
+ if text is not None and images is not None:
+ raise ValueError("Only one of text or images can be processed at a time")
+
+ if images is not None:
+ if is_valid_image(images):
+ images = [images]
+ elif isinstance(images, list) and is_valid_image(images[0]):
+ pass
+ elif not (isinstance(images, list) and isinstance(images[0], list) and is_valid_image(images[0][0])):
+ raise ValueError("images must be an image, list of images or list of list of images")
+
+ texts_doc = [self.visual_prompt_prefix] * len(images)
+ images = [image.convert("RGB") for image in images]
+
+ input_strings = [
+ build_string_from_input(
+ prompt=prompt,
+ bos_token=self.tokenizer.bos_token,
+ image_seq_len=self.image_seq_length,
+ image_token=IMAGE_TOKEN,
+ num_images=len(image_list) if isinstance(image_list, list) else 1,
+ )
+ for prompt, image_list in zip(texts_doc, images)
+ ]
+ images = make_batched_images(images)
+ pixel_values = self.image_processor(images, **output_kwargs["images_kwargs"])["pixel_values"]
+
+ # max_length has to account for the image tokens
+ if output_kwargs["text_kwargs"].get("max_length", None) is not None:
+ output_kwargs["text_kwargs"]["max_length"] += self.image_seq_length
+
+ inputs = self.tokenizer(
+ input_strings,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return_data = {**inputs, "pixel_values": pixel_values}
+
+ if return_token_type_ids:
+ labels = inputs["input_ids"].masked_fill(inputs["token_type_ids"] == 0, -100)
+ return_data.update({"labels": labels})
+
+ return BatchFeature(data=return_data)
+
+ elif text is not None:
+ if isinstance(text, str):
+ text = [text]
+ elif not (isinstance(text, list) and isinstance(text[0], str)):
+ raise ValueError("Text must be a string or a list of strings")
+
+ if suffix is None:
+ suffix = self.query_augmentation_token * 10
+ texts_query: List[str] = []
+
+ for query in text:
+ query = self.tokenizer.bos_token + self.query_prefix + query
+ query += suffix # add suffix (pad tokens)
+ query += "\n" # make input ISO to PaliGemma's processor
+ texts_query.append(query)
+
+ output_kwargs["text_kwargs"]["max_length"] = output_kwargs["text_kwargs"].get("max_length", 50)
+
+ batch_query = self.tokenizer(
+ texts_query,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return batch_query
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+ refer to the docstring of this method for more information.
+ """
+ return self.tokenizer.batch_decode(*args, **kwargs)
+
+ def decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+ the docstring of this method for more information.
+ """
+ return self.tokenizer.decode(*args, **kwargs)
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+ return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+ @property
+ def query_augmentation_token(self) -> str:
+ """
+ Return the query augmentation token.
+
+ Query augmentation buffers are used as reasoning buffers during inference.
+ """
+ return self.tokenizer.pad_token
+
+ def process_images(
+ self,
+ images: ImageInput = None,
+ **kwargs: Unpack[ColPaliProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Prepare for the model one or several image(s). This method is a wrapper around the `__call__` method of the ColPaliProcessor's
+ [`ColPaliProcessor.__call__`].
+
+ This method forwards the `images` and `kwargs` arguments to SiglipImageProcessor's [`~SiglipImageProcessor.__call__`].
+
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
+ number of channels, H and W are image height and width.
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+
+ - `'tf'`: Return TensorFlow `tf.constant` objects.
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+ - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ return self.__call__(images=images, **kwargs)
+
+ def process_queries(
+ self,
+ text: Union[TextInput, List[TextInput]],
+ **kwargs: Unpack[ColPaliProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColPaliProcessor's
+ [`ColPaliProcessor.__call__`].
+
+ This method forwards the `text` and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`].
+
+ Args:
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+
+ - `'tf'`: Return TensorFlow `tf.constant` objects.
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+ - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ """
+ return self.__call__(text=text, **kwargs)
+
+ def score_retrieval(
+ self,
+ query_embeddings: Union["torch.Tensor", List["torch.Tensor"]],
+ passage_embeddings: Union["torch.Tensor", List["torch.Tensor"]],
+ batch_size: int = 128,
+ output_dtype: Optional["torch.dtype"] = None,
+ output_device: Union["torch.device", str] = "cpu",
+ ) -> "torch.Tensor":
+ """
+ Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector
+ query embeddings (`qs`) and passage embeddings (`ps`). For ColPali, a passage is the
+ image of a document page.
+
+ Because the embedding tensors are multi-vector and can thus have different shapes, they
+ should be fed as:
+ (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim)
+ (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually
+ obtained by padding the list of tensors.
+
+ Args:
+ query_embeddings (`Union[torch.Tensor, List[torch.Tensor]`): Query embeddings.
+ passage_embeddings (`Union[torch.Tensor, List[torch.Tensor]`): Passage embeddings.
+ batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores.
+ output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor.
+ If `None`, the dtype of the input embeddings is used.
+ output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor.
+
+ Returns:
+ `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score
+ tensor is saved on the "cpu" device.
+ """
+
+ if len(query_embeddings) == 0:
+ raise ValueError("No queries provided")
+ if len(passage_embeddings) == 0:
+ raise ValueError("No passages provided")
+
+ if query_embeddings[0].device != passage_embeddings[0].device:
+ raise ValueError("Queries and passages must be on the same device")
+
+ if query_embeddings[0].dtype != passage_embeddings[0].dtype:
+ raise ValueError("Queries and passages must have the same dtype")
+
+ if output_dtype is None:
+ output_dtype = query_embeddings[0].dtype
+
+ scores: List[torch.Tensor] = []
+
+ for i in range(0, len(query_embeddings), batch_size):
+ batch_scores: List[torch.Tensor] = []
+ batch_queries = torch.nn.utils.rnn.pad_sequence(
+ query_embeddings[i : i + batch_size], batch_first=True, padding_value=0
+ )
+ for j in range(0, len(passage_embeddings), batch_size):
+ batch_passages = torch.nn.utils.rnn.pad_sequence(
+ passage_embeddings[j : j + batch_size], batch_first=True, padding_value=0
+ )
+ batch_scores.append(
+ torch.einsum("bnd,csd->bcns", batch_queries, batch_passages).max(dim=3)[0].sum(dim=2)
+ )
+ scores.append(torch.cat(batch_scores, dim=1).to(output_dtype).to(output_device))
+
+ return torch.cat(scores, dim=0)
+
+
+__all__ = ["ColPaliProcessor"]
diff --git a/src/transformers/models/data2vec/modeling_data2vec_vision.py b/src/transformers/models/data2vec/modeling_data2vec_vision.py
index 4d252ce1f19db7..770162285bf33b 100644
--- a/src/transformers/models/data2vec/modeling_data2vec_vision.py
+++ b/src/transformers/models/data2vec/modeling_data2vec_vision.py
@@ -362,6 +362,69 @@ def forward(
return outputs
+# Copied from transformers.models.beit.modeling_beit.BeitSdpaSelfAttention with Beit->Data2VecVision
+class Data2VecVisionSdpaSelfAttention(Data2VecVisionSelfAttention):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ head_mask: Optional[torch.Tensor] = None,
+ output_attentions: bool = False,
+ relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,
+ interpolate_pos_encoding: bool = False,
+ resolution: Optional[Tuple[int]] = None,
+ ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
+ if output_attentions or head_mask is not None:
+ logger.warning_once(
+ "`Data2VecVisionSdpaSelfAttention` is used but `torch.nn.functional.scaled_dot_product_attention` does not "
+ "support `output_attentions=True` or `head_mask`. Falling back to the manual attention implementation, "
+ "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ return super().forward(
+ hidden_states=hidden_states,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ relative_position_bias=relative_position_bias,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ resolution=resolution,
+ )
+
+ mixed_query_layer = self.query(hidden_states)
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+
+ attn_bias = None
+ if self.relative_position_bias is not None:
+ height, width = resolution
+ window_size = (height // self.config.patch_size, width // self.config.patch_size)
+ attn_bias = self.relative_position_bias(
+ window_size, interpolate_pos_encoding, dim_size=hidden_states.shape[1]
+ )
+
+ # Add shared relative position bias if provided.
+ if relative_position_bias is not None:
+ if attn_bias is None:
+ attn_bias = relative_position_bias
+ else:
+ attn_bias += relative_position_bias
+
+ scaling = 1 / math.sqrt(self.attention_head_size)
+ context_layer = torch.nn.functional.scaled_dot_product_attention(
+ query_layer,
+ key_layer,
+ value_layer,
+ attn_mask=attn_bias,
+ dropout_p=self.config.attention_probs_dropout_prob if self.training else 0.0,
+ is_causal=False,
+ scale=scaling,
+ )
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(*new_context_layer_shape)
+ return context_layer, None
+
+
# Copied from transformers.models.beit.modeling_beit.BeitSelfOutput with Beit->Data2VecVision
class Data2VecVisionSelfOutput(nn.Module):
"""
@@ -381,11 +444,19 @@ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor, gamma
return hidden_states
-# Copied from transformers.models.beit.modeling_beit.BeitAttention with Beit->Data2VecVision
+DATA2VEC_VISION_SELF_ATTENTION_CLASSES = {
+ "eager": Data2VecVisionSelfAttention,
+ "sdpa": Data2VecVisionSdpaSelfAttention,
+}
+
+
+# Copied from tests.models.beit.modeling_beit.BeitAttention with Beit->Data2VecVision, BEIT->DATA2VEC_VISION
class Data2VecVisionAttention(nn.Module):
def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:
super().__init__()
- self.attention = Data2VecVisionSelfAttention(config, window_size=window_size)
+ self.attention = DATA2VEC_VISION_SELF_ATTENTION_CLASSES[config._attn_implementation](
+ config, window_size=window_size
+ )
self.output = Data2VecVisionSelfOutput(config)
self.pruned_heads = set()
@@ -711,6 +782,7 @@ class Data2VecVisionPreTrainedModel(PreTrainedModel):
supports_gradient_checkpointing = True
_no_split_modules = ["Data2VecVisionLayer"]
_keys_to_ignore_on_load_unexpected = [r".*relative_position_index.*"]
+ _supports_sdpa = True
def _init_weights(self, module):
"""Initialize the weights"""
diff --git a/src/transformers/models/nougat/tokenization_nougat_fast.py b/src/transformers/models/nougat/tokenization_nougat_fast.py
index 0a7eec4ad98a4c..5d0a8934c05ee1 100644
--- a/src/transformers/models/nougat/tokenization_nougat_fast.py
+++ b/src/transformers/models/nougat/tokenization_nougat_fast.py
@@ -514,7 +514,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.M)
# most likely hallucinated titles
lines = generation.split("\n")
if lines[-1].startswith("#") and lines[-1].lstrip("#").startswith(" ") and len(lines) > 1:
diff --git a/src/transformers/models/sam/configuration_sam.py b/src/transformers/models/sam/configuration_sam.py
index b0045655d2066b..22a237615d1280 100644
--- a/src/transformers/models/sam/configuration_sam.py
+++ b/src/transformers/models/sam/configuration_sam.py
@@ -46,6 +46,8 @@ class SamPromptEncoderConfig(PretrainedConfig):
The non-linear activation function in the encoder and pooler.
"""
+ base_config_key = "prompt_encoder_config"
+
def __init__(
self,
hidden_size=256,
@@ -102,6 +104,8 @@ class SamMaskDecoderConfig(PretrainedConfig):
"""
+ base_config_key = "mask_decoder_config"
+
def __init__(
self,
hidden_size=256,
@@ -181,6 +185,8 @@ class SamVisionConfig(PretrainedConfig):
hidden_size`.
"""
+ base_config_key = "vision_config"
+
def __init__(
self,
hidden_size=768,
@@ -278,6 +284,11 @@ class SamConfig(PretrainedConfig):
```"""
model_type = "sam"
+ sub_configs = {
+ "prompt_encoder_config": SamPromptEncoderConfig,
+ "mask_decoder_config": SamMaskDecoderConfig,
+ "vision_config": SamVisionConfig,
+ }
def __init__(
self,
diff --git a/src/transformers/models/sam/modeling_sam.py b/src/transformers/models/sam/modeling_sam.py
index c99fb9d7e869f8..b935bc9e421e01 100644
--- a/src/transformers/models/sam/modeling_sam.py
+++ b/src/transformers/models/sam/modeling_sam.py
@@ -246,6 +246,47 @@ def forward(self, query: Tensor, key: Tensor, value: Tensor, attention_similarit
return out
+class SamSdpaAttention(SamAttention):
+ """
+ SAM's attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and
+ values. Using SDPA instead of the default attention.
+ """
+
+ def __init__(self, config, downsample_rate=None):
+ super().__init__(config, downsample_rate)
+
+ def forward(self, query: Tensor, key: Tensor, value: Tensor, attention_similarity: Tensor = None) -> Tensor:
+ # Input projections
+ query = self.q_proj(query)
+ key = self.k_proj(key)
+ value = self.v_proj(value)
+
+ point_batch_size = query.shape[1]
+ # Separate into heads
+ query = self._separate_heads(query, self.num_attention_heads)
+ key = self._separate_heads(key, self.num_attention_heads)
+ value = self._separate_heads(value, self.num_attention_heads)
+
+ # Scaled dot product attention
+ attn_mask = None
+ if attention_similarity is not None:
+ attn_mask = attention_similarity.unsqueeze(1).expand(-1, self.num_attention_heads, -1, -1)
+
+ out = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask)
+
+ # Get output
+ out = self._recombine_heads(out, point_batch_size)
+ out = self.out_proj(out)
+
+ return out
+
+
+SAM_ATTENTION_CLASSES = {
+ "eager": SamAttention,
+ "sdpa": SamSdpaAttention,
+}
+
+
class SamTwoWayAttentionBlock(nn.Module):
def __init__(self, config, attention_downsample_rate: int = 2, skip_first_layer_pe: bool = False):
"""
@@ -266,18 +307,21 @@ def __init__(self, config, attention_downsample_rate: int = 2, skip_first_layer_
self.hidden_size = config.hidden_size
self.layer_norm_eps = config.layer_norm_eps
- self.self_attn = SamAttention(config, downsample_rate=1)
+ self.self_attn = SAM_ATTENTION_CLASSES[config._attn_implementation](config, downsample_rate=1)
self.layer_norm1 = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
- self.cross_attn_token_to_image = SamAttention(config, downsample_rate=attention_downsample_rate)
+ self.cross_attn_token_to_image = SAM_ATTENTION_CLASSES[config._attn_implementation](
+ config, downsample_rate=attention_downsample_rate
+ )
self.layer_norm2 = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
self.mlp = SamMLPBlock(config)
self.layer_norm3 = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
self.layer_norm4 = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
- self.cross_attn_image_to_token = SamAttention(config, downsample_rate=attention_downsample_rate)
-
+ self.cross_attn_image_to_token = SAM_ATTENTION_CLASSES[config._attn_implementation](
+ config, downsample_rate=attention_downsample_rate
+ )
self.skip_first_layer_pe = skip_first_layer_pe
def forward(
@@ -344,7 +388,7 @@ def __init__(self, config: SamMaskDecoderConfig):
for i in range(self.num_hidden_layers):
self.layers.append(SamTwoWayAttentionBlock(config, skip_first_layer_pe=(i == 0)))
- self.final_attn_token_to_image = SamAttention(config)
+ self.final_attn_token_to_image = SAM_ATTENTION_CLASSES[config._attn_implementation](config)
self.layer_norm_final_attn = nn.LayerNorm(config.hidden_size)
def forward(
@@ -431,7 +475,7 @@ def forward(self, hidden_states):
class SamMaskDecoder(nn.Module):
def __init__(self, config: SamMaskDecoderConfig):
super().__init__()
-
+ self.config = config
self.hidden_size = config.hidden_size
self.num_multimask_outputs = config.num_multimask_outputs
@@ -856,11 +900,118 @@ def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch
return outputs
+class SamVisionSdpaAttention(SamVisionAttention):
+ """
+ Multi-head Attention block with relative position embeddings.
+ Using SDPA instead of the default attention.
+ """
+
+ def __init__(self, config, window_size):
+ super().__init__(config, window_size)
+
+ def add_decomposed_rel_pos(
+ self,
+ query: torch.Tensor,
+ rel_pos_h: torch.Tensor,
+ rel_pos_w: torch.Tensor,
+ q_size: Tuple[int, int],
+ k_size: Tuple[int, int],
+ ) -> torch.Tensor:
+ """
+ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+ https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
+ This method is reimplemented to follow the implementation in:
+ https://github.com/pytorch-labs/segment-anything-fast/blob/main/segment_anything_fast/modeling/image_encoder.py # noqa B950
+ This implementation is more memory efficient when using SDPA in the forward method.
+ Args:
+ q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
+ rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
+ rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
+ q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
+ k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
+
+ Returns:
+ attn (Tensor): attention map with added relative positional embeddings.
+ """
+ query_height, query_width = q_size
+ key_height, key_width = k_size
+ relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h)
+ relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w)
+
+ batch_size, _, dim = query.shape
+ reshaped_query = query.reshape(batch_size, query_height, query_width, dim)
+ rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height)
+ rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width)
+ rel_h = rel_h.unsqueeze(-1)
+ rel_w = rel_w.unsqueeze(-2)
+ rel_h = rel_h.reshape(batch_size, query_height * query_width, key_height, 1)
+ rel_w = rel_w.reshape(batch_size, query_height * query_width, 1, key_width)
+
+ return rel_h, rel_w
+
+ def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch.Tensor:
+ batch_size, height, width, _ = hidden_states.shape
+ # qkv with shape (3, B, nHead, H * W, C)
+ qkv = (
+ self.qkv(hidden_states)
+ .reshape(batch_size, height * width, 3, self.num_attention_heads, -1)
+ .permute(2, 0, 3, 1, 4)
+ )
+ # q, k, v with shape (B * nHead, H * W, C)
+ query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0)
+
+ rel_h, rel_w = None, None
+ if self.use_rel_pos:
+ rel_h, rel_w = self.add_decomposed_rel_pos(
+ query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width)
+ )
+
+ query = query.view(batch_size, self.num_attention_heads, height * width, -1)
+ key = key.view(batch_size, self.num_attention_heads, height * width, -1)
+ value = value.view(batch_size, self.num_attention_heads, height * width, -1)
+
+ if self.use_rel_pos:
+ rel_h = rel_h.view(batch_size, self.num_attention_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3))
+ rel_w = rel_w.view(batch_size, self.num_attention_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3))
+ attn_bias = (rel_h + rel_w).view(
+ batch_size, self.num_attention_heads, rel_h.size(2), rel_h.size(3) * rel_w.size(4)
+ )
+ attn_output = torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attn_bias)
+ else:
+ attn_output = torch.nn.functional.scaled_dot_product_attention(query, key, value)
+
+ attn_output = (
+ attn_output.view(batch_size, self.num_attention_heads, height, width, -1)
+ .permute(0, 2, 3, 1, 4)
+ .reshape(batch_size, height, width, -1)
+ )
+
+ attn_output = self.proj(attn_output)
+
+ if output_attentions:
+ # For output_attentions, calculate the attention weights
+ attn_weights = (query @ key.transpose(-2, -1)) * self.scale
+ if attn_bias is not None:
+ attn_weights = attn_weights + attn_bias
+ attn_weights = F.softmax(attn_weights, dim=-1)
+ outputs = (attn_output, attn_weights)
+ else:
+ outputs = (attn_output, None)
+
+ return outputs
+
+
+SAM_VISION_ATTENTION_CLASSES = {
+ "eager": SamVisionAttention,
+ "sdpa": SamVisionSdpaAttention,
+}
+
+
class SamVisionLayer(nn.Module):
def __init__(self, config, window_size):
super().__init__()
self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.attn = SamVisionAttention(config, window_size)
+ self.attn = SAM_VISION_ATTENTION_CLASSES[config._attn_implementation](config, window_size)
self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.mlp = SamMLPBlock(config)
self.window_size = window_size
@@ -1071,6 +1222,8 @@ class SamPreTrainedModel(PreTrainedModel):
base_model_prefix = "sam"
main_input_name = "pixel_values"
_no_split_modules = ["SamVisionAttention"]
+ supports_gradient_checkpointing = True
+ _supports_sdpa = True
def _init_weights(self, module):
std = self.config.initializer_range
diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py
index c6057088b7d506..823c51a290713d 100644
--- a/src/transformers/utils/dummy_pt_objects.py
+++ b/src/transformers/utils/dummy_pt_objects.py
@@ -813,6 +813,9 @@ def __init__(self, *args, **kwargs):
MODEL_FOR_QUESTION_ANSWERING_MAPPING = None
+MODEL_FOR_RETRIEVAL_MAPPING = None
+
+
MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = None
@@ -2258,6 +2261,20 @@ def __init__(self, *args, **kwargs):
requires_backends(self, ["torch"])
+class ColPaliForRetrieval(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+
+class ColPaliPreTrainedModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+
class ConditionalDetrForObjectDetection(metaclass=DummyObject):
_backends = ["torch"]
diff --git a/tests/models/beit/test_modeling_beit.py b/tests/models/beit/test_modeling_beit.py
index ac64f0fd3b0b11..e54273f7839965 100644
--- a/tests/models/beit/test_modeling_beit.py
+++ b/tests/models/beit/test_modeling_beit.py
@@ -14,18 +14,35 @@
# limitations under the License.
"""Testing suite for the PyTorch BEiT model."""
+import inspect
+import tempfile
import unittest
+import numpy as np
from datasets import load_dataset
from packaging import version
+from parameterized import parameterized
from transformers import BeitConfig
-from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
-from transformers.utils import cached_property, is_torch_available, is_vision_available
+from transformers.testing_utils import (
+ require_torch,
+ require_torch_multi_gpu,
+ require_torch_sdpa,
+ require_vision,
+ slow,
+ torch_device,
+)
+from transformers.utils import (
+ cached_property,
+ is_torch_available,
+ is_torch_bf16_available_on_device,
+ is_torch_fp16_available_on_device,
+ is_vision_available,
+)
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, _config_zero_init, floats_tensor, ids_tensor, sdpa_kernel
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -74,6 +91,8 @@ def __init__(
scope=None,
out_indices=[1, 2, 3, 4],
out_features=["stage1", "stage2", "stage3", "stage4"],
+ attn_implementation="eager",
+ mask_ratio=0.5,
):
self.parent = parent
self.vocab_size = vocab_size
@@ -100,6 +119,8 @@ def __init__(
# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
num_patches = (image_size // patch_size) ** 2
self.seq_length = num_patches + 1
+ self.num_masks = int(mask_ratio * self.seq_length)
+ self.attn_implementation = attn_implementation
def prepare_config_and_inputs(self):
pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
@@ -131,6 +152,7 @@ def get_config(self):
initializer_range=self.initializer_range,
out_indices=self.out_indices,
out_features=self.out_features,
+ attn_implementation=self.attn_implementation,
)
def create_and_check_model(self, config, pixel_values, labels, pixel_labels):
@@ -387,6 +409,193 @@ def test_model_from_pretrained(self):
model = BeitModel.from_pretrained(model_name)
self.assertIsNotNone(model)
+ @parameterized.expand([("float16",), ("bfloat16",), ("float32",)])
+ @require_torch_sdpa
+ def test_eager_matches_sdpa_inference(self, torch_dtype: str):
+ # The common test modifies the num_hidden_layers to be 1. However, for Beit we want to
+ # avoid that because the num_hidden_layers is generally assumed to be 4. Also, the code
+ # related to attention masks in the original common tests is not required as the Beit
+ # model does not handle attention masks. Furthermore, some extra code like modifying
+ # the norm layers eps values for specialized configs and checking for the 'noise'
+ # has been omitted to simply the test.
+ if not self.has_attentions:
+ self.skipTest(reason="Model architecture does not support attentions")
+
+ if not self.all_model_classes[0]._supports_sdpa:
+ self.skipTest(f"{self.all_model_classes[0].__name__} does not support SDPA")
+
+ if torch_dtype == "float16" and not is_torch_fp16_available_on_device(torch_device):
+ self.skipTest(f"float16 not supported on {torch_device} (on the specific device currently used)")
+
+ if torch_dtype == "bfloat16" and not is_torch_bf16_available_on_device(torch_device):
+ self.skipTest(
+ f"bfloat16 not supported on {torch_device} (on the specific device currently used, e.g. Nvidia T4 GPU)"
+ )
+
+ # Not sure whether it's fine to put torch.XXX in a decorator if torch is not available so hacking it here instead.
+ if torch_dtype == "float16":
+ torch_dtype = torch.float16
+ elif torch_dtype == "bfloat16":
+ torch_dtype = torch.bfloat16
+ elif torch_dtype == "float32":
+ torch_dtype = torch.float32
+
+ atols = {
+ ("cpu", False, torch.float32): 1e-6,
+ ("cpu", False, torch.float16): 5e-3,
+ ("cpu", False, torch.bfloat16): 1e-2,
+ ("cpu", True, torch.float32): 1e-6,
+ ("cpu", True, torch.float16): 5e-3,
+ ("cpu", True, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float32): 1e-6,
+ ("cuda", False, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float16): 5e-3,
+ ("cuda", True, torch.float32): 1e-6,
+ ("cuda", True, torch.bfloat16): 1e-2,
+ ("cuda", True, torch.float16): 5e-3,
+ }
+ rtols = {
+ ("cpu", False, torch.float32): 1e-4,
+ ("cpu", False, torch.float16): 5e-3,
+ ("cpu", False, torch.bfloat16): 1e-2,
+ ("cpu", True, torch.float32): 1e-4,
+ ("cpu", True, torch.float16): 5e-3,
+ ("cpu", True, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float32): 1e-4,
+ ("cuda", False, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float16): 5e-3,
+ ("cuda", True, torch.float32): 1e-4,
+ ("cuda", True, torch.bfloat16): 3e-2,
+ ("cuda", True, torch.float16): 5e-3,
+ }
+
+ def get_mean_reldiff(failcase, x, ref, atol, rtol):
+ return f"{failcase}: mean relative difference: {((x - ref).abs() / (ref.abs() + 1e-12)).mean():.3e}, torch atol = {atol}, torch rtol = {rtol}"
+
+ for model_class in self.all_model_classes:
+ config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+ config.rms_norm_eps = 1.0
+ config.layer_norm_eps = 1.0
+ config.norm_eps = 1.0
+ config.norm_epsilon = 1.0
+ config.layer_norm_epsilon = 1.0
+
+ model = model_class(config)
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ model.save_pretrained(tmpdirname)
+ model_sdpa = model_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype, use_mask_token=True)
+ model_sdpa = model_sdpa.eval().to(torch_device, dtype=torch_dtype)
+
+ model_eager = model_class.from_pretrained(
+ tmpdirname,
+ torch_dtype=torch_dtype,
+ attn_implementation="eager",
+ use_mask_token=True,
+ )
+ model_eager = model_eager.eval().to(torch_device, dtype=torch_dtype)
+
+ # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.)
+ for x in model_eager.modules():
+ if isinstance(x, (nn.LayerNorm, nn.GroupNorm)):
+ x.eps = 1.0
+ for x in model_sdpa.modules():
+ if isinstance(x, (nn.LayerNorm, nn.GroupNorm)):
+ x.eps = 1.0
+
+ # We use these for loops instead of parameterized.expand just for the interest of avoiding loading/saving 16 times the model,
+ # but it would be nicer to have an efficient way to use parameterized.expand
+ fail_cases = []
+ for padding_side in ["left", "right"]:
+ for use_mask in [False, True]:
+ for output_attentions in [True, False]:
+ can_output_attn = "output_attentions" in inspect.signature(model_sdpa.forward).parameters
+ if not (self.has_attentions and can_output_attn) and output_attentions:
+ continue
+ # TODO: if we can also check with `batch_size=1` without being flaky?
+ for batch_size in [7]:
+ dummy_input = inputs_dict[model.main_input_name]
+
+ if dummy_input.dtype in [torch.float32, torch.bfloat16, torch.float16]:
+ dummy_input = dummy_input.to(torch_dtype)
+
+ dummy_input = dummy_input[:batch_size]
+ for enable_kernels in [False, True]:
+ failcase = f"padding_side={padding_side}, use_mask={use_mask}, enable_kernels={enable_kernels}"
+ processed_inputs = {
+ model.main_input_name: dummy_input,
+ "output_hidden_states": True,
+ }
+
+ if (
+ self.has_attentions
+ and "output_attentions" in inspect.signature(model_sdpa.forward).parameters
+ ):
+ processed_inputs["output_attentions"] = output_attentions
+
+ if "bool_masked_pos" in inspect.signature(model_eager.forward).parameters:
+ dummy_mask = torch.ones((self.model_tester.num_masks,))
+ mask_length = self.model_tester.seq_length - 1 - dummy_mask.size(0)
+ dummy_mask = torch.cat([dummy_mask, torch.zeros(mask_length)])
+ dummy_bool_masked_pos = dummy_mask.expand(batch_size, -1).bool()
+ processed_inputs["bool_masked_pos"] = dummy_bool_masked_pos.to(torch_device)
+
+ with torch.no_grad():
+ with sdpa_kernel(
+ enable_flash=enable_kernels,
+ enable_math=True,
+ enable_mem_efficient=enable_kernels,
+ ):
+ prepared_inputs = self._prepare_for_class(processed_inputs, model_class)
+ outputs_eager = model_eager(**prepared_inputs)
+ outputs_sdpa = model_sdpa(**prepared_inputs)
+
+ logits_eager = outputs_eager.hidden_states[-1]
+ logits_sdpa = outputs_sdpa.hidden_states[-1]
+ if torch_device in ["cpu", "cuda"]:
+ atol = atols[torch_device, enable_kernels, torch_dtype]
+ rtol = rtols[torch_device, enable_kernels, torch_dtype]
+ elif torch_device == "xpu":
+ # As of PyTorch 2.5 XPU backend supports only torch.nn.attention.SDPBackend.MATH
+ # which is implemented on PyTorch level using aten operators and is
+ # device agnostic with respect to implementation of each aten operator.
+ atol = atols["cuda", False, torch_dtype]
+ rtol = rtols["cuda", False, torch_dtype]
+ else:
+ atol = 1e-7
+ rtol = 1e-4
+
+ # Masked tokens output slightly deviates - we don't mind that.
+ if use_mask:
+ _logits_sdpa = torch.zeros_like(input=logits_sdpa)
+ _logits_eager = torch.zeros_like(input=logits_eager)
+
+ _logits_sdpa[:-1] = logits_sdpa[:-1]
+ _logits_eager[:-1] = logits_eager[:-1]
+
+ if padding_side == "left":
+ _logits_sdpa[-1:, 2:] = logits_sdpa[-1:, 2:]
+ _logits_eager[-1:, 2:] = logits_eager[-1:, 2:]
+
+ elif padding_side == "right":
+ _logits_sdpa[-1:, 2:] = logits_sdpa[-1:, :-2]
+ _logits_eager[-1:, 2:] = logits_eager[-1:, :-2]
+
+ logits_sdpa = _logits_sdpa
+ logits_eager = _logits_eager
+
+ results = [
+ torch.allclose(_logits_sdpa, _logits_eager, atol=atol, rtol=rtol)
+ for (_logits_sdpa, _logits_eager) in zip(logits_sdpa, logits_eager)
+ ]
+ # If 80% batch elements have matched results, it's fine
+ if np.mean(results) < 0.8:
+ fail_cases.append(
+ get_mean_reldiff(failcase, logits_sdpa, logits_eager, atol, rtol)
+ )
+
+ self.assertTrue(len(fail_cases) == 0, "\n".join(fail_cases))
+
# We will verify our results on an image of cute cats
def prepare_img():
diff --git a/tests/models/colpali/__init__.py b/tests/models/colpali/__init__.py
new file mode 100644
index 00000000000000..e69de29bb2d1d6
diff --git a/tests/models/colpali/test_modeling_colpali.py b/tests/models/colpali/test_modeling_colpali.py
new file mode 100644
index 00000000000000..646726ac700ee5
--- /dev/null
+++ b/tests/models/colpali/test_modeling_colpali.py
@@ -0,0 +1,368 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Testing suite for the PyTorch ColPali model."""
+
+import gc
+import unittest
+from typing import ClassVar
+
+import torch
+from datasets import load_dataset
+from parameterized import parameterized
+
+from tests.test_configuration_common import ConfigTester
+from tests.test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
+from transformers import (
+ is_torch_available,
+ is_vision_available,
+)
+from transformers.models.colpali.configuration_colpali import ColPaliConfig
+from transformers.models.colpali.modeling_colpali import ColPaliForRetrieval, ColPaliForRetrievalOutput
+from transformers.models.colpali.processing_colpali import ColPaliProcessor
+from transformers.testing_utils import (
+ require_torch,
+ require_torch_sdpa,
+ require_vision,
+ slow,
+ torch_device,
+)
+
+
+if is_torch_available():
+ import torch
+
+if is_vision_available():
+ pass
+
+
+class ColPaliForRetrievalModelTester:
+ def __init__(
+ self,
+ parent,
+ ignore_index=-100,
+ image_token_index=0,
+ projector_hidden_act="gelu",
+ seq_length=25,
+ vision_feature_select_strategy="default",
+ vision_feature_layer=-1,
+ projection_dim=32,
+ text_config={
+ "model_type": "gemma",
+ "seq_length": 128,
+ "is_training": True,
+ "use_token_type_ids": False,
+ "use_labels": True,
+ "vocab_size": 99,
+ "hidden_size": 32,
+ "num_hidden_layers": 2,
+ "num_attention_heads": 4,
+ "num_key_value_heads": 1,
+ "head_dim": 8,
+ "intermediate_size": 37,
+ "hidden_activation": "gelu_pytorch_tanh",
+ "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,
+ "pad_token_id": 1,
+ },
+ is_training=False,
+ vision_config={
+ "use_labels": True,
+ "image_size": 20,
+ "patch_size": 5,
+ "num_image_tokens": 4,
+ "num_channels": 3,
+ "is_training": True,
+ "hidden_size": 32,
+ "projection_dim": 32,
+ "num_key_value_heads": 1,
+ "num_hidden_layers": 2,
+ "num_attention_heads": 4,
+ "intermediate_size": 37,
+ "dropout": 0.1,
+ "attention_dropout": 0.1,
+ "initializer_range": 0.02,
+ },
+ use_cache=False,
+ embedding_dim=128,
+ ):
+ self.parent = parent
+ self.ignore_index = ignore_index
+ # `image_token_index` is set to 0 to pass "resize_embeddings" test, do not modify
+ self.image_token_index = image_token_index
+ self.projector_hidden_act = projector_hidden_act
+ self.vision_feature_select_strategy = vision_feature_select_strategy
+ self.vision_feature_layer = vision_feature_layer
+ self.text_config = text_config
+ self.vision_config = vision_config
+ self.seq_length = seq_length
+ self.projection_dim = projection_dim
+ self.pad_token_id = text_config["pad_token_id"]
+
+ self.num_hidden_layers = text_config["num_hidden_layers"]
+ self.vocab_size = text_config["vocab_size"]
+ self.hidden_size = text_config["hidden_size"]
+ self.num_attention_heads = text_config["num_attention_heads"]
+ self.is_training = is_training
+
+ self.batch_size = 3
+ self.num_channels = vision_config["num_channels"]
+ self.image_size = vision_config["image_size"]
+ self.encoder_seq_length = seq_length
+ self.use_cache = use_cache
+
+ self.embedding_dim = embedding_dim
+ self.vlm_config = {
+ "model_type": "paligemma",
+ "text_config": self.text_config,
+ "vision_config": self.vision_config,
+ "ignore_index": self.ignore_index,
+ "image_token_index": self.image_token_index,
+ "projector_hidden_act": self.projector_hidden_act,
+ "projection_dim": self.projection_dim,
+ "vision_feature_select_strategy": self.vision_feature_select_strategy,
+ "vision_feature_layer": self.vision_feature_layer,
+ }
+
+ def get_config(self):
+ return ColPaliConfig(
+ vlm_config=self.vlm_config,
+ embedding_dim=self.embedding_dim,
+ )
+
+ def prepare_config_and_inputs(self):
+ pixel_values = floats_tensor(
+ [
+ self.batch_size,
+ self.vision_config["num_channels"],
+ self.vision_config["image_size"],
+ self.vision_config["image_size"],
+ ]
+ )
+ config = self.get_config()
+
+ return config, pixel_values
+
+ def prepare_config_and_inputs_for_common(self):
+ config_and_inputs = self.prepare_config_and_inputs()
+ config, pixel_values = config_and_inputs
+ input_ids = ids_tensor([self.batch_size, self.seq_length], config.vlm_config.text_config.vocab_size - 1) + 1
+ attention_mask = input_ids.ne(1).to(torch_device)
+ # set the 16 first tokens to be image, and ensure that no other tokens are image tokens
+ # do not change this unless you modified image size or patch size
+ input_ids[input_ids == config.vlm_config.image_token_index] = self.pad_token_id
+ input_ids[:, :16] = config.vlm_config.image_token_index
+ inputs_dict = {
+ "pixel_values": pixel_values,
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "labels": input_ids,
+ "token_type_ids": torch.zeros_like(input_ids),
+ }
+ return config, inputs_dict
+
+
+@require_torch
+class ColPaliForRetrievalModelTest(ModelTesterMixin, unittest.TestCase):
+ """
+ Model tester for `ColPaliForRetrieval`.
+ """
+
+ all_model_classes = (ColPaliForRetrieval,) if is_torch_available() else ()
+ fx_compatible = False
+ test_torchscript = False
+ test_pruning = False
+ test_resize_embeddings = True
+ test_head_masking = False
+
+ def setUp(self):
+ self.model_tester = ColPaliForRetrievalModelTester(self)
+ self.config_tester = ConfigTester(self, config_class=ColPaliConfig, has_text_modality=False)
+
+ # overwrite inputs_embeds tests because we need to delete "pixel values" for LVLMs
+
+ def test_inputs_embeds(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)
+ model.to(torch_device)
+ model.eval()
+
+ inputs = self._prepare_for_class(inputs_dict, model_class)
+
+ input_ids = inputs["input_ids"]
+ del inputs["input_ids"]
+ del inputs["pixel_values"]
+
+ wte = model.get_input_embeddings()
+ inputs["inputs_embeds"] = wte(input_ids)
+
+ with torch.no_grad():
+ model(**inputs)
+
+ # overwrite inputs_embeds tests because we need to delete "pixel values" for LVLMs
+ # while some other models require pixel_values to be present
+ def test_inputs_embeds_matches_input_ids(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)
+ model.to(torch_device)
+ model.eval()
+
+ inputs = self._prepare_for_class(inputs_dict, model_class)
+ input_ids = inputs["input_ids"]
+ del inputs["input_ids"]
+ del inputs["pixel_values"]
+
+ inputs_embeds = model.get_input_embeddings()(input_ids)
+
+ with torch.no_grad():
+ out_ids = model(input_ids=input_ids, **inputs)[0]
+ out_embeds = model(inputs_embeds=inputs_embeds, **inputs)[0]
+ self.assertTrue(torch.allclose(out_embeds, out_ids))
+
+ @slow
+ @require_vision
+ def test_colpali_forward_inputs(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)
+ model.to(torch_device)
+ model.eval()
+
+ inputs = self._prepare_for_class(inputs_dict, model_class)
+
+ with torch.no_grad():
+ outputs = model(**inputs, return_dict=True)
+
+ self.assertIsInstance(outputs, ColPaliForRetrievalOutput)
+
+ @unittest.skip(
+ reason="This architecure 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 architecure 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 architecure 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
+
+ @require_torch_sdpa
+ @slow
+ @parameterized.expand([("float16",), ("bfloat16",), ("float32",)])
+ def test_eager_matches_sdpa_inference(self, torch_dtype: str):
+ self.skipTest(
+ "Due to custom causal mask, there is a slightly too big difference between eager and sdpa in bfloat16."
+ )
+
+ @unittest.skip(
+ reason="From PaliGemma: Some undefined behavior encountered with test versions of this model. Skip for now."
+ )
+ def test_model_parallelism(self):
+ pass
+
+ @unittest.skip(
+ reason="PaliGemmma'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):
+ pass
+
+ @unittest.skip(reason="Pass because ColPali requires `attention_mask is not None`")
+ def test_sdpa_can_dispatch_on_flash(self):
+ pass
+
+ @unittest.skip(reason="Pass because ColPali requires `attention_mask is not None`")
+ def test_sdpa_can_compile_dynamic(self):
+ pass
+
+
+@require_torch
+class ColPaliModelIntegrationTest(unittest.TestCase):
+ model_name: ClassVar[str] = "vidore/colpali-v1.2-hf"
+
+ def setUp(self):
+ self.processor = ColPaliProcessor.from_pretrained(self.model_name)
+
+ def tearDown(self):
+ gc.collect()
+ torch.cuda.empty_cache()
+
+ @slow
+ def test_model_integration_test(self):
+ """
+ Test if the model is able to retrieve the correct pages for a small and easy dataset.
+ """
+ model = ColPaliForRetrieval.from_pretrained(
+ self.model_name,
+ torch_dtype=torch.bfloat16,
+ device_map=torch_device,
+ ).eval()
+
+ # Load the test dataset
+ ds = load_dataset("hf-internal-testing/document-visual-retrieval-test", split="test")
+
+ # Preprocess the examples
+ batch_images = self.processor(images=ds["image"]).to(torch_device)
+ batch_queries = self.processor(text=ds["query"]).to(torch_device)
+
+ # Run inference
+ with torch.inference_mode():
+ image_embeddings = model(**batch_images).embeddings
+ query_embeddings = model(**batch_queries).embeddings
+
+ # Compute retrieval scores
+ scores = self.processor.score_retrieval(
+ query_embeddings=query_embeddings,
+ passage_embeddings=image_embeddings,
+ ) # (len(qs), len(ps))
+
+ assert scores.ndim == 2, f"Expected 2D tensor, got {scores.ndim}"
+ assert scores.shape == (len(ds), len(ds)), f"Expected shape {(len(ds), len(ds))}, got {scores.shape}"
+
+ # Check if the maximum scores per row are in the diagonal of the matrix score
+ self.assertTrue((scores.argmax(axis=1) == torch.arange(len(ds), device=scores.device)).all())
+
+ # Further validation: fine-grained check, with a hardcoded score from the original implementation
+ expected_scores = torch.tensor(
+ [
+ [15.5625, 6.5938, 14.4375],
+ [12.2500, 16.2500, 11.0000],
+ [15.0625, 11.7500, 21.0000],
+ ],
+ dtype=scores.dtype,
+ )
+
+ assert torch.allclose(scores, expected_scores, atol=1), f"Expected scores {expected_scores}, got {scores}"
diff --git a/tests/models/colpali/test_processing_colpali.py b/tests/models/colpali/test_processing_colpali.py
new file mode 100644
index 00000000000000..42592460fa28ed
--- /dev/null
+++ b/tests/models/colpali/test_processing_colpali.py
@@ -0,0 +1,247 @@
+import shutil
+import tempfile
+import unittest
+
+import torch
+
+from transformers import GemmaTokenizer
+from transformers.models.colpali.processing_colpali import ColPaliProcessor
+from transformers.testing_utils import get_tests_dir, require_torch, require_vision
+from transformers.utils import is_vision_available
+from transformers.utils.dummy_vision_objects import SiglipImageProcessor
+
+from ...test_processing_common import ProcessorTesterMixin
+
+
+if is_vision_available():
+ from transformers import (
+ ColPaliProcessor,
+ PaliGemmaProcessor,
+ SiglipImageProcessor,
+ )
+
+SAMPLE_VOCAB = get_tests_dir("fixtures/test_sentencepiece.model")
+
+
+@require_vision
+class ColPaliProcessorTest(ProcessorTesterMixin, unittest.TestCase):
+ processor_class = ColPaliProcessor
+
+ def setUp(self):
+ self.tmpdirname = tempfile.mkdtemp()
+ image_processor = SiglipImageProcessor.from_pretrained("google/siglip-so400m-patch14-384")
+ image_processor.image_seq_length = 0
+ tokenizer = GemmaTokenizer(SAMPLE_VOCAB, keep_accents=True)
+ processor = PaliGemmaProcessor(image_processor=image_processor, tokenizer=tokenizer)
+ processor.save_pretrained(self.tmpdirname)
+
+ def tearDown(self):
+ shutil.rmtree(self.tmpdirname)
+
+ @require_torch
+ @require_vision
+ def test_process_images(self):
+ # Processor configuration
+ image_input = self.prepare_image_inputs()
+ image_processor = self.get_component("image_processor")
+ tokenizer = self.get_component("tokenizer", max_length=112, padding="max_length")
+ image_processor.image_seq_length = 14
+
+ # Get the processor
+ processor = self.processor_class(
+ tokenizer=tokenizer,
+ image_processor=image_processor,
+ )
+
+ # Process the image
+ batch_feature = processor.process_images(images=image_input, return_tensors="pt")
+
+ # Assertions
+ self.assertIn("pixel_values", batch_feature)
+ self.assertEqual(batch_feature["pixel_values"].shape, torch.Size([1, 3, 384, 384]))
+
+ @require_torch
+ @require_vision
+ def test_process_queries(self):
+ # Inputs
+ queries = [
+ "Is attention really all you need?",
+ "Are Benjamin, Antoine, Merve, and Jo best friends?",
+ ]
+
+ # Processor configuration
+ image_processor = self.get_component("image_processor")
+ tokenizer = self.get_component("tokenizer", max_length=112, padding="max_length")
+ image_processor.image_seq_length = 14
+
+ # Get the processor
+ processor = self.processor_class(
+ tokenizer=tokenizer,
+ image_processor=image_processor,
+ )
+
+ # Process the image
+ batch_feature = processor.process_queries(text=queries, return_tensors="pt")
+
+ # Assertions
+ self.assertIn("input_ids", batch_feature)
+ self.assertIsInstance(batch_feature["input_ids"], torch.Tensor)
+ self.assertEqual(batch_feature["input_ids"].shape[0], len(queries))
+
+ # The following tests are overwritten as ColPaliProcessor can only take one of images or text as input at a time
+
+ def test_tokenizer_defaults_preserved_by_kwargs(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor_components["tokenizer"] = self.get_component("tokenizer", max_length=117, padding="max_length")
+
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+ input_str = self.prepare_text_inputs()
+ inputs = processor(text=input_str, return_tensors="pt")
+ self.assertEqual(inputs[self.text_input_name].shape[-1], 117)
+
+ def test_image_processor_defaults_preserved_by_image_kwargs(self):
+ """
+ We use do_rescale=True, rescale_factor=-1 to ensure that image_processor kwargs are preserved in the processor.
+ We then check that the mean of the pixel_values is less than or equal to 0 after processing.
+ Since the original pixel_values are in [0, 255], this is a good indicator that the rescale_factor is indeed applied.
+ """
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor_components["image_processor"] = self.get_component(
+ "image_processor", do_rescale=True, rescale_factor=-1
+ )
+ processor_components["tokenizer"] = self.get_component("tokenizer", max_length=117, padding="max_length")
+
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ image_input = self.prepare_image_inputs()
+
+ inputs = processor(images=image_input, return_tensors="pt")
+ self.assertLessEqual(inputs[self.images_input_name][0][0].mean(), 0)
+
+ def test_kwargs_overrides_default_tokenizer_kwargs(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor_components["tokenizer"] = self.get_component("tokenizer", padding="longest")
+
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+ input_str = self.prepare_text_inputs()
+ inputs = processor(text=input_str, return_tensors="pt", max_length=112, padding="max_length")
+ self.assertEqual(inputs[self.text_input_name].shape[-1], 112)
+
+ def test_kwargs_overrides_default_image_processor_kwargs(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor_components["image_processor"] = self.get_component(
+ "image_processor", do_rescale=True, rescale_factor=1
+ )
+ processor_components["tokenizer"] = self.get_component("tokenizer", max_length=117, padding="max_length")
+
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ image_input = self.prepare_image_inputs()
+
+ inputs = processor(images=image_input, do_rescale=True, rescale_factor=-1, return_tensors="pt")
+ self.assertLessEqual(inputs[self.images_input_name][0][0].mean(), 0)
+
+ def test_unstructured_kwargs(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ input_str = self.prepare_text_inputs()
+ inputs = processor(
+ text=input_str,
+ return_tensors="pt",
+ do_rescale=True,
+ rescale_factor=-1,
+ padding="max_length",
+ max_length=76,
+ )
+
+ self.assertEqual(inputs[self.text_input_name].shape[-1], 76)
+
+ def test_unstructured_kwargs_batched(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ image_input = self.prepare_image_inputs(batch_size=2)
+ inputs = processor(
+ images=image_input,
+ return_tensors="pt",
+ do_rescale=True,
+ rescale_factor=-1,
+ padding="longest",
+ max_length=76,
+ )
+
+ self.assertLessEqual(inputs[self.images_input_name][0][0].mean(), 0)
+
+ def test_doubly_passed_kwargs(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ image_input = self.prepare_image_inputs()
+ with self.assertRaises(ValueError):
+ _ = processor(
+ images=image_input,
+ images_kwargs={"do_rescale": True, "rescale_factor": -1},
+ do_rescale=True,
+ return_tensors="pt",
+ )
+
+ def test_structured_kwargs_nested(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ input_str = self.prepare_text_inputs()
+
+ # Define the kwargs for each modality
+ all_kwargs = {
+ "common_kwargs": {"return_tensors": "pt"},
+ "images_kwargs": {"do_rescale": True, "rescale_factor": -1},
+ "text_kwargs": {"padding": "max_length", "max_length": 76},
+ }
+
+ inputs = processor(text=input_str, **all_kwargs)
+ self.skip_processor_without_typed_kwargs(processor)
+
+ self.assertEqual(inputs[self.text_input_name].shape[-1], 76)
+
+ def test_structured_kwargs_nested_from_dict(self):
+ if "image_processor" not in self.processor_class.attributes:
+ self.skipTest(f"image_processor attribute not present in {self.processor_class}")
+ processor_components = self.prepare_components()
+ processor = self.processor_class(**processor_components)
+ self.skip_processor_without_typed_kwargs(processor)
+ image_input = self.prepare_image_inputs()
+
+ # Define the kwargs for each modality
+ all_kwargs = {
+ "common_kwargs": {"return_tensors": "pt"},
+ "images_kwargs": {"do_rescale": True, "rescale_factor": -1},
+ "text_kwargs": {"padding": "max_length", "max_length": 76},
+ }
+
+ inputs = processor(images=image_input, **all_kwargs)
+ self.assertEqual(inputs[self.text_input_name].shape[-1], 76)
diff --git a/tests/models/data2vec/test_modeling_data2vec_vision.py b/tests/models/data2vec/test_modeling_data2vec_vision.py
index c729d88d614fbc..02276d905fa402 100644
--- a/tests/models/data2vec/test_modeling_data2vec_vision.py
+++ b/tests/models/data2vec/test_modeling_data2vec_vision.py
@@ -14,14 +14,32 @@
# limitations under the License.
"""Testing suite for the PyTorch Data2VecVision model."""
+import inspect
+import tempfile
import unittest
+import numpy as np
+from parameterized import parameterized
+
from transformers import Data2VecVisionConfig
-from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
-from transformers.utils import cached_property, is_torch_available, is_vision_available
+from transformers.testing_utils import (
+ require_torch,
+ require_torch_multi_gpu,
+ require_torch_sdpa,
+ require_vision,
+ slow,
+ torch_device,
+)
+from transformers.utils import (
+ cached_property,
+ is_torch_available,
+ is_torch_bf16_available_on_device,
+ is_torch_fp16_available_on_device,
+ 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, _config_zero_init, floats_tensor, ids_tensor, sdpa_kernel
from ...test_pipeline_mixin import PipelineTesterMixin
@@ -66,6 +84,8 @@ def __init__(
num_labels=3,
scope=None,
out_indices=[0, 1, 2, 3],
+ attn_implementation="eager",
+ mask_ratio=0.5,
):
self.parent = parent
self.vocab_size = 100
@@ -91,6 +111,8 @@ def __init__(
# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
num_patches = (image_size // patch_size) ** 2
self.seq_length = num_patches + 1
+ self.num_masks = int(mask_ratio * self.seq_length)
+ self.attn_implementation = attn_implementation
def prepare_config_and_inputs(self):
pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
@@ -121,6 +143,7 @@ def get_config(self):
is_decoder=False,
initializer_range=self.initializer_range,
out_indices=self.out_indices,
+ attn_implementation=self.attn_implementation,
)
def create_and_check_model(self, config, pixel_values, labels, pixel_labels):
@@ -300,6 +323,194 @@ def test_model_from_pretrained(self):
model = Data2VecVisionModel.from_pretrained(model_name)
self.assertIsNotNone(model)
+ @parameterized.expand([("float16",), ("bfloat16",), ("float32",)])
+ @require_torch_sdpa
+ # Copied from tests.models.beit.test_modeling_beit.BeitModelTest.test_eager_matches_sdpa_inference with Beit->Data2VecVision
+ def test_eager_matches_sdpa_inference(self, torch_dtype: str):
+ # The common test modifies the num_hidden_layers to be 1. However, for Data2VecVision we want to
+ # avoid that because the num_hidden_layers is generally assumed to be 4. Also, the code
+ # related to attention masks in the original common tests is not required as the Data2VecVision
+ # model does not handle attention masks. Furthermore, some extra code like modifying
+ # the norm layers eps values for specialized configs and checking for the 'noise'
+ # has been omitted to simply the test.
+ if not self.has_attentions:
+ self.skipTest(reason="Model architecture does not support attentions")
+
+ if not self.all_model_classes[0]._supports_sdpa:
+ self.skipTest(f"{self.all_model_classes[0].__name__} does not support SDPA")
+
+ if torch_dtype == "float16" and not is_torch_fp16_available_on_device(torch_device):
+ self.skipTest(f"float16 not supported on {torch_device} (on the specific device currently used)")
+
+ if torch_dtype == "bfloat16" and not is_torch_bf16_available_on_device(torch_device):
+ self.skipTest(
+ f"bfloat16 not supported on {torch_device} (on the specific device currently used, e.g. Nvidia T4 GPU)"
+ )
+
+ # Not sure whether it's fine to put torch.XXX in a decorator if torch is not available so hacking it here instead.
+ if torch_dtype == "float16":
+ torch_dtype = torch.float16
+ elif torch_dtype == "bfloat16":
+ torch_dtype = torch.bfloat16
+ elif torch_dtype == "float32":
+ torch_dtype = torch.float32
+
+ atols = {
+ ("cpu", False, torch.float32): 1e-6,
+ ("cpu", False, torch.float16): 5e-3,
+ ("cpu", False, torch.bfloat16): 1e-2,
+ ("cpu", True, torch.float32): 1e-6,
+ ("cpu", True, torch.float16): 5e-3,
+ ("cpu", True, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float32): 1e-6,
+ ("cuda", False, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float16): 5e-3,
+ ("cuda", True, torch.float32): 1e-6,
+ ("cuda", True, torch.bfloat16): 1e-2,
+ ("cuda", True, torch.float16): 5e-3,
+ }
+ rtols = {
+ ("cpu", False, torch.float32): 1e-4,
+ ("cpu", False, torch.float16): 5e-3,
+ ("cpu", False, torch.bfloat16): 1e-2,
+ ("cpu", True, torch.float32): 1e-4,
+ ("cpu", True, torch.float16): 5e-3,
+ ("cpu", True, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float32): 1e-4,
+ ("cuda", False, torch.bfloat16): 1e-2,
+ ("cuda", False, torch.float16): 5e-3,
+ ("cuda", True, torch.float32): 1e-4,
+ ("cuda", True, torch.bfloat16): 3e-2,
+ ("cuda", True, torch.float16): 5e-3,
+ }
+
+ def get_mean_reldiff(failcase, x, ref, atol, rtol):
+ return f"{failcase}: mean relative difference: {((x - ref).abs() / (ref.abs() + 1e-12)).mean():.3e}, torch atol = {atol}, torch rtol = {rtol}"
+
+ for model_class in self.all_model_classes:
+ config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+ config.rms_norm_eps = 1.0
+ config.layer_norm_eps = 1.0
+ config.norm_eps = 1.0
+ config.norm_epsilon = 1.0
+ config.layer_norm_epsilon = 1.0
+
+ model = model_class(config)
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ model.save_pretrained(tmpdirname)
+ model_sdpa = model_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype, use_mask_token=True)
+ model_sdpa = model_sdpa.eval().to(torch_device, dtype=torch_dtype)
+
+ model_eager = model_class.from_pretrained(
+ tmpdirname,
+ torch_dtype=torch_dtype,
+ attn_implementation="eager",
+ use_mask_token=True,
+ )
+ model_eager = model_eager.eval().to(torch_device, dtype=torch_dtype)
+
+ # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.)
+ for x in model_eager.modules():
+ if isinstance(x, (nn.LayerNorm, nn.GroupNorm)):
+ x.eps = 1.0
+ for x in model_sdpa.modules():
+ if isinstance(x, (nn.LayerNorm, nn.GroupNorm)):
+ x.eps = 1.0
+
+ # We use these for loops instead of parameterized.expand just for the interest of avoiding loading/saving 16 times the model,
+ # but it would be nicer to have an efficient way to use parameterized.expand
+ fail_cases = []
+ for padding_side in ["left", "right"]:
+ for use_mask in [False, True]:
+ for output_attentions in [True, False]:
+ can_output_attn = "output_attentions" in inspect.signature(model_sdpa.forward).parameters
+ if not (self.has_attentions and can_output_attn) and output_attentions:
+ continue
+ # TODO: if we can also check with `batch_size=1` without being flaky?
+ for batch_size in [7]:
+ dummy_input = inputs_dict[model.main_input_name]
+
+ if dummy_input.dtype in [torch.float32, torch.bfloat16, torch.float16]:
+ dummy_input = dummy_input.to(torch_dtype)
+
+ dummy_input = dummy_input[:batch_size]
+ for enable_kernels in [False, True]:
+ failcase = f"padding_side={padding_side}, use_mask={use_mask}, enable_kernels={enable_kernels}"
+ processed_inputs = {
+ model.main_input_name: dummy_input,
+ "output_hidden_states": True,
+ }
+
+ if (
+ self.has_attentions
+ and "output_attentions" in inspect.signature(model_sdpa.forward).parameters
+ ):
+ processed_inputs["output_attentions"] = output_attentions
+
+ if "bool_masked_pos" in inspect.signature(model_eager.forward).parameters:
+ dummy_mask = torch.ones((self.model_tester.num_masks,))
+ mask_length = self.model_tester.seq_length - 1 - dummy_mask.size(0)
+ dummy_mask = torch.cat([dummy_mask, torch.zeros(mask_length)])
+ dummy_bool_masked_pos = dummy_mask.expand(batch_size, -1).bool()
+ processed_inputs["bool_masked_pos"] = dummy_bool_masked_pos.to(torch_device)
+
+ with torch.no_grad():
+ with sdpa_kernel(
+ enable_flash=enable_kernels,
+ enable_math=True,
+ enable_mem_efficient=enable_kernels,
+ ):
+ prepared_inputs = self._prepare_for_class(processed_inputs, model_class)
+ outputs_eager = model_eager(**prepared_inputs)
+ outputs_sdpa = model_sdpa(**prepared_inputs)
+
+ logits_eager = outputs_eager.hidden_states[-1]
+ logits_sdpa = outputs_sdpa.hidden_states[-1]
+ if torch_device in ["cpu", "cuda"]:
+ atol = atols[torch_device, enable_kernels, torch_dtype]
+ rtol = rtols[torch_device, enable_kernels, torch_dtype]
+ elif torch_device == "xpu":
+ # As of PyTorch 2.5 XPU backend supports only torch.nn.attention.SDPBackend.MATH
+ # which is implemented on PyTorch level using aten operators and is
+ # device agnostic with respect to implementation of each aten operator.
+ atol = atols["cuda", False, torch_dtype]
+ rtol = rtols["cuda", False, torch_dtype]
+ else:
+ atol = 1e-7
+ rtol = 1e-4
+
+ # Masked tokens output slightly deviates - we don't mind that.
+ if use_mask:
+ _logits_sdpa = torch.zeros_like(input=logits_sdpa)
+ _logits_eager = torch.zeros_like(input=logits_eager)
+
+ _logits_sdpa[:-1] = logits_sdpa[:-1]
+ _logits_eager[:-1] = logits_eager[:-1]
+
+ if padding_side == "left":
+ _logits_sdpa[-1:, 2:] = logits_sdpa[-1:, 2:]
+ _logits_eager[-1:, 2:] = logits_eager[-1:, 2:]
+
+ elif padding_side == "right":
+ _logits_sdpa[-1:, 2:] = logits_sdpa[-1:, :-2]
+ _logits_eager[-1:, 2:] = logits_eager[-1:, :-2]
+
+ logits_sdpa = _logits_sdpa
+ logits_eager = _logits_eager
+
+ results = [
+ torch.allclose(_logits_sdpa, _logits_eager, atol=atol, rtol=rtol)
+ for (_logits_sdpa, _logits_eager) in zip(logits_sdpa, logits_eager)
+ ]
+ # If 80% batch elements have matched results, it's fine
+ if np.mean(results) < 0.8:
+ fail_cases.append(
+ get_mean_reldiff(failcase, logits_sdpa, logits_eager, atol, rtol)
+ )
+
+ self.assertTrue(len(fail_cases) == 0, "\n".join(fail_cases))
+
# We will verify our results on an image of cute cats
def prepare_img():
diff --git a/tests/models/sam/test_modeling_sam.py b/tests/models/sam/test_modeling_sam.py
index 7faace0096c8de..351016716a0cf1 100644
--- a/tests/models/sam/test_modeling_sam.py
+++ b/tests/models/sam/test_modeling_sam.py
@@ -14,12 +14,13 @@
# limitations under the License.
"""Testing suite for the PyTorch SAM model."""
+import tempfile
import unittest
import requests
from transformers import SamConfig, SamMaskDecoderConfig, SamPromptEncoderConfig, SamVisionConfig, pipeline
-from transformers.testing_utils import cleanup, require_torch, slow, torch_device
+from transformers.testing_utils import cleanup, require_torch, require_torch_sdpa, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available
from ...test_configuration_common import ConfigTester
@@ -295,6 +296,7 @@ class SamModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
test_resize_embeddings = False
test_head_masking = False
test_torchscript = False
+ _is_composite = True
# TODO: Fix me @Arthur: `run_batch_test` in `tests/test_pipeline_mixin.py` not working
def is_pipeline_test_to_skip(
@@ -311,22 +313,13 @@ def is_pipeline_test_to_skip(
def setUp(self):
self.model_tester = SamModelTester(self)
- self.vision_config_tester = ConfigTester(self, config_class=SamVisionConfig, has_text_modality=False)
- self.prompt_encoder_config_tester = ConfigTester(
- self,
- config_class=SamPromptEncoderConfig,
- has_text_modality=False,
- num_attention_heads=12,
- num_hidden_layers=2,
- )
- self.mask_decoder_config_tester = ConfigTester(
- self, config_class=SamMaskDecoderConfig, has_text_modality=False
+ common_properties = ["initializer_range"]
+ self.config_tester = ConfigTester(
+ self, config_class=SamConfig, has_text_modality=False, common_properties=common_properties
)
def test_config(self):
- self.vision_config_tester.run_common_tests()
- self.prompt_encoder_config_tester.run_common_tests()
- self.mask_decoder_config_tester.run_common_tests()
+ self.config_tester.run_common_tests()
@unittest.skip(reason="SAM's vision encoder does not use inputs_embeds")
def test_inputs_embeds(self):
@@ -450,6 +443,68 @@ def test_model_from_pretrained(self):
model = SamModel.from_pretrained(model_name)
self.assertIsNotNone(model)
+ @require_torch_sdpa
+ def test_sdpa_can_compile_dynamic(self):
+ self.skipTest(reason="SAM model can't be compiled dynamic yet")
+
+ @require_torch_sdpa
+ def test_sdpa_can_dispatch_composite_models(self):
+ """
+ Tests if composite models dispatch correctly on SDPA/eager when requested so when loading the model.
+ This tests only by looking at layer names, as usually SDPA layers are calles "SDPAAttention".
+ In contrast to the above test, this one checks if the "config._attn_implamentation" is a dict after the model
+ is loaded, because we manually replicate requested attn implementation on each sub-config when loading.
+ See https://github.com/huggingface/transformers/pull/32238 for more info
+
+ The test tries to cover most general cases of composite models, VLMs with vision and text configs. Any model
+ that has a different set of sub-configs has to overwrite this test.
+ """
+ if not self.has_attentions:
+ self.skipTest(reason="Model architecture does not support attentions")
+
+ if not self._is_composite:
+ self.skipTest(f"{self.all_model_classes[0].__name__} does not support SDPA")
+
+ for model_class in self.all_model_classes:
+ 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_sdpa = model_class.from_pretrained(tmpdirname, attn_implementation="sdpa")
+ model_sdpa = model_sdpa.eval().to(torch_device)
+
+ model_eager = model_class.from_pretrained(tmpdirname, attn_implementation="eager")
+ model_eager = model_eager.eval().to(torch_device)
+
+ # Root model determines SDPA support
+ attn_impl = "sdpa" if model._supports_sdpa else "eager"
+
+ # Check config propagation to submodels that support it
+ self.assertTrue(model_sdpa.config._attn_implementation == "sdpa")
+ self.assertTrue(model_sdpa.vision_encoder.config._attn_implementation == attn_impl)
+ self.assertTrue(model_sdpa.mask_decoder.config._attn_implementation == attn_impl)
+
+ self.assertTrue(model_eager.config._attn_implementation == "eager")
+ self.assertTrue(model_eager.vision_encoder.config._attn_implementation == "eager")
+ self.assertTrue(model_eager.mask_decoder.config._attn_implementation == "eager")
+
+ # Verify SDPA/eager layer presence
+ has_sdpa = False
+ for name, submodule in model_sdpa.named_modules():
+ class_name = submodule.__class__.__name__
+ if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name:
+ has_sdpa = True
+ break
+
+ if not has_sdpa and attn_impl == "sdpa":
+ raise ValueError("The SDPA model should have SDPA attention layers")
+
+ for name, submodule in model_eager.named_modules():
+ class_name = submodule.__class__.__name__
+ if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name:
+ raise ValueError("The eager model should not have SDPA attention layers")
+
def prepare_image():
img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index 13eacc4a596562..3aaf18c945451f 100755
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -4202,16 +4202,20 @@ def get_mean_reldiff(failcase, x, ref, atol, rtol):
outputs_eager = model_eager(**prepared_inputs)
outputs_sdpa = model_sdpa(**prepared_inputs)
- logits_eager = (
- outputs_eager.hidden_states[-1]
- if not is_encoder_decoder
- else outputs_eager.decoder_hidden_states[-1]
- )
- logits_sdpa = (
- outputs_sdpa.hidden_states[-1]
- if not is_encoder_decoder
- else outputs_sdpa.decoder_hidden_states[-1]
- )
+ if hasattr(outputs_eager, "vision_hidden_states"):
+ logits_eager = outputs_eager.vision_hidden_states[-1]
+ logits_sdpa = outputs_sdpa.vision_hidden_states[-1]
+ else:
+ logits_eager = (
+ outputs_eager.hidden_states[-1]
+ if not is_encoder_decoder
+ else outputs_eager.decoder_hidden_states[-1]
+ )
+ logits_sdpa = (
+ outputs_sdpa.hidden_states[-1]
+ if not is_encoder_decoder
+ else outputs_sdpa.decoder_hidden_states[-1]
+ )
if torch_device in ["cpu", "cuda"]:
atol = atols[torch_device, enable_kernels, torch_dtype]
@@ -4287,6 +4291,8 @@ def test_sdpa_can_dispatch_on_flash(self):
)
if config.model_type in ["idefics", "idefics2", "idefics3"]:
self.skipTest(reason="Idefics currently (transformers==4.39.1) requires an image_attention_mask input")
+ if config.model_type in ["sam"]:
+ self.skipTest(reason="SAM requires an attention_mask input for relative positional embeddings")
model = model_class(config)
with tempfile.TemporaryDirectory() as tmpdirname:
diff --git a/utils/check_table.py b/utils/check_table.py
index 5876818449558e..957bfd5af6af6f 100644
--- a/utils/check_table.py
+++ b/utils/check_table.py
@@ -87,7 +87,7 @@ def _find_text_in_file(filename: str, start_prompt: str, end_prompt: str) -> str
_re_tf_models = re.compile(r"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
_re_flax_models = re.compile(r"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch after the two previous regexes.
-_re_pt_models = re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
+_re_pt_models = re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration|ForRetrieval)")
# This is to make sure the transformers module imported is the one in the repo.
@@ -157,6 +157,7 @@ def _center_text(text: str, width: int) -> str:
"LayoutXLM": "LayoutLMv2",
"Llama2": "LLaMA",
"Llama3": "LLaMA",
+ "Falcon3": "LLaMA",
"MADLAD-400": "T5",
"MatCha": "Pix2Struct",
"mBART-50": "mBART",
diff --git a/utils/create_dependency_mapping.py b/utils/create_dependency_mapping.py
index f25a8fb5ca6ff1..0df782d1c21740 100644
--- a/utils/create_dependency_mapping.py
+++ b/utils/create_dependency_mapping.py
@@ -1,40 +1,48 @@
import ast
-from collections import defaultdict, deque
+from collections import defaultdict
# Function to perform topological sorting
def topological_sort(dependencies):
- # Create a graph and in-degree count for each node
+ new_dependencies = {}
graph = defaultdict(list)
- in_degree = defaultdict(int)
-
- # Build the graph
for node, deps in dependencies.items():
for dep in deps:
- graph[dep].append(node) # node depends on dep
- in_degree[node] += 1 # increase in-degree of node
+ if "example" not in node and "auto" not in dep:
+ graph[dep.split(".")[-2]].append(node.split("/")[-2])
+ new_dependencies[node.split("/")[-2]] = node
- # Add all nodes with zero in-degree to the queue
- zero_in_degree_queue = deque([node for node in dependencies if in_degree[node] == 0])
+ # Create a graph and in-degree count for each node
+ def filter_one_by_one(filtered_list, reverse):
+ if len(reverse) == 0:
+ return filtered_list
- sorted_list = []
- # Perform topological sorting
- while zero_in_degree_queue:
- current = zero_in_degree_queue.popleft()
- sorted_list.append(current)
+ graph = defaultdict(list)
+ # Build the graph
+ for node, deps in reverse.items():
+ for dep in deps:
+ graph[dep].append(node)
- # For each node that current points to, reduce its in-degree
- for neighbor in graph[current]:
- in_degree[neighbor] -= 1
- if in_degree[neighbor] == 0:
- zero_in_degree_queue.append(neighbor)
+ base_modules = set(reverse.keys()) - set(graph.keys())
+ if base_modules == reverse.keys():
+ # we are at the end
+ return filtered_list + list(graph.keys())
+ to_add = []
+ for k in graph.keys():
+ if len(graph[k]) == 1 and graph[k][0] in base_modules:
+ if graph[k][0] in reverse:
+ del reverse[graph[k][0]]
+ if k not in filtered_list:
+ to_add += [k]
+ for k in base_modules:
+ if k not in filtered_list:
+ to_add += [k]
+ filtered_list += list(to_add)
+ return filter_one_by_one(filtered_list, reverse)
- # Handle nodes that have no dependencies and were not initially part of the loop
- for node in dependencies:
- if node not in sorted_list:
- sorted_list.append(node)
+ final_order = filter_one_by_one([], graph)
- return sorted_list
+ return [new_dependencies.get(k) for k in final_order if k in new_dependencies]
# Function to extract class and import info from a file
@@ -46,7 +54,7 @@ def extract_classes_and_imports(file_path):
for node in ast.walk(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
module = node.module if isinstance(node, ast.ImportFrom) else None
- if module and "transformers" in module:
+ if module and (".modeling_" in module):
imports.add(module)
return imports
@@ -56,7 +64,7 @@ def map_dependencies(py_files):
dependencies = defaultdict(set)
# First pass: Extract all classes and map to files
for file_path in py_files:
- dependencies[file_path].add(None)
+ # dependencies[file_path].add(None)
class_to_file = extract_classes_and_imports(file_path)
for module in class_to_file:
dependencies[file_path].add(module)
@@ -66,4 +74,4 @@ def map_dependencies(py_files):
def find_priority_list(py_files):
dependencies = map_dependencies(py_files)
ordered_classes = topological_sort(dependencies)
- return ordered_classes[::-1]
+ return ordered_classes
diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py
index e8d117cd2af08f..28fcc4fc7b9e1a 100644
--- a/utils/modular_model_converter.py
+++ b/utils/modular_model_converter.py
@@ -1678,7 +1678,7 @@ def save_modeling_file(modular_file, converted_file):
parser = argparse.ArgumentParser()
parser.add_argument(
"--files_to_parse",
- default=["src/transformers/models/aria/modular_aria.py"],
+ default=["all"],
nargs="+",
help="A list of `modular_xxxx` files that should be converted to single model file",
)
diff --git a/utils/release.py b/utils/release.py
index b0349a80b49802..d5b74602e68c09 100644
--- a/utils/release.py
+++ b/utils/release.py
@@ -45,12 +45,14 @@
import argparse
import os
import re
+from pathlib import Path
import packaging.version
# All paths are defined with the intent that this script should be run from the root of the repo.
PATH_TO_EXAMPLES = "examples/"
+PATH_TO_MODELS = "src/transformers/models"
# This maps a type of file to the pattern to look for when searching where the version is defined, as well as the
# template to follow when replacing it with the new version.
REPLACE_PATTERNS = {
@@ -117,6 +119,17 @@ def global_version_update(version: str, patch: bool = False):
update_version_in_examples(version)
+def remove_conversion_scripts():
+ """
+ Delete the scripts that convert models from older, unsupported formats. We don't want to include these
+ in release wheels because they often have to open insecure file types (pickle, Torch .bin models). This results in
+ vulnerability scanners flagging us and can cause compliance issues for users with strict security policies.
+ """
+ model_dir = Path(PATH_TO_MODELS)
+ for conversion_script in list(model_dir.glob("**/convert*.py")):
+ conversion_script.unlink()
+
+
def get_version() -> packaging.version.Version:
"""
Reads the current version in the main __init__.
@@ -131,7 +144,7 @@ def pre_release_work(patch: bool = False):
"""
Do all the necessary pre-release steps:
- figure out the next minor release version and ask confirmation
- - update the version eveywhere
+ - update the version everywhere
- clean-up the model list in the main README
Args:
@@ -155,13 +168,15 @@ def pre_release_work(patch: bool = False):
print(f"Updating version to {version}.")
global_version_update(version, patch=patch)
+ print("Deleting conversion scripts.")
+ remove_conversion_scripts()
def post_release_work():
"""
- Do all the necesarry post-release steps:
+ Do all the necessary post-release steps:
- figure out the next dev version and ask confirmation
- - update the version eveywhere
+ - update the version everywhere
- clean-up the model list in the main README
"""
# First let's get the current version
diff --git a/utils/update_metadata.py b/utils/update_metadata.py
index b6ee1e7c8c13c2..8e4a7e3fe5340e 100755
--- a/utils/update_metadata.py
+++ b/utils/update_metadata.py
@@ -56,7 +56,7 @@
_re_tf_models = re.compile(r"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
_re_flax_models = re.compile(r"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
-_re_pt_models = re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
+_re_pt_models = re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration|ForRetrieval)")
# Fill this with tuples (pipeline_tag, model_mapping, auto_model)