Fix timestamp jumptimes indexOutofBoundsError based on Issue in origi…

…nal repo huggingface#12
erijung · Nov 27, 2024 · acd1651 · acd1651
1 parent 8eb900e
commit acd1651
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/transformers/models/whisper/generation_whisper.py b/src/transformers/models/whisper/generation_whisper.py
@@ -313,7 +313,11 @@ def _extract_token_timestamps(self, generate_outputs, alignment_heads, time_prec
                 # since jump_times is only extended in the next step, take index_ and not (index_ + 1)
                 next_element = jump_times[index_] if index_ < len(jump_times) else round((time_indices[-1] * time_precision), 4)
                 jump_times.insert(index_, next_element)
-            timestamps[batch_idx, 1:] = torch.tensor(jump_times)
+            # print("SURPRISE:", timestamps.shape, batch_idx, torch.tensor(jump_times).shape)
+            if timestamps.shape[-1] == len(jump_times):
+                timestamps[batch_idx, 0:] = torch.tensor(jump_times)
+            else:
+                timestamps[batch_idx, 1:] = torch.tensor(jump_times)
 
         return timestamps