VLMs: fix number of image tokens

src/transformers/models/chameleon/modeling_chameleon.py

@@ -1288,7 +1288,7 @@ def forward(
         if pixel_values is not None:
             image_tokens = self.get_image_tokens(pixel_values)
             n_image_tokens_in_text = (input_ids == self.vocabulary_mapping.image_token_id).sum().item()
-            n_image_features = image_tokens.shape[0]
+            n_image_features = image_tokens.shape[0] * image_tokens.shape[1]
             if n_image_tokens_in_text != n_image_features:
                 raise ValueError(
                     f"Image features and image tokens do not match: tokens: {n_image_tokens_in_text}, features {n_image_features}"

src/transformers/models/glm/modeling_glm.py

@@ -25,7 +25,6 @@
 import torch
 import torch.nn as nn
 import torch.utils.checkpoint
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache, StaticCache
@@ -921,6 +920,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
         device: torch.device,
         cache_position: torch.Tensor,
         batch_size: int,
+        **kwargs,
     ):
         """
         Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
@@ -1071,18 +1071,7 @@ def forward(
 
         loss = None
         if labels is not None:
-            # Upcast to float if we need to compute the loss to avoid potential precision issues
-            logits = logits.float()
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-            # Enable model parallelism
-            shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels)
+            loss = self.loss_function(logits, labels, self.vocab_size)
 
         if not return_dict:
             output = (logits,) + outputs[1:]
@@ -1186,27 +1175,8 @@ def forward(
 
         loss = None
         if labels is not None:
-            labels = labels.to(logits.device)
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = "regression"
-                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
-                    self.config.problem_type = "single_label_classification"
-                else:
-                    self.config.problem_type = "multi_label_classification"
-
-            if self.config.problem_type == "regression":
-                loss_fct = MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(pooled_logits, labels)
-            elif self.config.problem_type == "single_label_classification":
-                loss_fct = CrossEntropyLoss()
-                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
-            elif self.config.problem_type == "multi_label_classification":
-                loss_fct = BCEWithLogitsLoss()
-                loss = loss_fct(pooled_logits, labels)
+            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
         if not return_dict:
             output = (pooled_logits,) + transformer_outputs[1:]
             return ((loss,) + output) if loss is not None else output
@@ -1289,8 +1259,7 @@ def forward(
 
         loss = None
         if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            loss = self.loss_function(logits, labels, self.config)
 
         if not return_dict:
             output = (logits,) + outputs[2:]

src/transformers/models/llava/modeling_llava.py

@@ -518,8 +518,9 @@ def forward(
 
             # TODO: @raushan retain only the new behavior after v4.47
             else:
-                n_image_tokens = (input_ids == self.config.image_token_index).sum(dim=-1)[0].item()
-                n_image_features = image_features.shape[1]
+                n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
+                n_image_features = image_features.shape[0] * image_features.shape[1]
+
                 if n_image_tokens != n_image_features:
                     raise ValueError(
                         f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"

src/transformers/models/llava_next/modeling_llava_next.py

@@ -897,6 +897,7 @@ def forward(
             else:
                 n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
                 n_image_features = image_features.shape[0]
+
                 if n_image_tokens != n_image_features:
                     raise ValueError(
                         f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"

src/transformers/models/llava_next_video/modeling_llava_next_video.py

@@ -969,6 +969,7 @@ def forward(
             if image_features is not None:
                 n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
                 n_image_features = image_features.shape[0]
+
                 if n_image_tokens != n_image_features:
                     raise ValueError(
                         f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"

src/transformers/models/llava_next_video/modular_llava_next_video.py

@@ -484,6 +484,7 @@ def forward(
             if image_features is not None:
                 n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
                 n_image_features = image_features.shape[0]
+
                 if n_image_tokens != n_image_features:
                     raise ValueError(
                         f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"

src/transformers/models/llava_onevision/modeling_llava_onevision.py

@@ -621,6 +621,7 @@ def forward(
             )
             n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
             n_image_features = image_features.shape[0]
+
             if n_image_tokens != n_image_features:
                 raise ValueError(
                     f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
@@ -652,6 +653,7 @@ def forward(
             image_newline = self.image_newline[None, None, :].repeat(batch_size, 1, 1).to(video_features.device)
             video_features = torch.cat((video_features, image_newline), dim=1)
             video_features = video_features.flatten(0, 1)
+
             n_video_tokens = (input_ids == self.config.video_token_index).sum().item()
             n_video_features = video_features.shape[0]
             if n_video_tokens != n_video_features:

src/transformers/models/video_llava/modeling_video_llava.py

@@ -618,8 +618,9 @@ def forward(
             # TODO: @raushan retain only the new behavior after v4.47
             else:
                 if image_outputs is not None:
-                    n_image_tokens = (input_ids == self.config.image_token_index).sum(dim=-1)[0].item()
-                    n_image_features = image_features.shape[1]
+                    n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
+                    n_image_features = image_features.shape[0] * image_features.shape[1]
+
                     if n_image_tokens != n_image_features:
                         raise ValueError(
                             f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
@@ -633,8 +634,9 @@ def forward(
                     image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
                     inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
                 if video_outputs is not None:
-                    n_video_tokens = (input_ids == self.config.video_token_index).sum(dim=-1)[0].item()
-                    n_video_features = video_features.shape[1]
+                    n_video_tokens = (input_ids == self.config.video_token_index).sum().item()
+                    n_video_features = video_features.shape[0] * video_features.shape[1]
+
                     if n_video_tokens != n_video_features:
                         raise ValueError(
                             f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"

src/transformers/models/vipllava/modeling_vipllava.py

@@ -511,8 +511,9 @@ def forward(
 
             # TODO: @raushan retain only the new behavior after v4.47
             else:
-                n_image_tokens = (input_ids == self.config.image_token_index).sum(dim=-1)[0].item()
-                n_image_features = image_features.shape[1]
+                n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
+                n_image_features = image_features.shape[0] * image_features.shape[1]
+
                 if n_image_tokens != n_image_features:
                     raise ValueError(
                         f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"

tests/models/video_llava/test_modeling_video_llava.py

@@ -116,9 +116,9 @@ def __init__(
         self.batch_size = 5
         self.num_channels = 3
         self.image_size = 224
-        self.encoder_seq_length = 64
+        self.encoder_seq_length = 246
         self.num_image_tokens = 25
-        self.num_video_tokens = 26
+        self.num_video_tokens = 26 * self.num_frames
         self.seq_length = seq_length + self.num_image_tokens + self.num_video_tokens
 
     def get_config(self):