CI / generate: batch size computation compatible with all models (#29671

)

src/transformers/generation/utils.py

@@ -1949,11 +1949,9 @@ def _contrastive_search(
             )
 
         # keep track of which sequences are already finished
-        batch_size, cur_len = (
-            model_kwargs["attention_mask"].shape
-            if model_kwargs.get("attention_mask", None) is not None
-            else input_ids.shape
-        )
+        batch_size, cur_len = input_ids.shape
+        if "inputs_embeds" in model_kwargs:
+            cur_len = model_kwargs["inputs_embeds"].shape[1]
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
 
@@ -2398,12 +2396,10 @@ def _greedy_search(
             )
 
         # keep track of which sequences are already finished
+        batch_size, cur_len = input_ids.shape
+        if "inputs_embeds" in model_kwargs:
+            cur_len = model_kwargs["inputs_embeds"].shape[1]
         this_peer_finished = False
-        batch_size, cur_len = (
-            model_kwargs["attention_mask"].shape
-            if model_kwargs.get("attention_mask", None) is not None
-            else input_ids.shape
-        )
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
 
@@ -2686,12 +2682,10 @@ def _sample(
             )
 
         # keep track of which sequences are already finished
+        batch_size, cur_len = input_ids.shape
+        if "inputs_embeds" in model_kwargs:
+            cur_len = model_kwargs["inputs_embeds"].shape[1]
         this_peer_finished = False
-        batch_size, cur_len = (
-            model_kwargs["attention_mask"].shape
-            if model_kwargs.get("attention_mask", None) is not None
-            else input_ids.shape
-        )
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
 
@@ -4461,11 +4455,9 @@ def _assisted_decoding(
             )
 
         # keep track of which sequences are already finished
-        batch_size, cur_len = batch_size, cur_len = (
-            model_kwargs["attention_mask"].shape
-            if model_kwargs.get("attention_mask", None) is not None
-            else input_ids.shape
-        )
+        batch_size, cur_len = input_ids.shape
+        if "inputs_embeds" in model_kwargs:
+            cur_len = model_kwargs["inputs_embeds"].shape[1]
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)