update to transformers v5

README.md

@@ -27,7 +27,7 @@ The final evaluation on the test subset yielded a micro accuracy of 0.9722 and a
 Please see also [Google Colab Notebook](https://colab.research.google.com/drive/1jFMZENoTjjum3qlBxV0Q5dTxmpCvqlpf?usp=sharing).
 1. Install dependencies (Not required on Google Colab)
 ```sh
-python -m pip install einops torch torchvision transformers
+python -m pip install einops torch torchvision "transformers>=5.1.0"
 
 # Optional (This is also required on Google Colab if you want to use FlashAttention-2)
 pip install flash-attn --no-build-isolation
@@ -40,7 +40,7 @@ from io import BytesIO
 from PIL import Image
 import requests
 import torch
-from transformers import AutoModel, AutoProcessor
+from transformers import AutoImageProcessor, AutoModel
 
 repo_name = "SakanaAI/Metom"
 device = "cuda"
@@ -49,13 +49,21 @@ torch_dtype = torch.float32  # This can also set `torch.float16` or `torch.bfloa
 def get_image(image_url: str) -> Image.Image:
     return Image.open(BytesIO(requests.get(image_url).content)).convert("RGB")
 
-processor = AutoProcessor.from_pretrained(repo_name, trust_remote_code=True)
+processor = AutoImageProcessor.from_pretrained(repo_name)
 model = AutoModel.from_pretrained(
     repo_name,
-    torch_dtype=torch_dtype,
-    _attn_implementation="eager",  # This can also set `"sdpa"` or `"flash_attention_2"`
+    dtype=torch_dtype,
+    attn_implementation="sdpa",  # This can also set `"eager"`, `"flash_attention_2"` or other methods supported in transformers v5 (https://huggingface.co/docs/transformers/main/en/attention_interface)
     trust_remote_code=True
 ).to(device=device)
+# We still support transformers v4
+# model = AutoModel.from_pretrained(
+#     repo_name,
+#     torch_dtype=torch_dtype,
+#     _attn_implementation="sdpa",  # This can also set `"eager"` or `"flash_attention_2"`
+#     trust_remote_code=True,
+#     revision="transformers-v4",  # Use this revision
+# ).to(device=device)
 
 image1 = get_image("https://huggingface.co/SakanaAI/Metom/resolve/main/examples/example1_4E00.jpg")  # An example image
 image_array1 = processor(images=image1, return_tensors="pt")["pixel_values"].to(device=device, dtype=torch_dtype)
@@ -70,7 +78,7 @@ with torch.inference_mode():
     print(model.get_topk_labels(image_array2))  # Returns top-k prediction labels (label only)
     # [['定', '芝', '乏', '淀', '実'], ['倉', '衾', '斜', '会', '急']]
     print(model.get_topk_labels(image_array2, k=3, return_probs=True))  # Returns prediction top-k labels (label with probability)
-    # [[('定', 0.9979104399681091), ('芝', 0.0002953427319880575), ('乏', 0.00012814522779081017)], [('倉', 0.9862521290779114), ('衾', 0.0005956474924460053), ('斜', 0.00039981433656066656)]]
+    # [[('定', 0.9979110360145569), ('芝', 0.0002953446237370372), ('乏', 0.0001281465229112655)], [('倉', 0.9862518906593323), ('衾', 0.0005956498789601028), ('斜', 0.000399815384298563)]]
 ```
 
 ## Citation
@@ -98,7 +106,7 @@ with torch.inference_mode():
 [Google Colab Notebook](https://colab.research.google.com/drive/1jFMZENoTjjum3qlBxV0Q5dTxmpCvqlpf?usp=sharing)もご確認ください。
 1. 依存ライブラリをインストールする (Google Colabを使う場合は不要)
 ```sh
-python -m pip install einops torch torchvision transformers
+python -m pip install einops torch torchvision "transformers>=5.1.0"
 
 # 任意 (FlashAttention-2を使いたい場合はGoogle Colabを使う時でも必要)
 pip install flash-attn --no-build-isolation
@@ -111,7 +119,7 @@ from io import BytesIO
 from PIL import Image
 import requests
 import torch
-from transformers import AutoModel, AutoProcessor
+from transformers import AutoImageProcessor, AutoModel
 
 repo_name = "SakanaAI/Metom"
 device = "cuda"
@@ -120,13 +128,21 @@ torch_dtype = torch.float32  # `torch.float16` や `torch.bfloat16` も指定可
 def get_image(image_url: str) -> Image.Image:
     return Image.open(BytesIO(requests.get(image_url).content)).convert("RGB")
 
-processor = AutoProcessor.from_pretrained(repo_name, trust_remote_code=True)
+processor = AutoImageProcessor.from_pretrained(repo_name)
 model = AutoModel.from_pretrained(
     repo_name,
-    torch_dtype=torch_dtype,
-    _attn_implementation="eager",  # `"sdpa"` や `"flash_attention_2"` も指定可能
+    dtype=torch_dtype,
+    attn_implementation="sdpa",  # `"eager"`, `"flash_attention_2"` および transformers v5 でサポートされている Attention backends を指定可能 (https://huggingface.co/docs/transformers/main/en/attention_interface)
     trust_remote_code=True
 ).to(device=device)
+# transformers v4 もサポート
+# model = AutoModel.from_pretrained(
+#     repo_name,
+#     torch_dtype=torch_dtype,
+#     _attn_implementation="sdpa",  # `"eager"` や `"flash_attention_2"` も指定可能
+#     trust_remote_code=True,
+#     revision="transformers-v4",  # この revision を使用
+# ).to(device=device)
 
 image1 = get_image("https://huggingface.co/SakanaAI/Metom/resolve/main/examples/example1_4E00.jpg")  # 画像例
 image_array1 = processor(images=image1, return_tensors="pt")["pixel_values"].to(device=device, dtype=torch_dtype)
@@ -141,7 +157,7 @@ with torch.inference_mode():
     print(model.get_topk_labels(image_array2))  # 上位k件の予測ラベルを返す (ラベルのみ)
     # [['定', '芝', '乏', '淀', '実'], ['倉', '衾', '斜', '会', '急']]
     print(model.get_topk_labels(image_array2, k=3, return_probs=True))  # 上位k件の予測ラベルを返す (ラベルと確率)
-    # [[('定', 0.9979104399681091), ('芝', 0.0002953427319880575), ('乏', 0.00012814522779081017)], [('倉', 0.9862521290779114), ('衾', 0.0005956474924460053), ('斜', 0.00039981433656066656)]]
+    # [[('定', 0.9979110360145569), ('芝', 0.0002953446237370372), ('乏', 0.0001281465229112655)], [('倉', 0.9862518906593323), ('衾', 0.0005956498789601028), ('斜', 0.000399815384298563)]]
 ```
 
 ## 引用

config.json

@@ -2720,5 +2720,5 @@
   "model_type": "metom",
   "patch_size": 16,
   "pool": "cls",
-  "transformers_version": "4.46.2"
+  "transformers_version": "5.1.0"
 }

modeling_metom.py

@@ -1,27 +1,49 @@
 # This file is a modified version of the Vision Transformer - Pytorch implementation
 # https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
-from typing import List, Union, Tuple
+from collections.abc import Callable
 
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 import torch
 from torch import nn
 from transformers import PreTrainedModel
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.utils import logging
 
 from .configuration_metom import MetomConfig
 
-
-try:
-    from flash_attn import flash_attn_func
-    FLASH_ATTENTION_2_AVAILABLE = True
-except ImportError:
-    FLASH_ATTENTION_2_AVAILABLE = False
+logger = logging.get_logger(__name__)
 
 
 def size_pair(t):
     return t if isinstance(t, tuple) else (t, t)
 
 
+def metom_eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float | None = None,
+    dropout: float = 0.0,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if scaling is None:
+        scaling = query.size(-1) ** -0.5
+
+    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+    if attention_mask is not None:
+        attention_mask = attention_mask[:, :, :, : key.shape[-2]]
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
 class MetomFeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout):
         super().__init__()
@@ -31,7 +53,7 @@ class MetomFeedForward(nn.Module):
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(hidden_dim, dim),
-            nn.Dropout(dropout)
+            nn.Dropout(dropout),
         )
 
     def forward(self, x):
@@ -39,98 +61,99 @@ class MetomFeedForward(nn.Module):
 
 
 class MetomAttention(nn.Module):
-    def __init__(self, dim, heads, dim_head, dropout):
+    def __init__(self, config: MetomConfig):
         super().__init__()
-        inner_dim = dim_head *  heads
-        project_out = not (heads == 1 and dim_head == dim)
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-        self.norm = nn.LayerNorm(dim)
-        self.attend = nn.Softmax(dim = -1)
-        self.dropout = nn.Dropout(dropout)
-        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        inner_dim = config.dim_head * config.heads
+        project_out = not (config.heads == 1 and config.dim_head == config.dim)
+
+        self.config = config
+        self.heads = config.heads
+        self.scale = config.dim_head ** -0.5
+        self.norm = nn.LayerNorm(config.dim)
+        self.dropout = nn.Dropout(config.dropout)
+        self.to_qkv = nn.Linear(config.dim, inner_dim * 3, bias=False)
         self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, dim),
-            nn.Dropout(dropout)
+            nn.Linear(inner_dim, config.dim),
+            nn.Dropout(config.dropout),
         ) if project_out else nn.Identity()
 
-    def forward(self, x):
-        x = self.norm(x)
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h = self.heads), qkv)
-        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
-        out = torch.matmul(attn, v)
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class MetomSdpaAttention(MetomAttention):
-    def forward(self, x):
-        x = self.norm(x)
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h = self.heads), qkv)
-        out = nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout.p if self.training else 0.0)
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class MetomFlashAttention2(MetomAttention):
-    def forward(self, x):
+    def forward(self, x: torch.Tensor, **kwargs):
         x = self.norm(x)
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h = self.heads), qkv)
-        out = flash_attn_func(q, k, v, dropout_p=self.dropout.p if self.training else 0.0)
-        out = rearrange(out, "b h n d -> b n (h d)")
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv)
+
+        attn_implementation = self.config._attn_implementation or "eager"
+        if attn_implementation == "flex_attention":
+            if self.training and self.dropout.p > 0:
+                logger.warning_once(
+                    "`flex_attention` does not support attention dropout during training. Falling back to `eager`."
+                )
+                attn_implementation = "eager"
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            attn_implementation,
+            metom_eager_attention_forward,
+        )
+        out, _ = attention_interface(
+            self,
+            q,
+            k,
+            v,
+            None,
+            is_causal=False,
+            scaling=self.scale,
+            dropout=0.0 if not self.training else self.dropout.p,
+            **kwargs,
+        )
+        out = rearrange(out, "b n h d -> b n (h d)")
         return self.to_out(out)
 
 
 class MetomTransformer(nn.Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout, _attn_implementation = "eager"):
+    def __init__(self, config: MetomConfig):
         super().__init__()
-        if _attn_implementation == "flash_attention_2":
-            assert FLASH_ATTENTION_2_AVAILABLE, "FlashAttention-2 is not available. Please install `flash-attn`."
-        attn_cls = (
-            MetomAttention if _attn_implementation == "eager" else
-            MetomSdpaAttention if _attn_implementation == "sdpa" else
-            MetomFlashAttention2 if _attn_implementation == "flash_attention_2" else
-            MetomAttention
-        )
-        self.norm = nn.LayerNorm(dim)
+        self.norm = nn.LayerNorm(config.dim)
         self.layers = nn.ModuleList([])
-        for _ in range(depth):
-            self.layers.append(nn.ModuleList([
-                attn_cls(dim, heads = heads, dim_head = dim_head, dropout = dropout),
-                MetomFeedForward(dim, mlp_dim, dropout = dropout)
-            ]))
-
-    def forward(self, x):
+        for _ in range(config.depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        MetomAttention(config),
+                        MetomFeedForward(config.dim, config.mlp_dim, dropout=config.dropout),
+                    ]
+                )
+            )
+
+    def forward(self, x: torch.Tensor, **kwargs):
         for attn, ff in self.layers:
-            x = attn(x) + x
+            x = attn(x, **kwargs) + x
             x = ff(x) + x
         return self.norm(x)
 
 
 class MetomModel(PreTrainedModel):
     config_class = MetomConfig
+    main_input_name = "pixel_values"
+    _supports_attention_backend = True
+    _supports_flash_attn = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
+    _supports_flex_attn = True
 
     def __init__(self, config: MetomConfig):
         super().__init__(config)
         image_height, image_width = size_pair(config.image_size)
         patch_height, patch_width = size_pair(config.patch_size)
-        assert image_height % patch_height == 0 and image_width % patch_width == 0, "Image dimensions must be divisible by the patch size."
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, (
+            "Image dimensions must be divisible by the patch size."
+        )
 
         num_patches = (image_height // patch_height) * (image_width // patch_width)
         patch_dim = config.channels * patch_height * patch_width
         assert config.pool in {"cls", "mean"}, "pool type must be either cls (cls token) or mean (mean pooling)"
         assert len(config.labels) > 0, "labels must be composed of at least one label"
-        assert config._attn_implementation in {"eager", "sdpa", "flash_attention_2"}, "Attention implementation must be either eager, sdpa or flash_attention_2"
 
         self.to_patch_embedding = nn.Sequential(
-            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=patch_height, p2=patch_width),
             nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, config.dim),
             nn.LayerNorm(config.dim),
@@ -138,36 +161,35 @@ class MetomModel(PreTrainedModel):
         self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, config.dim))
         self.cls_token = nn.Parameter(torch.randn(1, 1, config.dim))
         self.dropout = nn.Dropout(config.emb_dropout)
-        self.transformer = MetomTransformer(
-            config.dim, config.depth, config.heads, config.dim_head, config.mlp_dim, config.dropout, config._attn_implementation
-        )
+        self.transformer = MetomTransformer(config)
         self.pool = config.pool
         self.to_latent = nn.Identity()
         self.mlp_head = nn.Linear(config.dim, len(config.labels))
         self.labels = config.labels
+        self.post_init()
 
-    def forward(self, processed_image):
-        x = self.to_patch_embedding(processed_image)
+    def forward(self, pixel_values: torch.Tensor, **kwargs):
+        x = self.to_patch_embedding(pixel_values)
         b, n, _ = x.shape
-        cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 d", b = b)
+        cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 d", b=b)
         x = torch.cat((cls_tokens, x), dim=1)
-        x += self.pos_embedding[:, :(n + 1)]
+        x += self.pos_embedding[:, : (n + 1)]
         x = self.dropout(x)
-        x = self.transformer(x)
-        x = x.mean(dim = 1) if self.pool == "mean" else x[:, 0]
+        x = self.transformer(x, **kwargs)
+        x = x.mean(dim=1) if self.pool == "mean" else x[:, 0]
         x = self.to_latent(x)
         return self.mlp_head(x)
 
-    def get_predictions(self, processed_image: torch.Tensor) -> List[str]:
-        logits = self(processed_image)
+    def get_predictions(self, pixel_values: torch.Tensor) -> list[str]:
+        logits = self(pixel_values=pixel_values)
         indices = torch.argmax(logits, dim=-1)
         return [self.labels[i] for i in indices]
 
     def get_topk_labels(
-        self, processed_image: torch.Tensor, k: int = 5, return_probs: bool = False
-    ) -> Union[List[List[str]], List[List[Tuple[str, float]]]]:
+        self, pixel_values: torch.Tensor, k: int = 5, return_probs: bool = False
+    ) -> list[list[str]] | list[list[tuple[str, float]]]:
         assert 0 < k <= len(self.labels), "k must be a positive integer less than or equal to the number of labels"
-        logits = self(processed_image)
+        logits = self(pixel_values=pixel_values)
         probs = torch.softmax(logits, dim=-1)
         topk_probs, topk_indices = torch.topk(probs, k, dim=-1)
         topk_labels = [[self.labels[i] for i in ti] for ti in topk_indices]